Rework coupled-buffer support: resolve fresh each work! call, support threading

docs/Project.toml

@@ -6,9 +6,9 @@ Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
 Ferrite = "c061ca5d-56c9-439f-9c0e-210fe06d3992"
 FerriteAssembly = "fd21fc07-c509-4fe1-9468-19963fd5935d"
+FerriteIGA = "e7b8d123-e02a-40ba-ad18-d3943ed54f1c"
 FerriteMeshParser = "0f8c756f-80dd-4a75-85c6-b0a5ab9d4620"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
-IGA = "e7b8d123-e02a-40ba-ad18-d3943ed54f1c"
 Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 LiveServer = "16fef848-5104-11e9-1b77-fb7a48bbb589"
 MaterialModelsBase = "af893363-701d-44dc-8b1e-d9a2c129bfc9"
@@ -20,7 +20,7 @@ WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
 
 [sources]
 FerriteAssembly = {path = ".."}
-IGA = {url = "https://github.com/lijas/IGA.jl"}
-MaterialModelsBase = {url = "https://github.com/KnutAM/MaterialModelsBase.jl.git"}
-MechanicalMaterialModels = {url = "https://github.com/KnutAM/MechanicalMaterialModels.jl.git"}
-Newton = {url = "https://github.com/KnutAM/Newton.jl.git"}
+FerriteIGA = {rev = "master", url = "https://github.com/Ferrite-FEM/FerriteIGA.jl"}
+MaterialModelsBase = {rev = "main", url = "https://github.com/KnutAM/MaterialModelsBase.jl.git"}
+MechanicalMaterialModels = {rev = "main", url = "https://github.com/KnutAM/MechanicalMaterialModels.jl.git"}
+Newton = {rev = "main", url = "https://github.com/KnutAM/Newton.jl.git"}

docs/src/DomainBuffers/Setup.md

@@ -31,9 +31,10 @@ especially in the case of `CoupledSimulations`.
 
 The idea behind the coupled simulation setup is to give access to values from a different simulation
 at the item level. For example, when solving two separate problems in parallel, and using staggered
-iterations. See the [Phase-field fracture tutorial](@ref Phase-field-fracture) for an example. 
+iterations. No separate setup-time linking is needed: pass the other simulation(s) directly via
+`CoupledSimulations` when calling [`work!`](@ref). See the
+[Phase-field fracture tutorial](@ref Phase-field-fracture) for an example.
 ```@docs
 Simulation
-couple_buffers
 CoupledSimulations
 ```

docs/src/literate_tutorials/iga.jl

@@ -11,7 +11,7 @@
 # [bottom of this page](@ref iga_plain_program).
 
 # Start by loading the necessary packages
-using Ferrite, IGA, LinearAlgebra, FerriteAssembly
+using Ferrite, FerriteIGA, LinearAlgebra, FerriteAssembly
 import FerriteAssembly.ExampleElements: ElasticPlaneStrain
 
 # ## Setup
@@ -107,7 +107,7 @@ work!(qe, buffer; a=a);
 # ```
 
 # Output results to VTK
-IGA.VTKIGAFile("plate_with_hole.vtu", grid) do vtk
+FerriteIGA.VTKIGAFile("plate_with_hole.vtu", grid) do vtk
     write_solution(vtk, dh, a)
 end
 

docs/src/literate_tutorials/phasefield_fracture.jl

@@ -146,20 +146,23 @@ function setup(m, grid, fieldname; qr_tri, qr_quad, ip_tri, ip_quad)
     return db, K, r, ndofs(dh)
 end
 
