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` 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} 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