feat: multi-field API — set!(fun, ϕ), op!, and grid/field compatibili…

DESIGN.md

@@ -603,7 +603,7 @@ and drifts by an ulp — invisible to a stencil, which reads only spacing, but e
 to make a coordinate-assembled RHS and therefore the Krylov iteration count depend
 on `nparts`. It is bit-for-bit a no-op for undistributed grids and forest leaf
 grids, where `first(local_range[d]) == 1`. The user-facing surface is
-`boundary_rhs(P)`, `set!(::MultiDeviceVector, P, fun)`, `assemble_rhs(P, f)`, and
+`boundary_rhs(P)`, `set!(fun, ::MultiDeviceVector, P)`, `assemble_rhs(P, f)`, and
 `local_grids(P)`. Only `prepare_distributed` carries a `distributed` qualifier,
 because only it shadows a single-device function; everything downstream dispatches
 on `P` and is named for what it computes, not for where it runs.

Project.toml

@@ -1,6 +1,6 @@
 name = "MatrixFreeOperators"
 uuid = "bf251007-91fc-422b-8727-f8d5215c1c76"
-version = "0.1.0"
+version = "0.2.0"
 authors = ["Kyle Beggs <beggskw@gmail.com>"]
 
 [deps]

README.md

@@ -101,13 +101,13 @@ g = CartesianGrid(((0.0, 1.0), (0.0, 1.0)), (32, 32))
 
 # σ(x,y) = 1 + xy as a coefficient field. Fields are device arrays plus
 # grid metadata; set! fills them from a function of position.
-σ = set!(scalar_field(g), x -> 1 + x[1] * x[2])
+σ = set!(x -> 1 + x[1] * x[2], scalar_field(g))
 
 # The operator is composed symbolically — no matrix is ever assembled.
 K = scaling(σ) - laplacian(g)
 
 # Manufactured right-hand side for u = sin(πx)sin(πy).
-f = set!(scalar_field(g), x -> (2π^2 + 1 + x[1] * x[2]) * sinpi(x[1]) * sinpi(x[2]))
+f = set!(x -> (2π^2 + 1 + x[1] * x[2]) * sinpi(x[1]) * sinpi(x[2]), scalar_field(g))
 
 # prepare walks the operator tree once and allocates all scratch buffers;
 # the result supports mul!/size/eltype with zero steady-state allocations,
@@ -116,7 +116,7 @@ P = prepare(K, scalar_field(g))
 u, stats = cg(P, flatten(f))
 
 # Compare against the exact solution on the interior DOFs.
