Skip to content

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

Open
jpsamaroo wants to merge 10 commits into
masterfrom
jps/dsparsematrix
Open

DArray: Add sparse matrix support and solvers and preconditioners#708
jpsamaroo wants to merge 10 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) 3.35 ± 3.8 ms 3.85 ± 1.3 ms 0.872 ± 1
array/dagger/N=1024 (block 512)/alloc (rand) 1.99 ± 0.68 ms 1.82 ± 0.54 ms 1.1 ± 0.5
array/dagger/N=1024 (block 512)/broadcast (X .+ 1) 2.15 ± 1.5 ms 1.87 ± 0.51 ms 1.15 ± 0.88
array/dagger/N=1024 (block 512)/map (sin.(X)) 6.24 ± 0.81 ms 7.06 ± 0.62 ms 0.884 ± 0.14
array/dagger/N=1024 (block 512)/norm 1.2 ± 0.074 ms 1.07 ± 0.56 ms 1.12 ± 0.59
array/dagger/N=1024 (block 512)/reduce (sum) 2.12 ± 2.2 ms 3.53 ± 1.9 ms 0.602 ± 0.7
array/dagger/N=1024 (block 512)/transpose (permutedims) 4.83 ± 0.38 ms 3.94 ± 0.66 ms 1.23 ± 0.23
array/dagger/N=256 (block 256)/add (X + X) 1.16 ± 0.014 ms 0.846 ± 0.14 ms 1.38 ± 0.23
array/dagger/N=256 (block 256)/alloc (rand) 0.679 ± 0.14 ms 0.743 ± 0.028 ms 0.913 ± 0.19
array/dagger/N=256 (block 256)/broadcast (X .+ 1) 0.428 ± 0.0042 ms 0.574 ± 0.3 ms 0.745 ± 0.39
array/dagger/N=256 (block 256)/map (sin.(X)) 1.04 ± 0.0095 ms 1.04 ± 0.028 ms 1 ± 0.029
array/dagger/N=256 (block 256)/norm 0.414 ± 0.029 ms 0.52 ± 0.032 ms 0.795 ± 0.073
array/dagger/N=256 (block 256)/reduce (sum) 0.783 ± 0.11 ms
array/dagger/N=256 (block 256)/transpose (permutedims) 0.778 ± 0.25 ms 0.83 ± 0.077 ms 0.938 ± 0.31
linalg/dagger/N=1024 (block 512)/cholesky 19.8 ± 1.5 ms 19.9 ± 0.86 ms 0.996 ± 0.087
linalg/dagger/N=1024 (block 512)/lu 0.0477 ± 0.013 s 0.0448 ± 0.0015 s 1.07 ± 0.3
linalg/dagger/N=1024 (block 512)/matmul (A*A) 0.0484 ± 0.018 s 0.0706 ± 0.00042 s 0.685 ± 0.25
linalg/dagger/N=1024 (block 512)/matvec (A*x) 4.2 ± 1.1 ms 3.53 ± 2.6 ms 1.19 ± 0.93
linalg/dagger/N=1024 (block 512)/qr 0.109 ± 0.0014 s 0.116 ± 0.0044 s 0.94 ± 0.037
linalg/dagger/N=1024 (block 512)/solve (A\b via lu) 0.0574 ± 0.0075 s 0.0488 ± 0.018 s 1.18 ± 0.46
linalg/dagger/N=1024 (block 512)/svd 0.0393 h 0.039 h 1.01
linalg/dagger/N=1024 (block 512)/syrk (A'*A) 0.0321 ± 0.0024 s 0.0643 ± 0.0015 s 0.499 ± 0.039
linalg/dagger/N=256 (block 256)/cholesky 2.35 ± 0.32 ms 4.54 ± 2.2 ms 0.517 ± 0.26
linalg/dagger/N=256 (block 256)/lu 6.67 ± 3.9 ms 3.17 ± 0.27 ms 2.11 ± 1.2
linalg/dagger/N=256 (block 256)/matmul (A*A) 2.38 ± 0.71 ms 2.71 ± 0.041 ms 0.88 ± 0.26
linalg/dagger/N=256 (block 256)/matvec (A*x) 1.39 ± 1.1 ms 0.891 ± 0.013 ms 1.57 ± 1.2
linalg/dagger/N=256 (block 256)/qr 4.49 ± 0.4 ms 5.2 ± 0.57 ms 0.865 ± 0.12
linalg/dagger/N=256 (block 256)/solve (A\b via lu) 8.09 ± 1.8 ms 11.3 ± 1.6 ms 0.717 ± 0.19
linalg/dagger/N=256 (block 256)/svd 0.637 ± 0.13 s 0.637 ± 0.026 s 1 ± 0.21
linalg/dagger/N=256 (block 256)/syrk (A'*A) 2.66 ± 1.7 ms 3.65 ± 0.37 ms 0.73 ± 0.47
stencil/dagger/N=1024 (block 512)/alloc (neighbors Wrap) 7.01 ± 0.79 ms 6.63 ± 1.3 ms 1.06 ± 0.24
stencil/dagger/N=1024 (block 512)/assign (const) 1.4 ± 0.069 ms 1.44 ± 0.89 ms 0.977 ± 0.61
stencil/dagger/N=1024 (block 512)/multi-expr 2.45 ± 0.34 ms 3.74 ± 2.6 ms 0.656 ± 0.46
stencil/dagger/N=1024 (block 512)/neighbors (Clamp) 5.98 ± 0.22 ms 5.99 ± 0.077 ms 0.999 ± 0.039
stencil/dagger/N=1024 (block 512)/neighbors (Pad) 5.85 ± 1.3 ms 6.19 ± 1.5 ms 0.945 ± 0.31
stencil/dagger/N=1024 (block 512)/neighbors (Reflect) 5.79 ± 1.7 ms 6.86 ± 1.9 ms 0.845 ± 0.34
stencil/dagger/N=1024 (block 512)/neighbors (Wrap) 5.36 ± 0.53 ms 6.6 ± 0.39 ms 0.813 ± 0.094
stencil/dagger/N=1024 (block 512)/update (+) 2.12 ± 1.3 ms 2.2 ± 1.6 ms 0.967 ± 0.91
stencil/dagger/N=256 (block 256)/alloc (neighbors Wrap) 1.39 ± 0.025 ms 1.44 ± 1.1 ms 0.964 ± 0.71
stencil/dagger/N=256 (block 256)/assign (const) 0.474 ± 0.011 ms 0.46 ± 0.039 ms 1.03 ± 0.091
stencil/dagger/N=256 (block 256)/multi-expr 1.08 ± 0.13 ms 2.01 ± 1.7 ms 0.536 ± 0.46
stencil/dagger/N=256 (block 256)/neighbors (Clamp) 1.2 ± 0.0036 ms 1.51 ± 1.2 ms 0.794 ± 0.61
stencil/dagger/N=256 (block 256)/neighbors (Pad) 1.16 ± 0.26 ms 1.48 ± 1.4 ms 0.781 ± 0.76
stencil/dagger/N=256 (block 256)/neighbors (Reflect) 1.24 ± 0.09 ms 1.21 ± 0.14 ms 1.02 ± 0.14
stencil/dagger/N=256 (block 256)/neighbors (Wrap) 1.22 ± 0.026 ms 1.24 ± 0.027 ms 0.986 ± 0.03
stencil/dagger/N=256 (block 256)/update (+) 0.553 ± 0.15 ms 0.915 ± 0.99 ms 0.605 ± 0.68
sparse/dagger/N=1024 (block 64)/cg solve (laplacian) 0.771 ± 0.019 s
sparse/dagger/N=256 (block 16)/cg solve (laplacian) 0.788 ± 0.032 s
sparse/dagger/N=256 (block 16)/spmv (S*x) 0.0445 ± 0.0066 s
sparse/dagger/N=1024 (block 64)/spgemm (S*S) 0.406 ± 0.016 s
sparse/dagger/N=256 (block 16)/spgemm (S*S) 0.419 ± 0.022 s
sparse/dagger/N=1024 (block 64)/spmv (S*x) 0.0447 ± 0.016 s
time_to_load 1.05 ± 0.0083 s 1.08 ± 0.0061 s 0.973 ± 0.0094

Allocations / memory

master dirty master / dirty
array/dagger/N=1024 (block 512)/add (X + X) 5.42 k allocs: 8.21 MB 5.51 k allocs: 8.22 MB 0.999
array/dagger/N=1024 (block 512)/alloc (rand) 2.71 k allocs: 8.09 MB 2.69 k allocs: 8.09 MB 1
array/dagger/N=1024 (block 512)/broadcast (X .+ 1) 2.36 k allocs: 8.08 MB 2.4 k allocs: 8.08 MB 1
array/dagger/N=1024 (block 512)/map (sin.(X)) 2.05 k allocs: 8.07 MB 2.08 k allocs: 8.07 MB 1
array/dagger/N=1024 (block 512)/norm 2.68 k allocs: 0.0899 MB 2.67 k allocs: 0.09 MB 0.999
array/dagger/N=1024 (block 512)/reduce (sum) 4.71 k allocs: 0.164 MB 4.5 k allocs: 0.156 MB 1.05
array/dagger/N=1024 (block 512)/transpose (permutedims) 3.29 k allocs: 8.14 MB 3.27 k allocs: 8.13 MB 1
array/dagger/N=256 (block 256)/add (X + X) 1.7 k allocs: 0.574 MB 1.72 k allocs: 0.575 MB 0.997
array/dagger/N=256 (block 256)/alloc (rand) 0.741 k allocs: 0.525 MB 0.745 k allocs: 0.526 MB 1
array/dagger/N=256 (block 256)/broadcast (X .+ 1) 0.638 k allocs: 0.522 MB 0.644 k allocs: 0.522 MB 1
array/dagger/N=256 (block 256)/map (sin.(X)) 0.561 k allocs: 0.519 MB 0.567 k allocs: 0.52 MB 1
array/dagger/N=256 (block 256)/norm 0.702 k allocs: 24.4 kB 0.708 k allocs: 24.6 kB 0.993
array/dagger/N=256 (block 256)/reduce (sum) 1.31 k allocs: 0.0472 MB
array/dagger/N=256 (block 256)/transpose (permutedims) 1.09 k allocs: 0.551 MB 1.09 k allocs: 0.551 MB 0.999
linalg/dagger/N=1024 (block 512)/cholesky 7.65 k allocs: 10.3 MB 7.75 k allocs: 10.3 MB 1
linalg/dagger/N=1024 (block 512)/lu 15.5 k allocs: 14.6 MB 15.5 k allocs: 14.6 MB 1
linalg/dagger/N=1024 (block 512)/matmul (A*A) 7.59 k allocs: 8.28 MB 7.73 k allocs: 8.29 MB 0.999
linalg/dagger/N=1024 (block 512)/matvec (A*x) 6.02 k allocs: 0.241 MB 6.02 k allocs: 0.242 MB 0.995
linalg/dagger/N=1024 (block 512)/qr 10.4 k allocs: 9.56 MB 13.1 k allocs: 9.65 MB 0.99
linalg/dagger/N=1024 (block 512)/solve (A\b via lu) 29 k allocs: 15.2 MB 29.3 k allocs: 15.2 MB 1
linalg/dagger/N=1024 (block 512)/svd 0.0481 M allocs: 0.197 GB 0.0517 M allocs: 0.198 GB 0.999
linalg/dagger/N=1024 (block 512)/syrk (A'*A) 8.74 k allocs: 24.4 MB 9.07 k allocs: 20.4 MB 1.2
linalg/dagger/N=256 (block 256)/cholesky 2.99 k allocs: 0.633 MB 3 k allocs: 0.633 MB 0.999
linalg/dagger/N=256 (block 256)/lu 5.34 k allocs: 1.24 MB 5.38 k allocs: 1.24 MB 0.998
linalg/dagger/N=256 (block 256)/matmul (A*A) 1.9 k allocs: 0.581 MB 1.93 k allocs: 0.582 MB 0.997
linalg/dagger/N=256 (block 256)/matvec (A*x) 2.33 k allocs: 0.0991 MB 2.43 k allocs: 0.104 MB 0.949
linalg/dagger/N=256 (block 256)/qr 3.4 k allocs: 0.778 MB 3.81 k allocs: 0.793 MB 0.981
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.997
linalg/dagger/N=256 (block 256)/svd 15.3 k allocs: 6.71 MB 15.8 k allocs: 6.73 MB 0.997
linalg/dagger/N=256 (block 256)/syrk (A'*A) 2.95 k allocs: 2.63 MB 3.06 k allocs: 2.14 MB 1.23
stencil/dagger/N=1024 (block 512)/alloc (neighbors Wrap) 4.65 k allocs: 8.2 MB 4.77 k allocs: 8.21 MB 0.999
stencil/dagger/N=1024 (block 512)/assign (const) 2.14 k allocs: 0.0947 MB 2.12 k allocs: 0.0941 MB 1.01
stencil/dagger/N=1024 (block 512)/multi-expr 4.71 k allocs: 0.21 MB 4.71 k allocs: 0.21 MB 0.998
stencil/dagger/N=1024 (block 512)/neighbors (Clamp) 3.7 k allocs: 0.23 MB 3.73 k allocs: 0.231 MB 0.995
stencil/dagger/N=1024 (block 512)/neighbors (Pad) 3.6 k allocs: 0.227 MB 3.57 k allocs: 0.226 MB 1
stencil/dagger/N=1024 (block 512)/neighbors (Reflect) 3.81 k allocs: 2.26 MB 3.71 k allocs: 0.262 MB 8.63
stencil/dagger/N=1024 (block 512)/neighbors (Wrap) 3.55 k allocs: 0.166 MB 3.5 k allocs: 0.164 MB 1.01
stencil/dagger/N=1024 (block 512)/update (+) 2.54 k allocs: 0.114 MB 2.59 k allocs: 0.116 MB 0.989
stencil/dagger/N=256 (block 256)/alloc (neighbors Wrap) 1.54 k allocs: 0.574 MB 1.53 k allocs: 0.575 MB 0.999
stencil/dagger/N=256 (block 256)/assign (const) 0.793 k allocs: 0.0405 MB 0.8 k allocs: 0.0408 MB 0.993
stencil/dagger/N=256 (block 256)/multi-expr 1.72 k allocs: 0.0871 MB 1.71 k allocs: 0.0869 MB 1
stencil/dagger/N=256 (block 256)/neighbors (Clamp) 1.24 k allocs: 0.0714 MB 1.23 k allocs: 0.0711 MB 1.01
stencil/dagger/N=256 (block 256)/neighbors (Pad) 1.2 k allocs: 0.0705 MB 1.18 k allocs: 0.0696 MB 1.01
stencil/dagger/N=256 (block 256)/neighbors (Reflect) 1.23 k allocs: 0.0793 MB 1.24 k allocs: 0.0795 MB 0.997
stencil/dagger/N=256 (block 256)/neighbors (Wrap) 1.17 k allocs: 0.0618 MB 1.18 k allocs: 0.0626 MB 0.988
stencil/dagger/N=256 (block 256)/update (+) 0.915 k allocs: 0.0462 MB 0.922 k allocs: 0.0464 MB 0.994
sparse/dagger/N=1024 (block 64)/cg solve (laplacian) 2.93 M allocs: 0.11 GB
sparse/dagger/N=256 (block 16)/cg solve (laplacian) 2.91 M allocs: 0.109 GB
sparse/dagger/N=256 (block 16)/spmv (S*x) 0.162 M allocs: 6.56 MB
sparse/dagger/N=1024 (block 64)/spgemm (S*S) 3.43 M allocs: 0.16 GB
sparse/dagger/N=256 (block 16)/spgemm (S*S) 3.32 M allocs: 0.121 GB
sparse/dagger/N=1024 (block 64)/spmv (S*x) 0.188 M allocs: 7.22 MB
time_to_load 0.147 k allocs: 10.8 kB 0.147 k allocs: 10.8 kB 1

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

  • linalg/dagger/N=1024 (block 512)/syrk (A'*A) (time): +100.3%
  • linalg/dagger/N=1024 (block 512)/matmul (A*A) (time): +45.9%
  • linalg/dagger/N=1024 (block 512)/qr (allocs): +26.0%
  • array/dagger/N=256 (block 256)/norm (time): +25.7%

Improvements

  • stencil/dagger/N=1024 (block 512)/neighbors (Reflect) (memory): -88.4%
  • array/dagger/N=256 (block 256)/add (X + X) (time): -27.3%
Within noise (12 metric(s) past threshold but inside the ±spread; not counted)
  • linalg/dagger/N=256 (block 256)/cholesky (time): 93.5%
  • stencil/dagger/N=256 (block 256)/multi-expr (time): 86.4%
  • array/dagger/N=1024 (block 512)/reduce (sum) (time): 66.1%
  • stencil/dagger/N=256 (block 256)/update (+) (time): 65.4%
  • stencil/dagger/N=1024 (block 512)/multi-expr (time): 52.5%
  • linalg/dagger/N=256 (block 256)/solve (A\b via lu) (time): 39.4%
  • linalg/dagger/N=256 (block 256)/syrk (A'*A) (time): 37.0%
  • array/dagger/N=256 (block 256)/broadcast (X .+ 1) (time): 34.2%
  • stencil/dagger/N=256 (block 256)/neighbors (Pad) (time): 28.0%
  • stencil/dagger/N=256 (block 256)/neighbors (Clamp) (time): 25.9%
  • linalg/dagger/N=256 (block 256)/matvec (A*x) (time): -36.1%
  • linalg/dagger/N=256 (block 256)/lu (time): -52.5%

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 and others added 10 commits August 31, 2026 17:50
@jpsamaroo
DArray: Add proper sparse matrix support
5a0f604
@jpsamaroo
DArray/sparse: Add iterative solvers via Krylov.jl
f51946f
@jpsamaroo
DArray/sparse: Add AMG and ILU preconditioner integrations
0f324ae
@jpsamaroo @cursoragent
DArray/sparse: Add sparse direct solvers (KLU/UMFPACK)
c22e4f2 
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
0840a27
@jpsamaroo @cursoragent
DArray/sparse: Add GPU sparse tile support (CUDA/ROC/OpenCL/Metal/one…
84c9c1e 
…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
4f57696 
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
afa5319
@jpsamaroo
DArray: Add fill!
821cab7
@jpsamaroo @cursoragent
AGENTS.md: Record the acceleration x backend and extension-isolation …
c283be0 
…lessons

Two bugs in whole-object sparse tiles (a 4-argument `move!` that is ambiguous
with every GPU extension's `AbstractArray` method, and a `collect` gather that
scalar-indexes device storage because it runs under the caller's GPU scope)
passed both the single-process GPU suite and the MPI-CPU suite, and only showed
up once MPI x GPU ran. Record what makes that cell distinct, and record why
shared generics for cross-cutting interop belong in core rather than reached at
across two extensions of the same package.

Co-authored-by: Cursor <cursoragent@cursor.com>
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