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

Abierto
  • ¿Algo no está bien?
#51 1 comentario 0 reacciones 0 asignados Ver en GitHub

Nadie ha tomado este issue todavía.

Evaluación

Dificultad
5/5
Tiempo estimado
Más de una semana
Aptitud para principiantes
35/100
Tipo de issue
Nueva funcionalidad
Claridad
Bastante claro
Estado de actividad
Tranquilo
Stack tecnológico
julia
Área
performance

Línea de trabajo

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.

Escrito por el modelo de indexación a partir del texto del issue.

Descripción

performance

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:

interior MFO apply! -t 1 MFO apply! -t auto (5 threads)
126² 9.5 µs/step 9.6 µs/step
510² 148.5 µs 144.3 µs
1022² 569.4 µs 569.2 µs
2046² 2653.9 µs 2625.3 µs

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

Lenguaje dominante
Julia
Estrellas
3
Forks
1
Merge medio
4 d 19 h
PR fusionados (30 d)
5

Preparar el entorno

Este proyecto no incluye contenedor de desarrollo, Dockerfile ni guía de contribución, así que la configuración corre por tu cuenta: empieza por su README y consulta nuestra guía para la primera contribución para los pasos generales.

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de RallypointOne/MatrixFreeOperators.jl

Todos los issues de RallypointOne/MatrixFreeOperators.jl

Issues similares

Más issues de Julia

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.