-u_exact = flatten(set!(scalar_field(g), x -> sinpi(x[1]) * sinpi(x[2])))
+u_exact = flatten(set!(x -> sinpi(x[1]) * sinpi(x[2]), scalar_field(g)))
 maximum(abs, u .- u_exact)   # ~1e-3, second-order accurate
 ```
 

benchmark/Project.toml

@@ -6,5 +6,5 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [compat]
 BenchmarkTools = "1.5"
-MatrixFreeOperators = "0.1"
+MatrixFreeOperators = "0.2"
 julia = "1.10"

benchmark/benchmarks.jl

@@ -16,15 +16,15 @@ g3 = CartesianGrid(((0.0, 1.0), (0.0, 1.0), (0.0, 1.0)), (64, 64, 64))
 for (dim, g) in (("2D 256²", g2), ("3D 64³", g3))
     L = laplacian(g)
     P = prepare(L, scalar_field(g))
-    x = flatten(set!(scalar_field(g), p -> sin(4p[1]) + cos(3p[end])))
+    x = flatten(set!(p -> sin(4p[1]) + cos(3p[end]), scalar_field(g)))
     y = similar(x)
     SUITE["grid"][dim]["laplacian prepare"] = @benchmarkable prepare($L, $(scalar_field(g)))
     SUITE["grid"][dim]["laplacian mul!"] = @benchmarkable mul!($y, $P, $x)
 end
 
-xs2 = flatten(set!(scalar_field(g2), p -> sin(4p[1]) + cos(3p[2])))
+xs2 = flatten(set!(p -> sin(4p[1]) + cos(3p[2]), scalar_field(g2)))
 ys2 = similar(xs2)
-xv2 = flatten(set!(vector_field(g2), p -> SVector(sin(p[2]), cos(p[1]))))
+xv2 = flatten(set!(p -> SVector(sin(p[2]), cos(p[1])), vector_field(g2)))
 yv2 = similar(xv2)
 
 PG = prepare(gradient(g2), scalar_field(g2))
@@ -33,11 +33,11 @@ SUITE["grid"]["2D 256²"]["gradient mul!"] = @benchmarkable mul!($yv2, $PG, $xs2
 PD = prepare(divergence(g2), vector_field(g2))
 SUITE["grid"]["2D 256²"]["divergence mul!"] = @benchmarkable mul!($ys2, $PD, $xv2)
 
-vel = set!(vector_field(g2), p -> SVector(sin(p[2]), cos(p[1])))
+vel = set!(p -> SVector(sin(p[2]), cos(p[1])), vector_field(g2))
 PA = prepare(advection(g2, vel), scalar_field(g2))
 SUITE["grid"]["2D 256²"]["advection mul!"] = @benchmarkable mul!($ys2, $PA, $xs2)
 
-κ = set!(scalar_field(g2), p -> 1 + 0.5 * sin(p[1]))
+κ = set!(p -> 1 + 0.5 * sin(p[1]), scalar_field(g2))
 PS = prepare(2.0 * laplacian(g2) + scaling(κ), scalar_field(g2))
 SUITE["grid"]["2D 256²"]["2λ + κ·I mul!"] = @benchmarkable mul!($ys2, $PS, $xs2)
 
@@ -50,15 +50,15 @@ SUITE["grid"]["2D 256²"]["∇·(κ∇u) mul!"] = @benchmarkable mul!($ys2, $PK,
 # benchpkg runs this file on the PR's base branch too, where Diffusion does not exist.
 if isdefined(MatrixFreeOperators, :Diffusion)
     for (dim, g, xf, yf) in (("2D 256²", g2, xs2, ys2),)
-        D = diffusion(g, set!(scalar_field(g), p -> 1 + 0.5 * sin(p[1])))
+        D = diffusion(g, set!(p -> 1 + 0.5 * sin(p[1]), scalar_field(g)))
         PDf = prepare(D, scalar_field(g))
         SUITE["grid"][dim]["diffusion prepare"] = @benchmarkable prepare(
             $D, $(scalar_field(g))
         )
         SUITE["grid"][dim]["diffusion mul!"] = @benchmarkable mul!($yf, $PDf, $xf)
     end
-    D3 = diffusion(g3, set!(scalar_field(g3), p -> 1 + 0.5 * sin(p[1])))
-    x3 = flatten(set!(scalar_field(g3), p -> sin(4p[1]) + cos(3p[3])))
+    D3 = diffusion(g3, set!(p -> 1 + 0.5 * sin(p[1]), scalar_field(g3)))
+    x3 = flatten(set!(p -> sin(4p[1]) + cos(3p[3]), scalar_field(g3)))
     y3 = similar(x3)
     PD3 = prepare(D3, scalar_field(g3))
     SUITE["grid"]["3D 64³"]["diffusion mul!"] = @benchmarkable mul!($y3, $PD3, $x3)
@@ -73,10 +73,10 @@ end
 # revisions without the distributed seam have no `_slab_op`.
 if isdefined(MatrixFreeOperators, :Diffusion) && isdefined(MatrixFreeOperators, :_slab_op)
     for (dim, g) in (("2D 256²", g2), ("3D 64³", g3))
-        D = diffusion(g, set!(scalar_field(g), p -> 1 + 0.5 * sin(p[1])))
+        D = diffusion(g, set!(p -> 1 + 0.5 * sin(p[1]), scalar_field(g)))
         lg = partition_grid(g, 2)[1]
         Ds = MatrixFreeOperators._slab_op(D, lg)
-        ȳs = set!(scalar_field(lg), p -> sin(4p[1]) + cos(3p[end]))
+        ȳs = set!(p -> sin(4p[1]) + cos(3p[end]), scalar_field(lg))
         x̄s = scalar_field(lg)
         SUITE["grid"][dim]["diffusion slab apply!"] =
             @benchmarkable apply!($x̄s, $Ds, $ȳs, $lg, 1.0, 0.0)
@@ -102,8 +102,8 @@ end
 # rank-changers gather unconditionally; a Laplacian would shortcut to its forward
 # action on this all-physical grid and measure nothing.
 
-sadj = set!(scalar_field(g2), p -> sin(4p[1]) + cos(3p[2]))
-vadj = set!(vector_field(g2), p -> SVector(sin(p[2]), cos(p[1])))
+sadj = set!(p -> sin(4p[1]) + cos(3p[2]), scalar_field(g2))
+vadj = set!(p -> SVector(sin(p[2]), cos(p[1])), vector_field(g2))
 sadj_out = scalar_field(g2)
 vadj_out = vector_field(g2)
 Dx, Dy = derivative(g2, 1), derivative(g2, 2)
@@ -132,7 +132,7 @@ SUITE["grid"]["2D 256²"]["adjoint(∂x + ∂y) mul!"] = @benchmarkable mul!($ys
 # 8×8 root tiling of 32² blocks (64 uniform leaves) — same DOFs as the 2D grid above,
 # so the forest overhead (halo exchange + per-leaf dispatch) is directly comparable.
 bf = BlockForest(g2; blocksize=(32, 32), maxlevel=2)
-xb = set!(scalar_field(bf), p -> sin(4p[1]) + cos(3p[2]))
+xb = set!(p -> sin(4p[1]) + cos(3p[2]), scalar_field(bf))
 Lf = laplacian(bf)
 Pf = prepare(Lf, scalar_field(bf))
 xf = flatten(xb)
@@ -159,7 +159,7 @@ end
 # Same 32² blocks, so the extra cost over the uniform leg is coarse–fine work.
 bfr = BlockForest(g2; blocksize=(32, 32), maxlevel=2)
 refine!(bfr, p -> p[1] < 0.5)
-xbr = set!(scalar_field(bfr), p -> sin(4p[1]) + cos(3p[2]))
+xbr = set!(p -> sin(4p[1]) + cos(3p[2]), scalar_field(bfr))
 Lr = laplacian(bfr)
 
 SUITE["forest"]["2D refined"]["halo_update!"] = @benchmarkable halo_update!($xbr, $bfr)
@@ -184,7 +184,7 @@ SUITE["forest"]["2D refined"]["laplacian apply_adjoint!"] =
 g3r = CartesianGrid(((0.0, 1.0), (0.0, 1.0), (0.0, 1.0)), (32, 32, 32))
 bfr3 = BlockForest(g3r; blocksize=(8, 8, 8), maxlevel=2)
 refine!(bfr3, p -> p[1] < 0.5)
-xbr3 = set!(scalar_field(bfr3), p -> sin(4p[1]) + cos(3p[3]))
+xbr3 = set!(p -> sin(4p[1]) + cos(3p[3]), scalar_field(bfr3))
 
 SUITE["forest"]["3D refined"]["halo_update!"] = @benchmarkable halo_update!($xbr3, $bfr3)
 SUITE["forest"]["3D refined"]["halo_update_adjoint!"] =

benchmark/gpu.jl

@@ -33,7 +33,7 @@ println("laplacian mul! — CPU per-leaf / GPU per-leaf / GPU packed (min time)"
     "forest", "nleaves", "DOFs", "CPU/leaf", "GPU/leaf", "GPU packed", "packed speedup", "host alloc")
 for (n, refined) in ((256, false), (512, false), (1024, false), (2048, false), (512, true))
     bf = make_forest(n; refined)
-    u = set!(scalar_field(bf), x -> sin(4x[1]) + cos(3x[2]))
+    u = set!(x -> sin(4x[1]) + cos(3x[2]), scalar_field(bf))
     L = laplacian(bf)
     x = flatten(u)
     y = similar(x)
@@ -70,7 +70,7 @@ println("\nhalo_update! on the device packed field — batched kernels vs per-de
     "forest", "nleaves", "kernels", "loop", "kernel speedup")
 for (n, refined) in ((512, false), (2048, false), (512, true))
     bf = make_forest(n; refined)
-    u = set!(scalar_field(bf), x -> sin(4x[1]) + cos(3x[2]))
+    u = set!(x -> sin(4x[1]) + cos(3x[2]), scalar_field(bf))
     pg = Adapt.adapt(CuArray, pack(u))
     sched = MFO._exchange_schedule(bf)
     halo_update!(pg, pg.grid)                        # warm the _device_schedule cache

docs/index.qmd

@@ -36,12 +36,12 @@ using MatrixFreeOperators, Krylov, LinearAlgebra
 
 # -Δu + σu = f on (0,1)², homogeneous Dirichlet, manufactured solution
 g = CartesianGrid(((0.0, 1.0), (0.0, 1.0)), (32, 32))
-σ = set!(scalar_field(g), x -> 1 + x[1] * x[2])
+σ = set!(x -> 1 + x[1] * x[2], scalar_field(g))
 K = scaling(σ) - laplacian(g)
 
-f = set!(scalar_field(g), x -> (2π^2 + 1 + x[1] * x[2]) * sinpi(x[1]) * sinpi(x[2]))
+f = set!(x -> (2π^2 + 1 + x[1] * x[2]) * sinpi(x[1]) * sinpi(x[2]), scalar_field(g))
 P = prepare(K, scalar_field(g))
 u, stats = cg(P, flatten(f))
 
-maximum(abs, u .- flatten(set!(scalar_field(g), x -> sinpi(x[1]) * sinpi(x[2]))))
+maximum(abs, u .- flatten(set!(x -> sinpi(x[1]) * sinpi(x[2]), scalar_field(g))))
 ```

