Skip to content

DArray: Add sparse matrix support and solvers and preconditioners - #708

Merged
jpsamaroo merged 15 commits into
masterfrom
jps/dsparsematrix
Sep 4, 2026
Merged

DArray: Add sparse matrix support and solvers and preconditioners#708
jpsamaroo merged 15 commits into
masterfrom
jps/dsparsematrix

Conversation

@jpsamaroo

@jpsamaroo jpsamaroo commented Jun 10, 2026

Copy link
Copy Markdown
Member

Adds a DSparseMatrix type for sparse DArray tiles, that allows in-place modifications, necessary for Datadeps algorithms. Then this builds algorithms like GEMM, GEMV, iterative solvers, and preconditioners on top that allow the DArray to work with sparse matrix and vector data effectively.

Written by Claude Opus

@jpsamaroo
jpsamaroo force-pushed the jps/dsparsematrix branch from 1269c08 to a212025 Compare July 24, 2026 03:06
@github-actions

github-actions Bot commented Jul 24, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Dagger benchmarks: dirty vs master

Median time

master dirty master / dirty
array/dagger/N=1024 (block 512)/add (X + X) 5.07 ± 1.2 ms 2.69 ± 0.82 ms 1.89 ± 0.72
array/dagger/N=1024 (block 512)/alloc (rand) 2.25 ± 0.39 ms 2.49 ± 2.7 ms 0.903 ± 0.99
array/dagger/N=1024 (block 512)/broadcast (X .+ 1) 2.46 ± 1.5 ms 2.03 ± 2.8 ms 1.22 ± 1.8
array/dagger/N=1024 (block 512)/map (sin.(X)) 5.84 ± 0.98 ms 4.04 ± 1.1 ms 1.45 ± 0.46
array/dagger/N=1024 (block 512)/norm 1.16 ± 0.15 ms 1.15 ± 0.44 ms 1.01 ± 0.41
array/dagger/N=1024 (block 512)/reduce (sum) 1.91 ± 0.68 ms 2.06 ± 0.28 ms 0.925 ± 0.35
array/dagger/N=1024 (block 512)/transpose (permutedims) 3.24 ± 0.33 ms 2.58 ± 0.7 ms 1.26 ± 0.36
array/dagger/N=256 (block 256)/add (X + X) 0.732 ± 0.077 ms 1.27 ± 1.1 ms 0.576 ± 0.49
array/dagger/N=256 (block 256)/alloc (rand) 0.67 ± 0.29 ms 0.679 ± 0.068 ms 0.987 ± 0.44
array/dagger/N=256 (block 256)/broadcast (X .+ 1) 0.538 ± 0.14 ms 0.387 ± 0.036 ms 1.39 ± 0.38
array/dagger/N=256 (block 256)/map (sin.(X)) 0.893 ± 0.17 ms 0.868 ± 0.1 ms 1.03 ± 0.23
array/dagger/N=256 (block 256)/norm 0.378 ± 0.1 ms 0.424 ± 0.088 ms 0.893 ± 0.31
array/dagger/N=256 (block 256)/reduce (sum) 0.761 ± 0.36 ms 0.519 ± 0.045 ms 1.47 ± 0.71
array/dagger/N=256 (block 256)/transpose (permutedims) 0.544 ± 0.021 ms 0.698 ± 0.037 ms 0.779 ± 0.051
linalg/dagger/N=1024 (block 512)/cholesky 17.7 ± 3 ms 18.5 ± 7.8 ms 0.954 ± 0.43
linalg/dagger/N=1024 (block 512)/lu 31.4 ± 3.1 ms 29.8 ± 6.8 ms 1.06 ± 0.26
linalg/dagger/N=1024 (block 512)/matmul (A*A) 25.6 ± 1.4 ms 27.6 ± 3.7 ms 0.928 ± 0.13
linalg/dagger/N=1024 (block 512)/matvec (A*x) 3.81 ± 1.3 ms 3.32 ± 1.3 ms 1.15 ± 0.6
linalg/dagger/N=1024 (block 512)/qr 0.0851 ± 0.0064 s 0.0865 ± 0.006 s 0.984 ± 0.1
linalg/dagger/N=1024 (block 512)/solve (A\b via lu) 0.0385 ± 0.003 s 0.0416 ± 0.0071 s 0.923 ± 0.17
linalg/dagger/N=1024 (block 512)/svd 48.1 s 49.4 s 0.972
linalg/dagger/N=1024 (block 512)/syrk (A'*A) 25 ± 2.1 ms 22.9 ± 0.51 ms 1.09 ± 0.097
linalg/dagger/N=256 (block 256)/cholesky 4.64 ± 2.9 ms 4.52 ± 0.93 ms 1.03 ± 0.68
linalg/dagger/N=256 (block 256)/lu 3.39 ± 0.11 ms 5.87 ± 0.087 ms 0.577 ± 0.02
linalg/dagger/N=256 (block 256)/matmul (A*A) 1.51 ± 0.23 ms 1.55 ± 0.38 ms 0.975 ± 0.28
linalg/dagger/N=256 (block 256)/matvec (A*x) 1.94 ± 0.6 ms 1.21 ± 0.59 ms 1.6 ± 0.93
linalg/dagger/N=256 (block 256)/qr 4.98 ± 0.36 ms 3.57 ± 0.46 ms 1.4 ± 0.21
linalg/dagger/N=256 (block 256)/solve (A\b via lu) 8.55 ± 6.6 ms 12.7 ± 3.2 ms 0.673 ± 0.55
linalg/dagger/N=256 (block 256)/svd 0.39 ± 0.0061 s 0.387 ± 0.0069 s 1.01 ± 0.024
linalg/dagger/N=256 (block 256)/syrk (A'*A) 3.19 ± 2.2 ms 4.52 ± 2.3 ms 0.707 ± 0.6
stencil/dagger/N=1024 (block 512)/alloc (neighbors Wrap) 5.26 ± 1.8 ms 5.59 ± 2.3 ms 0.941 ± 0.5
stencil/dagger/N=1024 (block 512)/assign (const) 1.08 ± 0.41 ms 1.19 ± 0.19 ms 0.911 ± 0.38
stencil/dagger/N=1024 (block 512)/multi-expr 4.2 ± 1.8 ms 2.75 ± 2.7 ms 1.53 ± 1.6
stencil/dagger/N=1024 (block 512)/neighbors (Clamp) 5.38 ± 0.86 ms 4.41 ± 0.29 ms 1.22 ± 0.21
stencil/dagger/N=1024 (block 512)/neighbors (Pad) 5.29 ± 0.76 ms 5.18 ± 0.34 ms 1.02 ± 0.16
stencil/dagger/N=1024 (block 512)/neighbors (Reflect) 5.73 ± 0.81 ms 5.62 ± 0.87 ms 1.02 ± 0.21
stencil/dagger/N=1024 (block 512)/neighbors (Wrap) 5.37 ± 0.58 ms 5.44 ± 0.6 ms 0.986 ± 0.15
stencil/dagger/N=1024 (block 512)/update (+) 1.67 ± 0.3 ms 1.53 ± 0.071 ms 1.09 ± 0.2
stencil/dagger/N=256 (block 256)/alloc (neighbors Wrap) 1.58 ± 0.071 ms 1.69 ± 0.072 ms 0.932 ± 0.058
stencil/dagger/N=256 (block 256)/assign (const) 0.488 ± 0.11 ms 0.439 ± 0.058 ms 1.11 ± 0.29
stencil/dagger/N=256 (block 256)/multi-expr 1.09 ± 0.88 ms 2.05 ± 1.3 ms 0.532 ± 0.55
stencil/dagger/N=256 (block 256)/neighbors (Clamp) 1.12 ± 0.47 ms 1.06 ± 0.045 ms 1.05 ± 0.44
stencil/dagger/N=256 (block 256)/neighbors (Pad) 1.59 ± 0.79 ms 1.4 ± 0.33 ms 1.13 ± 0.62
stencil/dagger/N=256 (block 256)/neighbors (Reflect) 1.05 ± 0.1 ms 1.22 ± 0.061 ms 0.857 ± 0.093
stencil/dagger/N=256 (block 256)/neighbors (Wrap) 1.09 ± 0.44 ms 1.18 ± 0.024 ms 0.927 ± 0.38
stencil/dagger/N=256 (block 256)/update (+) 1.56 ± 2.3 ms 0.698 ± 0.033 ms 2.23 ± 3.4
sparse/dagger/N=1024 (block 64)/cg solve (laplacian) 0.562 ± 0.018 s
sparse/dagger/N=256 (block 16)/cg solve (laplacian) 0.596 ± 0.018 s
sparse/dagger/N=256 (block 16)/spmv (S*x) 29.9 ± 2.6 ms
sparse/dagger/N=1024 (block 64)/spgemm (S*S) 0.307 ± 0.017 s
sparse/dagger/N=256 (block 16)/spgemm (S*S) 0.308 ± 0.018 s
sparse/dagger/N=1024 (block 64)/spmv (S*x) 0.0362 ± 4.7e-05 s
time_to_load 0.948 ± 0.022 s 0.938 ± 0.0023 s 1.01 ± 0.023

Allocations / memory

master dirty master / dirty
array/dagger/N=1024 (block 512)/add (X + X) 5.36 k allocs: 8.21 MB 5.56 k allocs: 8.22 MB 0.999
array/dagger/N=1024 (block 512)/alloc (rand) 2.66 k allocs: 8.09 MB 2.77 k allocs: 8.09 MB 1
array/dagger/N=1024 (block 512)/broadcast (X .+ 1) 2.35 k allocs: 8.08 MB 2.38 k allocs: 8.08 MB 1
array/dagger/N=1024 (block 512)/map (sin.(X)) 2.04 k allocs: 8.07 MB 2.09 k allocs: 8.07 MB 1
array/dagger/N=1024 (block 512)/norm 2.64 k allocs: 0.0886 MB 2.7 k allocs: 0.09 MB 0.985
array/dagger/N=1024 (block 512)/reduce (sum) 3.79 k allocs: 0.131 MB 4.07 k allocs: 0.139 MB 0.939
array/dagger/N=1024 (block 512)/transpose (permutedims) 3.22 k allocs: 8.13 MB 3.32 k allocs: 8.14 MB 1
array/dagger/N=256 (block 256)/add (X + X) 1.69 k allocs: 0.574 MB 1.72 k allocs: 0.575 MB 1
array/dagger/N=256 (block 256)/alloc (rand) 0.744 k allocs: 0.526 MB 0.754 k allocs: 0.526 MB 1
array/dagger/N=256 (block 256)/broadcast (X .+ 1) 0.612 k allocs: 0.521 MB 0.651 k allocs: 0.522 MB 0.998
array/dagger/N=256 (block 256)/map (sin.(X)) 0.561 k allocs: 0.519 MB 0.574 k allocs: 0.52 MB 1
array/dagger/N=256 (block 256)/norm 0.702 k allocs: 24.4 kB 0.715 k allocs: 24.6 kB 0.992
array/dagger/N=256 (block 256)/reduce (sum) 1.05 k allocs: 0.0381 MB 0.816 k allocs: 30.5 kB 1.28
array/dagger/N=256 (block 256)/transpose (permutedims) 1.09 k allocs: 0.551 MB 1.11 k allocs: 0.551 MB 0.999
linalg/dagger/N=1024 (block 512)/cholesky 7.64 k allocs: 10.3 MB 7.9 k allocs: 10.3 MB 0.999
linalg/dagger/N=1024 (block 512)/lu 15.2 k allocs: 14.6 MB 15.6 k allocs: 14.6 MB 0.999
linalg/dagger/N=1024 (block 512)/matmul (A*A) 7.51 k allocs: 8.28 MB 7.75 k allocs: 8.29 MB 0.999
linalg/dagger/N=1024 (block 512)/matvec (A*x) 5.71 k allocs: 0.231 MB 6.38 k allocs: 0.253 MB 0.913
linalg/dagger/N=1024 (block 512)/qr 10.3 k allocs: 9.55 MB 11.7 k allocs: 9.6 MB 0.995
linalg/dagger/N=1024 (block 512)/solve (A\b via lu) 28.8 k allocs: 15.2 MB 29.6 k allocs: 15.2 MB 0.998
linalg/dagger/N=1024 (block 512)/svd 0.119 M allocs: 0.2 GB 0.0502 M allocs: 0.197 GB 1.01
linalg/dagger/N=1024 (block 512)/syrk (A'*A) 8.73 k allocs: 24.4 MB 9.16 k allocs: 20.4 MB 1.2
linalg/dagger/N=256 (block 256)/cholesky 2.99 k allocs: 0.632 MB 3.03 k allocs: 0.634 MB 0.998
linalg/dagger/N=256 (block 256)/lu 5.34 k allocs: 1.24 MB 5.41 k allocs: 1.24 MB 0.999
linalg/dagger/N=256 (block 256)/matmul (A*A) 1.9 k allocs: 0.581 MB 1.94 k allocs: 0.582 MB 0.998
linalg/dagger/N=256 (block 256)/matvec (A*x) 2.33 k allocs: 0.0991 MB 2.38 k allocs: 0.1 MB 0.988
linalg/dagger/N=256 (block 256)/qr 3.4 k allocs: 0.778 MB 3.46 k allocs: 0.78 MB 0.998
linalg/dagger/N=256 (block 256)/solve (A\b via lu) 11.8 k allocs: 1.53 MB 11.9 k allocs: 1.53 MB 0.998
linalg/dagger/N=256 (block 256)/svd 15.3 k allocs: 6.71 MB 15.5 k allocs: 6.71 MB 0.999
linalg/dagger/N=256 (block 256)/syrk (A'*A) 2.96 k allocs: 2.63 MB 3.06 k allocs: 2.14 MB 1.23
stencil/dagger/N=1024 (block 512)/alloc (neighbors Wrap) 4.6 k allocs: 8.2 MB 4.79 k allocs: 8.2 MB 1
stencil/dagger/N=1024 (block 512)/assign (const) 2.04 k allocs: 0.091 MB 2.13 k allocs: 0.0934 MB 0.974
stencil/dagger/N=1024 (block 512)/multi-expr 4.55 k allocs: 0.204 MB 4.72 k allocs: 0.209 MB 0.974
stencil/dagger/N=1024 (block 512)/neighbors (Clamp) 3.67 k allocs: 0.229 MB 3.79 k allocs: 0.231 MB 0.991
stencil/dagger/N=1024 (block 512)/neighbors (Pad) 3.53 k allocs: 0.225 MB 3.68 k allocs: 0.228 MB 0.985
stencil/dagger/N=1024 (block 512)/neighbors (Reflect) 3.65 k allocs: 0.26 MB 3.8 k allocs: 0.263 MB 0.989
stencil/dagger/N=1024 (block 512)/neighbors (Wrap) 3.51 k allocs: 0.165 MB 3.65 k allocs: 0.168 MB 0.983
stencil/dagger/N=1024 (block 512)/update (+) 2.52 k allocs: 0.113 MB 2.58 k allocs: 0.115 MB 0.989
stencil/dagger/N=256 (block 256)/alloc (neighbors Wrap) 1.52 k allocs: 0.573 MB 1.56 k allocs: 0.575 MB 0.998
stencil/dagger/N=256 (block 256)/assign (const) 0.798 k allocs: 0.0405 MB 0.803 k allocs: 0.0408 MB 0.994
stencil/dagger/N=256 (block 256)/multi-expr 1.71 k allocs: 0.0864 MB 1.74 k allocs: 0.0874 MB 0.989
stencil/dagger/N=256 (block 256)/neighbors (Clamp) 1.23 k allocs: 0.0708 MB 1.26 k allocs: 0.0715 MB 0.99
stencil/dagger/N=256 (block 256)/neighbors (Pad) 1.18 k allocs: 0.0696 MB 1.2 k allocs: 0.0696 MB 1
stencil/dagger/N=256 (block 256)/neighbors (Reflect) 1.24 k allocs: 0.0789 MB 1.27 k allocs: 0.0796 MB 0.991
stencil/dagger/N=256 (block 256)/neighbors (Wrap) 1.18 k allocs: 0.0625 MB 1.2 k allocs: 0.0625 MB 1
stencil/dagger/N=256 (block 256)/update (+) 0.912 k allocs: 0.0459 MB 0.913 k allocs: 0.046 MB 0.998
sparse/dagger/N=1024 (block 64)/cg solve (laplacian) 2.82 M allocs: 0.105 GB
sparse/dagger/N=256 (block 16)/cg solve (laplacian) 2.89 M allocs: 0.107 GB
sparse/dagger/N=256 (block 16)/spmv (S*x) 0.182 M allocs: 6.92 MB
sparse/dagger/N=1024 (block 64)/spgemm (S*S) 3.14 M allocs: 0.15 GB
sparse/dagger/N=256 (block 16)/spgemm (S*S) 3.11 M allocs: 0.112 GB
sparse/dagger/N=1024 (block 64)/spmv (S*x) 0.178 M allocs: 6.8 MB
time_to_load 0.147 k allocs: 10.8 kB 0.147 k allocs: 10.8 kB 1

Plots

⚠️ Regressions (time > 25.0% and outside the reported ±spread; allocs/memory > 25.0%)

  • linalg/dagger/N=256 (block 256)/lu (time): +73.2%
  • array/dagger/N=256 (block 256)/transpose (permutedims) (time): +28.4%

Improvements

  • linalg/dagger/N=1024 (block 512)/svd (allocs): -58.0%
  • array/dagger/N=1024 (block 512)/add (X + X) (time): -47.0%
  • linalg/dagger/N=256 (block 256)/qr (time): -28.4%
Within noise (10 metric(s) past threshold but inside the ±spread; not counted)
  • stencil/dagger/N=256 (block 256)/multi-expr (time): 88.1%
  • array/dagger/N=256 (block 256)/add (X + X) (time): 73.6%
  • linalg/dagger/N=256 (block 256)/solve (A\b via lu) (time): 48.5%
  • linalg/dagger/N=256 (block 256)/syrk (A'*A) (time): 41.5%
  • array/dagger/N=256 (block 256)/broadcast (X .+ 1) (time): -28.1%
  • array/dagger/N=1024 (block 512)/map (sin.(X)) (time): -30.8%
  • array/dagger/N=256 (block 256)/reduce (sum) (time): -31.8%
  • stencil/dagger/N=1024 (block 512)/multi-expr (time): -34.6%
  • linalg/dagger/N=256 (block 256)/matvec (A*x) (time): -37.5%
  • stencil/dagger/N=256 (block 256)/update (+) (time): -55.2%

Full results and plots (download the benchmark-results artifact).

@jpsamaroo
jpsamaroo force-pushed the jps/dsparsematrix branch 3 times, most recently from c9a194b to 03a0842 Compare July 27, 2026 23:43
@jpsamaroo
jpsamaroo force-pushed the jps/dsparsematrix branch 2 times, most recently from e0b3eb7 to 0a5eee0 Compare September 3, 2026 19:04
jpsamaroo and others added 12 commits September 3, 2026 20:23
@jpsamaroo
DArray: Add proper sparse matrix support
86d8339
@jpsamaroo @cursoragent
DArray/sparse: Add iterative solvers via Krylov.jl
b292650 
Co-authored-by: Cursor <cursoragent@cursor.com>
@jpsamaroo
DArray/sparse: Add AMG and ILU preconditioner integrations
16b54d8
@jpsamaroo @cursoragent
DArray/sparse: Add sparse direct solvers (KLU/UMFPACK)
331a41a 
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>
@jpsamaroo
docs: Document sparse DArrays
794a426
@jpsamaroo @cursoragent
DArray/sparse: Add GPU sparse tile support (CUDA/ROC/OpenCL/Metal/one…
f2725c7 
…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>
@jpsamaroo @cursoragent
DArray/sparse: Add GPU sparse Krylov solver tests
35ef351 
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>
@jpsamaroo
DArray: Base.fetch(::DArray) just converts DTask to Chunk
6f406a0
@jpsamaroo
DArray: Add fill!
6fe8ed0
@jpsamaroo @cursoragent
compat: Require DataStructures 0.19
01ed280 
Sch.jl uses popfirst!(::PriorityQueue), which exists only in 0.19. 0.18
resolves and loads, then the scheduler throws on the first pop — so a
downstream pin should be bumped rather than re-adding 0.18 here.

Co-authored-by: Cursor <cursoragent@cursor.com>
@jpsamaroo @cursoragent
GPU: Move device processor types and with_context into Dagger
9392ac4 
GPU×SparseArrays extensions were reaching into CUDAExt/ROCExt/… via
get_extension for the processor type and with_context, which is a coin
flip on extension load order. The processor types now live in Dagger
with only owner+device fields, and with_context is a Dagger generic
that each GPU extension specializes.

Co-authored-by: Cursor <cursoragent@cursor.com>
@jpsamaroo @cursoragent
CI: Exercise Julia pre-releases in GHA matrices
4c119f9 
Catch incompatibilities on the next RC/beta/alpha instead of only after it becomes stable.

Co-authored-by: Cursor <cursoragent@cursor.com>
jpsamaroo and others added 3 commits September 4, 2026 09:43
@jpsamaroo @cursoragent
GPU: Fix vendor sparse CI failures on CUDA, ROCm, oneAPI, and Metal
393689e 
CUSPARSE stores nzVal not nzval, rocSPARSE CSC*CSC falls into scalar generic
mul, oneAPI 2.x types zeMemOpenIpcHandle's out-param as PtrOrZePtr, and
MetalExt must not overwrite the core show during precompilation.

Co-authored-by: Cursor <cursoragent@cursor.com>
@jpsamaroo @cursoragent
MPI/GPU: Do not IPC-export host Arrays stamped with a VRAM space
54f4130 
Collect's MPI cat tree densifies tiles to Array under the GPU compute
scope; labeling that result with the processor space sent the next hop
through device IPC and failed on Matrix.

Co-authored-by: Cursor <cursoragent@cursor.com>
@jpsamaroo @cursoragent
compat: Pin test AMD below 0.5.4 to unbreak LinearSolve precompile
c3326ac 
AMD 0.5.4 dropped SS_Int, which SparseColumnPivotedQR 2.1.7 still reads, so AlgebraicMultigrid cannot load.

Co-authored-by: Cursor <cursoragent@cursor.com>
@jpsamaroo
jpsamaroo merged commit 23837ab into master Sep 4, 2026
27 of 39 checks passed
@jpsamaroo
jpsamaroo deleted the jps/dsparsematrix branch September 4, 2026 20:37
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

linear algebra sparse

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@jpsamaroo