From 2befa6c92ee96935d7e5a425456b22392454ca35 Mon Sep 17 00:00:00 2001 From: djukic14 Date: Tue, 29 Sep 2026 10:22:06 +0200 Subject: [PATCH] Reuse assembly scratch buffers to reduce block-assembly allocations Add thread-safe scratch-pool management for `zlocal`, `tadjq`, and quadrature buffers, update block assembly to borrow and return scratch bundles safely. --- src/integralop.jl | 239 ++++++++++++++++++++++++++++++------ src/operator.jl | 2 +- test/runtests.jl | 1 + test/test_blockassembler.jl | 64 ++++++---- 4 files changed, 248 insertions(+), 58 deletions(-) diff --git a/src/integralop.jl b/src/integralop.jl index cfd1478f7..a36756ab8 100644 --- a/src/integralop.jl +++ b/src/integralop.jl @@ -28,6 +28,165 @@ for `I` and `J` permutations of `1:numfunctions(test_space)` and """ function blockassembler end +""" + AssemblyScratch{Z,T,Q} + +Reusable local buffers for one block-assembly operation. + +An `AssemblyScratch` must be owned by at most one assembly task at a time. +""" +struct AssemblyScratch{Z,T,Q} + zlocal::Z + tadjq::T + qbuffer::Q +end + +""" + AbstractAssemblyScratchSource + +Abstract selector for the source of block-assembly scratch buffers. +""" +abstract type AbstractAssemblyScratchSource end + +""" + NewAssemblyScratch() + +Select fresh buffers for an assembly operation. + +`assemblyscratch` allocates a new `AssemblyScratch` when this source is used. +""" +struct NewAssemblyScratch <: AbstractAssemblyScratchSource end + +""" + StoredAssemblyScratch(scratch) + +Select an existing scratch bundle for an assembly operation. + +The referenced bundle is returned without allocation. It must not be used by +another task until the current operation has finished. +""" +struct StoredAssemblyScratch{S} <: AbstractAssemblyScratchSource + scratch::S +end + +""" + AssemblyScratchPool(available, available_condition) + +Thread-safe pool of reusable scratch bundles for block assembly. + +Each bundle is exclusively borrowed by one task. The condition protects the +available-bundle list and lets tasks wait when all bundles are in use. +""" +struct AssemblyScratchPool{S} + available::Vector{S} + available_condition::Threads.Condition +end + +""" + acquirescratch(pool::AssemblyScratchPool) + +Acquire an available scratch bundle from `pool`. + +The call waits if every bundle is currently in use. The caller owns the +returned bundle until it passes it to [`releasescratch!`](@ref). +""" +function acquirescratch(pool::AssemblyScratchPool) + lock(pool.available_condition) + try + while isempty(pool.available) + wait(pool.available_condition) + end + return pop!(pool.available) + finally + unlock(pool.available_condition) + end +end + +""" + releasescratch!(pool::AssemblyScratchPool, scratch) + +Return `scratch` to `pool` and notify one waiting task. + +Every successful call to [`acquirescratch`](@ref) must be paired with one +call to this function, including when assembly exits through an exception. +""" +function releasescratch!(pool::AssemblyScratchPool, scratch) + lock(pool.available_condition) + try + push!(pool.available, scratch) + notify(pool.available_condition) + finally + unlock(pool.available_condition) + end + return nothing +end + +""" + assemblyscratch( + ::NewAssemblyScratch, + biop, + test_space, + trial_space, + test_assembly_data, + trial_assembly_data, + quadstrat, + ) + +Allocate an `AssemblyScratch` for the supplied assembly configuration. + +The local matrix, trial-adjoint workspace, and quadrature buffer are sized +from the assembly data, spaces, and quadrature strategy. +""" +function assemblyscratch( + ::NewAssemblyScratch, + biop, + test_space, + trial_space, + test_assembly_data, + trial_assembly_data, + quadstrat, +) + return AssemblyScratch( + zeros( + scalartype(biop, test_space, trial_space), + size(test_assembly_data.data, 2), + size(trial_assembly_data.data, 2), + ), + Vector{eltype(trial_assembly_data.data)}( + undef, size(trial_assembly_data.data, 1) + ), + quadraturebuffer(quadstrat, test_space, trial_space), + ) +end + +""" + assemblyscratch( + source::StoredAssemblyScratch, + biop, + test_space, + trial_space, + test_assembly_data, + trial_assembly_data, + quadstrat, + ) + +Return the stored scratch bundle without allocating new buffers. + +The caller is responsible for ensuring that the bundle is not used +concurrently by another assembly task. +""" +function assemblyscratch( + source::StoredAssemblyScratch, + biop, + test_space, + trial_space, + test_assembly_data, + trial_assembly_data, + quadstrat, +) + return source.scratch +end + # """ # quadrule(operator,test_refspace,trial_refspace,p,test_element,q_trial_element, qd) @@ -139,9 +298,6 @@ function assemblechunk!(biop::IntegralOperator, tfs::Space, bfs::Space, store; tshapes = refspace(tfs); #num_tshapes = numfunctions(tshapes, tdom) bshapes = refspace(bfs); #num_bshapes = numfunctions(bshapes, bdom) - num_tshapes = size(tad.data, 2) - num_bshapes = size(bad.data, 2) - qs = if CompScienceMeshes.refines(tgeo, bgeo) TestRefinesTrialQStrat(quadstrat) elseif CompScienceMeshes.refines(bgeo, tgeo) @@ -151,7 +307,6 @@ function assemblechunk!(biop::IntegralOperator, tfs::Space, bfs::Space, store; end qd = quaddata(biop, tshapes, bshapes, test_elements, bsis_elements, qs) - zlocal = zeros(scalartype(biop, tfs, bfs), 2num_tshapes, 2num_bshapes) # @show "after" qs # assemblechunk_body!(biop, # tfs, test_elements, tad, tcells, @@ -163,7 +318,7 @@ function assemblechunk!(biop::IntegralOperator, tfs::Space, bfs::Space, store; assemblechunk_body!(biop, tfs, bfs, test_elements, tcells, tad, eachindex(tcells), bsis_elements, bcells, bad, eachindex(bcells), - qd, zlocal, store; quadstrat=qs) + qd, store; quadstrat=qs) end @@ -198,19 +353,28 @@ end function assemblechunk_body!(biop, test_space, trial_space, test_elements, test_element_ptrs, test_assembly_data, active_test_els, trial_elements, trial_element_ptrs, trial_assembly_data, active_trial_els, - qd, zlocal, store; quadstrat, scheduler=:serial) + qd, store; + quadstrat, + scheduler=:serial, + scratchsource=NewAssemblyScratch()) num_tshapes = size(test_assembly_data.data,2) num_bshapes = size(trial_assembly_data.data,2) @tasks for p in eachindex(active_test_els) @set scheduler = scheduler - @local begin - zlocal = zeros(scalartype(biop, test_space, trial_space), num_tshapes, num_bshapes) - tadjq = Vector{eltype(trial_assembly_data.data)}(undef, size(trial_assembly_data.data,1)) - qbuffer = quadraturebuffer(quadstrat, test_space, trial_space) - end - qaction = ApplyIntegrate(qbuffer) + @local scratch = assemblyscratch( + scratchsource, + biop, + test_space, + trial_space, + test_assembly_data, + trial_assembly_data, + quadstrat, + ) + zlocal = scratch.zlocal + tadjq = scratch.tadjq + qaction = ApplyIntegrate(scratch.qbuffer) P = active_test_els[p] tcell = test_elements[P] tptr = test_element_ptrs[P] @@ -288,7 +452,7 @@ struct AssembleblockbodyFunctor{B,T1,T2,T3,T4,T5,T6,T7,T8,T9} trialelements::T5 trialassemblydata::T6 quadraturedata::T7 - zlocals::T8 + scratchpool::T8 quadstrat::T9 end @@ -308,18 +472,25 @@ function (f::AssembleblockbodyFunctor)(testids, trialids, store) test_element_ptrs = eachindex(f.testelements) trial_element_ptrs = eachindex(f.trialelements) - assemblechunk_body!(f.biop, f.tfs, f.bfs, - f.testelements, test_element_ptrs, tad1, active_test_els, - f.trialelements, trial_element_ptrs, bad1, active_trial_els, - f.quadraturedata, nothing, store; - quadstrat=f.quadstrat, scheduler=:serial) + scratch = acquirescratch(f.scratchpool) + try + assemblechunk_body!(f.biop, f.tfs, f.bfs, + f.testelements, test_element_ptrs, tad1, active_test_els, + f.trialelements, trial_element_ptrs, bad1, active_trial_els, + f.quadraturedata, store; + quadstrat=f.quadstrat, + scheduler=:serial, + scratchsource=StoredAssemblyScratch(scratch)) + finally + releasescratch!(f.scratchpool, scratch) + end # assembleblock_body!(f.biop, f.tfs, testids, f.testelements, f.testassemblydata, # f.bfs, trialids, f.trialelements, f.trialassemblydata, # f.quadraturedata, f.zlocals, store; quadstrat=f.quadstrat) end function blockassembler(biop::IntegralOperator, tfs::Space, bfs::Space; - quadstrat=defaultquadstrat(biop, tfs, bfs)) + quadstrat=defaultquadstrat(biop, tfs, bfs), primer=nothing) tgeo = geometry(tfs) bgeo = geometry(bfs) @@ -332,9 +503,13 @@ function blockassembler(biop::IntegralOperator, tfs::Space, bfs::Space; quadstrat end + if primer === nothing + primer = assembleblock_primer(biop, tfs, bfs; quadstrat=qs) + end + test_elements, test_assembly_data, trial_elements, trial_assembly_data, - quadrature_data, zlocals = assembleblock_primer(biop, tfs, bfs; quadstrat=qs) + quadrature_data, scratchpool = primer return AssembleblockbodyFunctor( biop, @@ -345,7 +520,7 @@ function blockassembler(biop::IntegralOperator, tfs::Space, bfs::Space; trial_elements, trial_assembly_data, quadrature_data, - zlocals, + scratchpool, qs, ) @@ -405,24 +580,18 @@ function assembleblock_primer(biop, tfs, bfs; test_elements, tad = assemblydata(tfs; onlyactives=false) bsis_elements, bad = assemblydata(bfs; onlyactives=false) - tgeo = geometry(tfs) - bgeo = geometry(bfs) - - tdom = domain(chart(tgeo, first(tgeo))) - bdom = domain(chart(bgeo, first(bgeo))) - - tshapes = refspace(tfs); num_tshapes = numfunctions(tshapes, tdom) - bshapes = refspace(bfs); num_bshapes = numfunctions(bshapes, bdom) + tshapes = refspace(tfs) + bshapes = refspace(bfs) qd = quaddata(biop, tshapes, bshapes, test_elements, bsis_elements, quadstrat) - zlocals = Channel{Matrix{scalartype(biop, tfs, bfs)}}(2*Threads.nthreads()) - - for _ in 1:2*Threads.nthreads() - put!(zlocals, zeros(scalartype(biop, tfs, bfs), num_tshapes, num_bshapes)) - end + scratches = [ + assemblyscratch(NewAssemblyScratch(), biop, tfs, bfs, tad, bad, quadstrat) + for _ in 1:(2 * Threads.nthreads()) + ] + scratchpool = AssemblyScratchPool(scratches, Threads.Condition()) - return test_elements, tad, bsis_elements, bad, qd, zlocals + return test_elements, tad, bsis_elements, bad, qd, scratchpool end # function assembleblock_body!(biop::IntegralOperator, diff --git a/src/operator.jl b/src/operator.jl index f26dd0fec..0ca1d9ed9 100644 --- a/src/operator.jl +++ b/src/operator.jl @@ -274,7 +274,7 @@ function assemble!(operator::Operator, testfunctions::Space, trialfunctions::Spa assemblechunk_body!(operator, testfunctions, trialfunctions, testelements, eachindex(testelements), testad1, coloredtestelements, trialelements, eachindex(trialelements), trialad, coloredtrialelements, - qdata, nothing, store; quadstrat=qs, scheduler) + qdata, store; quadstrat=qs, scheduler) next!(pbar; step = length(testelementcolors[i]) * length(trialelementcolors[j])) end end finish!(pbar) diff --git a/test/runtests.jl b/test/runtests.jl index feac5ac38..b442b4ea1 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -49,6 +49,7 @@ include("test_local_storage.jl") include("test_embedding.jl") include("test_assemblerow.jl") +include("test_blockassembler.jl") # include("test_mixed_blkassm.jl") include("test_local_assembly.jl") include("test_assemble_refinements.jl") diff --git a/test/test_blockassembler.jl b/test/test_blockassembler.jl index 91fdb4fc1..ab0663cbe 100644 --- a/test/test_blockassembler.jl +++ b/test/test_blockassembler.jl @@ -3,26 +3,46 @@ using CompScienceMeshes using LinearAlgebra using Test -r = 10.0 -λ = 20 * r -k = 2 * π / λ - -sphere = readmesh(joinpath(dirname(@__FILE__),"assets","sphere5.in"), T=Float64) - -D = Maxwell3D.doublelayer(wavenumber=k) -X = raviartthomas(sphere) -Y = buffachristiansen(sphere) - -A = assemble(D, X, X) - -@views blkasm = BEAST.blockassembler(D, X, X) - -@views function assembler(Z, tdata, sdata) - @views store(v,m,n) = (Z[m,n] += v) - blkasm(tdata,sdata,store) +@testset "blockassembler" begin + r = 10.0 + lambda = 20 * r + wavenumber = 2 * pi / lambda + + sphere = readmesh(joinpath(@__DIR__, "assets", "sphere5.in"); T=Float64) + operator = Maxwell3D.doublelayer(; wavenumber) + space = raviartthomas(sphere) + + matrix = assemble(operator, space, space) + blockassembler = BEAST.blockassembler(operator, space, space) + ids = collect(eachindex(space.fns)) + + function assembleblock(assembler, rows=ids, columns=ids) + out = zeros(ComplexF64, length(rows), length(columns)) + store(v, m, n) = (out[m, n] += v) + assembler(rows, columns, store) + return out + end + + expected = assembleblock(blockassembler) + @test matrix ≈ expected atol=eps(Float64) + + poolsize = length(blockassembler.scratchpool.available) + scratchids = Set(objectid.(blockassembler.scratchpool.available)) + @test assembleblock(blockassembler) ≈ expected + @test length(blockassembler.scratchpool.available) == poolsize + @test Set(objectid.(blockassembler.scratchpool.available)) == scratchids + + failingstore(args...) = error("store failed") + @test_throws ErrorException blockassembler( + [first(ids)], [first(ids)], failingstore + ) + @test length(blockassembler.scratchpool.available) == poolsize + + results = fetch.([ + Threads.@spawn assembleblock(blockassembler) for + _ in 1:(2 * Threads.nthreads() + 1) + ]) + @test all(result -> result ≈ expected, results) + @test length(blockassembler.scratchpool.available) == poolsize + @test Set(objectid.(blockassembler.scratchpool.available)) == scratchids end - -A_blk = zeros(ComplexF64, length(X.fns), length(Y.fns)) -assembler(A_blk, [1:length(X.fns);], [1:length(Y.fns);]) - -@test norm(A - A_blk) ≈ 0 atol=eps(Float64)