docs/pages/amr.qmd

@@ -171,8 +171,8 @@ different, block-by-block order), so any order-independent reduction matches:
 g  = CartesianGrid(((0.0, 1.0), (0.0, 1.0)), (8, 8))
 bf = BlockForest(g; blocksize = (4, 4), maxlevel = 3)   # level 0 ⇒ same 8×8
 
-u  = set!(scalar_field(g),  x -> sinpi(x[1]) * sinpi(x[2]))
-uf = set!(scalar_field(bf), x -> sinpi(x[1]) * sinpi(x[2]))
+u  = set!(x -> sinpi(x[1]) * sinpi(x[2]),  scalar_field(g))
+uf = set!(x -> sinpi(x[1]) * sinpi(x[2]), scalar_field(bf))
 
 norm(flatten(laplacian(g) * u)) ≈ norm(flatten(laplacian(bf) * uf))
 ```
@@ -188,7 +188,7 @@ conflicts), re-establishes 2:1 balance, and returns freshly allocated fields wit
 the data carried across:
 
 ```{julia}
-bump = set!(scalar_field(forest), x -> exp(-((x[1] - 0.5)^2 + (x[2] - 0.5)^2) / 0.01))
+bump = set!(x -> exp(-((x[1] - 0.5)^2 + (x[2] - 0.5)^2) / 0.01), scalar_field(forest))
 bump = regrid!(bump; refine = b -> maximum(abs, interior(b)) > 0.5)
 length(collect(leaves(forest)))                  # resolution now follows the bump
 ```

docs/pages/autodiff.qmd

@@ -45,7 +45,7 @@ function diffusion_loss(κdata, udata, w, gg)
 end
 
 rng = MersenneTwister(1)
-u = set!(scalar_field(g), x -> sinpi(x[1]) * sinpi(x[2]))
+u = set!(x -> sinpi(x[1]) * sinpi(x[2]), scalar_field(g))
 w = rand(rng, local_size(g)...)
 κ = 1.0 .+ rand(rng, padded_size(g)...)
 
@@ -156,12 +156,12 @@ using Enzyme, StaticArrays, LinearAlgebra
 
 g1 = CartesianGrid(((0.0, 2π),), (32,); bc=((Periodic(), Periodic()),))
 F  = advection(g1, SelfAdvection())          # nonlinear u·∇u
-u0 = set!(vector_field(g1), x -> SVector(2 + sin(x[1])))
+u0 = set!(x -> SVector(2 + sin(x[1])), vector_field(g1))
 
 J  = linearize(F, u0)                        # FiniteDifferenceJVP — the default
 Ja = linearize(F, u0, EnzymeJVP())           # exact, and has a transpose
 
-v = set!(vector_field(g1), x -> SVector(cos(2x[1])))
+v = set!(x -> SVector(cos(2x[1])), vector_field(g1))
 Jv  = collect(interior(apply(J,  copy(v))))
 Jav = collect(interior(apply(Ja, copy(v))))
 
@@ -174,7 +174,7 @@ Only the Enzyme-backed one has a transpose:
 # A ramp rather than a trig mode: on this periodic grid J·v is orthogonal to the
 # low harmonics, so a sin/cos test vector would make both sides vanish and the
 # identity would hold vacuously.
-w = set!(vector_field(g1), x -> SVector(x[1] / 2π))
+w = set!(x -> SVector(x[1] / 2π), vector_field(g1))
 lhs = sum(dot.(collect(interior(apply(Ja, copy(v)))), collect(interior(w))))
 rhs = sum(dot.(collect(interior(v)), collect(interior(apply(adjoint(Ja), copy(w))))))
 (lhs, rhs, abs(lhs - rhs))                   # ⟨Jv, w⟩ = ⟨v, Jᵀw⟩

docs/pages/distributed.qmd

@@ -229,7 +229,7 @@ assembles the lift of its own slab and the result comes back as a
 b = assemble_rhs(P, x -> sinpi(x[1]) * sinpi(x[2]))   # f - boundary_rhs(P)
 
 # ...or in two steps, if you want the pieces
-src = set!(MultiDeviceVector{Float64}(undef, P.spec), P, x -> sinpi(x[1]))
+src = set!(x -> sinpi(x[1]), MultiDeviceVector{Float64}(undef, P.spec), P)
 b   = src .- boundary_rhs(P)
 ```
 
