DArray: Add sparse matrix support and solvers and preconditioners

.buildkite/pipeline.yml

@@ -111,5 +111,17 @@ steps:
       julia --project=test/metalenv -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()'
       julia --project=test/metalenv test/run_mpi.jl 2 2 test/mpi_metal.jl
 
+  - label: Julia 1.11 (Finch)
+    timeout_in_minutes: 90
+    if: build.message !~ /\[skip tests\]/
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "1.11"
+      - JuliaCI/julia-test#v1: ~
+      - JuliaCI/julia-coverage#v1:
+          codecov: true
+    env:
+      CI_TEST_FINCH: "1"
+
 # env:
 #   SECRET_CODECOV_TOKEN: ""

.github/workflows/CI.yml

@@ -33,6 +33,7 @@ jobs:
           - {version: '1', os: ubuntu-latest, arch: x64}
           - {version: '1', os: windows-latest, arch: x64}
           - {version: '1', os: macos-26, arch: arm64}
+          - {version: 'pre', os: ubuntu-latest, arch: x64}
           - {version: 'nightly', os: ubuntu-latest, arch: x64}
     steps:
       - uses: actions/checkout@v7
@@ -69,6 +70,7 @@ jobs:
           - {version: '1', os: ubuntu-latest, arch: x64}
           - {version: '1', os: windows-latest, arch: x64}
           - {version: '1', os: macos-26, arch: arm64}
+          - {version: 'pre', os: ubuntu-latest, arch: x64}
           - {version: 'nightly', os: ubuntu-latest, arch: x64}
     steps:
       - uses: actions/checkout@v7
@@ -244,6 +246,7 @@ jobs:
           - '1.10'   # LTS
           - '1.11'
           - '1'      # latest stable
+          - 'pre'    # latest prerelease
         ranks:
           - 2
           - 4

AGENTS.md

@@ -156,3 +156,118 @@ lesson.
    term to one candidate and asserts the *other* one wins is what catches it.
    The behavioral check is cheaper still: 40 sleeping tasks over 4 workers
    land `[1 => 40]` when the term is dead and `10/10/10/10` when it is live.
