meta(cpu): single-grid apply! throughput — threading (AK now in core) and stencil fusion

Start with DESIGN.md and the leaf operators in src/operators/laplacian.jl, gradient.jl, divergence.jl, and scaling.jl, then inspect src/operators/forest_packed.jl and src/transfer_kernels.jl. Compare the proposed executor paths, including their Enzyme and view-backed Field constraints, and check #83 for the shared forest requirement. Done means the chosen path scales in comparisons/chmy_diffusion_2d/run.sh, preserves the stated invariants, and includes the required view test when applicable.

Status (2026-09-30): since PR #105 (57e2d8a) AcceleratedKernels is a core dependency and set! / op! run threaded on CPU, but apply! stencil sweeps are still serial — the root cause below is unchanged. This issue is now the meta tracker for single-grid CPU apply! throughput: threading here, stencil fusion in #52, the forest path in #83.

Context

The new Chmy.jl comparison benchmark (benchmark/comparisons/chmy_diffusion_2d/, Apple M5 Pro, Julia 1.12.6) shows MFO's CPU operator application is completely insensitive to -t — the thread pool sits idle:

Chmy's KernelAbstractions CPU backend does scale with threads, and overtakes us at DRAM-resident sizes despite moving ~40% more memory per step (7 array passes for its materialized flux form vs our fused 5): at 2046² with 5 threads, Chmy runs 2193 µs/step vs our 2654 µs. Single-threaded we win every size 2.5–16.9×, so the loss is purely about parallelization, not kernel design.

Root cause

Every leaf apply! body is a plain serial Base broadcast over a SubArray (src/operators/laplacian.jl, gradient.jl, divergence.jl, scaling.jl). The KA kernels that do exist (src/operators/forest_packed.jl, src/transfer_kernels.jl) only launch when backend isa KernelAbstractions.GPU; on CPU they fall back to the serial host loop. There is no Threads.@threads anywhere in src/; the only CPU-threaded code is set! (AK.map!) and op! (AK.foreachindex) in src/Fields.jl, added by #105, neither of which is on the apply! path.

Why it costs so much on modern hardware

One core cannot reach the package's aggregate memory bandwidth: STREAM triad on this M5 Pro measures 129 GB/s single-thread vs 246 GB/s with 5 threads. A perfect serial stencil sweep therefore leaves ≥2× bandwidth unused — this gap only grows on higher-core-count machines, and it is exactly where Chmy beats us.

Options

1. Launch leaf kernels through the KA CPU backend. The grid already carries device::KernelAbstractions.CPU, and DESIGN.md names KA @kernel as the per-operator escape hatch — this is Chmy's model, keeps device-agnosticism, and gives threading for free. Cost: leaf bodies stop being array-level broadcasts, which is a structural decision DESIGN.md must record (AD-friendliness of the leaves is a locked invariant).
2. A threaded executor at the apply! seam. Keep leaf bodies as broadcast expressions but materialize them through a chunked Threads.@threads sweep over interior rows. Preserves the array-level-leaf invariant (Enzyme rules, Reactant tracing, Adapt), smaller blast radius, CPU-only.
3. AK.foreachindex as the executor. AcceleratedKernels is already a core dependency (#105) and op! already uses AK.foreachindex(..., AK.get_backend(g)) for a threaded per-cell sweep — the same shape an apply! sweep needs. On CPU it runs on Julia tasks rather than the KA CPU backend, and on GPU it dispatches through the same backend, so one call site covers both. Cost: like option 1, the leaf body becomes a per-index function rather than a whole-array broadcast, and it has to be checked under Enzyme/Reactant (threaded tasks inside a differentiated leaf body).
4. User-level threading only (what examples/niederer_benchmark.jl does over forest leaves) — doesn't help the single-grid case, which is where we lose today.

Whichever route: the Enzyme rule constraints still apply (no allocation in rule bodies, no mixed GC-pointer/inline-float rule arguments), so the executor must not disturb the _exchange_storage! / _bc_storage! seams.

Sub-issues

• #83 — thread the reference BlockForest path (leaf sweeps and exchange descriptor loops); pick the executor once for both.
• #52 — constant-fold affine time-step operators (α·I + β·Laplacian) into one stencil pass; same Chmy benchmark and the same 2046² / 5-thread target numbers, and it compounds with threading here.

Acceptance

• If option 1 or 3 (leaf kernels through the KA CPU backend or AK) is chosen, apply! must stay safe when a Field wraps a SubArray view — KA's CPU() backend is known to race on views (the downstream case in #64, closed as not-an-MFO-bug because today's CPU path is plain broadcast), and op! already does threaded writes into possibly view-backed fields via AK. A test with view-backed input and output fields gates this.
• The forest reference path is tracked separately in #83 (sub-issue of this one); pick the executor once for both.

benchmark/comparisons/chmy_diffusion_2d/run.sh shows apply! scaling with -t, and MFO ≥ Chmy at every sweep size at -t auto. With 5/7 of Chmy's per-step traffic there is no structural reason to lose any threaded row.

—
🤖 Claude updated this now that AcceleratedKernels has moved into core.
_{claude-opus-5-5[1m] (high) · claude-code 2.1.285 · unreviewed}