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: ""

AGENTS.md

@@ -156,3 +156,36 @@ 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`, and the
+   `mpi_device_direct`/`mpi_remap_space` hooks in `src/gpu.jl`). 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.

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"
 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"

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,76 @@ 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
+
+Dagger provides distributed, matrix-free Krylov solvers (`cg`, `minres`,
+`gmres`, `bicgstab`) for `A x = b`. Load `Krylov` to enable them.
+
+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
+b  = distribute(rand(n), Blocks(250))
+
+x, stats = Dagger.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 = Dagger.cg(DA, b; M = P)      # pass as `M`
+```
+
+Other preconditioners: `Dagger.JacobiPreconditioner`,
+`Dagger.BlockJacobiPreconditioner` (core), and `Dagger.BlockILUPreconditioner`
+(load `IncompleteLU`). 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.