stencils: Performance enhancements - #731
Merged
Merged
Conversation
Stencil throughput was dominated by two costs unrelated to the actual arithmetic. `@neighbors` returned a `SubArray` view of a `HaloArray`, so every element access paid `HaloArray`'s region-code dispatch plus the view's index translation, which prevented the sweep loop from vectorizing at all. The sweep now splits each chunk into an interior -- where every neighborhood access provably lands in the center array, so it can run against plain arrays under `@inbounds @simd` with the kernel force-inlined -- and a thin boundary shell that keeps the general path. `@neighbors` returns a direct-offset wrapper rather than a `SubArray`, since the latter's index translation was itself enough to block vectorization. Halos were built by a second task per chunk per expression that copied the whole center chunk into a cached `HaloArray`. For a double-buffered stencil that copy is as much memory traffic as the sweep itself. Halos are now assembled inside the sweeping task from the neighbor chunks passed as read dependencies, wrapping the center in place; for boundary conditions whose halo regions are plain slices of a neighbor (e.g. `Wrap`) they are views rather than materialized copies. Expressions that read and write the same chunks take their neighbors from a snapshot instead. This drops the per-chunk halo cache, its finalizer, and the in-place `load_*_region_into!` fill helpers, which only existed to serve it. Co-authored-by: Cursor <cursoragent@cursor.com>
Every GPU backend launched one kernel over the whole chunk, with all `@neighbors` accesses going through the general `HaloArray` path: a region code recomputed per access, a branch on it, and an index remap. On a GPU that scalar bookkeeping runs for all (2w+1)^N accesses of every element. Apply the same interior/boundary split the CPU sweep already uses. The interior sweeps against `HaloInterior` stand-ins, which index the center array directly, leaving the general path to the 2N boundary slabs. The five backends launched byte-identical kernels, so this lives in core as `gpu_stencil_sweep!` and each extension now delegates to it. A 4096^2 sweep over a single GPU chunk goes from 192 GB/s to 222 GB/s, against 298 GB/s for an equivalent hand-written AMDGPU kernel. Co-authored-by: Cursor <cursoragent@cursor.com>
`with_context` recorded the stream to restore by calling `AMDGPU.stream()`, which creates a stream when the running task doesn't have one. Every Dagger task that touches a ROC processor gets a fresh Julia task, so each one manufactured a HIP stream purely to be saved and put back. Those streams accumulate until finalized, and a run gets through enough tasks that stream creation itself stalls inside the driver, after which unrelated HIP calls start reporting illegal addresses. Read the task-local slot directly so a task without a stream reports `nothing` instead of being given one. A 4096^2 stencil sweep over a single GPU chunk goes from 20.4 GB/s to 192 GB/s, and multi-chunk sweeps stop faulting on every run. Co-authored-by: Cursor <cursoragent@cursor.com>
`@stencil` spawns a task per chunk per expression, so the block size decides how many times the scheduler's per-task cost is paid. That cost is large enough to dominate: a 4096x4096 Jacobi sweep on 12 threads is ~1.7x slower at `Blocks(512, 512)` than at `AutoBlocks()`, purely from running 64 tasks per sweep instead of 12. Nothing in the docs pointed at this, and partitioning much finer than the processor count is an easy default to reach for. Co-authored-by: Cursor <cursoragent@cursor.com>
Contributor
Dagger benchmarks:
|
| master | dirty | master / dirty | |
|---|---|---|---|
| array/dagger/N=1024 (block 512)/add (X + X) | 4.27 ± 0.23 ms | 4.41 ± 0.79 ms | 0.968 ± 0.18 |
| array/dagger/N=1024 (block 512)/alloc (rand) | 3.64 ± 0.42 ms | 3.73 ± 0.067 ms | 0.976 ± 0.11 |
| array/dagger/N=1024 (block 512)/broadcast (X .+ 1) | 2.95 ± 0.19 ms | 2.98 ± 0.025 ms | 0.988 ± 0.065 |
| array/dagger/N=1024 (block 512)/map (sin.(X)) | 6.85 ± 2.3 ms | 8.08 ± 0.41 ms | 0.848 ± 0.29 |
| array/dagger/N=1024 (block 512)/norm | 1.17 ± 0.023 ms | 1.24 ± 0.36 ms | 0.939 ± 0.27 |
| array/dagger/N=1024 (block 512)/reduce (sum) | 3.28 ± 1.3 ms | 2.72 ± 0.47 ms | 1.21 ± 0.52 |
| array/dagger/N=1024 (block 512)/transpose (permutedims) | 7.65 ± 0.46 ms | 7.92 ± 0.41 ms | 0.966 ± 0.077 |
| array/dagger/N=256 (block 256)/add (X + X) | 1.39 ± 0.28 ms | 1.8 ± 0.63 ms | 0.776 ± 0.32 |
| array/dagger/N=256 (block 256)/alloc (rand) | 1.24 ± 0.04 ms | 1.09 ± 0.12 ms | 1.13 ± 0.13 |
| array/dagger/N=256 (block 256)/broadcast (X .+ 1) | 0.533 ± 0.016 ms | 0.593 ± 0.24 ms | 0.899 ± 0.36 |
| array/dagger/N=256 (block 256)/map (sin.(X)) | 1.26 ± 0.044 ms | 1.66 ± 2 ms | 0.762 ± 0.92 |
| array/dagger/N=256 (block 256)/norm | 0.591 ± 0.14 ms | 0.602 ± 0.018 ms | 0.982 ± 0.23 |
| array/dagger/N=256 (block 256)/reduce (sum) | 1.15 ± 0.3 ms | 1.13 ± 0.21 ms | 1.02 ± 0.32 |
| array/dagger/N=256 (block 256)/transpose (permutedims) | 1.18 ± 0.26 ms | 1.13 ± 0.17 ms | 1.04 ± 0.28 |
| linalg/dagger/N=1024 (block 512)/cholesky | 20.9 ± 2.6 ms | 22.5 ± 0.56 ms | 0.926 ± 0.12 |
| linalg/dagger/N=1024 (block 512)/lu | 0.0524 ± 0.00087 s | 0.0464 ± 0.0024 s | 1.13 ± 0.061 |
| linalg/dagger/N=1024 (block 512)/matmul (A*A) | 0.056 ± 0.011 s | 0.0572 ± 0.009 s | 0.98 ± 0.24 |
| linalg/dagger/N=1024 (block 512)/matvec (A*x) | 3.72 ± 0.51 ms | 3.17 ± 0.14 ms | 1.17 ± 0.17 |
| linalg/dagger/N=1024 (block 512)/qr | 0.121 ± 0.0052 s | 0.12 ± 0.0039 s | 1 ± 0.054 |
| linalg/dagger/N=1024 (block 512)/solve (A\b via lu) | 0.0583 ± 0.0061 s | 0.0606 ± 0.002 s | 0.963 ± 0.11 |
| linalg/dagger/N=1024 (block 512)/syrk (A'*A) | 0.0381 ± 0.0058 s | 0.0367 ± 0.003 s | 1.04 ± 0.18 |
| linalg/dagger/N=256 (block 256)/cholesky | 3.03 ± 0.79 ms | 3.97 ± 0.38 ms | 0.765 ± 0.21 |
| linalg/dagger/N=256 (block 256)/lu | 5.27 ± 0.51 ms | 6.02 ± 1.2 ms | 0.875 ± 0.19 |
| linalg/dagger/N=256 (block 256)/matmul (A*A) | 2.68 ± 0.35 ms | 2.39 ± 0.39 ms | 1.12 ± 0.23 |
| linalg/dagger/N=256 (block 256)/matvec (A*x) | 1.57 ± 0.096 ms | 1.64 ± 0.073 ms | 0.957 ± 0.072 |
| linalg/dagger/N=256 (block 256)/qr | 5.05 ± 0.43 ms | 6 ± 0.57 ms | 0.842 ± 0.11 |
| linalg/dagger/N=256 (block 256)/solve (A\b via lu) | 9.68 ± 1.1 ms | 9.95 ± 2 ms | 0.973 ± 0.23 |
| linalg/dagger/N=256 (block 256)/syrk (A'*A) | 3.84 ± 1.2 ms | 3.93 ± 0.72 ms | 0.976 ± 0.36 |
| stencil/dagger/N=1024 (block 512)/alloc (neighbors Wrap) | 8.38 ± 0.73 ms | 8.43 ± 0.48 ms | 0.994 ± 0.1 |
| stencil/dagger/N=1024 (block 512)/assign (const) | 1.43 ± 0.073 ms | 1.42 ± 0.064 ms | 1.01 ± 0.069 |
| stencil/dagger/N=1024 (block 512)/multi-expr | 3.14 ± 0.17 ms | 3.3 ± 0.14 ms | 0.949 ± 0.066 |
| stencil/dagger/N=1024 (block 512)/neighbors (Clamp) | 7.45 ± 1.6 ms | 9.18 ± 3.1 ms | 0.812 ± 0.32 |
| stencil/dagger/N=1024 (block 512)/neighbors (Pad) | 7.99 ± 0.43 ms | 7.29 ± 0.43 ms | 1.1 ± 0.088 |
| stencil/dagger/N=1024 (block 512)/neighbors (Reflect) | 7.42 ± 2.6 ms | 7.49 ± 1.3 ms | 0.99 ± 0.39 |
| stencil/dagger/N=1024 (block 512)/neighbors (Wrap) | 6.58 ± 1.2 ms | 8.57 ± 0.47 ms | 0.767 ± 0.15 |
| stencil/dagger/N=1024 (block 512)/update (+) | 3.37 ± 1.1 ms | 2.51 ± 0.11 ms | 1.35 ± 0.44 |
| stencil/dagger/N=256 (block 256)/alloc (neighbors Wrap) | 2.41 ± 0.29 ms | 2.41 ± 0.19 ms | 1 ± 0.14 |
| stencil/dagger/N=256 (block 256)/assign (const) | 0.778 ± 0.058 ms | 0.843 ± 0.093 ms | 0.923 ± 0.12 |
| stencil/dagger/N=256 (block 256)/multi-expr | 1.39 ± 0.12 ms | 1.47 ± 0.11 ms | 0.949 ± 0.11 |
| stencil/dagger/N=256 (block 256)/neighbors (Clamp) | 1.53 ± 0.046 ms | 1.64 ± 0.067 ms | 0.932 ± 0.048 |
| stencil/dagger/N=256 (block 256)/neighbors (Pad) | 1.6 ± 0.067 ms | 1.65 ± 0.11 ms | 0.964 ± 0.074 |
| stencil/dagger/N=256 (block 256)/neighbors (Reflect) | 1.7 ± 0.11 ms | 1.67 ± 0.14 ms | 1.02 ± 0.11 |
| stencil/dagger/N=256 (block 256)/neighbors (Wrap) | 1.6 ± 0.089 ms | 1.73 ± 0.0092 ms | 0.925 ± 0.051 |
| stencil/dagger/N=256 (block 256)/update (+) | 0.921 ± 0.17 ms | 0.91 ± 0.12 ms | 1.01 ± 0.23 |
| time_to_load | 1.07 ± 0.0079 s | 1.06 ± 0.0025 s | 1.01 ± 0.0079 |
Plots
⚠️ Regressions (> 25.0%)
array/dagger/N=256 (block 256)/map (sin.(X)): +31.3%linalg/dagger/N=256 (block 256)/cholesky: +30.8%stencil/dagger/N=1024 (block 512)/neighbors (Wrap): +30.4%array/dagger/N=256 (block 256)/add (X + X): +28.8%
Improvements (> 25.0% faster)
stencil/dagger/N=1024 (block 512)/update (+): -25.7%
Full results and plots (download the benchmark-results artifact).
jpsamaroo
force-pushed
the
jps/stencil-fix-mpi
branch
from
August 7, 2026 13:38
a3ec45f to
6c6bab0
Compare
jpsamaroo
force-pushed
the
jps/stencil-fix-mpi
branch
from
August 7, 2026 21:47
8910dfd to
ad993c9
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Also adds tests for stencils over MPI.
Written by Claude