+
+16. **Only the acceleration x backend cross product exercises the data-movement
+   paths, and one of the axes alone will pass while they are broken.** Two
+   independent bugs in whole-object (`aliases_as_whole`) sparse tiles survived a
+   green single-process-GPU suite *and* a green MPI-CPU suite, and both fell out
+   the first time MPI x GPU ran. First: a host-to-device `move!` of a whole
+   object only happens when tiles are *built* on the host and then placed on a
+   device, which no single-axis configuration does. And such a hook must be
+   defined at the `dep_mod` (5-argument) arity, because every GPU extension
+   claims `move!(::TheirVRAMSpace, ::OtherSpace, ::AbstractArray{T,N},
+   ::AbstractArray{T,N})` for the 4-argument form — a container method with
+   unconstrained spaces is genuinely ambiguous with all of them (more specific
+   in the value arguments, less specific in the space arguments), so it would
+   need a tie-breaker per backend per space pair; nothing outside core defines
+   the 5-argument form for arrays, and that is what every Datadeps copy path
+   calls anyway. Second: `collect`'s gather tasks run wherever the caller's
+   compute scope puts them, so under a GPU scope the `cat` tree runs *on the
+   device* — and generic `cat` fills its output element by element, which is
+   scalar indexing. Keep shared test bodies in a `test/array/*_defs.jl` file
+   (as `stencil_defs.jl` and `sparse_defs.jl` do) and call them from all four
+   entry points; `test/mpi_opencl.jl` makes the fourth cell cheap to run
+   locally.
+
+17. **Extensions of the same package must not reach into each other.** Load
+   order between two extensions of one package is unspecified, so
+   `Base.get_extension(Dagger, :MPIExt)` from another extension is a coin flip.
+   When `AExt` and `A×BExt` both need to extend the same generic, declare that
+   generic in core Dagger and let each extension add methods to it (see
+   `inplace_mpi_parts` in `src/memory-spaces.jl`, the
+   `mpi_device_direct`/`mpi_remap_space` hooks in `src/gpu.jl`, and the GPU
+   processor types / `with_context` in `src/gpu.jl`). GPU×SparseArrays
+   extensions import `CuArrayDeviceProc` (etc.) from Dagger and call
+   `Dagger.with_context` — never `Base.get_extension(Dagger, :CUDAExt)`. The tempting
+   shortcut — adding `B` to `AExt`'s trigger list — is worse than it looks: it
+   makes `AExt` refuse to load at all until `B` is loaded, so MPI acceleration
+   would have silently required SparseArrays.
+
+18. **The scheduler needs DataStructures 0.19.** `Sch.jl` does
+   `popfirst!(::PriorityQueue)`. That method exists only in DataStructures
+   0.19; 0.18 resolves, loads, and then the scheduler throws on the first
+   pop. Compat is `0.19` only — do not re-add `0.18` to satisfy a downstream
+   pin. If a demo package pins 0.18, bump *that* package's compat (as the
+   Jutul clone patch does), not Dagger's.
+
+19. **Per-tile AMG can report `stats.solved` while `‖Ax−b‖` is O(1)–O(100).**
+   `AMGPreconditioner` is block-diagonal: one V-cycle per diagonal tile,
+   applied as Krylov `M` (`ldiv=false`). Left-preconditioned GMRES/BiCGStab
+   then converge in the *preconditioned* residual. With many tiles that
+   V-cycle is a weak additive-Schwarz operator, so Krylov stops while the
+   true residual is huge (seen on Chan, VoronoiFVM penalty rows, and Jutul
+   heat). BlockJacobi (exact tile LU) on the same layout is fine. Global AMG
+   is `Blocks(n, n)` (one tile). Do not treat `stats.solved` as `Ax≈b` for
+   block AMG; check the un-preconditioned residual. CG will also reject AMG
+   as non-SPD — use GMRES.
+
+21. **Reassigning a local to a value of a different type deoptimizes the
+   *whole* body, not the assignment.** `arg = adopt_sparse_arg!(state, arg,
+   deps)` in `_populate_one_arg!` looks like a cheap normalization, but
+   inference must pick one type for `arg` over the entire method, so it
+   widens to the join and every later use — `type_may_alias(typeof(arg))`,
+   `supports_inplace_move`, the `ArgumentWrapper` construction,
+   `get_or_make_arg_chunk!` — becomes a dynamic dispatch on a boxed value.
+   The cost lands on the *common* path (the branch that never fires) and is
+   invisible in a diff of the branch that does. Argument-processing code is
+   per-argument per-task, so this is the worst possible place for it. Pass
+   the new value forward into a function barrier instead of assigning it
+   back; the callee then specializes per concrete type and the conditional
+   is confined to choosing which call to make. `get_or_make_arg_chunk!`
+   already exists for exactly this reason — extend the pattern rather than
+   reintroducing the reassignment next to it.
+
+22. **A re-tiling fallback needs the tile *backend*, and needs to outlive the
+   call.** Making non-square-tiled operands work instead of erroring is two
+   traps in one. First, allocate the destination through a tile-type-dispatched
+   allocator (`allocate_tiled`): `DArray{T}(undef, part, dims)` gives dense
+   tiles, so "repartition this sparse operator" silently becomes "densify this
+   sparse operator" — an out-of-memory multiplier that a correctness test
+   passes. A `DArray`'s type parameters do not record its tile type, so read it
+   off a chunk (`chunktype(first(A.chunks))`). Second, `maybe_copy_buffered`
+   frees its buffers when its body returns, which is only right when nothing
+   escapes; a block preconditioner builds per-tile operators *from* the re-tiled
+   tiles in tasks it never awaits, so it needs an ordinary array
+   (`repartition`) whose lifetime is its own. And note that a *device* sparse
+   tile supports neither end of a sub-range copy — reading it is scalar
+   indexing, writing it would insert nonzeros into a CSC in place — so
+   `copyto_view!` has to stage the whole thing on the host and re-upload
+   through `move`, which is the one hook every GPU sparse extension already
+   defines.
+
+23. **A memory space must be keyed on the device, never on a handle that
+   records current *ownership*.** `OpenCLExt.memory_space(::CLArray)` looked the
+   buffer's `Managed.queue` up in Dagger's registered `QUEUES`. That field is
+   OpenCL.jl's ownership tracking, not provenance: `convert(::CLPtr, ::Managed)`
+   synchronizes and then re-stamps it with `cl.queue()` whenever the accessing
+   task's queue differs, and `cl.queue()` is task-local *and lazily created*, so
+   the first touch from a task that never ran `with_context!` re-stamps the
+   buffer with a queue Dagger has never seen. The lookup then yields `nothing`
+   and `CLMemorySpace(myid(), nothing)` does not even construct — a
+   `MethodError` deep inside `aliasing`, arbitrarily far from the access that
+   moved the ownership, on an array that was allocated perfectly correctly.
+   Nothing about the memory changed; only a mutable bookkeeping field did. Key
+   on `queue.device` (matched against `DEVICES`) instead. Suspect this shape
+   whenever a space lookup fails for a value that a `Chunk` already carries a
+   valid space for: the chunk recorded the space once, the value is being asked
+   to re-derive it, and only the second one goes through the mutable field.
+
+24. **`similar(::DArray)` must not fetch the source tiles.** Spawning
+   `similar(chunk, T, sz)` per result tile looks like the way to preserve
+   sparse/GPU backends, but it is a false data-dependency: `A * A` then moves
+   every tile of `A` into the allocation tasks (and those tasks have no
+   concrete `return_type`). At the small sizes the dense GEMM bench uses, that
+   is a measured ~2×. Use `allocate_tiled` instead — dense stays
+   `DArray(undef)` (GPU processors still override `AllocateUndef`), sparse
+   stays sparse zeros. `similar(chunk)` is only right when you actually need
+   the source value.