-db_u_uc, Ku, ru, ndofs_u = setup(PhaseFieldFracture{:u}(mbase), grid, :u; 
+db_u, Ku, ru, ndofs_u = setup(PhaseFieldFracture{:u}(mbase), grid, :u;
     qr_tri, qr_quad, 
     ip_tri = Lagrange{RefTriangle, 1}()^2, 
     ip_quad = Lagrange{RefQuadrilateral, 1}()^2
     )
 
-db_d_uc, Kd, rd, ndofs_d = setup(PhaseFieldFracture{:d}(mbase), grid, :d; 
+db_d, Kd, rd, ndofs_d = setup(PhaseFieldFracture{:d}(mbase), grid, :d;
     qr_tri, qr_quad,
     ip_tri = Lagrange{RefTriangle, 2}(), 
     ip_quad = Lagrange{RefQuadrilateral, 2}()
     )
 
-sim_u = Simulation(couple_buffers(db_u_uc; d = db_d_uc), zeros(ndofs_u), zeros(ndofs_u))
-sim_d = Simulation(couple_buffers(db_d_uc; u = db_u_uc), zeros(ndofs_d), zeros(ndofs_d));
+# No separate setup-time coupling step is needed: `db_u` and `db_d` are used directly, and the
+# coupling below (`CoupledSimulations(d = sim_d)` / `CoupledSimulations(u = sim_u)`) links them
+# fresh on each `work!` call.
+sim_u = Simulation(db_u, zeros(ndofs_u), zeros(ndofs_u))
+sim_d = Simulation(db_d, zeros(ndofs_d), zeros(ndofs_d));
 
 # Setup loading and boundary conditions
 load_function(t) = 1e-4 * t

src/Autodiff/autodiff.jl

@@ -55,16 +55,42 @@ reinit_buffer!(cb::AutoDiffCellBuffer, args...; kwargs...) = reinit_buffer!(cb.c
 set_time_increment!(c::AutoDiffCellBuffer, Δt) = set_time_increment!(c.cb, Δt)
 
 function _replace_material_with(ad_cb::AutoDiffCellBuffer{CB}, new_material) where CB
-    cb = setproperties(ad_cb.cb; material = new_material)
+    # A fresh `coupled_cache`: reusing `ad_cb.cb`'s would let a stale wrapper - built with the old
+    # material - be returned from `couple_itembuffers` on the new buffer.
+    cb = setproperties(ad_cb.cb; material = new_material, coupled_cache = Ref{Any}(nothing))
     if isa(cb, CB) # If type didn't change, no need to recalculate autodiff buffers
         return setproperties(ad_cb; cb)
     else
         return AutoDiffCellBuffer(cb)
     end
 end
 
-function couple_buffers(cb::AutoDiffCellBuffer; kwargs...)
-    return AutoDiffCellBuffer(couple_buffers(cb.cb; kwargs...))
+"""
+    couple_itembuffers(ad_cb::AutoDiffCellBuffer, coupled::NamedTuple)
+
+Link `ad_cb` to `coupled`. `ad_cb.cb`'s `coupled_buffers` is one of `CellBuffer`'s type
+parameters, so if `coupled` differs in structure from what `ad_cb.cb` currently holds, the whole
+`AutoDiffCellBuffer` - including a new `ForwardDiff.JacobianConfig` - must be rebuilt (the same
+"rebuild only if the type changed" contract [`_replace_material_with`](@ref) already uses).
+
+`ad_cb` is always the domain's own persistent buffer (the same object every `work!` call, never
+itself replaced by this function - see [`couple_itembuffers(::CellBuffer, ...)`](@ref) for why),
+so `ad_cb.cb`'s `coupled_cache` field is used to cache the *result* of this function: repeated
+calls linking to the same buffer objects (the common case, across staggered iterations reusing
+the same partner) return the previously-built `AutoDiffCellBuffer` - JacobianConfig included -
+instead of rebuilding it every `work!` call.
+"""
+function couple_itembuffers(ad_cb::AutoDiffCellBuffer{CB}, coupled::NT) where {T,CC,CV,DR,MT,ST,UD,UC,CB<:CellBuffer{T,CC,CV,DR,MT,ST,UD,UC},NT<:NamedTuple}
+    get_coupled_buffers(ad_cb.cb) === coupled && return ad_cb
+    cache = ad_cb.cb.coupled_cache
+    cached = cache[]
+    if cached !== nothing && cached[1] === coupled
+        return cached[2]::AutoDiffCellBuffer{<:CellBuffer{T,CC,CV,DR,MT,ST,UD,UC,NT}}
+    end
+    cb = setproperties(ad_cb.cb; coupled_buffers = coupled)::CellBuffer{T,CC,CV,DR,MT,ST,UD,UC,NT}
+    result = isa(cb, CB) ? setproperties(ad_cb; cb) : AutoDiffCellBuffer(cb)
+    cache[] = (coupled, result)
+    return result
 end
 
 function create_local(c::AutoDiffCellBuffer)

src/DomainBuffers.jl

@@ -153,25 +153,6 @@ function replace_material(dbs::DomainBuffers, replacement_function)
     return Dict(key=>replace_material(db, replacement_function) for (key,db) in dbs)
 end
 
-"""
-    couple_buffers(dbs::Dict{String, <:AbstractDomainBuffer}; kwargs::Dict{String, <:AbstractDomainBuffer}...)
-    couple_buffers(db::AbstractDomainBuffer; kwargs::AbstractDomainBuffer...)
-
-Return new buffer(s) that are coupled with the buffers provided as keyword arguments. The key is used in 
-[`get_coupled_buffer`](@ref) to get the coupled itembuffer, such that its values may be queried. 
-
-!!! note
-    This functionality assumes that each setup has the same grid, and in case of multiple domains, these should also 
-    match.
-"""
-function couple_buffers(dbs::DomainBuffers; kwargs...)
-    return Dict(
-        key => (all(haskey(v, key) for (_, v) in kwargs) ? 
-            couple_buffers(db; (k => v[key] for (k, v) in kwargs)...) :
-            db) for (key, db) in dbs)
-    #return Dict(key => couple_buffers(db; (k => v[key] for (k, v) in kwargs)...) for (key, db) in dbs)
-end
-
 """
     getset(dbs::Dict{String,AbstractDomainBuffer}, domain::String)
     getset(db::AbstractDomainBuffer)
@@ -205,6 +186,7 @@ function ThreadedDomainBuffer(set, itembuffer::AbstractItemBuffer, states::State
 end
 
 get_num_tasks(db::ThreadedDomainBuffer) = db.num_tasks
+get_num_tasks(::DomainBuffer) = 1 # A single, shared buffer - conceptually one "task".
 get_num_tasks(dbs::DomainBuffers) = maximum(get_num_tasks, values(dbs))
 
 get_chunks(db::ThreadedDomainBuffer) = db.chunks
@@ -243,19 +225,6 @@ function replace_material(db::ThreadedDomainBuffer, replacement_function)
     return setproperties(db; itembuffer = TaskLocals(base_ibuf, task_ibuf))
 end
 
-function couple_buffers(db::DomainBuffer; kwargs...)
-    itembuffer = couple_buffers(db.itembuffer; (k => v.itembuffer for (k, v) in kwargs)...)
-    return setproperties(db; itembuffer)
-end
-
-function couple_buffers(db::ThreadedDomainBuffer; kwargs...)
-    base_ibuf = couple_buffers(get_base(db.itembuffer); (k => get_base(v.itembuffer) for (k, v) in kwargs)...)
-    task_ibuf = map(enumerate(get_locals(db.itembuffer))) do (i, ibuf)
-        couple_buffers(ibuf; (k => get_local(v.itembuffer, i) for (k, v) in kwargs)...)
-    end
-    return setproperties(db; itembuffer = TaskLocals(base_ibuf, task_ibuf))
-end
-
 # Experimental: Insert new states, allows reusing the buffer for multiple simulations with same 
 # initial state (grid, dh, etc.), but which experience different loading. Typically for RVE simulations. 
 function replace_states!(dbs::Dict{String, <:AbstractDomainBuffer}, states::Dict{String, <:StateVariables})

src/FerriteAssembly.jl

@@ -29,7 +29,7 @@ include("LoadHandler/LoadHandler.jl")
 
 # Setup 
 export DomainSpec, setup_domainbuffer, setup_domainbuffers
-export Simulation, CoupledSimulations, couple_buffers
+export Simulation, CoupledSimulations
 # Main functions to use during simulations
 export work!, update_states!, revert_states!, set_time_increment!
 # Workers

src/ItemBuffers/AbstractItemBuffer.jl

@@ -58,12 +58,32 @@ function get_user_cache end
 """
     get_coupled_buffer(b::AbstractItemBuffer, key::Symbol)
 
-Get the coupled buffer `key` from `b`. To enable this, use [`couple_buffers`](@ref) on the 
-domain buffers. The coupled buffer can be queried just like a normal item buffer,
+Get the coupled buffer `key` from `b`. To enable this, supply a `coupled_simulations` to
+[`work!`](@ref); the coupled itembuffer can be queried just like a normal item buffer,
 e.g. by calling `get_state(coupled_buffer)`.
-""" 
+"""
 @inline get_coupled_buffer(b::AbstractItemBuffer, key::Symbol) = getfield(get_coupled_buffers(b), key)
 
+get_coupled_buffers(::AbstractItemBuffer) = NamedTuple() # Default: buffer types that don't support coupling
+
+"""
+    couple_itembuffers(itembuffer::AbstractItemBuffer, coupled::NamedTuple)
+
+Return an `itembuffer`-like buffer linked to `coupled` (a `NamedTuple` of the actual buffer
+objects to link to, e.g. from [`work_domain_sequential!`](@ref)/[`work_domain_threaded!`](@ref)).
+Buffer types that support coupling (e.g. `CellBuffer`) override this to skip reconstruction
+entirely when nothing has changed since last time - the common case across repeated `work!` calls
+reusing the same partner buffers - so steady-state coupled work costs nothing extra. This default
+is for buffer types that don't support coupling at all (e.g. `FacetBuffer`): a no-op when
+`coupled` is empty, or an error if actually asked to couple.
+"""
+function couple_itembuffers(itembuffer::AbstractItemBuffer, coupled::NamedTuple)
+    get_coupled_buffers(itembuffer) === coupled && return itembuffer
+    isempty(coupled) && return itembuffer
+    throw(ArgumentError("$(typeof(itembuffer)) does not support coupled simulations"))
+end
+
+
 """
     Ferrite.celldofs(::AbstractItemBuffer)
 

src/ItemBuffers/CellBuffer.jl

@@ -12,7 +12,7 @@ function work_single_cell! end
 mutable struct CellBuffer{T,CC,CV,DR,MT,ST,UD,UC,CB} <: AbstractCellBuffer
     const ae_old::Vector{T}           # Old element dof values
     const ae::Vector{T}               # Current element dof values
-    const re::Vector{T}               # Residual/force vector 
+    const re::Vector{T}               # Residual/force vector
     const Ke::Matrix{T}               # Element stiffness matrix
     const dofs::Vector{Int}           # celldofs
     const coords::CC                  # cellcoords (or what is required to reinit cellvalues)
@@ -26,7 +26,11 @@ mutable struct CellBuffer{T,CC,CV,DR,MT,ST,UD,UC,CB} <: AbstractCellBuffer
     old_state::ST                     # Old state variables for the cell (updated in reinit!)
     const user_data::UD               # User data for the cell (used for additional information)
     const user_cache::UC              # Cache for the cell (user type) (deepcopy for each thread)
-    const coupled_buffers::CB         # nothing or NamedTuple with staggered coupled `CellBuffer`s. 
+    const coupled_buffers::CB         # NamedTuple with coupled `CellBuffer`s, empty if not coupled.
+    # Cache of the last (coupled_buffers key, couple_itembuffers result) pair, so that repeated
+    # `work!` calls linking to the *same* coupled buffer objects (the common case) skip
+    # reconstruction entirely instead of paying for it every call - see `couple_itembuffers`.
+    const coupled_cache::Base.RefValue{Any}
 end
 
 """
@@ -49,9 +53,10 @@ function CellBuffer(numdofs::Int, coords, cellvalues, material, state, dofrange,
     cellid = -1
     cache = allocate_cell_cache(material, cellvalues)
     return CellBuffer(
-        zeros(numdofs), zeros(numdofs), zeros(numdofs), zeros(numdofs,numdofs), 
-        zeros(Int, numdofs), coords, 
-        cellvalues, Δt, cellid, dofrange, material, state, state, user_data, cache, nothing)
+        zeros(numdofs), zeros(numdofs), zeros(numdofs), zeros(numdofs,numdofs),
+        zeros(Int, numdofs), coords,
+        cellvalues, Δt, cellid, dofrange, material, state, state, user_data, cache, NamedTuple(),
+        Ref{Any}(nothing))
 end
 
 setup_cellbuffer(ad::Bool, args...; kwargs...) = setup_cellbuffer(Val(ad), args...; kwargs...)
@@ -61,10 +66,6 @@ function setup_cellbuffer(::Val{false}, sdh, cv, material, cell_state, dofrange,
     return CellBuffer(numdofs, coords, cv, material, cell_state, dofrange, user_data)
 end
 
-function couple_buffers(cb::CellBuffer; kwargs...)
-    return setproperties(cb; coupled_buffers = NamedTuple{keys(kwargs)}(values(kwargs)))
-end
-
 function setup_cellbuffer(::Val{true}, args...)
     return AutoDiffCellBuffer(setup_cellbuffer(Val(false), args...))
 end
@@ -73,7 +74,7 @@ end
 # TaskLocals interface (only `create_local` required for other `AbstractCellBuffer`s) (unless gather! is req.)
 function create_local(cb::CellBuffer)
     dcpy = map(deepcopy, (cb.ae_old, cb.ae, cb.re, cb.Ke, cb.dofs, cb.coords, cb.cellvalues, cb.Δt, cb.cellid, cb.dofrange, cb.material, cb.state, cb.old_state))
-    return CellBuffer(dcpy..., cb.user_data, deepcopy(cb.user_cache), create_local(cb.coupled_buffers))
+    return CellBuffer(dcpy..., cb.user_data, deepcopy(cb.user_cache), NamedTuple(), Ref{Any}(nothing))
 end
 
 set_time_increment!(cb::CellBuffer, Δt) = (cb.Δt=Δt)
@@ -143,17 +144,44 @@ function reinit_buffer!(cb::CellBuffer, sim::Simulation, coupled, cellnum::Int)
     return nothing  # Ferrite's reinit! doesn't return 
 end
 
-# No coupled buffer, no coupled simulation
-reinit_coupled!(::Nothing, coupled::CoupledSimulations{@NamedTuple{}}, cellnum::Int) = nothing
-
 function reinit_coupled!(coupled_buffers::NamedTuple, coupled::CoupledSimulations, cellnum::Int)
     if length(coupled_buffers) != length(coupled.sims)
         throw(ArgumentError("When using coupled simulations, the coupled buffers must match the coupled simulations"))
     end
-    tuple((reinit_buffer!(cb, coupled.sims[k], CoupledSimulations(), cellnum) for (k, cb) in pairs(coupled_buffers))...)
+    for (k, cb) in pairs(coupled_buffers)
+        reinit_buffer!(cb, coupled.sims[k], CoupledSimulations(), cellnum)
+    end
     return nothing
 end
 
 function _replace_material_with(cb::CellBuffer, new_material)
-    return setproperties(cb; material = new_material)
+    # A fresh `coupled_cache`: reusing `cb`'s would let a stale wrapper - built with the old
+    # material - be returned from `couple_itembuffers` on the new buffer.
+    return setproperties(cb; material = new_material, coupled_cache = Ref{Any}(nothing))
 end
+
+"""
+    couple_itembuffers(cb::CellBuffer, coupled_buffers::NamedTuple)
+
+Link `cb` directly to `coupled_buffers` (already the correct buffer objects to link to - e.g.
+partner base buffers for sequential work, or this task's own per-task partner buffers for
+threaded work - no fetching needed here). Since `coupled_buffers` is one of `CellBuffer`'s type
+parameters, actually changing it requires constructing a new `CellBuffer` (all other fields keep
+the same references as `cb`, so this is cheap - no arrays are copied).
+
+`cb` is always the domain's own persistent buffer (the same object every `work!` call), so its
+[`coupled_cache`](@ref) lets repeated calls linking to the *same* buffer objects (the common case,
+across staggered iterations reusing the same partner) return the previously-built result instead
+of reconstructing - which matters far more for [`AutoDiffCellBuffer`](@ref), where reconstruction
+means rebuilding a `ForwardDiff.JacobianConfig`.
+"""
+function couple_itembuffers(cb::CellBuffer{T,CC,CV,DR,MT,ST,UD,UC}, coupled_buffers::NT) where {T,CC,CV,DR,MT,ST,UD,UC,NT<:NamedTuple}
+    cb.coupled_buffers === coupled_buffers && return cb
+    cached = cb.coupled_cache[]
+    if cached !== nothing && cached[1] === coupled_buffers
+        return cached[2]::CellBuffer{T,CC,CV,DR,MT,ST,UD,UC,NT}
+    end
+    result = setproperties(cb; coupled_buffers)::CellBuffer{T,CC,CV,DR,MT,ST,UD,UC,NT}
+    cb.coupled_cache[] = (coupled_buffers, result)
+    return result
+end

src/Multithreading/TaskLocals.jl

@@ -25,6 +25,8 @@ end
 Get the `i`th `local` variable from `tl`
 """
 get_local(tl::TaskLocals, i::Int) = tl.locals[i]
+get_local(nt::NamedTuple, i::Int) = map(tl -> get_local(tl, i), nt) # e.g. for coupled buffers
+get_local(x, ::Int) = x # e.g. a sequential (non-TaskLocals) coupled partner's single buffer
 
 """
     get_locals(tl::TaskLocals)

src/Simulation.jl

@@ -53,16 +53,25 @@ Base.iterate(sim::Simulation{<:DomainBuffers}, iter) = _iterate(sim, iterate(sim
 """
     CoupledSimulations(; key1 = sim1::Simulation, key2 = sim2::Simulation, ...)
 
-Setup the collection of coupled simulations to allow values (such as state variables and 
-local dof-values from these simulations to be available when `work!`ing another simulation, 
-if the buffers have been coupled with [`couple_buffers`](@ref). 
-The coupled itembuffer on the local level is accessed with [`get_coupled_buffer`](@ref). 
+Setup the collection of coupled simulations to allow values (such as state variables and
+local dof-values) from these simulations to be available when `work!`ing another simulation.
+The coupled itembuffer on the local level is accessed with [`get_coupled_buffer`](@ref).
 """
 struct CoupledSimulations{NT <: NamedTuple{<:Any, <:NTuple{<:Any, Simulation}}}
     sims::NT
 end
 CoupledSimulations(; kwargs...) = CoupledSimulations(NamedTuple{keys(kwargs)}(values(kwargs)))
 
+"""
+    get_itembuffer(coupled::CoupledSimulations, num_tasks::Int)
+
+Return a `NamedTuple` (same keys as `coupled.sims`) of each coupled partner's own itembuffer
+(its `TaskLocals`, if threaded, or its single buffer, if sequential)
+"""
+function get_itembuffer(coupled::CoupledSimulations)
+    return map(get_itembuffer, coupled.sims)
+end
+
 function get_domain_simulation(cs::CoupledSimulations, name::String)
     # Need to return a named tuple with only the simulations that have a domain called `name`
     sims = Pair{Symbol, Simulation}[]