From 4dc2760e8754234d9ed8ed329ff17991ea8d0883 Mon Sep 17 00:00:00 2001 From: Kyle Beggs Date: Thu, 10 Sep 2026 07:17:45 -0400 Subject: [PATCH 1/3] =?UTF-8?q?perf(amr):=20fuse=20each=20coarse=E2=80=93f?= =?UTF-8?q?ine=20ghost=20fill=20into=20one=20gather=20pass?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The host fill sweeps ran one broadcast per CSR term — 7 per fill in 2D, 19 in 3D — over a handful of dst cells each, so broadcast setup, not arithmetic, dominated the exchange on a refined forest. Evaluate the whole weighted sum per cell inside a single broadcast instead: the K term views are built once per fill by recursion over Val(K) down the row (closure-free, so the two-array walk keeps inlining) and ride an immutable Ref subtype as a broadcast scalar, which a mutable RefValue cannot do without escaping the un-inlined gather body and heap-allocating ~2 KB per 3D fill inside a rule seam that must not allocate. halo_update! on a refined 32³ forest: 853 → 205 µs packed and 751 → 229 µs per-block at 8³ blocks, 3.0 → 0.67 ms at 4³ (min-of-N, alternating processes, three rounds, ±4%). Term and accumulation order are unchanged, so the result is bit-identical to the retired loop whenever the field eltype is the schedule's weight type, and — unlike that loop, which rounded its partial sum into the field K−1 times — bit-identical to the device CSR kernel in mixed precision too. _close_fill now asserts the two row invariants the fusion rests on: every term window has the dst slab's shape, and every term window is disjoint from the dst slab (the views ride a non-AbstractArray scalar, so Base's aliasing machinery never sees them). The adjoint keeps its per-term scatter: a fill's term windows collide on shared source cells, so a dst-centric single pass would reorder the accumulation. The bit-identical fused route is a source-centric transposed CSR, scoped to #94. Claude-Session: https://claude.ai/code/session_016F4h1y22x3ohpRdGVCpjHP --- src/schedule.jl | 15 ++++ src/transfer.jl | 188 ++++++++++++++++++++++++++++++++++++++++-------- 2 files changed, 175 insertions(+), 28 deletions(-) diff --git a/src/schedule.jl b/src/schedule.jl index 16b214f..ba529ff 100644 --- a/src/schedule.jl +++ b/src/schedule.jl @@ -46,6 +46,21 @@ box plus a row index pair; a term is one slab window plus a weight) and the host sweeps walk exactly the layout the device schedule uploads (`_DevFill` / `_DevTerm`). Row widths are dimension-fixed — see [`_ninterp_terms`](@ref) / [`_nrestrict_terms`](@ref) — which the emitters assert at close. + +Two shape invariants hold of every row, asserted at close by `_close_fill` and +relied on by both the host sweep and the device kernel: + + - **Every term window has the dst slab's shape** — cell `I` of the dst reads cell + `I` of every term — which is what lets the host sweep run one fused gather + `dst[I] = Σₖ wₖ · srcₖ[I]` per fill (`_run_fills!`) and the device kernel index + every term by the dst cell's decoded offset (`_fill_kernel!`). + - **Every term window is disjoint from the dst slab.** dst boxes are ghost layers + and terms read interiors, so no emitter produces an overlap; the fused gather + makes it load-bearing, because its term views ride an immutable `Ref` wrapper + that Base's `broadcast_unalias` never sees. + +A mis-shaped or self-reading emitter therefore fails at schedule build, never as an +out-of-bounds read or a silently wrong ghost value in a sweep. """ struct GhostFill{N} dst_block::Int diff --git a/src/transfer.jl b/src/transfer.jl index eeb310d..ea7c4f6 100644 --- a/src/transfer.jl +++ b/src/transfer.jl @@ -144,30 +144,71 @@ function _emit_interp!( push!( interp, _close_fill( - i, _cf_box(Val(N), d, g_n:1:g_n, tdims, fine_t), - tfirst, length(terms), _ninterp_terms(Val(N)), + terms, i, _cf_box(Val(N), d, g_n:1:g_n, tdims, fine_t), + tfirst, _ninterp_terms(Val(N)), ), ) end return nothing end -# Close one fill over the terms pushed onto the phase buffer since `tfirst`. The -# row width is a function of N alone (schedule.jl); a mismatch is an emitter bug, -# not a topology case, so it throws rather than recording a ragged row. +# Close one fill over the terms pushed onto the phase buffer since `tfirst`, and +# check the three invariants the sweeps and the device kernel index by (see the +# `GhostFill` docstring). All three are emitter bugs rather than topology cases, so +# they throw at schedule build — never as an out-of-bounds read, a shape error +# inside a broadcast, or (for the disjointness one) a silently wrong answer: +# +# 1. the row is `nexpected` terms wide (a function of N alone, schedule.jl), +# 2. every term window has the dst slab's shape, cell for cell, +# 3. every term window is disjoint from the dst slab. +# +# (3) is what lets the fused gather read all K windows while writing the dst in one +# broadcast: the term views ride an immutable `Ref` wrapper, which takes them out of +# Base's `broadcast_unalias` machinery, so nothing else would notice a dst that +# aliased a source. No emitter produces one — dst boxes are ghost layers and the +# terms read interiors — but the fusion is what makes it load-bearing. function _close_fill( - dst_block::Int, dst_ranges::NTuple{N,StepRange{Int,Int}}, - tfirst::Int, tlast::Int, nexpected::Int, -) where {N} + terms::Vector{SlabTerm{N,T}}, dst_block::Int, + dst_ranges::NTuple{N,StepRange{Int,Int}}, tfirst::Int, nexpected::Int, +) where {N,T} + tlast = length(terms) nterms = tlast - tfirst + 1 nterms == nexpected || throw( AssertionError( "coarse–fine fill emitted $nterms terms, expected $nexpected for N = $N", ), ) + len = length.(dst_ranges) + for k in tfirst:tlast + t = terms[k] + length.(t.ranges) == len || throw( + AssertionError( + "coarse–fine fill term $(k - tfirst + 1) reads a $(length.(t.ranges)) " * + "window but the dst slab is $len; every term must match it cell for cell", + ), + ) + _boxes_overlap(dst_block, dst_ranges, t.block, t.ranges) && throw( + AssertionError( + "coarse–fine fill term $(k - tfirst + 1) reads block $(t.block) " * + "$(t.ranges), which overlaps the dst slab $dst_block $dst_ranges; a " * + "fill's sources must be disjoint from the cells it writes", + ), + ) + end return GhostFill{N}(dst_block, dst_ranges, tfirst, tlast) end +# Do two index boxes of the same storage share a cell? Build-time only. +function _boxes_overlap( + b1::Int, r1::NTuple{N,StepRange{Int,Int}}, b2::Int, r2::NTuple{N,StepRange{Int,Int}} +) where {N} + b1 == b2 || return false + for d in 1:N + isempty(intersect(r1[d], r2[d])) && return false + end + return true +end + # Fine→coarse (flux-matching restriction): fill the coarse leaf K's ghost layer # on face (d, side) so its stencil flux through the interface equals the mean of # the fine-grid fluxes — g = u₁ + 2/2^(N−1) · Σ (u_f1 − g_f) over the fine @@ -210,8 +251,8 @@ function _emit_restrict!( push!( restrict, _close_fill( - i, _cf_box(Val(N), d, gC_n:1:gC_n, tdims, dst_t), - tfirst, length(terms), _nrestrict_terms(Val(N)), + terms, i, _cf_box(Val(N), d, gC_n:1:gC_n, tdims, dst_t), + tfirst, _nrestrict_terms(Val(N)), ), ) # The same child walk also records the weight-free coarse–fine-face topology @@ -337,10 +378,10 @@ Julia 1.12 on (EnzymeAD/Enzyme.jl#2707). Keeping the seam here rather than at Restriction runs last because it reads the interpolation-filled fine ghosts; no other cross-phase dependency exists. """ -function _exchange_storage!(store, lay::BlockLayout, sched::ExchangeSchedule) +function _exchange_storage!(store, lay::BlockLayout, sched::ExchangeSchedule{N}) where {N} _run_copies!(store, lay, sched.copies) - _run_fills!(store, lay, sched.interp, sched.interp_terms) - _run_fills!(store, lay, sched.restrict, sched.restrict_terms) + _run_fills!(store, lay, sched.interp, sched.interp_terms, _interp_width(Val(N))) + _run_fills!(store, lay, sched.restrict, sched.restrict_terms, _restrict_width(Val(N))) return nothing end @@ -354,28 +395,111 @@ function _run_copies!(store, lay::BlockLayout, copies::Vector{CopyDescriptor{N}} return nothing end -# `fills` and `terms` are the phase's flat isbits buffers: each fill names its CSR -# row `tfirst:tlast` into `terms`, so no per-fill record wider than a dst box is -# ever copied (96 bytes in 3D against 1752 for a 19-term row held inline). That -# is a structural choice, not a measured one — this loop times the same either -# way, the broadcasts dominating the descriptor walk. Seed with term 1, -# accumulate in term order — the arithmetic the device CSR kernel reproduces -# term-for-term. +# The phase row widths as compile-time values. Both are functions of N alone +# (schedule.jl), so `Val(N)` in, `Val(K)` out: dispatching the gather on K +# specializes it over two values per dimension, not over a dynamic loop bound. +@inline _interp_width(::Val{N}) where {N} = Val(_ninterp_terms(Val(N))) +@inline _restrict_width(::Val{N}) where {N} = Val(_nrestrict_terms(Val(N))) + +# One fused gather pass per fill: `dst[I] = Σₖ wₖ · srcₖ[I]` over the fill's K term +# windows, which all have the dst slab's shape (`_close_fill` enforces it). The +# retired form ran one `dst .= / .+=` broadcast per term — 7 per fill in 2D, 19 in +# 3D — each re-reading and re-writing the same handful of dst cells and each paying +# broadcast setup for them, so the fill phases dominated the exchange on a refined +# forest (#49/#84). Here the K term views are built once per fill and the weighted +# sum is evaluated per cell inside a single broadcast over the dst indices, the same +# `f.(CartesianIndices(dst), Ref(arrays))` shape the stencil leaves use, so the body +# stays array-level and device-agnostic. +# +# `fills`/`terms` stay the phase's flat CSR buffers — the fill record is its dst box +# plus a `[tfirst, tlast]` row (96 bytes in 3D against 1752 for a 19-term row held +# inline) — and the K views are built by recursion over `Val(K)` off `terms[tfirst + +# k - 1]`, closure-free: an `ntuple(Val(K)) do k` over two arrays stops inlining and +# loses the vectorization (`_diff_axes` in operators/diffusion.jl is the same trap). +# +# Term and accumulation order are unchanged: seed with term 1, add term k to the +# running sum, plain `*`/`+` and no fma, so `((w₁s₁ + w₂s₂) + w₃s₃) + …` is what the +# per-term loop computed and what the device CSR kernel (`_fill_kernel!`) reproduces +# term-for-term. It is therefore bit-identical to the retired loop whenever the +# field eltype IS the schedule's weight type `T`. When it is not — a Float32 field +# on a Float64 forest — the retired loop rounded its partial sum into the field's +# eltype K−1 times where this keeps the promoted accumulator and rounds once: about +# one ulp on the affected cells, and in the device kernel's direction, since +# `_fill_kernel!` accumulates the same promoted way. The fusion removes that +# pre-existing host/device divergence rather than introducing one (#84). function _run_fills!( - store, lay::BlockLayout, fills::Vector{GhostFill{N}}, terms::Vector{SlabTerm{N,T}} -) where {N,T} + store, lay::BlockLayout, fills::Vector{GhostFill{N}}, terms::Vector{SlabTerm{N,T}}, + ::Val{K}, +) where {N,T,K} for f in fills + f.tlast - f.tfirst + 1 == K && checkbounds(Bool, terms, f.tfirst:f.tlast) || + _throw_row_width(f, K, length(terms)) dst = _leaf_view(store, lay, f.dst_block, f.dst_ranges) - t1 = terms[f.tfirst] - dst .= t1.weight .* _leaf_view(store, lay, t1.block, t1.ranges) - for k in (f.tfirst + 1):f.tlast - tk = terms[k] - dst .+= tk.weight .* _leaf_view(store, lay, tk.block, tk.ranges) - end + srcs = _term_views(store, lay, terms, f.tfirst, Val(K)) + ws = _term_weights(terms, f.tfirst, Val(K)) + dst .= _gather_at.(CartesianIndices(dst), _AsScalar(srcs), _AsScalar(ws)) end return nothing end +# The row width licenses the `@inbounds` walk down `terms`, so it is checked per +# fill rather than trusted — once, outside the broadcast. +@noinline _throw_row_width(f::GhostFill, K::Int, nterms::Int) = throw( + AssertionError( + "ghost fill row $(f.tfirst):$(f.tlast) is not $K terms inside a $nterms-term " * + "buffer — the phase's row width is fixed by the dimension", + ), +) + +# Immutable stand-in for `Ref` as a broadcast scalar. `Ref(x)` is a mutable +# `RefValue`: once the gather body is too big to inline into `copyto!` (it is, at 19 +# term views) the Ref escapes and is heap-allocated per fill — the whole tuple of +# views, on every fill of every exchange, inside the Enzyme rule seam that must not +# allocate. An immutable `Ref` subtype rides every 0-dimensional `Ref` +# specialization of Base.Broadcast and stays on the stack. +# +# The Adapt rule keeps the wrapped views convertible for a device broadcast: a +# refined `BlockField` of GPU arrays is the one path that reaches this host body on +# a GPU (packed fields take the batched CSR kernels in transfer_kernels.jl, and a +# vector of device arrays is not something a kernel can index). It ships K device +# `SubArray`s as kernel parameters — ~2.4 KB for a 19-term 3D interpolation fill +# against CUDA's 4 KB parameter budget on pre-sm_90 hardware, so it fits, but with +# under 2× of headroom and none to spare for a wider stencil. Untested on GPU here. +struct _AsScalar{X} <: Ref{X} + x::X +end +@inline Base.getindex(s::_AsScalar) = s.x +Adapt.adapt_structure(to, s::_AsScalar) = _AsScalar(Adapt.adapt(to, s.x)) + +# The K source windows of one fill as a tuple of concrete SubArrays, and its K +# weights as a tuple — both by recursion down the CSR row from `i` (closure-free; +# see the note above). +@inline _term_views(store, lay::BlockLayout, terms, i::Int, ::Val{0}) = () +@inline function _term_views(store, lay::BlockLayout, terms, i::Int, ::Val{K}) where {K} + t = @inbounds terms[i] + return ( + _leaf_view(store, lay, t.block, t.ranges), + _term_views(store, lay, terms, i + 1, Val(K - 1))..., + ) +end + +@inline _term_weights(terms, i::Int, ::Val{0}) = () +@inline _term_weights(terms, i::Int, ::Val{K}) where {K} = + (@inbounds(terms[i].weight), _term_weights(terms, i + 1, Val(K - 1))...) + +# Per-cell body of the fused gather: the K-term weighted sum at the dst-local index +# I, left-associated in term order. `@inbounds` is licensed by the shape invariant — +# every term window has the dst slab's shape, and I ranges over that slab. +@inline function _gather_at( + I::CartesianIndex, srcs::Tuple{Vararg{AbstractArray,K}}, ws::Tuple{Vararg{Number,K}} +) where {K} + return _gather_terms(I, srcs, ws, Val(K)) +end +@inline _gather_terms(I::CartesianIndex, srcs::Tuple, ws::Tuple, ::Val{1}) = + @inbounds ws[1] * srcs[1][I] +@inline _gather_terms(I::CartesianIndex, srcs::Tuple, ws::Tuple, ::Val{k}) where {k} = + @inbounds _gather_terms(I, srcs, ws, Val(k - 1)) + ws[k] * srcs[k][I] + """ halo_update_adjoint!(x::AbstractBlockField, g::BlockForest) -> x @@ -429,6 +553,14 @@ end # Fills in reverse, terms within a fill in forward order — the scatter-adds of # one fill collide on shared source cells, so the term order is part of the # bit-exact contract. +# +# Deliberately still one scatter-add broadcast per term, unlike the fused forward +# gather, and the slower half of the exchange because of it. Those collisions are +# exactly what blocks the transposition: a dst-centric single pass would accumulate +# into a shared source cell in dst-cell order instead of term order and change the +# roundoff. The bit-identical fused form is a source-centric transposed CSR — a +# second descriptor set built at schedule time, deferred to issue #94. This runs off +# the mul! hot path, in apply_adjoint! and the reverse rule only. function _run_fills_adjoint!( store, lay::BlockLayout, fills::Vector{GhostFill{N}}, terms::Vector{SlabTerm{N,T}} ) where {N,T} From 3e8abaf9c533c83e0cbab694ac24f8fa52b6df8c Mon Sep 17 00:00:00 2001 From: Kyle Beggs Date: Thu, 10 Sep 2026 07:17:58 -0400 Subject: [PATCH 2/3] test(amr): guard the fused ghost gather structurally, not just numerically A bit-parity testset alone cannot see a revert: its reference is a re-written copy of the per-term loop the fusion replaces, and the zero-allocation testset passes either way. Redefining _run_fills! back to the loop left the suite green. So count what the sweep does. ProbeArray is block storage that tallies every element write and every broadcast materialized into a view of it; one phase run over it must issue exactly one broadcast per fill and write every dst cell exactly once. A per-term loop over a K-wide row issues K and writes K times, and now fails. Also added: bit-parity against the retired loop in Float32 and Float64, 2D and 3D, both storage layouts, scalar and SVector fields, forward and adjoint; a mixed-precision case (Float32 field on a Float64 forest) checked against the device CSR kernel on the CPU backend, where the fused gather is exact and the retired loop differed on 73 cells in 2D and 4420 in 3D; the row invariants _close_fill asserts, on hand-built rows and on every row the emitters produce; and zero allocation for both layouts, forward and adjoint, on refined 2D and 3D forests. Claude-Session: https://claude.ai/code/session_016F4h1y22x3ohpRdGVCpjHP --- test/exchange_schedule.jl | 339 +++++++++++++++++++++++++++++++++++++- 1 file changed, 335 insertions(+), 4 deletions(-) diff --git a/test/exchange_schedule.jl b/test/exchange_schedule.jl index 7beaf64..0dd2be2 100644 --- a/test/exchange_schedule.jl +++ b/test/exchange_schedule.jl @@ -1,3 +1,36 @@ +# Adaptor that doubles every Array it reaches — proves an Adapt rule recursed into a +# wrapper's contents (the _AsScalar broadcast scalar of the fused ghost gather). +struct DoubleAdaptor end +Adapt.adapt_storage(::DoubleAdaptor, x::Array) = 2 .* x + +# Block storage that counts what a sweep does to it: every element write, and every +# broadcast materialized into a view of it. This is the structural guard on the +# fused gather — a per-term loop over a K-wide row issues K broadcasts and writes +# each dst cell K times, the fused gather issues one and writes each cell once, and +# no numerical test can tell them apart. +struct ProbeArray{T,N,A<:AbstractArray{T,N}} <: AbstractArray{T,N} + parent::A + writes::Base.RefValue{Int} + bcasts::Base.RefValue{Int} +end +ProbeArray(a::AbstractArray, w::Base.RefValue{Int}, b::Base.RefValue{Int}) = + ProbeArray{eltype(a),ndims(a),typeof(a)}(a, w, b) +Base.size(p::ProbeArray) = size(p.parent) +Base.IndexStyle(::Type{<:ProbeArray}) = IndexCartesian() +Base.@propagate_inbounds Base.getindex(p::ProbeArray{T,N}, I::Vararg{Int,N}) where {T,N} = + p.parent[I...] +Base.@propagate_inbounds function Base.setindex!( + p::ProbeArray{T,N}, v, I::Vararg{Int,N} +) where {T,N} + p.writes[] += 1 + return p.parent[I...] = v +end +const ProbeView{T,N} = SubArray{T,N,<:ProbeArray} +function Base.copyto!(dst::ProbeView, bc::Base.Broadcast.Broadcasted{Nothing}) + parent(dst).bcasts[] += 1 + return invoke(copyto!, Tuple{AbstractArray,Base.Broadcast.Broadcasted{Nothing}}, dst, bc) +end + @testset "Exchange schedule" begin MFO = MatrixFreeOperators @@ -159,10 +192,46 @@ @test all(ts -> sum(t.weight for t in ts) ≈ 1, rterms) end end - # A row of the wrong width is an emitter bug and must throw, never be recorded. - box = (1:1:1, 2:1:3) - @test_throws AssertionError MFO._close_fill(1, box, 1, 2, 3) - @test MFO._close_fill(1, box, 4, 6, 3) === MFO.GhostFill{2}(1, box, 4, 6) + # `_close_fill` asserts the three row invariants the sweeps index by: width, + # every term window shaped like the dst slab, and every term window disjoint + # from it. All three are emitter bugs — a ragged row, an out-of-bounds read + # in the fused gather, and a fill that reads the cells it writes (which the + # gather could not see, its term views riding a non-array broadcast scalar). + box = (1:1:1, 2:1:3) # 1×2 dst on block 1 + good(b, r) = MFO.SlabTerm{2,Float64}(b, r, 0.5) + row = [good(2, (4:1:4, 2:1:3)), good(2, (5:1:5, 2:1:3))] + @test MFO._close_fill(row, 1, box, 1, 2) === MFO.GhostFill{2}(1, box, 1, 2) + @test_throws AssertionError MFO._close_fill(row, 1, box, 1, 3) # width + @test_throws AssertionError MFO._close_fill( # width + [row; good(2, (6:1:6, 2:1:3))], 1, box, 1, 2 + ) + @test_throws AssertionError MFO._close_fill( # shape + [row[1], good(2, (5:1:5, 2:1:4))], 1, box, 1, 2 + ) + @test_throws AssertionError MFO._close_fill( # reads its dst + [row[1], good(1, (1:1:1, 2:1:3))], 1, box, 1, 2 + ) + # a same-block term that misses the dst cells is fine (step-2 windows + # interleave all over the coarse–fine descriptors) + @test MFO._close_fill([row[1], good(1, (1:1:1, 4:1:5))], 1, box, 1, 2) isa + MFO.GhostFill{2} + # and both invariants hold of every row a real emitter produces + for N in (2, 3) + ext = ntuple(_ -> (0.0, 1.0), N) + g = CartesianGrid(ext, ntuple(_ -> 16, N); bc=ntuple(_ -> (Dirichlet(), Neumann()), N)) + bfe = BlockForest(g; blocksize=ntuple(_ -> 4, N), maxlevel=3) + refine!(bfe, x -> x[1] < 0.5) + refine!(bfe, x -> x[1] < 0.25) + sched = MFO._exchange_schedule(bfe) + for (fills, terms) in + ((sched.interp, sched.interp_terms), (sched.restrict, sched.restrict_terms)) + @test !isempty(fills) + for f in fills, t in MFO._fill_terms(terms, f) + @test length.(t.ranges) == length.(f.dst_ranges) + @test !MFO._boxes_overlap(f.dst_block, f.dst_ranges, t.block, t.ranges) + end + end + end end @testset "bcfaces: physical-face lists per (dim, side)" begin @@ -300,6 +369,239 @@ end end + @testset "fused gather: one pass per fill, and what it matches bit for bit" begin + # `_run_fills!` evaluates dst[I] = Σₖ wₖ·srcₖ[I] in ONE broadcast per fill, + # over the CSR row `tfirst:tlast` of its phase. Three claims, each checked by + # something that fails when the fusion is reverted or reassociated: + # (1) same term and accumulation order as the retired per-term loop ⇒ + # bit-identical to it whenever the field eltype IS the weight type, + # (2) it really is one pass — a structural count, because (1) passes either + # way (its reference is a re-written copy of the loop it replaces) and + # so does the zero-allocation testset, + # (3) in mixed precision it is bit-identical to the device CSR kernel, which + # the retired loop was not. + MFOL = MatrixFreeOperators + # (1)'s reference: the retired per-term loop, over the same CSR buffers. + function ref_fills!(store, lay, fills, terms) + for f in fills + dst = MFOL._leaf_view(store, lay, f.dst_block, f.dst_ranges) + t1 = terms[f.tfirst] + dst .= t1.weight .* MFOL._leaf_view(store, lay, t1.block, t1.ranges) + for k in (f.tfirst + 1):f.tlast + tk = terms[k] + dst .+= tk.weight .* MFOL._leaf_view(store, lay, tk.block, tk.ranges) + end + end + return nothing + end + function ref_fills_adjoint!(store, lay, fills, terms) + for f in Iterators.reverse(fills) + dst = MFOL._leaf_view(store, lay, f.dst_block, f.dst_ranges) + for k in f.tfirst:f.tlast + tk = terms[k] + MFOL._leaf_view(store, lay, tk.block, tk.ranges) .+= tk.weight .* dst + end + fill!(dst, zero(eltype(dst))) + end + return nothing + end + function ref_exchange!(x, sched) + store, lay = MFOL._storage(x), MFOL._layout(x) + MFOL._run_copies!(store, lay, sched.copies) + ref_fills!(store, lay, sched.interp, sched.interp_terms) + ref_fills!(store, lay, sched.restrict, sched.restrict_terms) + return x + end + function ref_exchange_adjoint!(x, sched) + store, lay = MFOL._storage(x), MFOL._layout(x) + ref_fills_adjoint!(store, lay, sched.restrict, sched.restrict_terms) + ref_fills_adjoint!(store, lay, sched.interp, sched.interp_terms) + MFOL._run_copies_adjoint!(store, lay, sched.copies) + return x + end + # exact per-cell comparison over the full padded storage, with a diagnostic + # on mismatch (how many cells, how far apart, on which leaf) + function bit_equal(a, b, what) + for i in 1:MFO.nleaves(a.grid) + A, B = MFO._block_array(a, i), MFO._block_array(b, i) + A == B && continue + bad = count(!iszero, A .- B) + @info "fused gather mismatch" what leaf = i ncells = bad worst = + maximum(maximum.(abs.(A .- B))) + return false + end + return true + end + # dimension-fixed row widths reach the sweep as compile-time values + @test @inferred(MFO._interp_width(Val(2))) === Val(7) + @test @inferred(MFO._interp_width(Val(3))) === Val(19) + @test @inferred(MFO._restrict_width(Val(2))) === Val(5) + @test @inferred(MFO._restrict_width(Val(3))) === Val(9) + + rng = Random.MersenneTwister(31) + for T in (Float64, Float32), N in (2, 3) + ext = ntuple(_ -> (T(0), T(1)), N) + bc = ntuple(d -> d == 1 ? (Dirichlet(), Neumann()) : (Periodic(), Periodic()), N) + g = CartesianGrid(ext, ntuple(_ -> 16, N); bc=bc) + bf = BlockForest(g; blocksize=ntuple(_ -> 4, N), maxlevel=3) + refine!(bf, x -> x[1] < 0.5) + refine!(bf, x -> x[1] < 0.25 && x[2] < 0.5) + sched = MFO._exchange_schedule(bf) + @test eltype(sched.interp_terms) === MFO.SlabTerm{N,T} + @test !isempty(sched.interp) && !isempty(sched.restrict) + x = scalar_field(bf) + @test eltype(x) === T + for i in 1:MFO.nleaves(bf) + rand!(rng, x.blocks[i]) + end + what = (T=T, N=N) + # (1) forward, BlockField and PackedBlockField (both run the host gather) + @test bit_equal(halo_update!(copy(x), bf), ref_exchange!(copy(x), sched), (what..., :fwd)) + @test bit_equal( + halo_update!(pack(copy(x)), bf), ref_exchange!(pack(copy(x)), sched), + (what..., :fwd_packed), + ) + # adjoint: the transposed per-term scatter, unchanged by the fusion, but + # its fill/term order is part of the same contract + @test bit_equal( + MFO.halo_update_adjoint!(copy(x), bf), ref_exchange_adjoint!(copy(x), sched), + (what..., :adj), + ) + @test bit_equal( + MFO.halo_update_adjoint!(pack(copy(x)), bf), + ref_exchange_adjoint!(pack(copy(x)), sched), (what..., :adj_packed), + ) + # SVector elements ride the same fills (the weight is a scalar T) + v = vector_field(bf) + for i in 1:MFO.nleaves(bf) + v.blocks[i] .= SVector{N,T}.(ntuple(_ -> rand(rng, T, size(v.blocks[i])), N)...) + end + @test bit_equal(halo_update!(copy(v), bf), ref_exchange!(copy(v), sched), (what..., :fwd_vec)) + @test bit_equal( + MFO.halo_update_adjoint!(copy(v), bf), ref_exchange_adjoint!(copy(v), sched), + (what..., :adj_vec), + ) + + # (2) one broadcast per fill, one write per dst cell. A per-term loop over + # a K-wide row issues K broadcasts and writes every dst cell K times. + for (fills, terms, K) in ( + (sched.interp, sched.interp_terms, MFO._ninterp_terms(Val(N))), + (sched.restrict, sched.restrict_terms, MFO._nrestrict_terms(Val(N))), + ) + nw, nb = Ref(0), Ref(0) + store = [ProbeArray(copy(b), nw, nb) for b in x.blocks] + MFO._run_fills!(store, MFO.BlocksLayout(), fills, terms, Val(K)) + @test nb[] == length(fills) + @test nw[] == sum(f -> prod(length.(f.dst_ranges)), fills) + # and the counted run really did the work: same cells as the reference + ref = [copy(b) for b in x.blocks] + ref_fills!(ref, MFO.BlocksLayout(), fills, terms) + @test all(i -> store[i].parent == ref[i], eachindex(ref)) + # a row that is not K wide is an emitter bug, not an out-of-bounds read + @test_throws AssertionError MFO._run_fills!( + store, MFO.BlocksLayout(), fills, terms, Val(K + 1) + ) + end + end + + # (3) mixed precision — a Float32 field on a Float64 forest, which the API + # allows (`scalar_field(bf, Float32)`). The retired loop rounded its partial + # sum into the Float32 slab K−1 times; the fused gather keeps the promoted + # accumulator and rounds once, which is what the device CSR kernel + # `_fill_kernel!` has always done. So the fusion does not merely differ from + # the retired loop here — it removes a host/device divergence. Compared with + # the host copy phase run on both sides, to isolate the fills from the + # documented corner divergence of the batched copy kernel. + cpu = KernelAbstractions.CPU() + for N in (2, 3) + ext = ntuple(_ -> (0.0, 1.0), N) + bc = ntuple(d -> d == 1 ? (Dirichlet(), Neumann()) : (Periodic(), Periodic()), N) + g = CartesianGrid(ext, ntuple(_ -> 16, N); bc=bc) + bf = BlockForest(g; blocksize=ntuple(_ -> 4, N), maxlevel=3) + refine!(bf, x -> x[1] < 0.5) + refine!(bf, x -> x[1] < 0.25 && x[2] < 0.5) + sched = MFO._exchange_schedule(bf) + @test eltype(sched.interp_terms) === MFO.SlabTerm{N,Float64} + ds = MFO._flatten_schedule(sched, bf, cpu) + rng32 = Random.MersenneTwister(17) + x32 = scalar_field(bf, Float32) + @test eltype(x32) === Float32 + for i in 1:MFO.nleaves(bf) + rand!(rng32, x32.blocks[i]) + end + copies!(y) = MFO._run_copies!(MFO._storage(y), MFO._layout(y), sched.copies) + fused = pack(copy(x32)) + copies!(fused) + MFO._run_fills!( + MFO._storage(fused), MFO._layout(fused), sched.interp, sched.interp_terms, + MFO._interp_width(Val(N)), + ) + MFO._run_fills!( + MFO._storage(fused), MFO._layout(fused), sched.restrict, + sched.restrict_terms, MFO._restrict_width(Val(N)), + ) + retired = pack(copy(x32)) + copies!(retired) + ref_fills!(MFO._storage(retired), MFO._layout(retired), sched.interp, sched.interp_terms) + ref_fills!( + MFO._storage(retired), MFO._layout(retired), sched.restrict, sched.restrict_terms + ) + device = pack(copy(x32)) + copies!(device) + MFO._run_fills_device!( + device.data, ds.interp, ds.interp_terms, ds.interp_maxcells, cpu + ) + MFO._run_fills_device!( + device.data, ds.restrict, ds.restrict_terms, ds.restrict_maxcells, cpu + ) + @test fused.data == device.data # fused ≡ device kernel + ndiff = count(!iszero, retired.data .- device.data) # retired loop was not + @test ndiff > 0 + @test maximum(abs, retired.data .- device.data) <= 4 * eps(Float32) + @info "mixed precision (Float32 field, Float64 weights)" N ndiff + end + + # The pieces: term views alias storage cell-for-cell from the CSR row, the + # weights come out in term order, and the per-cell body is the LEFT-associated + # sum — checked with values whose reassociation is visible in Float32: + # (1 + 1e8) + (−1e8) = 0 but 1 + (1e8 − 1e8) = 1. + ws = (1.0f0, 1.0f0, 1.0f0) + blocks = [fill(1.0f0, 2, 2), fill(1.0f8, 2, 2), fill(-1.0f8, 2, 2)] + row = MFO.SlabTerm{2,Float32}[ + MFO.SlabTerm{2,Float32}(k, (1:1:2, 2:1:2), ws[k]) for k in 1:3 + ] + srcs = MFO._term_views(blocks, MFO.BlocksLayout(), row, 1, Val(3)) + @test srcs === ntuple(k -> view(blocks[k], 1:1:2, 2:1:2), 3) + @test MFO._term_weights(row, 1, Val(3)) === ws + @test MFO._term_weights([row; row], 4, Val(3)) === ws # rows start anywhere + I = CartesianIndex(2, 1) + @test (ws[1] * srcs[1][I] + ws[2] * srcs[2][I]) + ws[3] * srcs[3][I] === 0.0f0 + @test ws[1] * srcs[1][I] + (ws[2] * srcs[2][I] + ws[3] * srcs[3][I]) === 1.0f0 + @test MFO._gather_at(I, srcs, ws) === 0.0f0 + packed = cat(blocks...; dims=3) + psrcs = MFO._term_views(packed, MFO.PackedLayout(), row, 1, Val(3)) + @test all(k -> psrcs[k] == srcs[k], 1:3) + @test MFO._gather_at(I, psrcs, ws) === 0.0f0 + # and a mis-weighted probe: the seed really is term 1 (not a zero accumulator, + # which would also make (0 + 1) + 1e8 − 1e8 = 0) + @test MFO._gather_at(I, srcs, (2.0f0, 1.0f0, 1.0f0)) === 0.0f0 + @test MFO._gather_at(I, srcs, (1.0f0, 1.0f0, 0.0f0)) === 1.0f8 + + # _AsScalar: an immutable Ref that broadcasts as a scalar (a mutable Ref of + # the term views escapes into copyto! and heap-allocates per fill), whose + # contents Adapt still converts for a device broadcast. + sc = MFO._AsScalar((1, 2)) + @test sc isa Ref + @test isbitstype(typeof(sc)) + @test sc[] === (1, 2) + @test ((i, t) -> i + t[1] * t[2]).(1:3, sc) == [3, 4, 5] + @test ((i, t) -> i + t[1] * t[2]).(1:3, MFO._AsScalar((1, 2))) == [3, 4, 5] + a = Float32[1 2; 3 4] + adapted = Adapt.adapt(DoubleAdaptor(), MFO._AsScalar((a, view(a, 1:1:1, 1:1:2)))) + @test adapted isa MFO._AsScalar + @test adapted[] == (2 .* a, view(2 .* a, 1:1:1, 1:1:2)) + end + @testset "forest apply_bc! matches the per-leaf physical fill" begin # Reference: per-leaf single-grid apply_bc! on hand-built leaf grids carrying # the physical/Interface mix leaf grids themselves no longer encode. @@ -369,6 +671,35 @@ halo_update!(f, g) return @allocated halo_update!(f, g) end + # (qualified through the module constant: `MFO` is a testset-local binding, + # and a captured local makes the call a dynamic lookup that boxes its + # arguments — an allocation of the scaffolding, not of the sweep) + function alloc_halo_adjoint(f, g) + MatrixFreeOperators.halo_update_adjoint!(f, g) + MatrixFreeOperators.halo_update_adjoint!(f, g) + return @allocated MatrixFreeOperators.halo_update_adjoint!(f, g) + end @test alloc_halo(uf, bf) == 0 + # Refined forests run the coarse–fine fills — 7 terms per interpolation fill + # in 2D, 19 in 3D. The fused gather's term views ride an immutable broadcast + # scalar precisely so this stays at zero: a mutable `Ref` escapes the + # un-inlined gather body and heap-allocates the whole tuple of views on every + # fill (~2 KB per 3D fill), inside the `_exchange_storage!` rule seam that + # must not allocate. + for N in (2, 3) + ext = ntuple(_ -> (0.0, 1.0), N) + bc = ntuple(d -> d == 1 ? (Dirichlet(), Neumann()) : (Periodic(), Periodic()), N) + g = CartesianGrid(ext, ntuple(_ -> 16, N); bc=bc) + bfr = BlockForest(g; blocksize=ntuple(_ -> 4, N), maxlevel=3) + refine!(bfr, x -> x[1] < 0.5) + sched = MFO._exchange_schedule(bfr) + @test !isempty(sched.interp) && !isempty(sched.restrict) + for xr in (scalar_field(bfr), pack(scalar_field(bfr))) + @inferred halo_update!(xr, bfr) + @inferred MFO.halo_update_adjoint!(xr, bfr) + @test alloc_halo(xr, bfr) == 0 + @test alloc_halo_adjoint(xr, bfr) == 0 + end + end end end From ab2664fc431334b64e0e99f423fe3315a8b98a7f Mon Sep 17 00:00:00 2001 From: Kyle Beggs Date: Thu, 10 Sep 2026 07:17:58 -0400 Subject: [PATCH 3/3] docs(amr): record the fused ghost gather and the deferred fused adjoint What the fusion buys, what it is bit-identical to (the retired loop at matched eltypes; the device kernel always, which the loop was not), the two row invariants that make it legal, the compile-latency price, and why the adjoint stays per-term until #94. Claude-Session: https://claude.ai/code/session_016F4h1y22x3ohpRdGVCpjHP --- DESIGN.md | 25 ++++++++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/DESIGN.md b/DESIGN.md index 0bc1d6f..3ee2f06 100644 --- a/DESIGN.md +++ b/DESIGN.md @@ -823,7 +823,30 @@ pre-built. host sweeps (both within run-to-run noise on refined 2D 256²/32² and 3D 32³/8³ forests) — so the choice was made on the smaller record and on the host walking the same flat buffers the batched device gather consumes, not - on a timing. The restriction's conservation guarantee is for *unweighted* + on a timing. Over that layout the host sweep runs each fill as **one fused + gather** — `dst[I] = Σₖ wₖ·srcₖ[I]` in a single broadcast over the dst slab, + the K term views built by recursion over `Val(K)` off the CSR row and carried + as an *immutable* broadcast scalar so the rule seam stays allocation-free + (#84). The retired form issued one broadcast per term (7 in 2D, 19 in 3D) over + a handful of cells each, and broadcast setup, not arithmetic, was the cost: on + a refined 32³ forest `halo_update!` goes 853 → 205 µs packed and 751 → 229 µs + per-block at 8³ blocks, and 3.0 → 0.67 ms at 4³ (min-of-N, alternating + processes, three rounds, ±4% spread). Term and accumulation order are + unchanged, so the result is bit-identical to the retired loop whenever the + field eltype is the schedule's weight type `T`, and — unlike that loop, which + rounded its partial sum into the field K−1 times — bit-identical to the device + CSR kernel in mixed precision too, which removes a pre-existing host/device + divergence. Two row invariants make the fusion legal and are asserted at + emission: every term window has the dst slab's shape, and every term window is + disjoint from the dst slab (the views ride a non-`AbstractArray` scalar, so + Base's `broadcast_unalias` never sees them). The price is compile latency — the + first `halo_update!` on a refined 3D forest roughly doubles, 0.55 → 1.06 s + per-block and 0.27 → 0.68 s packed — accepted for a 4× runtime win. The + adjoint's transposed scatter keeps its per-term form and is now the slower half + of the exchange (3–5× the forward): a fill's term windows collide on shared + source cells, so a dst-centric single pass would reorder the accumulation into + those cells; the bit-identical fused route is a source-centric transposed CSR + built at schedule time, scoped to issue #94. The restriction's conservation guarantee is for *unweighted* differences: a variable-coefficient flux weights each side of a coarse–fine face by an independently formed face κ, so `Diffusion` owns a κ-weighted coarse-ghost rewrite (issue #58) fed by weight-free `cfflux`