Project.toml

@@ -23,7 +23,6 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 ScopedValues = "7e506255-f358-4e82-b7e4-beb19740aa63"
 Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
-SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 TaskLocalValues = "ed4db957-447d-4319-bfb6-7fa9ae7ecf34"
@@ -33,63 +32,92 @@ UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 [weakdeps]
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
+AlgebraicMultigrid = "2169fc97-5a83-5252-b627-83903c6c433c"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+Finch = "9177782c-1635-4eb9-9bfb-d9dfa25e6bce"
+IncompleteLU = "40713840-3770-5561-ab4c-a76e7d0d7895"
 GraphViz = "f526b714-d49f-11e8-06ff-31ed36ee7ee0"
 JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
+Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7"
 LinuxPerf = "b4c46c6c-4fb0-484d-a11a-41bc3392d094"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
+Metis = "2679e427-3c69-5b7f-982b-ece356f1e94b"
 OpenCL = "08131aa3-fb12-5dee-8b74-c09406e224a2"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+PureKLU = "0c0d3e7f-3a8b-4f7e-b6f1-9a4d2e7c1f01"
+PureUMFPACK = "b7e1f0a2-3c4d-4e5f-9a0b-1c2d3e4f5a6b"
 PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b"
 
 [extensions]
 AbstractFFTsExt = "AbstractFFTs"
+AlgebraicMultigridExt = ["AlgebraicMultigrid", "SparseArrays"]
 CUDAExt = "CUDA"
+CUDASparseArraysExt = ["CUDA", "SparseArrays"]
 DistributionsExt = "Distributions"
+FinchExt = "Finch"
 GraphVizExt = "GraphViz"
 GraphVizSimpleExt = "Colors"
+IncompleteLUExt = ["IncompleteLU", "SparseArrays"]
 IntelExt = "oneAPI"
+IntelSparseArraysExt = ["oneAPI", "SparseArrays"]
 JSON3Ext = "JSON3"
+KrylovExt = "Krylov"
 LinuxPerfExt = "LinuxPerf"
 MPIExt = "MPI"
+MPISparseExt = ["MPI", "SparseArrays"]
 MetalExt = "Metal"
+MetalSparseArraysExt = ["Metal", "SparseArrays"]
+MetisExt = ["Metis", "SparseArrays"]
 OpenCLExt = "OpenCL"
+OpenCLSparseArraysExt = ["OpenCL", "SparseArrays"]
 PlotsExt = ["DataFrames", "Plots"]
+PureKLUExt = ["PureKLU", "SparseArrays"]
+PureUMFPACKExt = ["PureUMFPACK", "SparseArrays"]
 PythonExt = "PythonCall"
 ROCExt = "AMDGPU"
+ROCSparseArraysExt = ["AMDGPU", "SparseArrays"]
+SparseArraysExt = "SparseArrays"
 
 [compat]
 AMDGPU = "1, 2"
 AbstractFFTs = "1.5.0"
 Adapt = "4"
+AlgebraicMultigrid = "1"
 CUDA = "3, 4, 5"
 Colors = "0.12, 0.13"
 DataFrames = "1"
-DataStructures = "0.18, 0.19"
+DataStructures = "0.19"
 DistributedNext = "1.0.0"
 Distributions = "0.25"
 FillArrays = "1.13.0"
+Finch = "1.2.9"
 GPUArraysCore = "0.2.0"
 GraphViz = "0.2"
 Graphs = "1"
+IncompleteLU = "0.2"
 JSON3 = "1"
 KernelAbstractions = "0.9"
+Krylov = "0.10"
 LinuxPerf = "0.4.2"
 MPI = "0.20.22"
 MacroTools = "0.5"
 MemPool = "0.4.18"
 Metal = "1.1"