@@ -247,7 +247,7 @@ Bring your own source data with [`local_grids`](../pages/api.qmd), which hands
 back each partition's slab grid, on the host:
 
 ```julia
-fields = [set!(scalar_field(lg), myfun) for lg in local_grids(P)]
+fields = [set!(myfun, scalar_field(lg)) for lg in local_grids(P)]
 b = assemble_rhs(P, fields)
 ```
 

docs/pages/internals.qmd

@@ -1084,7 +1084,7 @@ g = CartesianGrid(((0.0, 1.0), (0.0, 1.0)), (256, 256))  # homogeneous Dirichlet
 L = -laplacian(g)
 P = prepare(L)
 M = MultigridPreconditioner(L)
-b = flatten(set!(scalar_field(g), x -> 2π^2 * sinpi(x[1]) * sinpi(x[2])))
+b = flatten(set!(x -> 2π^2 * sinpi(x[1]) * sinpi(x[2]), scalar_field(g)))
 u, stats = Krylov.cg(P, b; M)
 ```
 

examples/adaptive_poisson.jl

@@ -25,15 +25,15 @@ function l1_error(u, bf)
     e = 0.0
     for i in 1:MFO.nleaves(bf)
         sp = MFO._leaf_spacing(bf, bf.forest.leaves[i].level)
-        ref = set!(similar(MFO.block(u, i)), uexact)
+        ref = set!(uexact, similar(MFO.block(u, i)))
         e += sum(abs, interior(MFO.block(u, i)) .- interior(ref)) * prod(sp)
     end
     return e
 end
 
 for cycle in 1:4
     P = prepare(laplacian(bf), u)
-    rhs = .-flatten(set!(scalar_field(bf), f))          # Δu = -f; Dirichlet lift is zero
+    rhs = .-flatten(set!(f, scalar_field(bf)))          # Δu = -f; Dirichlet lift is zero
     sol, stats = Krylov.gmres(P, rhs; rtol=1e-10)       # nonsymmetric on an adapted forest
     flat_to_interior!(u, sol)
     lo, hi = extrema(k -> k.level, bf.forest.leaves)

examples/heat_equation.jl

@@ -20,8 +20,8 @@ nsteps = nframes * steps_per_frame
 dt = (tspan[2] - tspan[1]) / nsteps
 # 20 Gaussian blobs at random centers; diameter ≈ 4σ drawn uniform in [0.1, 0.3]
 blobs = [(rand(), rand(), 0.2 * (0.5 + rand())) for _ in 1:20]
-u0 = set!(scalar_field(g), x ->
-    sum(exp(-((x[1] - cx)^2 + (x[2] - cy)^2) / (s^2 / 8)) for (cx, cy, s) in blobs))
+u0 = set!(x ->
+    sum(exp(-((x[1] - cx)^2 + (x[2] - cy)^2) / (s^2 / 8)) for (cx, cy, s) in blobs), scalar_field(g))
 interior(u0) ./= maximum(interior(u0))
 
 function step!(u, du, L, αdt)

examples/inverse_diffusion.jl

@@ -62,10 +62,10 @@ end
 
 # This example uses one drive. The data are sensitive to κ through the flux κ∇u,
 # so recovery remains weak where this excitation's gradient is small.
-u = set!(scalar_field(g), x -> sinpi(x[1]) * sinpi(x[2]))
+u = set!(x -> sinpi(x[1]) * sinpi(x[2]), scalar_field(g))
 
 rng = MersenneTwister(20260731)
-κ★ = set!(scalar_field(g), κ_true)
+κ★ = set!(κ_true, scalar_field(g))
 clean = collect(response(κ★.data, u.data, g))
 obs = clean .+ 0.01 * maximum(abs, clean) .* randn(rng, size(clean))
 

examples/monodomain_amr.jl

@@ -100,7 +100,7 @@ function simulate(; snap_times=Float64[], on_frame=nothing, frame_every=250)
                          bc=ntuple(_ -> (Neumann(), Neumann()), 2))
     bf = BlockForest(base; blocksize=(BS, BS), maxlevel=MAXLEV)
 
-    V = set!(scalar_field(bf), x -> x[1] < 4.0 ? 1.0 : 0.0)   # S1 along the left edge
+    V = set!(x -> x[1] < 4.0 ? 1.0 : 0.0, scalar_field(bf))   # S1 along the left edge
     W = scalar_field(bf)
     dv, dw = zeros(BS, BS), zeros(BS, BS)
 

examples/multigrid_poisson.jl

@@ -16,8 +16,8 @@ A = prepare(L)
 uexact(x) = 1 + exp(x[1]) * sinpi(x[1]) * sinpi(x[2])
 f(x) = -exp(x[1]) * sinpi(x[2]) * ((1 - 2 * pi^2) * sinpi(x[1]) + 2 * pi * cospi(x[1]))
 # inhomogeneous boundary data enters the RHS through the affine lift
-b = flatten(set!(scalar_field(g), f)) .- flatten(boundary_rhs(L, g))
-uex = flatten(set!(scalar_field(g), uexact))
+b = flatten(set!(f, scalar_field(g))) .- flatten(boundary_rhs(L, g))
+uex = flatten(set!(uexact, scalar_field(g)))
 
 mg = MultigridPreconditioner(L)                          # Jacobi(2/3), levels=:auto
 println(mg)

examples/niederer_benchmark.jl

@@ -403,8 +403,8 @@ function simulate(; adaptive::Bool)
         balance!(bf)
     end
 
-    V = set!(scalar_field(bf), _ -> initial_voltage())
-    S = set!(state_field(bf), _ -> initial_state())
+    V = set!(_ -> initial_voltage(), scalar_field(bf))
+    S = set!(_ -> initial_state(), state_field(bf))
     LV = scalar_field(bf)
     Lop = diffusion_operator(bf)
 
@@ -419,8 +419,8 @@ function simulate(; adaptive::Bool)
     diffuse!(V, LV, Lop, bf)
     react!(V, S, bf, 0.0)
     probe_all(V, bf, probes)
-    V = set!(scalar_field(bf), _ -> initial_voltage())
-    S = set!(state_field(bf), _ -> initial_state())
+    V = set!(_ -> initial_voltage(), scalar_field(bf))
+    S = set!(_ -> initial_state(), state_field(bf))
     LV = scalar_field(bf)
 
     snaps, hist = Any[], Tuple{Float64,Int}[]

ext/MatrixFreeOperatorsMDLAExt.jl

@@ -293,7 +293,7 @@ call it concurrently with a solve on the same prepared operator.
 
 ```julia
 P = prepare_distributed(laplacian(g), 2)
