DO NOT MERGE: Sparse+Hierarchical and faster stencils - #727
Draft
jpsamaroo wants to merge 43 commits into
Draft
Conversation
Contributor
Dagger benchmarks:
|
| master | dirty | master / dirty | |
|---|---|---|---|
| array/dagger/N=1024 (block 512)/add (X + X) | 4.45 ± 0.082 ms | 4.49 ± 0.99 ms | 0.991 ± 0.22 |
| array/dagger/N=1024 (block 512)/alloc (rand) | 3.55 ± 0.69 ms | 3.62 ± 0.18 ms | 0.98 ± 0.2 |
| array/dagger/N=1024 (block 512)/broadcast (X .+ 1) | 3 ± 0.067 ms | 2.88 ± 0.53 ms | 1.04 ± 0.19 |
| array/dagger/N=1024 (block 512)/map (sin.(X)) | 7.23 ± 1.8 ms | 7.35 ± 0.7 ms | 0.983 ± 0.26 |
| array/dagger/N=1024 (block 512)/norm | 1.43 ± 0.069 ms | 1.33 ± 0.054 ms | 1.08 ± 0.068 |
| array/dagger/N=1024 (block 512)/reduce (sum) | 2.71 ± 0.31 ms | 2.51 ± 0.9 ms | 1.08 ± 0.41 |
| array/dagger/N=1024 (block 512)/transpose (permutedims) | 7.38 ± 0.29 ms | 7.16 ± 0.18 ms | 1.03 ± 0.049 |
| array/dagger/N=256 (block 256)/add (X + X) | 1.39 ± 0.11 ms | 1.1 ± 0.051 ms | 1.26 ± 0.11 |
| array/dagger/N=256 (block 256)/alloc (rand) | 1.13 ± 0.041 ms | 1.18 ± 0.008 ms | 0.959 ± 0.036 |
| array/dagger/N=256 (block 256)/broadcast (X .+ 1) | 0.727 ± 0.21 ms | 0.659 ± 0.26 ms | 1.1 ± 0.54 |
| array/dagger/N=256 (block 256)/map (sin.(X)) | 1.05 ± 0.17 ms | 1.16 ± 0.095 ms | 0.908 ± 0.16 |
| array/dagger/N=256 (block 256)/norm | 0.604 ± 0.047 ms | 0.583 ± 0.088 ms | 1.04 ± 0.18 |
| array/dagger/N=256 (block 256)/reduce (sum) | 1.35 ± 2.3 ms | 0.722 ± 0.032 ms | 1.87 ± 3.2 |
| array/dagger/N=256 (block 256)/transpose (permutedims) | 1.15 ± 0.12 ms | 1.32 ± 0.3 ms | 0.875 ± 0.22 |
| linalg/dagger/N=1024 (block 512)/cholesky | 21.1 ± 2.3 ms | 23.1 ± 3.1 ms | 0.916 ± 0.16 |
| linalg/dagger/N=1024 (block 512)/lu | 0.0454 ± 0.003 s | 0.0476 ± 0.0033 s | 0.953 ± 0.092 |
| linalg/dagger/N=1024 (block 512)/matmul (A*A) | 0.0462 ± 0.015 s | 0.084 ± 0.0029 s | 0.55 ± 0.18 |
| linalg/dagger/N=1024 (block 512)/matvec (A*x) | 2.97 ± 0.18 ms | 3.56 ± 0.54 ms | 0.835 ± 0.14 |
| linalg/dagger/N=1024 (block 512)/qr | 0.117 ± 0.0049 s | 0.117 ± 0.0016 s | 0.995 ± 0.044 |
| linalg/dagger/N=1024 (block 512)/solve (A\b via lu) | 0.0525 ± 0.0022 s | 0.0561 ± 0.0034 s | 0.936 ± 0.07 |
| linalg/dagger/N=1024 (block 512)/syrk (A'*A) | 0.0358 ± 0.0018 s | 0.0725 ± 0.0016 s | 0.494 ± 0.027 |
| linalg/dagger/N=256 (block 256)/cholesky | 3.79 ± 0.48 ms | 2.77 ± 0.17 ms | 1.37 ± 0.19 |
| linalg/dagger/N=256 (block 256)/lu | 4.93 ± 0.34 ms | 4.8 ± 0.099 ms | 1.03 ± 0.075 |
| linalg/dagger/N=256 (block 256)/matmul (A*A) | 2.69 ± 0.22 ms | 4.11 ± 0.52 ms | 0.655 ± 0.099 |
| linalg/dagger/N=256 (block 256)/matvec (A*x) | 2.63 ± 2 ms | 1.95 ± 0.23 ms | 1.35 ± 1 |
| linalg/dagger/N=256 (block 256)/qr | 6.46 ± 0.99 ms | 5.36 ± 1.1 ms | 1.2 ± 0.31 |
| linalg/dagger/N=256 (block 256)/solve (A\b via lu) | 9.67 ± 2.9 ms | 11.1 ± 3.4 ms | 0.872 ± 0.38 |
| linalg/dagger/N=256 (block 256)/syrk (A'*A) | 4.77 ± 1.4 ms | 5.65 ± 0.39 ms | 0.844 ± 0.26 |
| stencil/dagger/N=1024 (block 512)/alloc (neighbors Wrap) | 26.6 ± 0.79 ms | 9.52 ± 0.48 ms | 2.79 ± 0.16 |
| stencil/dagger/N=1024 (block 512)/assign (const) | 1.47 ± 0.075 ms | 1.49 ± 1.6 ms | 0.986 ± 1.1 |
| stencil/dagger/N=1024 (block 512)/multi-expr | 4.18 ± 0.48 ms | 4.44 ± 0.85 ms | 0.942 ± 0.21 |
| stencil/dagger/N=1024 (block 512)/neighbors (Clamp) | 24.2 ± 2.2 ms | 7.35 ± 0.55 ms | 3.29 ± 0.39 |
| stencil/dagger/N=1024 (block 512)/neighbors (Pad) | 23.5 ± 1.2 ms | 7.97 ± 0.84 ms | 2.95 ± 0.35 |
| stencil/dagger/N=1024 (block 512)/neighbors (Reflect) | 24.5 ± 0.82 ms | 7.19 ± 0.84 ms | 3.4 ± 0.41 |
| stencil/dagger/N=1024 (block 512)/neighbors (Wrap) | 24.3 ± 0.6 ms | 8.18 ± 1.5 ms | 2.97 ± 0.53 |
| stencil/dagger/N=1024 (block 512)/update (+) | 2.97 ± 0.38 ms | 2.93 ± 0.29 ms | 1.01 ± 0.16 |
| stencil/dagger/N=256 (block 256)/alloc (neighbors Wrap) | 5.12 ± 0.081 ms | 2.45 ± 0.77 ms | 2.09 ± 0.66 |
| stencil/dagger/N=256 (block 256)/assign (const) | 0.778 ± 0.076 ms | 0.765 ± 0.1 ms | 1.02 ± 0.17 |
| stencil/dagger/N=256 (block 256)/multi-expr | 1.49 ± 0.088 ms | 1.52 ± 0.19 ms | 0.985 ± 0.14 |
| stencil/dagger/N=256 (block 256)/neighbors (Clamp) | 5.11 ± 0.22 ms | 1.8 ± 0.15 ms | 2.85 ± 0.27 |
| stencil/dagger/N=256 (block 256)/neighbors (Pad) | 4.81 ± 0.56 ms | 2.06 ± 0.84 ms | 2.34 ± 0.99 |
| stencil/dagger/N=256 (block 256)/neighbors (Reflect) | 4.84 ± 0.24 ms | 2.35 ± 0.16 ms | 2.06 ± 0.17 |
| stencil/dagger/N=256 (block 256)/neighbors (Wrap) | 4.69 ± 0.15 ms | 1.78 ± 0.81 ms | 2.64 ± 1.2 |
| stencil/dagger/N=256 (block 256)/update (+) | 0.863 ± 0.024 ms | 0.827 ± 0.11 ms | 1.04 ± 0.14 |
| sparse/dagger/N=1024 (block 64)/cg solve (laplacian) | 1.52 ± 0.074 s | ||
| sparse/dagger/N=256 (block 16)/cg solve (laplacian) | 1.47 ± 0.011 s | ||
| sparse/dagger/N=256 (block 16)/spmv (S*x) | 0.0942 ± 0.002 s | ||
| sparse/dagger/N=1024 (block 64)/spgemm (S*S) | 1.28 ± 0.02 s | ||
| sparse/dagger/N=256 (block 16)/spgemm (S*S) | 1.35 ± 0.039 s | ||
| sparse/dagger/N=1024 (block 64)/spmv (S*x) | 0.0966 ± 0.0033 s | ||
| time_to_load | 1.06 ± 0.0074 s | 1.38 ± 0.012 s | 0.767 ± 0.0087 |
Plots
⚠️ Regressions (> 25.0%)
linalg/dagger/N=1024 (block 512)/syrk (A'*A): +102.3%linalg/dagger/N=1024 (block 512)/matmul (A*A): +81.7%linalg/dagger/N=256 (block 256)/matmul (A*A): +52.7%
Improvements (> 25.0% faster)
stencil/dagger/N=1024 (block 512)/neighbors (Reflect): -70.6%stencil/dagger/N=1024 (block 512)/neighbors (Clamp): -69.6%stencil/dagger/N=1024 (block 512)/neighbors (Wrap): -66.4%stencil/dagger/N=1024 (block 512)/neighbors (Pad): -66.1%stencil/dagger/N=256 (block 256)/neighbors (Clamp): -64.9%stencil/dagger/N=1024 (block 512)/alloc (neighbors Wrap): -64.2%stencil/dagger/N=256 (block 256)/neighbors (Wrap): -62.1%stencil/dagger/N=256 (block 256)/neighbors (Pad): -57.2%stencil/dagger/N=256 (block 256)/alloc (neighbors Wrap): -52.2%stencil/dagger/N=256 (block 256)/neighbors (Reflect): -51.5%array/dagger/N=256 (block 256)/reduce (sum): -46.6%linalg/dagger/N=256 (block 256)/cholesky: -26.8%linalg/dagger/N=256 (block 256)/matvec (A*x): -25.8%
Full results and plots (download the benchmark-results artifact).
Add `Dagger.klu` (PureKLU) and `Dagger.splu` (PureUMFPACK) whole-matrix direct solves for sparse `DMatrix`, plus block direct preconditioners (`BlockKLUPreconditioner`/`BlockUMFPACKPreconditioner`). The pure-Julia factorizations are movable, so the factor is gathered/factored on the worker owning the most tiles and pinned there; solves move only O(n) vectors. `Dagger.splu` also exposes two opt-in parallel variants (PureUMFPACK): - `distributed=true, method=:trsv`: re-tile the L/U factors as sparse DMatrices and run a blocked datadeps forward/backward substitution. - `distributed=true, method=:schur`: single-level METIS vertex-separator domain decomposition (`ext/MetisExt.jl`) that factors interior blocks in parallel across workers and reduces/factors the Schur complement. Supporting changes: `_gather_sparse`/`_sparse_copy_of` hooks in SparseArraysExt, Metis/PureKLU/PureUMFPACK weak deps + extensions, and a multi-worker `array/linalg/sparsedirect` test suite (314 tests). Co-authored-by: Cursor <cursoragent@cursor.com>
Always compute region-end write-back from history. The MPI SPMD arg_current shortcut could skip needed copies and leave uninitialized Krylov workspace tiles under multi-worker Distributed. Co-authored-by: Cursor <cursoragent@cursor.com>
Remainder copies previously re-resolved syncdeps via live ainfos_owner / get_read_deps!, which can miss the true producer after ownership moves. Store the writer DTask on each HistoryEntry and wait on it directly. Co-authored-by: Cursor <cursoragent@cursor.com>
MPIExt imports SparseMatrixCSC; Julia 1.12 requires SparseArrays in the extension trigger list so the extension can precompile. Co-authored-by: Cursor <cursoragent@cursor.com>
MPIExt now declares SparseArrays as an extension trigger; the MPI test environment must provide it. Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
…API) Preserve sparsity across GPU moves and Datadeps for SpGEMM/SpMV, using vendor sparse libraries where available and DeviceSparseMatrixCSC otherwise. Co-authored-by: Cursor <cursoragent@cursor.com>
Exercise cg/minres/gmres/bicgstab and Jacobi preconditioning on GPU sparse tiles, with host-safe diagonal/factorize hooks for device CSC. Co-authored-by: Cursor <cursoragent@cursor.com>
Hierarchical Datadeps previously keyed multi-owner partitions on root_worker_id, which is always myid() under MPI. Use partition_affinity_id (rank under MPIExt) so tasks spread across ranks, run Phase 4 with a shared DataDepsState for SPMD-coherent planning, and skip the unused shared-chunk registry on that path so write-back stays correct for Cholesky/LU. Co-authored-by: Cursor <cursoragent@cursor.com>
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>
Clearing `sticky` is not enough to make a task migratable. A hand-built `@task` has no thread pool assigned, and `Base.enq_work` routes such a task onto the *current* thread's work queue rather than the shared multi-queue, so it cannot be picked up until the thread that scheduled it yields. `Threads.@spawn` sets both; our task pools only cleared `sticky`. The effect was that fire tasks piled up behind whichever thread was submitting work. A submitter that runs a long stretch without yielding -- Datadeps planning a whole region, say -- held off every fire task it had just created, so nothing began executing until planning had finished. Setting the thread pool as well is what allows execution to overlap submission. Co-authored-by: Cursor <cursoragent@cursor.com>
`estimate_task_costs!` scanned every chunk once per candidate processor, but chunks are located per worker, so the result depends only on the processor's parent. With `procs` commonly being every thread of a single worker, that scan (and the transfer-rate lock acquisition alongside it) was repeated identically once per thread. Both are now done once per parent. Also skip the final sort when every processor costs the same, which is the usual single-worker case; the shuffle just above already gives the arbitrary-but-fair order that sorting would leave us with. Co-authored-by: Cursor <cursoragent@cursor.com>
`generate_slot!` always routed through `move_rewrap`, even for a `Chunk` that already lives in the destination space, allocating a second Chunk and DRef over the very same memory. That is a MemPool round-trip per argument per region for no benefit. Only leaves qualify for the fast path: `move_rewrap` does more than move bytes, rebuilding wrappers and resolving handles so that what reaches the task is a plain destination-space value (a `ChunkView` becomes a `Chunk` over a real `SubArray`, a nested `Chunk` is flattened). Those results differ from the input even when nothing moves. Co-authored-by: Cursor <cursoragent@cursor.com>
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>
Select halo regions without a runtime tuple index Fusing halos out of views made the halo tuple heterogeneous on GPU backends, which return a plain device array for a contiguous `view` and a `SubArray` for a strided one. `halos[idx]` with a runtime `idx` then infers to a union wide enough that the load after it is a dynamic dispatch, and GPU backends cannot compile a dynamic dispatch at all -- `@stencil` on AMDGPU crashed inside GPUCompiler's IR validation, which segfaults trying to name the offending method. Unroll the selection so each branch sees one concrete array type. Homogeneous tuples, which is what the CPU path always builds, keep the plain indexed load. 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>
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>
jpsamaroo
force-pushed
the
jps/hier-stencils
branch
from
July 27, 2026 23:51
60080d3 to
ec75eb8
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.
Just opening this to get CI and benchmarks - these commits will go into master as their own PRs.