+Metis = "1"
 NextLA = "0.2.2"
 OnlineStats = "1"
 OpenCL = "0.10"
 Plots = "1"
 PrecompileTools = "1.2"
 Preferences = "1.4.3"
+PureKLU = "1"
+PureUMFPACK = "0.1"
 PythonCall = "0.9"
 Requires = "1"
 ScopedValues = "1.1"

contrib/mpi/mpi_transfer_bench.jl

@@ -99,8 +99,7 @@ function bench_macro_cpu()
 end
 
 function bench_macro_gpu()
-    CUDAExt = Base.get_extension(Dagger, :CUDAExt)
-    CuProc = CUDAExt.CuArrayDeviceProc
+    CuProc = Dagger.CuArrayDeviceProc
     procs = sort(collect(Dagger.get_processors(Dagger.MPIClusterProc(comm)));
                  by=p->(p.rank, Dagger.short_name(p)))
     gpu_scope = Dagger.UnionScope([Dagger.ExactScope(p) for p in procs if p.innerProc isa CuProc]...)

docs/make.jl

@@ -23,6 +23,8 @@ makedocs(;
         "Task Affinity" => "task-affinity.md",
         "Data Management" => "data-management.md",
         "Distributed Arrays" => "darray.md",
+        "Sparse Arrays" => "sparse-arrays.md",
+        "Iterative Solvers" => "iterative-solving.md",
         "Streaming Tasks" => "streaming.md",
         "Scopes" => "scopes.md",
         "Processors" => "processors.md",

docs/src/index.md

@@ -363,6 +363,77 @@ collect(DA) # returns a `Matrix{Float64}`
 
 -----
 
+## Quickstart: Sparse Arrays
+
+A `DArray` can hold sparse tiles, giving a distributed, tiled sparse matrix.
+Load a sparse backend (`SparseArrays`) to enable it.
+
+For more details: [Sparse Distributed Arrays](@ref)
+
+### Distribute or allocate a sparse `DArray`
+
+```julia
+using SparseArrays
+
+# Partition an existing sparse matrix into sparse tiles
+DA = distribute(sprand(1000, 1000, 0.01), Blocks(250, 250))
+
+# Or allocate sparse DArrays directly
+Z = spzeros(Blocks(250, 250), Float64, 1000, 1000)
+R = sprand(Blocks(250, 250), Float64, (1000, 1000), 0.01)
+```
+
+### Multiply sparse arrays
+
+```julia
+using LinearAlgebra
+x = distribute(rand(1000), Blocks(250))
+y = DA * x        # distributed sparse matrix-vector multiply (dense result)
+C = DA * DA       # distributed sparse-sparse matmul (sparse result)
+```
+
+`collect(DA)` returns a dense `Array`; operate on the `DArray` to stay sparse.
+
+-----
+
+## Quickstart: Iterative Solvers
+
+Load `Krylov` and its own solvers (`cg`, `minres`, `gmres`, `bicgstab`, ...)
+work directly on Dagger arrays, distributed and matrix-free.
+
+For more details: [Iterative Solvers](@ref)
+
+### Solve a sparse system
+
+```julia
+using SparseArrays, Krylov, LinearAlgebra
+
+n = 1000
+A = spdiagm(-1 => fill(-1.0, n-1), 0 => fill(2.0, n), 1 => fill(-1.0, n-1))
+DA = distribute(A, Blocks(250, 250))   # square tiles are fastest
+b  = distribute(rand(n), Blocks(250))
+
+x, stats = Krylov.cg(DA, b)
+@show stats.solved, stats.niter
+```
+
+### Add a preconditioner
+
+```julia
+using AlgebraicMultigrid   # enables Dagger.AMGPreconditioner
+
+P = Dagger.AMGPreconditioner(DA)        # build once
+x, stats = Krylov.cg(DA, b; M = P)      # pass as `M`
+```
+
+Other preconditioners: `Dagger.JacobiPreconditioner`,
+`Dagger.BlockJacobiPreconditioner` (core), `Dagger.BlockILUPreconditioner`
+(load `IncompleteLU`), and `Dagger.BlockPreconditioner(A, build)` to plug in any
+third-party per-tile factory. Any object implementing `mul!(y, A, x)` over
+`DVector`s can be used as a matrix-free operator `A`.
+
+-----
+
 ## Quickstart: Stencil Operations
 
 Dagger's `@stencil` macro allows for easy specification of stencil operations on `DArray`s, often used in simulations and image processing. These operations typically involve updating an element based on the values of its neighbors.