-b = set!(MultiDeviceVector{Float64}(undef, P.spec), P, x -> sin(x[1]))
+b = set!(x -> sin(x[1]), MultiDeviceVector{Float64}(undef, P.spec), P)
 b .-= boundary_rhs(P)
 u, stats = Krylov.cg(P, b)
 ```
@@ -333,12 +333,12 @@ MatrixFreeOperators.local_grids(P::MDLAPreparedOperator) =
     [Adapt.adapt(Array, p.grid) for p in P.parts]
 
 """
-    set!(x::MultiDeviceVector, P::MDLAPreparedOperator, fun) -> x
+    set!(fun, x::MultiDeviceVector, P::MDLAPreparedOperator) -> x
 
 Fill `x` with `fun(coords)` evaluated slab-locally on each partition.
 
-The distributed twin of `set!(::Field, fun)`, and exact: `cell_center` evaluates
-at the global cell index, so this is bit-for-bit `MultiDeviceVector(flatten(set!(scalar_field(g), fun)), P.spec)`
+The distributed twin of `set!(fun, ::Field)`, and exact: `cell_center` evaluates
+at the global cell index, so this is bit-for-bit `MultiDeviceVector(flatten(set!(fun, scalar_field(g))), P.spec)`
 without ever building the global field. Uses `P`'s input scratch, so the same
 concurrency caveat as `mul!` applies.
 
@@ -349,7 +349,7 @@ captured host arrays. For anything heavier, build the fields yourself on
 them.
 """
 function MatrixFreeOperators.set!(
-    x::MultiDeviceVector{T}, P::MDLAPreparedOperator{T}, fun
+    fun, x::MultiDeviceVector{T}, P::MDLAPreparedOperator{T}
 ) where {T}
     _dist_set!(P.xpads, fun, P.ctx)
     _dist_map!(P.ctx) do d
@@ -388,7 +388,7 @@ function MatrixFreeOperators.assemble_rhs(P::MDLAPreparedOperator{T}, f) where {
     return x
 end
 
-_source!(x, P::MDLAPreparedOperator, fun) = MatrixFreeOperators.set!(x, P, fun)
+_source!(x, P::MDLAPreparedOperator, fun) = MatrixFreeOperators.set!(fun, x, P)
 function _source!(x, P::MDLAPreparedOperator, fields::AbstractVector)
     length(fields) == length(P.parts) || throw(
         ArgumentError(