diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index b3f72a3c9..b61fbc9b4 100644 --- a/.buildkite/pipeline.yml +++ b/.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: "" diff --git a/Project.toml b/Project.toml index 21eb4a00d..3f64360ae 100644 --- a/Project.toml +++ b/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,40 +32,62 @@ 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" +MPIExt = ["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" @@ -74,22 +95,28 @@ 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.17" 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" diff --git a/docs/make.jl b/docs/make.jl index 611f7301d..d7f6f0168 100644 --- a/docs/make.jl +++ b/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", diff --git a/docs/src/index.md b/docs/src/index.md index 4b79aaa72..3e4b8ca59 100644 --- a/docs/src/index.md +++ b/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. diff --git a/docs/src/iterative-solving.md b/docs/src/iterative-solving.md new file mode 100644 index 000000000..31b62d755 --- /dev/null +++ b/docs/src/iterative-solving.md @@ -0,0 +1,226 @@ +# Iterative Solvers + +Dagger provides distributed, **matrix-free** Krylov solvers for linear systems +`A x = b`, built on top of [Krylov.jl](https://github.com/JuliaSmoothOptimizers/Krylov.jl). +They run entirely over `DArray`s: the operator `A` may be a dense or +[sparse](@ref "Sparse Distributed Arrays") `DMatrix`, or any object that +implements distributed `mul!(y, A, x)` over `DVector`s. This makes them a natural +fit for large, sparse systems arising from ODE/PDE discretizations. + +Solvers are provided through a package extension; load `Krylov` to enable them: + +```julia +using Distributed +addprocs(4) +using Dagger, SparseArrays, Krylov +``` + +If you call a solver without `Krylov` loaded, you'll get a clear error telling +you to `using Krylov`. + +## Solving a system + +```julia +using SparseArrays, Krylov + +# A symmetric positive-definite operator (1-D Laplacian) and a right-hand side +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 (see below) +b = distribute(rand(n), Blocks(250)) + +x, stats = Dagger.cg(DA, b) +@show stats.solved, stats.niter +``` + +Each solver returns a tuple `(x, stats)`, where `x` is a `DVector` and `stats` +is Krylov's statistics object (`stats.solved`, `stats.niter`, residual history, +etc.). Keyword arguments such as `atol`, `rtol`, and `itmax` are forwarded +to Krylov. + +### Available methods + +| Function | Operator class | Notes | +|---------------------|-------------------------------|------------------------------------| +| [`Dagger.cg`](@ref) | symmetric positive-definite | cheapest; the PDE workhorse | +| [`Dagger.minres`](@ref) | symmetric (indefinite ok) | saddle-point / indefinite systems | +| [`Dagger.gmres`](@ref) | general nonsymmetric | robust; `restart`/`memory` to bound memory | +| [`Dagger.bicgstab`](@ref) | general nonsymmetric | short recurrence, low memory | + +[`Dagger.krylov_solve`](@ref) is a generic entry point taking the method as a +symbol (`:cg`, `:minres`, `:gmres`, `:bicgstab`): + +```julia +x, stats = Dagger.krylov_solve(:gmres, DA, b; memory=50) +``` + +## Matrix-free operators + +The solvers never form `A⁻¹`, and they never require `A` to be a materialized +matrix. They only need `mul!(y, A, x)` (and `mul!(y, A', x)` for two-sided +methods like GMRES/BiCGStab) to work over `DVector`s. So you can pass any custom +operator type that implements distributed `mul!`: + +```julia +struct MyStencilOperator + # ...your distributed state... +end +function LinearAlgebra.mul!(y::Dagger.DVector, A::MyStencilOperator, x::Dagger.DVector) + # fill y with A*x using Dagger tasks / datadeps + return y +end + +x, stats = Dagger.cg(MyStencilOperator(...), b) +``` + +Workspace vectors are allocated via `similar(b)`, so every internal vector +inherits `b`'s element type **and** partitioning. The distributed BLAS-1 building +blocks the solvers rely on — `dot`, `norm`, `axpy!`, `axpby!`, `rmul!`, +`copyto!`, `fill!`, broadcasting — are all implemented for `DVector` and align +mismatched partitionings automatically. + +!!! note "Square tiles for `DMatrix` operators" + When the operator is a `DMatrix`, use **square tiles** (`Blocks(k, k)`). The + solver's workspace vectors are all allocated as `similar(b)` (one + partitioning), and each must serve as both the length-`n` input and the + length-`n` output of `mul!(y, A, x)`. The distributed SpMV requires the input + to match `A`'s column blocks and the output to match `A`'s row blocks, which + is only simultaneously possible when those block sizes are equal. Any uniform + square tile size works (a ragged final block is fine). This constraint is on + how *you* partition the operator — generic code calling `mul!`/`dot` never + needs to know about it. + +## Preconditioners + +A preconditioner accelerates convergence by approximating `A⁻¹`. Dagger's +preconditioners follow Krylov's `ldiv=false` convention: the object **represents +the inverse operator `M⁻¹`** and is applied via `mul!(y, P, x)` (computing +`y = M⁻¹ x`). Pass one as the `M` keyword: + +```julia +P = Dagger.JacobiPreconditioner(DA) +x, stats = Dagger.cg(DA, b; M = P) +``` + +The built-in preconditioners, from cheapest to strongest: + +| Preconditioner | Needs | Idea | +|--------------------------------------|------------------------|--------------------------------------------| +| [`Dagger.JacobiPreconditioner`](@ref) | (core) | scale by `1 ./ diag(A)` | +| [`Dagger.BlockJacobiPreconditioner`](@ref) | (core) | exact `lu` solve per diagonal tile | +| [`Dagger.BlockILUPreconditioner`](@ref) | `IncompleteLU` | incomplete-LU (drop tol `τ`) per tile | +| [`Dagger.AMGPreconditioner`](@ref) | `AlgebraicMultigrid` | AMG V-cycle per tile | + +```julia +using AlgebraicMultigrid, IncompleteLU + +x, _ = Dagger.cg(DA, b; M = Dagger.BlockJacobiPreconditioner(DA)) +x, _ = Dagger.cg(DA, b; M = Dagger.BlockILUPreconditioner(DA; τ = 0.01)) +x, _ = Dagger.gmres(DA, b; M = Dagger.AMGPreconditioner(DA; method = :ruge_stuben)) +``` + +### How block preconditioners are distributed + +`BlockJacobiPreconditioner`, `BlockILUPreconditioner`, and `AMGPreconditioner` +are all **block-diagonal** preconditioners: they build one operator per diagonal +tile of `A` and apply them independently per block. They share a common +mechanism ([`Dagger.AbstractBlockPreconditioner`](@ref)): + +- The per-tile operator (an `lu`/ILU factorization, or an AMG hierarchy) is built + **once**, at construction. +- A factorization/hierarchy generally cannot be moved between workers (sparse + `lu` factors and AMG hierarchies hold process-bound resources). So each + operator is **pinned** to the worker owning its tile, and every apply for that + block is scheduled there — only the (small, movable) vector chunks are + transferred. + +A useful consequence of the per-tile design: with a **single tile** +(`Blocks(n, n)`), any of these becomes a *global* preconditioner over the whole +matrix (e.g. a global AMG, or an exact direct solve for block-Jacobi). With many +tiles, it becomes a scalable block-Jacobi / additive-Schwarz variant that trades +some convergence for parallelism. As with the operator, these require square +diagonal tiles. + +### Choosing a preconditioner + +- **SPD elliptic (Poisson-like) problems:** `AMGPreconditioner` gives near + mesh-independent convergence and is usually the best choice; `cg` as the + solver. +- **General sparse systems:** `BlockILUPreconditioner` is a solid, cheap-setup + general-purpose option; pair with `gmres` or `bicgstab`. +- **Quick baseline / very well-conditioned systems:** `JacobiPreconditioner` (or + none) may suffice. +- **Strong per-subdomain coupling:** `BlockJacobiPreconditioner` (exact tile + solves) is stronger than diagonal Jacobi. + +## Sparse direct solvers + +For systems that fit on a single worker, Dagger also offers **direct** sparse +solves via pure-Julia factorization backends. Unlike the C-bound `UmfpackLU`, +these factorizations are plain Julia data, so Dagger can move and schedule them +freely. + +Load `PureKLU` (KLU; good for unsymmetric/circuit systems) or `PureUMFPACK` +(UMFPACK-style multifrontal LU): + +```julia +using SparseArrays, PureKLU, PureUMFPACK + +A = distribute(sprand(2000, 2000, 0.005) + 10I, Blocks(500, 500)) +b = distribute(rand(2000), Blocks(500)) + +F = Dagger.klu(A) # or Dagger.splu(A) +x = F \ b # returns a DVector partitioned like b +``` + +`Dagger.klu`/`Dagger.splu` gather the sparse `DMatrix` into one +`SparseMatrixCSC` (without densifying), factor it once, and return a +[`Dagger.DaggerSparseLU`](@ref) supporting `F \ b` and `ldiv!(x, F, b)`. Factor +once, solve many right-hand sides cheaply. + +There are also **block direct preconditioners** that factor each diagonal tile +exactly (`Dagger.BlockKLUPreconditioner`, `Dagger.BlockUMFPACKPreconditioner`), +usable like the other block preconditioners. With a single tile they are exact +whole-matrix solves; with many tiles they are exact-block-Jacobi preconditioners +for the iterative solvers. + +## A worked example: implicit time stepping + +Implicit ODE/PDE integrators repeatedly solve systems with the same operator +`(I - Δt·L)` and changing right-hand sides. Build the preconditioner once and +reuse it across steps: + +```julia +using SparseArrays, Krylov, AlgebraicMultigrid + +L = distribute(laplacian, Blocks(k, k)) # discretized operator (square tiles) +A = I - Δt * L # or a custom matrix-free operator +P = Dagger.AMGPreconditioner(L) # build hierarchy once + +u = distribute(u0, Blocks(k)) +for step in 1:nsteps + rhs = ... # depends on current state + u, stats = Dagger.cg(A, rhs; M = P, rtol = 1e-8) +end +``` + +## API + +```@docs +Dagger.cg +Dagger.minres +Dagger.gmres +Dagger.bicgstab +Dagger.krylov_solve +Dagger.AbstractDaggerPreconditioner +Dagger.JacobiPreconditioner +Dagger.AbstractBlockPreconditioner +Dagger.BlockJacobiPreconditioner +Dagger.BlockILUPreconditioner +Dagger.AMGPreconditioner +Dagger.BlockKLUPreconditioner +Dagger.BlockUMFPACKPreconditioner +Dagger.klu +Dagger.splu +Dagger.DaggerSparseLU +``` diff --git a/docs/src/sparse-arrays.md b/docs/src/sparse-arrays.md new file mode 100644 index 000000000..9935b033c --- /dev/null +++ b/docs/src/sparse-arrays.md @@ -0,0 +1,177 @@ +# Sparse Distributed Arrays + +Dagger's [`DArray`](@ref) can hold **sparse** tiles, giving you a distributed, +tiled sparse matrix (or vector) that participates in the same scheduling, +[Datadeps](@ref), and linear-algebra machinery as dense `DArray`s. This is the +foundation for distributed sparse matrix multiplication and the +[matrix-free iterative solvers](@ref "Iterative Solvers"). + +Sparse support is provided through package extensions, so you opt in by loading a +sparse backend: + +- **`SparseArrays`** (the standard library) — tiles are `SparseMatrixCSC` / + `SparseVector` on the CPU. This is the default, well-supported backend. +- **GPU + `SparseArrays`** — under a GPU compute scope (`cuda_gpu`, `rocm_gpu`, + `cl_device`, `metal_gpu`, `intel_gpu`), tiles use the vendor sparse type when + available (CUDA cuSPARSE / AMDGPU rocSPARSE) or else + `Dagger.DeviceSparseMatrixCSC` (OpenCL / Metal / oneAPI) with host SpGEMM/SpMV + fallbacks. Load the GPU package together with `SparseArrays`. +- **`Finch`** — tiles are `Finch.Tensor`s, enabling a wider range of sparse + formats. This backend is more experimental (CPU only for now). + +```julia +using Distributed +addprocs(4) +using Dagger, SparseArrays +``` + +!!! note "Load order with workers" + As with all Dagger usage, add your workers *before* `using Dagger` and the + backend package, so the packages load on every worker. See the note at the + top of the [home page](@ref "Dagger: A framework for out-of-core and parallel execution"). + +## Creating a sparse `DArray` + +### From an existing sparse array + +`distribute` accepts a sparse matrix or vector and partitions it into sparse +tiles according to a `Blocks` specification: + +```julia +using SparseArrays +A = sprand(1000, 1000, 0.01) # a SparseMatrixCSC +DA = distribute(A, Blocks(250, 250)) # a 4×4 grid of sparse tiles +``` + +Each tile is a sparse matrix in its own right, stored on one of the workers. + +### Allocating directly + +You can also allocate a sparse `DArray` without first building a local sparse +array, using the `Blocks`-aware methods of `spzeros` and `sprand`: + +```julia +using SparseArrays + +# All-zeros sparse DArray, Float64, 1000×1000 in 250×250 tiles +Z = spzeros(Blocks(250, 250), Float64, 1000, 1000) + +# Random sparse DArray with ~1% nonzeros per tile +R = sprand(Blocks(250, 250), Float64, (1000, 1000), 0.01) +``` + +These run the per-tile allocation on the owning worker, so no large sparse array +is ever materialized on a single process. + +### Converting back to a dense array + +`collect` gathers the tiles and returns a **dense** `Array`: + +```julia +M = collect(DA) # dense Matrix{Float64} +``` + +To keep data sparse and distributed, operate on the `DArray` directly rather +than collecting. + +## How it works + +### The `DSparseArray` wrapper + +Internally, each sparse tile is wrapped in a [`Dagger.DSparseArray`](@ref) — a +small mutable container holding the actual sparse storage (`mat`): + +```julia +mutable struct DSparseArray{T,N} <: AbstractArray{T,N} + mat # e.g. a SparseMatrixCSC, SparseVector, or Finch.Tensor +end +``` + +`DSparseVector{T}` and `DSparseMatrix{T}` are the 1- and 2-dimensional aliases. + +The wrapper exists because **sparse storage is reallocated on writes**. Many +sparse operations (e.g. `A*B`, or anything that changes the sparsity pattern) +cannot update their result in place — they produce a brand-new sparse array of a +different size. Datadeps, however, tracks data dependencies by the *identity* of +the objects it manages, and it does not support objects that grow or shrink. The +`DSparseArray` wrapper solves this: its identity is stable, and a write simply +swaps the inner `mat` for the new storage: + +```julia +# Conceptually, how an in-place sparse update is modeled: +tile.mat = tile.mat * other # identity of `tile` is unchanged +``` + +### Aliasing as a whole + +Because the inner storage may move, it is never safe to alias *part* of a sparse +tile (e.g. via a `view` or a strided sub-region). Dagger therefore treats a +`DSparseArray` as an **indivisible aliasing unit**: any access — including +through `view`, `transpose`, `adjoint`, or `reshape` — resolves to the +container's stable whole-object aliasing. This is what keeps Datadeps correct +when sparse writes reallocate storage. (For the curious: the type opts in via +`Dagger.aliases_as_whole`, and Datadeps' `aliasing_root` unwraps any wrapper of a +`DSparseArray` before computing aliasing. Calling `pointer` on a `DSparseArray` +intentionally errors, to catch any code path that tries to treat it as raw +strided memory.) + +The practical upshot: you can pass sparse tiles, or views of sparse `DArray`s, +into `Dagger.spawn_datadeps` regions and trust that read/write ordering is +tracked correctly. + +## Operations + +Sparse `DArray`s support the array operations that have distributed +implementations, including: + +- **Matrix–matrix multiply** (`A * B`, `mul!`), sparse × sparse, producing a + sparse result. +- **Sparse matrix–vector multiply** (SpMV: `A * x`, `mul!(y, A, x)`) with a + sparse matrix and dense vectors — the workhorse of iterative solvers. +- **Transpose/adjoint**, **`collect`**, and elementwise/`norm` operations. + +```julia +using SparseArrays, LinearAlgebra +A = distribute(sprand(1000, 1000, 0.01), Blocks(250, 250)) +x = distribute(rand(1000), Blocks(250)) + +y = A * x # distributed SpMV -> dense DVector +C = A * A # distributed sparse-sparse matmul -> sparse DArray +``` + +### Partitioning guidance + +- Choose tile sizes so each tile comfortably fits on a worker, and so the number + of tiles is at least the number of workers (for parallelism). +- For **square operators used with the iterative solvers**, use **square tiles** + (`Blocks(k, k)`); see [Iterative Solvers](@ref) for why. +- Operands with mismatched partitionings are aligned automatically (by buffered + copy) where needed, but matching partitionings avoid that overhead. + +## Backends + +### `SparseArrays` (recommended) + +Tiles are `SparseMatrixCSC` (matrices) or `SparseVector` (vectors). This backend +provides efficient SpMV (including transposed/adjoint operands) and uses +`SparseArrays`' own `*` for sparse–sparse products. + +### `Finch` (experimental) + +Loading `Finch` makes tiles `Finch.Tensor`s, supporting a broader set of sparse +and structured formats. Finch support is newer and exercised by a dedicated test +suite; prefer `SparseArrays` unless you specifically need a Finch format. + +## Limitations + +- `collect` densifies; there is no sparse-preserving global gather. +- A sparse tile is aliased as a whole — Datadeps cannot track independent writes + to disjoint sub-regions of a single sparse tile (use finer tiling instead). +- Not every dense `DArray` operation has a sparse counterpart yet; sparse support + focuses on multiplication and the building blocks needed for iterative solving. + +## API + +```@docs +Dagger.DSparseArray +``` diff --git a/ext/AlgebraicMultigridExt.jl b/ext/AlgebraicMultigridExt.jl new file mode 100644 index 000000000..dd68bbcfa --- /dev/null +++ b/ext/AlgebraicMultigridExt.jl @@ -0,0 +1,45 @@ +module AlgebraicMultigridExt + +import AlgebraicMultigrid +import SparseArrays +import Dagger +import Dagger: DMatrix +import LinearAlgebra + +# Algebraic-multigrid preconditioner, built per diagonal tile on top of Dagger's +# block-preconditioner machinery (see `src/array/iterativesolvers.jl`). The AMG +# hierarchy's expensive setup runs once per tile at construction; each apply runs +# one V-cycle (a sequence of SpMVs, smoother sweeps, and a coarse solve). + +_as_sparse(A::SparseArrays.SparseMatrixCSC) = A +_as_sparse(A::AbstractMatrix) = SparseArrays.sparse(A) + +function _amg_operator(tile, method::Symbol; kwargs...) + S = _as_sparse(Dagger._tile_matrix(tile)) + ml = if method === :ruge_stuben + AlgebraicMultigrid.ruge_stuben(S; kwargs...) + elseif method === :smoothed_aggregation + AlgebraicMultigrid.smoothed_aggregation(S; kwargs...) + else + throw(ArgumentError("AMGPreconditioner: unknown method $(method); use \ + :ruge_stuben or :smoothed_aggregation")) + end + return AlgebraicMultigrid.aspreconditioner(ml) +end + +function Dagger.AMGPreconditioner(A::DMatrix; method::Symbol=:ruge_stuben, kwargs...) + build = tile -> _amg_operator(tile, method; kwargs...) + return Dagger._build_block_preconditioner(Dagger.AMGPreconditioner, A, + "AMGPreconditioner", build) +end + +# An AMG preconditioner applies a V-cycle via `ldiv!` (`\` is not defined). +Dagger._block_apply!(y, p::AlgebraicMultigrid.Preconditioner, x) = + (LinearAlgebra.ldiv!(y, p, x); return nothing) + +# The hierarchy is read-only and pinned to its worker; it is never written and +# its internals (a `Vector{Level}`) aren't introspectable by datadeps, so treat +# it as non-aliasing rather than recursing into it. +Dagger.aliasing(::AlgebraicMultigrid.Preconditioner) = Dagger.NoAliasing() + +end # module AlgebraicMultigridExt diff --git a/ext/CUDASparseArraysExt.jl b/ext/CUDASparseArraysExt.jl new file mode 100644 index 000000000..072577ff3 --- /dev/null +++ b/ext/CUDASparseArraysExt.jl @@ -0,0 +1,177 @@ +module CUDASparseArraysExt + +# Sparse DArray support for CUDA / cuSPARSE. Loaded when both CUDA and +# SparseArrays are available (see Project.toml combo extension). + +import Dagger +import CUDA +import SparseArrays +import SparseArrays: SparseMatrixCSC, SparseVector +import LinearAlgebra +import CUDA: CuArray +import CUDA.CUSPARSE: CuSparseMatrixCSC, CuSparseMatrixCSR, CuSparseVector + +const CPUProc = Union{Dagger.OSProc,Dagger.ThreadProc} + +const CUDAExt = Base.get_extension(Dagger, :CUDAExt)::Module +using .CUDAExt: CuArrayDeviceProc +_with_context(f, proc) = CUDAExt.with_context(f, proc) + +#----- Memory / aliasing ------------------------------------------------------- + +Dagger.value_memory_space(x::CuSparseMatrixCSC) = Dagger.memory_space(x.nzVal) +Dagger.value_memory_space(x::CuSparseMatrixCSR) = Dagger.memory_space(x.nzVal) +Dagger.value_memory_space(x::CuSparseVector) = Dagger.memory_space(x.nzVal) +Dagger.memory_space(x::CuSparseMatrixCSC) = Dagger.value_memory_space(x) +Dagger.memory_space(x::CuSparseMatrixCSR) = Dagger.value_memory_space(x) +Dagger.memory_space(x::CuSparseVector) = Dagger.value_memory_space(x) + +function Dagger.aliasing(x::Union{CuSparseMatrixCSC,CuSparseMatrixCSR,CuSparseVector}, _=identity) + space = Dagger.value_memory_space(x) + ptr = Dagger.RemotePtr{Cvoid}(UInt(pointer_from_objref(x)), space) + return Dagger.ObjectAliasing(ptr, sizeof(typeof(x))) +end +Dagger.aliases_as_whole(::Union{CuSparseMatrixCSC,CuSparseMatrixCSR,CuSparseVector}) = true +Dagger.maybe_wrap_tile(x::Union{CuSparseMatrixCSC,CuSparseMatrixCSR,CuSparseVector}) = + Dagger.DSparseArray(x) + +#----- Allocation -------------------------------------------------------------- + +Dagger.allocate_sparse_zeros(::CuArrayDeviceProc, ::Type{T}, dims::Dims{2}) where T = + CuSparseMatrixCSC(SparseArrays.spzeros(T, dims...)) +Dagger.allocate_sparse_zeros(::CuArrayDeviceProc, ::Type{T}, dims::Dims{1}) where T = + CuSparseVector(SparseArrays.spzeros(T, dims...)) +Dagger.allocate_sparse_rand(::CuArrayDeviceProc, ::Type{T}, dims::Dims{2}, sparsity::AbstractFloat) where T = + CuSparseMatrixCSC(SparseArrays.sprand(T, dims..., sparsity)) +Dagger.allocate_sparse_rand(::CuArrayDeviceProc, ::Type{T}, dims::Dims{1}, sparsity::AbstractFloat) where T = + CuSparseVector(SparseArrays.sprand(T, dims..., sparsity)) + +#----- Similar / collect ------------------------------------------------------- + +Dagger._sparse_similar(::CuSparseMatrixCSC, ::Type{T}, dims::Dims{2}) where T = + CuSparseMatrixCSC(SparseArrays.spzeros(T, dims...)) +Dagger._sparse_similar(::CuSparseVector, ::Type{T}, dims::Dims{1}) where T = + CuSparseVector(SparseArrays.spzeros(T, dims...)) +Dagger._sparse_collect(A::CuSparseMatrixCSC) = SparseMatrixCSC(A) +Dagger._sparse_collect(A::CuSparseMatrixCSR) = SparseMatrixCSC(CuSparseMatrixCSC(A)) +Dagger._sparse_collect(A::CuSparseVector) = SparseVector(A) +Dagger._sparse_copy(A::Union{CuSparseMatrixCSC,CuSparseMatrixCSR,CuSparseVector}) = copy(A) + +#----- Host ↔ device helpers --------------------------------------------------- + +_to_cu_sparse(x::SparseMatrixCSC) = CuSparseMatrixCSC(x) +_to_cu_sparse(x::SparseVector) = CuSparseVector(x) +_to_cu_sparse(x::CuSparseMatrixCSC) = x +_to_cu_sparse(x::CuSparseMatrixCSR) = CuSparseMatrixCSC(x) +_to_cu_sparse(x::CuSparseVector) = x +_to_cu_sparse(x::Dagger.DeviceSparseMatrixCSC) = CuSparseMatrixCSC(SparseMatrixCSC(x)) +_to_cu_dsparse(x::Dagger.DSparseArray) = Dagger.DSparseArray(_to_cu_sparse(x.mat)) + +_to_host_sparse(x::CuSparseMatrixCSC) = SparseMatrixCSC(x) +_to_host_sparse(x::CuSparseMatrixCSR) = SparseMatrixCSC(CuSparseMatrixCSC(x)) +_to_host_sparse(x::CuSparseVector) = SparseVector(x) +_to_host_sparse(x::SparseMatrixCSC) = x +_to_host_sparse(x::SparseVector) = x +_to_host_dsparse(x::Dagger.DSparseArray) = Dagger.DSparseArray(_to_host_sparse(x.mat)) + +_to_csc(A::CuSparseMatrixCSC) = A +_to_csc(A::CuSparseMatrixCSR) = CuSparseMatrixCSC(A) +_to_csc(A) = CuSparseMatrixCSC(A) + +function _op_csc(X, t::Char) + Xc = _to_csc(X) + t == 'N' && return Xc + Sh = SparseMatrixCSC(Xc) + t == 'T' && return CuSparseMatrixCSC(SparseArrays.sparse(transpose(Sh))) + t == 'C' && return CuSparseMatrixCSC(SparseArrays.sparse(adjoint(Sh))) + throw(ArgumentError("Invalid trans char: $t")) +end + +_apply_trans(X, t::Char) = + t == 'N' ? X : t == 'T' ? transpose(X) : t == 'C' ? adjoint(X) : + throw(ArgumentError("Invalid trans char: $t")) + +#----- Move -------------------------------------------------------------------- + +function Dagger.move(from_proc::CPUProc, to_proc::CuArrayDeviceProc, x::SparseMatrixCSC) + _with_context(to_proc) do + return Dagger.DSparseArray(CuSparseMatrixCSC(x)) + end +end +function Dagger.move(from_proc::CPUProc, to_proc::CuArrayDeviceProc, x::SparseVector) + _with_context(to_proc) do + return Dagger.DSparseArray(CuSparseVector(x)) + end +end +function Dagger.move(from_proc::CPUProc, to_proc::CuArrayDeviceProc, x::Dagger.DSparseArray) + _with_context(to_proc) do + return _to_cu_dsparse(x) + end +end +function Dagger.move(from_proc::CuArrayDeviceProc, to_proc::CPUProc, x::Dagger.DSparseArray) + _with_context(from_proc) do + CUDA.synchronize() + return _to_host_dsparse(x) + end +end +function Dagger.move(from_proc::CuArrayDeviceProc, to_proc::CPUProc, + x::Union{CuSparseMatrixCSC,CuSparseMatrixCSR,CuSparseVector}) + _with_context(from_proc) do + CUDA.synchronize() + return Dagger.DSparseArray(_to_host_sparse(x)) + end +end +function Dagger.move(from_proc::CuArrayDeviceProc, to_proc::CuArrayDeviceProc, x::Dagger.DSparseArray) + if from_proc == to_proc + _with_context(CUDA.synchronize, from_proc) + return x + end + _with_context(to_proc) do + return _to_cu_dsparse(_to_host_dsparse(x)) + end +end + +#----- Tile kernels ------------------------------------------------------------ + +function Dagger.matmatmul!( + C::Dagger.DSparseMatrix, + transA::Char, transB::Char, + A::Union{CuSparseMatrixCSC,CuSparseMatrixCSR}, + B::Union{CuSparseMatrixCSC,CuSparseMatrixCSR}, + alpha, beta +) + AB = _to_csc(_op_csc(A, transA) * _op_csc(B, transB)) + prod = _to_csc(isone(alpha) ? AB : alpha * AB) + if iszero(beta) + C.mat = prod + elseif isone(beta) + C.mat = _to_csc(prod + _to_csc(C.mat)) + else + C.mat = _to_csc(prod + beta * _to_csc(C.mat)) + end + return C +end + +function Dagger.matvecmul!(C::CuArray, transA::Char, + A::Union{CuSparseMatrixCSC,CuSparseMatrixCSR}, + B::CuArray, alpha, beta) + LinearAlgebra.mul!(C, _apply_trans(A, transA), B, alpha, beta) + return C +end + +function Dagger.transpose_tile(B::CuSparseMatrixCSC) + return CuSparseMatrixCSC(SparseArrays.sparse(SparseMatrixCSC(B)')) +end +function Dagger.transpose_tile(B::CuSparseMatrixCSC, uplo::Char) + Bh = SparseMatrixCSC(B) + Bt = uplo == 'U' ? SparseArrays.triu(Bh) : + uplo == 'L' ? SparseArrays.tril(Bh) : + throw(ArgumentError("uplo must be 'U' or 'L', got $uplo")) + Ct = Bt + Bt' + for i in 1:LinearAlgebra.checksquare(Bh) + Ct[i, i] = Bh[i, i] + end + return CuSparseMatrixCSC(Ct) +end + +end # module CUDASparseArraysExt diff --git a/ext/FinchExt.jl b/ext/FinchExt.jl new file mode 100644 index 000000000..be2f0e356 --- /dev/null +++ b/ext/FinchExt.jl @@ -0,0 +1,188 @@ +module FinchExt + +import Finch +import LinearAlgebra +import Dagger +import Dagger: Blocks, AutoBlocks, BlocksOrAuto, AssignmentType, DSparseArray, DSparseMatrix + +Dagger._sparse_collect(A::Finch.Tensor) = Array(A) +# Finch's generic `similar` produces tensor formats that destabilize later +# `@finch`/`@einsum` kernels (observed as segfaults during tile moves), so +# allocate an empty COO-backed tile explicitly instead. +Dagger._sparse_similar(::Finch.Tensor, ::Type{T}, dims::Dims) where {T} = + Finch.fspzeros(T, dims...) +Dagger.maybe_wrap_tile(x::Finch.Tensor) = DSparseArray(x) + +function Finch.fspzeros(p::Blocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) + d = Dagger.ArrayDomain(map(x->1:x, dims)) + N = length(dims) + a = Dagger.AllocateArray(T, (T, _dims) -> DSparseArray(Finch.fspzeros(T, _dims...)), false, d, Dagger.partition(p, d), p, assignment; + return_type=DSparseArray{T,N}) + return Dagger._to_darray(a) +end +Finch.fspzeros(p::BlocksOrAuto, T::Type, dims::Integer...; assignment::AssignmentType = :arbitrary) = + Finch.fspzeros(p, T, dims; assignment) +Finch.fspzeros(p::BlocksOrAuto, dims::Integer...; assignment::AssignmentType = :arbitrary) = + Finch.fspzeros(p, Float64, dims; assignment) +Finch.fspzeros(p::BlocksOrAuto, dims::Dims; assignment::AssignmentType = :arbitrary) = + Finch.fspzeros(p, Float64, dims; assignment) +Finch.fspzeros(::AutoBlocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) = + Finch.fspzeros(Dagger.auto_blocks(dims), T, dims; assignment) + +function Finch.fsprand(p::Blocks, T::Type, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) + d = Dagger.ArrayDomain(map(x->1:x, dims)) + N = length(dims) + a = Dagger.AllocateArray(T, (T, _dims) -> DSparseArray(Finch.fsprand(T, _dims..., sparsity)), false, d, Dagger.partition(p, d), p, assignment; + return_type=DSparseArray{T,N}) + return Dagger._to_darray(a) +end +Finch.fsprand(p::BlocksOrAuto, T::Type, dims_and_sparsity::Real...; assignment::AssignmentType = :arbitrary) = + Finch.fsprand(p, T, dims_and_sparsity[1:end-1], dims_and_sparsity[end]; assignment) +Finch.fsprand(p::BlocksOrAuto, dims_and_sparsity::Real...; assignment::AssignmentType = :arbitrary) = + Finch.fsprand(p, Float64, dims_and_sparsity[1:end-1], dims_and_sparsity[end]; assignment) +Finch.fsprand(p::BlocksOrAuto, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) = + Finch.fsprand(p, Float64, dims, sparsity; assignment) +Finch.fsprand(::AutoBlocks, T::Type, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) = + Finch.fsprand(Dagger.auto_blocks(dims), T, dims, sparsity; assignment) + +# Materialize a (possibly lazy/`SwizzleArray`) result into a concrete sparse +# `Tensor`. `@einsum` returns a lazy `SwizzleArray` for transposed-output +# patterns, which is not itself a `Finch.Tensor` and corrupts subsequent +# accumulation steps; copying through `@finch` produces a clean `SparseCOO`. +function _finch_materialize(X, ::Type{T}, dims) where {T} + out = Finch.fspzeros(T, dims...) + Finch.@finch begin + out .= 0 + for j = _, i = _ + if X[i, j] != 0 + out[i, j] = X[i, j] + end + end + end + return out +end + +# Finch tensors do not define `Base.copy`; build a fresh equivalent tensor. +Dagger._sparse_copy(mat::Finch.Tensor) = _finch_materialize(mat, eltype(mat), size(mat)) + +# Finch tensors do not support `setindex!`, so copy-buffering writeback into a +# Finch tile densifies, applies the (partial-range) copy, then re-sparsifies. +function Dagger._sparse_copyto_view!(mat::Finch.Tensor, Brange, src) + dense = Array(mat) + copyto!(view(dense, Brange), src) + return _finch_materialize(dense, eltype(mat), size(mat)) +end + +# Compute `C_out += alpha * op(A) * op(B)` via Finch's `@einsum`, choosing the +# right index pattern for each transpose option. We deliberately use *explicit* +# index patterns (`A[k,i]`) rather than `Finch.swizzle`, because feeding a +# swizzled tensor into `@einsum` makes it return a lazy `SwizzleArray` (rather +# than a materialized `Tensor`), which breaks subsequent accumulation steps. +# Adjoint ('C') additionally applies `conj` inside the einsum. +function _finch_einsum_matmul!(C_out, A, B, transA::Char, transB::Char, alpha) + aT = transA == 'T' || transA == 'C' + aC = transA == 'C' + bT = transB == 'T' || transB == 'C' + bC = transB == 'C' + if !aT && !bT + bC ? Finch.@einsum(C_out[i,j] += alpha * (A[i,k] * conj(B[k,j]))) : + Finch.@einsum(C_out[i,j] += alpha * (A[i,k] * B[k,j])) + elseif !aT && bT + bC ? Finch.@einsum(C_out[i,j] += alpha * (A[i,k] * conj(B[j,k]))) : + Finch.@einsum(C_out[i,j] += alpha * (A[i,k] * B[j,k])) + elseif aT && !bT + aC ? Finch.@einsum(C_out[i,j] += alpha * (conj(A[k,i]) * B[k,j])) : + Finch.@einsum(C_out[i,j] += alpha * (A[k,i] * B[k,j])) + else # aT && bT + if aC && bC + Finch.@einsum C_out[i,j] += alpha * (conj(A[k,i]) * conj(B[j,k])) + elseif aC + Finch.@einsum C_out[i,j] += alpha * (conj(A[k,i]) * B[j,k]) + elseif bC + Finch.@einsum C_out[i,j] += alpha * (A[k,i] * conj(B[j,k])) + else + Finch.@einsum C_out[i,j] += alpha * (A[k,i] * B[j,k]) + end + end + return C_out +end + +function Dagger.matmatmul!( + C::DSparseMatrix, + transA::Char, + transB::Char, + A::Finch.Tensor, + B::Finch.Tensor, + alpha, + beta +) + if !isa(C.mat, Finch.Tensor) + # Not supported here, forward to generic matmatmul! + return Dagger.matmatmul!(C.mat, transA, transB, A, B, alpha, beta) + end + transA in ('N', 'T', 'C') || throw(ArgumentError("Invalid transA: $transA")) + transB in ('N', 'T', 'C') || throw(ArgumentError("Invalid transB: $transB")) + + # `@einsum` cannot write into an existing sparse tensor in place, so we build + # a fresh output tensor and (re)assign it to the wrapper. `DSparseMatrix` + # hides this reallocation from Datadeps. + # + # N.B. BEWARE: `@einsum Cm[i,j] = ...` (assignment into an existing tensor) + # does not work; use a fresh `C_out` with `+=`. + Cm = C.mat + C_out = Finch.fspzeros(eltype(Cm), size(Cm)...) + C_out = _finch_einsum_matmul!(C_out, A, B, transA, transB, alpha) + # `@einsum` may return a lazy `SwizzleArray`; materialize to a concrete + # sparse `Tensor` so accumulation and later iterations stay well-typed. + C_out = _finch_materialize(C_out, eltype(Cm), size(Cm)) + if beta != 0 + C_out = C_out + beta * Cm + end + C.mat = C_out + + return C +end + +# Sparse matrix-vector multiply tile kernel: `C = alpha*op(A)*B + beta*C` with a +# Finch tensor `A` and dense vectors `B`/`C`. Finch's `@einsum` does not reliably +# accumulate into a dense output, so we densify the single sparse tile and use a +# dense BLAS-backed `mul!`; this is a per-tile densification, not the whole matrix. +function Dagger.matvecmul!(C::AbstractVector, transA::Char, A::Finch.Tensor, B::AbstractVector, alpha, beta) + Ad = Array(A) + opA = transA == 'N' ? Ad : + transA == 'C' ? adjoint(Ad) : + transA == 'T' ? transpose(Ad) : + throw(ArgumentError("Invalid transA: $transA")) + LinearAlgebra.mul!(C, opA, B, alpha, beta) + return C +end + +# Tile transpose/symmetrization used by `copytri!`. Finch tensors don't support +# in-place `setindex!`, so we densify, build the result, and re-sparsify via +# `_finch_materialize`. +# - `uplo === nothing`: off-diagonal tile, return the conjugate transpose. +# - `uplo == 'U'`/`'L'`: diagonal tile, build the full Hermitian tile from the +# given triangle (matching the dense `copydiagtile!` semantics). +function Dagger.transpose_tile(B::Finch.Tensor, uplo=nothing) + _B = Array(B) + if uplo === nothing + C = adjoint(_B) + return _finch_materialize(C, eltype(_B), size(C)) + end + n = LinearAlgebra.checksquare(_B) + C = zeros(eltype(_B), n, n) + if uplo == 'U' + for j in 1:n, i in 1:n + C[i, j] = i <= j ? _B[i, j] : conj(_B[j, i]) + end + elseif uplo == 'L' + for j in 1:n, i in 1:n + C[i, j] = i >= j ? _B[i, j] : conj(_B[j, i]) + end + else + throw(ArgumentError("uplo must be 'U' or 'L', got $uplo")) + end + return _finch_materialize(C, eltype(_B), size(C)) +end + +end # module FinchExt \ No newline at end of file diff --git a/ext/IncompleteLUExt.jl b/ext/IncompleteLUExt.jl new file mode 100644 index 000000000..06ecab522 --- /dev/null +++ b/ext/IncompleteLUExt.jl @@ -0,0 +1,36 @@ +module IncompleteLUExt + +import IncompleteLU +import SparseArrays +import Dagger +import Dagger: DMatrix +import LinearAlgebra + +# Block incomplete-LU preconditioner, built per diagonal tile on top of Dagger's +# block-preconditioner machinery (see `src/array/iterativesolvers.jl`). Each tile +# gets an ILU factorization (drop tolerance `τ`) once at construction; the apply +# is a forward/backward substitution per block. + +_as_sparse(A::SparseArrays.SparseMatrixCSC) = A +_as_sparse(A::AbstractMatrix) = SparseArrays.sparse(A) + +function _ilu_operator(tile; τ=0.001, kwargs...) + S = _as_sparse(Dagger._tile_matrix(tile)) + return IncompleteLU.ilu(S; τ=τ, kwargs...) +end + +function Dagger.BlockILUPreconditioner(A::DMatrix; kwargs...) + build = tile -> _ilu_operator(tile; kwargs...) + return Dagger._build_block_preconditioner(Dagger.BlockILUPreconditioner, A, + "BlockILUPreconditioner", build) +end + +# An ILU factorization applies via `ldiv!` (`\` is not defined). +Dagger._block_apply!(y, F::IncompleteLU.ILUFactorization, x) = + (LinearAlgebra.ldiv!(y, F, x); return nothing) + +# The factorization is read-only and pinned to its worker; treat it as +# non-aliasing rather than recursing into its internals. +Dagger.aliasing(::IncompleteLU.ILUFactorization) = Dagger.NoAliasing() + +end # module IncompleteLUExt diff --git a/ext/IntelSparseArraysExt.jl b/ext/IntelSparseArraysExt.jl new file mode 100644 index 000000000..9cf2f2c01 --- /dev/null +++ b/ext/IntelSparseArraysExt.jl @@ -0,0 +1,68 @@ +module IntelSparseArraysExt + +# Sparse DArray support for oneAPI via DeviceSparseMatrixCSC (no vendor sparse +# library assumed). SpGEMM/SpMV fall back to host SparseArrays. + +import Dagger +import oneAPI +import Adapt +import SparseArrays +import SparseArrays: SparseMatrixCSC, SparseVector +import LinearAlgebra +import oneAPI: oneArray + +const CPUProc = Union{Dagger.OSProc,Dagger.ThreadProc} + +const IntelExt = Base.get_extension(Dagger, :IntelExt)::Module +using .IntelExt: oneArrayDeviceProc +_with_context(f, proc) = IntelExt.with_context(f, proc) + +Adapt.adapt_structure(::Type{<:oneArray}, S::SparseMatrixCSC) = + Dagger.device_sparse_from_host(oneArray, S) + +Dagger.allocate_sparse_zeros(::oneArrayDeviceProc, ::Type{T}, dims::Dims{2}) where T = + Dagger.device_sparse_from_host(oneArray, SparseArrays.spzeros(T, dims...)) +Dagger.allocate_sparse_zeros(::oneArrayDeviceProc, ::Type{T}, dims::Dims{1}) where T = + SparseArrays.spzeros(T, dims...) +Dagger.allocate_sparse_rand(::oneArrayDeviceProc, ::Type{T}, dims::Dims{2}, sparsity::AbstractFloat) where T = + Dagger.device_sparse_from_host(oneArray, SparseArrays.sprand(T, dims..., sparsity)) +Dagger.allocate_sparse_rand(::oneArrayDeviceProc, ::Type{T}, dims::Dims{1}, sparsity::AbstractFloat) where T = + SparseArrays.sprand(T, dims..., sparsity) + +function Dagger.move(from_proc::CPUProc, to_proc::oneArrayDeviceProc, x::SparseMatrixCSC) + _with_context(to_proc) do + return Dagger.DSparseArray(Dagger.device_sparse_from_host(oneArray, x)) + end +end +function Dagger.move(from_proc::CPUProc, to_proc::oneArrayDeviceProc, x::SparseVector) + return Dagger.DSparseArray(copy(x)) +end +function Dagger.move(from_proc::CPUProc, to_proc::oneArrayDeviceProc, x::Dagger.DSparseArray) + _with_context(to_proc) do + mat = x.mat + S = mat isa SparseMatrixCSC ? mat : SparseMatrixCSC(mat) + return Dagger.DSparseArray(Dagger.device_sparse_from_host(oneArray, S)) + end +end +function Dagger.move(from_proc::oneArrayDeviceProc, to_proc::CPUProc, x::Dagger.DSparseArray) + _with_context(from_proc) do + oneAPI.synchronize() + mat = x.mat + if mat isa Dagger.DeviceSparseMatrixCSC + return Dagger.DSparseArray(SparseMatrixCSC(mat)) + else + return Dagger.DSparseArray(copy(mat)) + end + end +end +function Dagger.move(from_proc::oneArrayDeviceProc, to_proc::oneArrayDeviceProc, x::Dagger.DSparseArray) + if from_proc == to_proc + return x + end + _with_context(to_proc) do + S = x.mat isa Dagger.DeviceSparseMatrixCSC ? SparseMatrixCSC(x.mat) : SparseMatrixCSC(x.mat) + return Dagger.DSparseArray(Dagger.device_sparse_from_host(oneArray, S)) + end +end + +end # module IntelSparseArraysExt diff --git a/ext/KrylovExt.jl b/ext/KrylovExt.jl new file mode 100644 index 000000000..5d36ba2fa --- /dev/null +++ b/ext/KrylovExt.jl @@ -0,0 +1,36 @@ +module KrylovExt + +import Krylov +import Dagger +import Dagger: DVector, DMatrix +import LinearAlgebra + +# Distributed iterative solvers, implemented on top of Krylov.jl. +# +# We never form `A⁻¹`: Krylov only needs `mul!(y, A, x)` (and `mul!(y, A', x)` +# for two-sided methods), which Dagger provides over `DVector`s via its +# distributed SpMV/`gemv_dagger!` path. The remaining requirements -- `dot`, +# `norm`, `axpy!`, `axpby!`, `rmul!`, `copyto!`, `fill!`, broadcasting -- are the +# distributed BLAS-1 ops in `Dagger`. +# +# Workspace vectors are allocated through a `KrylovConstructor` built from +# `similar(b)`, so every internal vector inherits `b`'s element type *and* +# partitioning (`similar` preserves both). This keeps all `mul!`/`dot`/`axpy!` +# operands on a compatible chunk layout, which the distributed kernels require. + +function _dagger_krylov(method::Symbol, A, b::DVector; kwargs...) + kc = Krylov.KrylovConstructor(similar(b)) + workspace = Krylov.krylov_workspace(Val(method), kc) + Krylov.krylov_solve!(workspace, A, b; kwargs...) + return Krylov.solution(workspace), Krylov.statistics(workspace) +end + +Dagger.krylov_solve(method::Symbol, A, b::DVector; kwargs...) = + _dagger_krylov(method, A, b; kwargs...) + +Dagger.cg(A, b::DVector; kwargs...) = _dagger_krylov(:cg, A, b; kwargs...) +Dagger.minres(A, b::DVector; kwargs...) = _dagger_krylov(:minres, A, b; kwargs...) +Dagger.gmres(A, b::DVector; kwargs...) = _dagger_krylov(:gmres, A, b; kwargs...) +Dagger.bicgstab(A, b::DVector; kwargs...) = _dagger_krylov(:bicgstab, A, b; kwargs...) + +end # module KrylovExt diff --git a/ext/MPIExt.jl b/ext/MPIExt.jl index 9274496b4..f7521bd6e 100644 --- a/ext/MPIExt.jl +++ b/ext/MPIExt.jl @@ -8,7 +8,8 @@ import Dagger: @dagdebug, @opcounter # extends with new methods (for MPI-specific types) or calls directly. import Dagger: AbstractAliasing, accelerate!, accel_matches_proc, aliased_object!, - AliasedObjectCache, AliasedObjectCacheStore, aliasing, bind_moved_argument, + AliasedObjectCache, AliasedObjectCacheStore, aliasing, aliasing_unwrapped, + bind_moved_argument, chunktype, ChunkView, check_uniform, check_uniformity!, CHECK_UNIFORMITY, cleanup_tasks_accel!, constrain, CPURAMMemorySpace, current_acceleration, CyclicProcGrid, datasize, default_enabled, @@ -1757,7 +1758,9 @@ function aliasing(accel::MPIAcceleration, x::ChunkView, dep_mod) if handle.rank == rank ainfo = _with_default_acceleration() do v = view(unwrap(x.chunk), x.slices...) - aliasing(v, dep_mod) + # Resolve whole-object containers (e.g. `DSparseArray`) where `v` + # lives; see `aliasing_unwrapped`. + aliasing_unwrapped(v, dep_mod) end ainfo = mpi_remap_ainfo(ainfo, handle.rank) @opcounter :aliasing_bcast_send_yield @@ -2104,7 +2107,11 @@ function mpi_propagate_chunk_types!(tasks, accel::MPIAcceleration, expected_type for t in tasks if t isa Thunk if t.options !== nothing - t.options.return_type = expected_type + # Respect a more specific tile type already set by the spawner + # (e.g. `DSparseArray` from sparse allocators). + if t.options.return_type === nothing + t.options.return_type = expected_type + end else t.options = Options(return_type=expected_type) end diff --git a/ext/MetalSparseArraysExt.jl b/ext/MetalSparseArraysExt.jl new file mode 100644 index 000000000..0d47f97b2 --- /dev/null +++ b/ext/MetalSparseArraysExt.jl @@ -0,0 +1,68 @@ +module MetalSparseArraysExt + +# Sparse DArray support for Metal via DeviceSparseMatrixCSC (no vendor sparse +# library). SpGEMM/SpMV fall back to host SparseArrays. + +import Dagger +import Metal +import Adapt +import SparseArrays +import SparseArrays: SparseMatrixCSC, SparseVector +import LinearAlgebra +import Metal: MtlArray + +const CPUProc = Union{Dagger.OSProc,Dagger.ThreadProc} + +const MetalExt = Base.get_extension(Dagger, :MetalExt)::Module +using .MetalExt: MtlArrayDeviceProc +_with_context(f, proc) = MetalExt.with_context(f, proc) + +Adapt.adapt_structure(::Type{<:MtlArray}, S::SparseMatrixCSC) = + Dagger.device_sparse_from_host(MtlArray, S) + +Dagger.allocate_sparse_zeros(::MtlArrayDeviceProc, ::Type{T}, dims::Dims{2}) where T = + Dagger.device_sparse_from_host(MtlArray, SparseArrays.spzeros(T, dims...)) +Dagger.allocate_sparse_zeros(::MtlArrayDeviceProc, ::Type{T}, dims::Dims{1}) where T = + SparseArrays.spzeros(T, dims...) +Dagger.allocate_sparse_rand(::MtlArrayDeviceProc, ::Type{T}, dims::Dims{2}, sparsity::AbstractFloat) where T = + Dagger.device_sparse_from_host(MtlArray, SparseArrays.sprand(T, dims..., sparsity)) +Dagger.allocate_sparse_rand(::MtlArrayDeviceProc, ::Type{T}, dims::Dims{1}, sparsity::AbstractFloat) where T = + SparseArrays.sprand(T, dims..., sparsity) + +function Dagger.move(from_proc::CPUProc, to_proc::MtlArrayDeviceProc, x::SparseMatrixCSC) + _with_context(to_proc) do + return Dagger.DSparseArray(Dagger.device_sparse_from_host(MtlArray, x)) + end +end +function Dagger.move(from_proc::CPUProc, to_proc::MtlArrayDeviceProc, x::SparseVector) + return Dagger.DSparseArray(copy(x)) +end +function Dagger.move(from_proc::CPUProc, to_proc::MtlArrayDeviceProc, x::Dagger.DSparseArray) + _with_context(to_proc) do + mat = x.mat + S = mat isa SparseMatrixCSC ? mat : SparseMatrixCSC(mat) + return Dagger.DSparseArray(Dagger.device_sparse_from_host(MtlArray, S)) + end +end +function Dagger.move(from_proc::MtlArrayDeviceProc, to_proc::CPUProc, x::Dagger.DSparseArray) + _with_context(from_proc) do + Metal.synchronize() + mat = x.mat + if mat isa Dagger.DeviceSparseMatrixCSC + return Dagger.DSparseArray(SparseMatrixCSC(mat)) + else + return Dagger.DSparseArray(copy(mat)) + end + end +end +function Dagger.move(from_proc::MtlArrayDeviceProc, to_proc::MtlArrayDeviceProc, x::Dagger.DSparseArray) + if from_proc == to_proc + return x + end + _with_context(to_proc) do + S = x.mat isa Dagger.DeviceSparseMatrixCSC ? SparseMatrixCSC(x.mat) : SparseMatrixCSC(x.mat) + return Dagger.DSparseArray(Dagger.device_sparse_from_host(MtlArray, S)) + end +end + +end # module MetalSparseArraysExt diff --git a/ext/MetisExt.jl b/ext/MetisExt.jl new file mode 100644 index 000000000..9f4a2623b --- /dev/null +++ b/ext/MetisExt.jl @@ -0,0 +1,75 @@ +module MetisExt + +import Metis +import SparseArrays +import SparseArrays: SparseMatrixCSC, nnz, rowvals, nzrange +import LinearAlgebra +import Dagger + +# Single-level (non-recursive) METIS k-way partition + vertex-separator +# construction for Stage 4c Schur-complement domain decomposition. +# Recursive / nested dissection is future work. + +# Symmetric adjacency pattern of `A + Aᵀ` (structure only, unit weights). +function _symmetrize_pattern(A::SparseMatrixCSC{Tv,Ti}) where {Tv,Ti} + n = size(A, 1) + S = A + transpose(A) + return SparseMatrixCSC{Tv,Ti}(n, n, S.colptr, S.rowval, ones(Tv, nnz(S))) +end + +""" + Dagger._nested_dissection_partition(A, nparts) -> (; perm, interiors, separator) + +K-way METIS partition of the undirected adjacency of `A + Aᵀ`, then a vertex +separator `Γ` = every vertex with a neighbor in a different part. Interiors are +the remaining vertices grouped by part (empty parts dropped). The permutation is +`[I₁; …; Iₖ; Γ]`. + +For tiny graphs (or `nparts == 1`), returns a trivial partition: one interior +covering `1:n` and an empty separator (caller should fall back to a serial +factorization). +""" +function Dagger._nested_dissection_partition(A::SparseMatrixCSC, nparts::Integer) + n = LinearAlgebra.checksquare(A) + nparts = Int(nparts) + nparts >= 1 || throw(ArgumentError("nparts must be ≥ 1, got $nparts")) + + # Too small to partition usefully → trivial (serial) partition. + if nparts == 1 || n < 4 || n <= nparts + return (; perm=collect(1:n), interiors=[collect(1:n)], separator=Int[]) + end + + pattern = _symmetrize_pattern(A) + g = Metis.graph(pattern) + part = Int.(Metis.partition(g, nparts)) + + is_sep = falses(n) + rows = rowvals(pattern) + @inbounds for v in 1:n + for k in nzrange(pattern, v) + u = rows[k] + u == v && continue + if part[u] != part[v] + is_sep[v] = true + break + end + end + end + + separator = findall(is_sep) + interiors = Vector{Vector{Int}}() + for p in 1:nparts + Ii = findall(i -> !is_sep[i] && part[i] == p, 1:n) + isempty(Ii) || push!(interiors, Ii) + end + + # Everything landed in the separator (or no interiors) → serial fallback. + if isempty(interiors) + return (; perm=collect(1:n), interiors=[collect(1:n)], separator=Int[]) + end + + perm = vcat(interiors..., separator) + return (; perm, interiors, separator) +end + +end # module MetisExt diff --git a/ext/OpenCLSparseArraysExt.jl b/ext/OpenCLSparseArraysExt.jl new file mode 100644 index 000000000..2a2f52772 --- /dev/null +++ b/ext/OpenCLSparseArraysExt.jl @@ -0,0 +1,82 @@ +module OpenCLSparseArraysExt + +# Sparse DArray support for OpenCL via DeviceSparseMatrixCSC (no vendor sparse +# library). SpGEMM/SpMV fall back to host SparseArrays. + +import Dagger +import OpenCL +import Adapt +import SparseArrays +import SparseArrays: SparseMatrixCSC, SparseVector +import LinearAlgebra +import OpenCL: CLArray + +const CPUProc = Union{Dagger.OSProc,Dagger.ThreadProc} + +const OpenCLExt = Base.get_extension(Dagger, :OpenCLExt)::Module +using .OpenCLExt: CLArrayDeviceProc +_with_context(f, proc) = OpenCLExt.with_context(f, proc) + +# Do not densify via Adapt (generic `adapt(CLArray, SparseMatrixCSC)` would). +Adapt.adapt_structure(::Type{<:CLArray}, S::SparseMatrixCSC) = + Dagger.device_sparse_from_host(CLArray, S) + +_empty_device_sparse(::Type{T}, dims::Dims{2}) where T = + Dagger.device_sparse_from_host(CLArray, SparseArrays.spzeros(T, dims...)) +_empty_device_sparse(::Type{T}, dims::Dims{1}) where T = + # Store 1-D sparse as a 1-column DeviceSparseMatrixCSC + Dagger.device_sparse_from_host(CLArray, SparseMatrixCSC(SparseArrays.spzeros(T, dims...))) + +Dagger.allocate_sparse_zeros(::CLArrayDeviceProc, ::Type{T}, dims::Dims{2}) where T = + _empty_device_sparse(T, dims) +Dagger.allocate_sparse_zeros(::CLArrayDeviceProc, ::Type{T}, dims::Dims{1}) where T = + SparseArrays.spzeros(T, dims...) # vectors stay host CSC; wrap still applies +Dagger.allocate_sparse_rand(::CLArrayDeviceProc, ::Type{T}, dims::Dims{2}, sparsity::AbstractFloat) where T = + Dagger.device_sparse_from_host(CLArray, SparseArrays.sprand(T, dims..., sparsity)) +Dagger.allocate_sparse_rand(::CLArrayDeviceProc, ::Type{T}, dims::Dims{1}, sparsity::AbstractFloat) where T = + SparseArrays.sprand(T, dims..., sparsity) + +function Dagger.move(from_proc::CPUProc, to_proc::CLArrayDeviceProc, x::SparseMatrixCSC) + _with_context(to_proc) do + return Dagger.DSparseArray(Dagger.device_sparse_from_host(CLArray, x)) + end +end +function Dagger.move(from_proc::CPUProc, to_proc::CLArrayDeviceProc, x::SparseVector) + # Keep SparseVector on host (wrapped); OpenCL SpMV host-falls-back anyway. + return Dagger.DSparseArray(copy(x)) +end +function Dagger.move(from_proc::CPUProc, to_proc::CLArrayDeviceProc, x::Dagger.DSparseArray) + _with_context(to_proc) do + mat = x.mat + if mat isa SparseMatrixCSC + return Dagger.DSparseArray(Dagger.device_sparse_from_host(CLArray, mat)) + elseif mat isa Dagger.DeviceSparseMatrixCSC + return Dagger.DSparseArray(Dagger.device_sparse_from_host(CLArray, SparseMatrixCSC(mat))) + else + return Dagger.DSparseArray(Dagger.device_sparse_from_host(CLArray, SparseMatrixCSC(mat))) + end + end +end +function Dagger.move(from_proc::CLArrayDeviceProc, to_proc::CPUProc, x::Dagger.DSparseArray) + _with_context(from_proc) do + OpenCL.cl.finish(OpenCL.cl.queue()) + mat = x.mat + if mat isa Dagger.DeviceSparseMatrixCSC + return Dagger.DSparseArray(SparseMatrixCSC(mat)) + else + return Dagger.DSparseArray(copy(mat)) + end + end +end +function Dagger.move(from_proc::CLArrayDeviceProc, to_proc::CLArrayDeviceProc, x::Dagger.DSparseArray) + if from_proc == to_proc + return x + end + _with_context(to_proc) do + mat = x.mat + S = mat isa Dagger.DeviceSparseMatrixCSC ? SparseMatrixCSC(mat) : SparseMatrixCSC(mat) + return Dagger.DSparseArray(Dagger.device_sparse_from_host(CLArray, S)) + end +end + +end # module OpenCLSparseArraysExt diff --git a/ext/PureKLUExt.jl b/ext/PureKLUExt.jl new file mode 100644 index 000000000..c2bbf22f5 --- /dev/null +++ b/ext/PureKLUExt.jl @@ -0,0 +1,38 @@ +module PureKLUExt + +import PureKLU +import SparseArrays +import SparseArrays: SparseMatrixCSC +import Dagger +import Dagger: DMatrix +import LinearAlgebra + +# Sparse direct (KLU) integration. KLU factorizations are pure-Julia data +# (serializable/movable), so they back both a whole-matrix direct solve +# (`Dagger.klu`) and per-tile block direct solves (`BlockKLUPreconditioner`), +# the latter via Dagger's block-preconditioner machinery. + +_as_sparse(A::SparseMatrixCSC) = A +_as_sparse(A::AbstractMatrix) = SparseArrays.sparse(A) + +# Whole-matrix direct solve: gather+factor on one worker, pin the factor there. +function Dagger.klu(A::DMatrix; kwargs...) + return Dagger._spawn_direct_factorization(A, S -> PureKLU.klu(S; kwargs...)) +end + +# Per-tile block direct preconditioner. +function Dagger.BlockKLUPreconditioner(A::DMatrix; kwargs...) + build = tile -> PureKLU.klu(_as_sparse(Dagger._tile_matrix(tile)); kwargs...) + return Dagger._build_block_preconditioner(Dagger.BlockKLUPreconditioner, A, + "BlockKLUPreconditioner", build) +end + +# KLU supports in-place `ldiv!`, so apply each block solve without allocating. +Dagger._block_apply!(y, F::PureKLU.KLUFactorization, x) = + (LinearAlgebra.ldiv!(y, F, x); return nothing) + +# The factorization is read-only and pinned to its worker; treat it as +# non-aliasing rather than recursing into its internals. +Dagger.aliasing(::PureKLU.KLUFactorization) = Dagger.NoAliasing() + +end # module PureKLUExt diff --git a/ext/PureUMFPACKExt.jl b/ext/PureUMFPACKExt.jl new file mode 100644 index 000000000..1817a5010 --- /dev/null +++ b/ext/PureUMFPACKExt.jl @@ -0,0 +1,63 @@ +module PureUMFPACKExt + +import PureUMFPACK +import SparseArrays +import SparseArrays: SparseMatrixCSC +import Dagger +import Dagger: DMatrix +import LinearAlgebra + +# Sparse direct (UMFPACK-style, pure-Julia) integration. Like KLU, these +# factorizations are plain Julia data (movable), backing both a whole-matrix +# direct solve (`Dagger.splu`) and per-tile block direct solves +# (`BlockUMFPACKPreconditioner`). With `distributed=true`: +# - `method=:trsv` (default) — Stage 4b tiled triangular solves +# - `method=:schur` — Stage 4c Schur-complement domain decomposition + +_as_sparse(A::SparseMatrixCSC) = A +_as_sparse(A::AbstractMatrix) = SparseArrays.sparse(A) + +function _umfpack_factorize(; kwargs...) + return S -> PureUMFPACK.splu(_as_sparse(S); kwargs...) +end + +# Whole-matrix direct solve: gather+factor on one worker, pin the factor there. +# Opt-in distributed paths select Stage 4b (`:trsv`) or Stage 4c (`:schur`). +function Dagger.splu(A::DMatrix; distributed::Bool=false, method::Symbol=:trsv, + blocksize=nothing, nparts=nothing, kwargs...) + factorize = _umfpack_factorize(; kwargs...) + if !distributed + return Dagger._spawn_direct_factorization(A, factorize) + end + if method === :trsv + F = Dagger._spawn_direct_factorization(A, factorize) + return Dagger._make_distributed_sparse_lu(F, A; blocksize) + elseif method === :schur + return Dagger._make_schur_sparse_lu(A, factorize; nparts) + else + throw(ArgumentError( + "Dagger.splu distributed method must be :trsv or :schur, got $(repr(method))")) + end +end + +# `(Rs .* A)[p, q] == L * U` with unit-lower `L` (explicit stored ones on the +# diagonal; solve treats it as unit diagonal via `UnitLowerTriangular`). +function Dagger._extract_lu_factors(F::PureUMFPACK.PureLU) + return (F.L, F.U, copy(F.p), copy(F.q), copy(F.Rs)) +end + +# Per-tile block direct preconditioner. +function Dagger.BlockUMFPACKPreconditioner(A::DMatrix; kwargs...) + build = tile -> PureUMFPACK.splu(_as_sparse(Dagger._tile_matrix(tile)); kwargs...) + return Dagger._build_block_preconditioner(Dagger.BlockUMFPACKPreconditioner, A, + "BlockUMFPACKPreconditioner", build) +end + +# `PureLU` provides `\` (no `ldiv!`); the default `_block_apply!` (`y .= op \ x`) +# already uses it, so no apply override is needed. + +# The factorization is read-only and pinned to its worker; treat it as +# non-aliasing rather than recursing into its internals. +Dagger.aliasing(::PureUMFPACK.PureLU) = Dagger.NoAliasing() + +end # module PureUMFPACKExt diff --git a/ext/ROCSparseArraysExt.jl b/ext/ROCSparseArraysExt.jl new file mode 100644 index 000000000..7340109d5 --- /dev/null +++ b/ext/ROCSparseArraysExt.jl @@ -0,0 +1,185 @@ +module ROCSparseArraysExt + +# Sparse DArray support for AMDGPU / rocSPARSE. Loaded when both AMDGPU and +# SparseArrays are available (see Project.toml combo extension). + +import Dagger +import AMDGPU +import SparseArrays +import SparseArrays: SparseMatrixCSC, SparseVector +import LinearAlgebra +import AMDGPU: ROCArray +import AMDGPU.rocSPARSE: ROCSparseMatrixCSC, ROCSparseMatrixCSR, ROCSparseVector + +const CPUProc = Union{Dagger.OSProc,Dagger.ThreadProc} + +# ROCExt is triggered by AMDGPU and provides the processor + context helpers. +const ROCExt = Base.get_extension(Dagger, :ROCExt)::Module +using .ROCExt: ROCArrayDeviceProc +_with_context(f, proc) = ROCExt.with_context(f, proc) + +#----- Memory / aliasing ------------------------------------------------------- + +Dagger.value_memory_space(x::ROCSparseMatrixCSC) = Dagger.memory_space(x.nzVal) +Dagger.value_memory_space(x::ROCSparseMatrixCSR) = Dagger.memory_space(x.nzVal) +Dagger.value_memory_space(x::ROCSparseVector) = Dagger.memory_space(x.nzVal) +# Datadeps `aliased_object!` keys slots by `memory_space(x)` (see DSparseArray). +Dagger.memory_space(x::ROCSparseMatrixCSC) = Dagger.value_memory_space(x) +Dagger.memory_space(x::ROCSparseMatrixCSR) = Dagger.value_memory_space(x) +Dagger.memory_space(x::ROCSparseVector) = Dagger.value_memory_space(x) + +function Dagger.aliasing(x::Union{ROCSparseMatrixCSC,ROCSparseMatrixCSR,ROCSparseVector}, _=identity) + space = Dagger.value_memory_space(x) + ptr = Dagger.RemotePtr{Cvoid}(UInt(pointer_from_objref(x)), space) + return Dagger.ObjectAliasing(ptr, sizeof(typeof(x))) +end +Dagger.aliases_as_whole(::Union{ROCSparseMatrixCSC,ROCSparseMatrixCSR,ROCSparseVector}) = true +Dagger.maybe_wrap_tile(x::Union{ROCSparseMatrixCSC,ROCSparseMatrixCSR,ROCSparseVector}) = + Dagger.DSparseArray(x) + +#----- Allocation -------------------------------------------------------------- + +Dagger.allocate_sparse_zeros(::ROCArrayDeviceProc, ::Type{T}, dims::Dims{2}) where T = + ROCSparseMatrixCSC(SparseArrays.spzeros(T, dims...)) +Dagger.allocate_sparse_zeros(::ROCArrayDeviceProc, ::Type{T}, dims::Dims{1}) where T = + ROCSparseVector(SparseArrays.spzeros(T, dims...)) +Dagger.allocate_sparse_rand(::ROCArrayDeviceProc, ::Type{T}, dims::Dims{2}, sparsity::AbstractFloat) where T = + ROCSparseMatrixCSC(SparseArrays.sprand(T, dims..., sparsity)) +Dagger.allocate_sparse_rand(::ROCArrayDeviceProc, ::Type{T}, dims::Dims{1}, sparsity::AbstractFloat) where T = + ROCSparseVector(SparseArrays.sprand(T, dims..., sparsity)) + +#----- Similar / collect ------------------------------------------------------- + +Dagger._sparse_similar(::ROCSparseMatrixCSC, ::Type{T}, dims::Dims{2}) where T = + ROCSparseMatrixCSC(SparseArrays.spzeros(T, dims...)) +Dagger._sparse_similar(::ROCSparseVector, ::Type{T}, dims::Dims{1}) where T = + ROCSparseVector(SparseArrays.spzeros(T, dims...)) +Dagger._sparse_collect(A::ROCSparseMatrixCSC) = SparseMatrixCSC(A) +Dagger._sparse_collect(A::ROCSparseMatrixCSR) = SparseMatrixCSC(ROCSparseMatrixCSC(A)) +Dagger._sparse_collect(A::ROCSparseVector) = SparseVector(A) +Dagger._sparse_copy(A::Union{ROCSparseMatrixCSC,ROCSparseMatrixCSR,ROCSparseVector}) = copy(A) + +#----- Host ↔ device helpers --------------------------------------------------- + +_to_roc_sparse(x::SparseMatrixCSC) = ROCSparseMatrixCSC(x) +_to_roc_sparse(x::SparseVector) = ROCSparseVector(x) +_to_roc_sparse(x::ROCSparseMatrixCSC) = x +_to_roc_sparse(x::ROCSparseMatrixCSR) = ROCSparseMatrixCSC(x) +_to_roc_sparse(x::ROCSparseVector) = x +_to_roc_sparse(x::Dagger.DeviceSparseMatrixCSC) = ROCSparseMatrixCSC(SparseMatrixCSC(x)) +_to_roc_dsparse(x::Dagger.DSparseArray) = Dagger.DSparseArray(_to_roc_sparse(x.mat)) + +_to_host_sparse(x::ROCSparseMatrixCSC) = SparseMatrixCSC(x) +_to_host_sparse(x::ROCSparseMatrixCSR) = SparseMatrixCSC(ROCSparseMatrixCSC(x)) +_to_host_sparse(x::ROCSparseVector) = SparseVector(x) +_to_host_sparse(x::SparseMatrixCSC) = x +_to_host_sparse(x::SparseVector) = x +_to_host_dsparse(x::Dagger.DSparseArray) = Dagger.DSparseArray(_to_host_sparse(x.mat)) + +_to_csc(A::ROCSparseMatrixCSC) = A +_to_csc(A::ROCSparseMatrixCSR) = ROCSparseMatrixCSC(A) +_to_csc(A) = ROCSparseMatrixCSC(A) + +# Materialize transpose/adjoint as CSC so SpGEMM stays on rocSPARSE / host CSC +# paths; `A * transpose(B)` would otherwise fall into scalar-indexing generic mul. +function _op_csc(X, t::Char) + Xc = _to_csc(X) + t == 'N' && return Xc + Sh = SparseMatrixCSC(Xc) + t == 'T' && return ROCSparseMatrixCSC(SparseArrays.sparse(transpose(Sh))) + t == 'C' && return ROCSparseMatrixCSC(SparseArrays.sparse(adjoint(Sh))) + throw(ArgumentError("Invalid trans char: $t")) +end + +#----- Move (preserve sparsity; wrap in DSparseArray) -------------------------- + +function Dagger.move(from_proc::CPUProc, to_proc::ROCArrayDeviceProc, x::SparseMatrixCSC) + _with_context(to_proc) do + return Dagger.DSparseArray(ROCSparseMatrixCSC(x)) + end +end +function Dagger.move(from_proc::CPUProc, to_proc::ROCArrayDeviceProc, x::SparseVector) + _with_context(to_proc) do + return Dagger.DSparseArray(ROCSparseVector(x)) + end +end +function Dagger.move(from_proc::CPUProc, to_proc::ROCArrayDeviceProc, x::Dagger.DSparseArray) + _with_context(to_proc) do + return _to_roc_dsparse(x) + end +end +function Dagger.move(from_proc::ROCArrayDeviceProc, to_proc::CPUProc, x::Dagger.DSparseArray) + _with_context(from_proc) do + AMDGPU.synchronize() + return _to_host_dsparse(x) + end +end +function Dagger.move(from_proc::ROCArrayDeviceProc, to_proc::CPUProc, + x::Union{ROCSparseMatrixCSC,ROCSparseMatrixCSR,ROCSparseVector}) + _with_context(from_proc) do + AMDGPU.synchronize() + return Dagger.DSparseArray(_to_host_sparse(x)) + end +end +function Dagger.move(from_proc::ROCArrayDeviceProc, to_proc::ROCArrayDeviceProc, x::Dagger.DSparseArray) + # Same device: identity (like dense ROCArray). A copy here would discard + # in-place SpGEMM writes under Datadeps/Sch argument moves. + if from_proc == to_proc + _with_context(AMDGPU.synchronize, from_proc) + return x + end + # Distinct devices: stage through host (no peer-copy helper yet). + _with_context(to_proc) do + return _to_roc_dsparse(_to_host_dsparse(x)) + end +end + +#----- Tile kernels ------------------------------------------------------------ + +function Dagger.matmatmul!( + C::Dagger.DSparseMatrix, + transA::Char, transB::Char, + A::Union{ROCSparseMatrixCSC,ROCSparseMatrixCSR}, + B::Union{ROCSparseMatrixCSC,ROCSparseMatrixCSR}, + alpha, beta +) + AB = _to_csc(_op_csc(A, transA) * _op_csc(B, transB)) + prod = _to_csc(isone(alpha) ? AB : alpha * AB) + if iszero(beta) + C.mat = prod + elseif isone(beta) + C.mat = _to_csc(prod + _to_csc(C.mat)) + else + C.mat = _to_csc(prod + beta * _to_csc(C.mat)) + end + return C +end + +_apply_trans(X, t::Char) = + t == 'N' ? X : t == 'T' ? transpose(X) : t == 'C' ? adjoint(X) : + throw(ArgumentError("Invalid trans char: $t")) + +function Dagger.matvecmul!(C::ROCArray, transA::Char, + A::Union{ROCSparseMatrixCSC,ROCSparseMatrixCSR}, + B::ROCArray, alpha, beta) + # rocSPARSE SpMV supports transpose/adjoint wrappers. + LinearAlgebra.mul!(C, _apply_trans(A, transA), B, alpha, beta) + return C +end + +function Dagger.transpose_tile(B::ROCSparseMatrixCSC) + return ROCSparseMatrixCSC(SparseArrays.sparse(SparseMatrixCSC(B)')) +end +function Dagger.transpose_tile(B::ROCSparseMatrixCSC, uplo::Char) + Bh = SparseMatrixCSC(B) + Bt = uplo == 'U' ? SparseArrays.triu(Bh) : + uplo == 'L' ? SparseArrays.tril(Bh) : + throw(ArgumentError("uplo must be 'U' or 'L', got $uplo")) + Ct = Bt + Bt' + for i in 1:LinearAlgebra.checksquare(Bh) + Ct[i, i] = Bh[i, i] + end + return ROCSparseMatrixCSC(Ct) +end + +end # module ROCSparseArraysExt diff --git a/ext/SparseArraysExt.jl b/ext/SparseArraysExt.jl new file mode 100644 index 000000000..8820bd45a --- /dev/null +++ b/ext/SparseArraysExt.jl @@ -0,0 +1,238 @@ +module SparseArraysExt + +import SparseArrays +import SparseArrays: SparseMatrixCSC, SparseVector +import LinearAlgebra +import Dagger +import Dagger: Blocks, AutoBlocks, BlocksOrAuto, AssignmentType, DSparseArray, DSparseMatrix + +# Keep tiles sparse through `collect`/`cat`; the outer `collect` densifies. +Dagger._sparse_collect(M::SparseMatrixCSC) = copy(M) + +# Assemble already-local tiles into one global `SparseMatrixCSC` without +# densifying: unwrap each tile, offset its (i,j) by the precomputed subdomain +# offsets, and build from triplets. Intended to run inside a worker-scoped task +# (the scheduler moves the tile chunks there); used by `Dagger.klu` / `Dagger.splu`. +function Dagger._gather_sparse(::Type{T}, tiles, row_offsets, col_offsets, m, n) where T + Is = Int[]; Js = Int[]; Vs = T[] + for k in 1:length(tiles) + tile = SparseMatrixCSC(Dagger._tile_matrix(tiles[k])) + ti, tj, tv = SparseArrays.findnz(tile) + append!(Is, ti .+ row_offsets[k]) + append!(Js, tj .+ col_offsets[k]) + append!(Vs, tv) + end + return SparseArrays.sparse(Is, Js, Vs, m, n) +end + +# Dense → sparse for Stage-4c Schur complements (fill-in is expected). +Dagger._sparse_copy_of(S::AbstractMatrix) = SparseArrays.sparse(S) +Dagger._sparse_copy_of(S::SparseMatrixCSC) = S + +# Wrap bare sparse tiles (e.g. from `distribute`) so Datadeps sees a stable container. +Dagger.maybe_wrap_tile(x::SparseMatrixCSC) = DSparseArray(x) +Dagger.maybe_wrap_tile(x::SparseVector) = DSparseArray(x) +Dagger.maybe_wrap_tile(x::Dagger.DeviceSparseMatrixCSC) = DSparseArray(x) + +# Host defaults for scoped sparse allocation (GPU Exts override per-processor). +Dagger.allocate_sparse_zeros_default(::Dagger.Processor, ::Type{T}, dims::Dims{2}) where T = + SparseArrays.spzeros(T, dims...) +Dagger.allocate_sparse_zeros_default(::Dagger.Processor, ::Type{T}, dims::Dims{1}) where T = + SparseArrays.spzeros(T, dims...) +Dagger.allocate_sparse_rand_default(::Dagger.Processor, ::Type{T}, dims::Dims{2}, sparsity::AbstractFloat) where T = + SparseArrays.sprand(T, dims..., sparsity) +Dagger.allocate_sparse_rand_default(::Dagger.Processor, ::Type{T}, dims::Dims{1}, sparsity::AbstractFloat) where T = + SparseArrays.sprand(T, dims..., sparsity) + +# DeviceSparseMatrixCSC ↔ SparseMatrixCSC +function SparseArrays.SparseMatrixCSC(A::Dagger.DeviceSparseMatrixCSC{Tv,Ti}) where {Tv,Ti} + return SparseMatrixCSC{Tv,Ti}(A.m, A.n, Array(A.colptr), Array(A.rowval), Array(A.nzval)) +end +""" + device_sparse_from_host(Arr, S::SparseMatrixCSC) -> DeviceSparseMatrixCSC + +Upload a host CSC onto device vectors of type `Arr` (e.g. `CLArray`, `MtlArray`, +`oneArray`). Indices are converted to `Int32` for device friendliness. +""" +function Dagger.device_sparse_from_host(::Type{Arr}, S::SparseMatrixCSC{Tv}) where {Tv,Arr} + colptr = Arr(Int32.(S.colptr)) + rowval = Arr(Int32.(S.rowval)) + nzval = Arr(Array(S.nzval)) + return Dagger.DeviceSparseMatrixCSC(S.m, S.n, colptr, rowval, nzval) +end +Base.copy(A::Dagger.DeviceSparseMatrixCSC) = + Dagger.DeviceSparseMatrixCSC(A.m, A.n, copy(A.colptr), copy(A.rowval), copy(A.nzval)) +Dagger._sparse_copy(A::Dagger.DeviceSparseMatrixCSC) = copy(A) +Dagger._sparse_collect(A::Dagger.DeviceSparseMatrixCSC) = SparseMatrixCSC(A) +function Dagger._sparse_similar(A::Dagger.DeviceSparseMatrixCSC{Tv,Ti}, ::Type{T}, dims::Dims{2}) where {Tv,Ti,T} + n = dims[2] + colptr = similar(A.colptr, Ti, n + 1) + fill!(colptr, one(Ti)) + rowval = similar(A.rowval, Ti, 0) + nzval = similar(A.nzval, T, 0) + return Dagger.DeviceSparseMatrixCSC(dims[1], n, colptr, rowval, nzval) +end + +# Rebuild a DeviceSparseMatrixCSC on the same device array type as `like`. +function _to_device_sparse(like::Dagger.DeviceSparseMatrixCSC, S::SparseMatrixCSC) + colptr = similar(like.colptr, eltype(like.colptr), length(S.colptr)) + rowval = similar(like.rowval, eltype(like.rowval), length(S.rowval)) + nzval = similar(like.nzval, eltype(S), length(S.nzval)) + copyto!(colptr, eltype(colptr).(S.colptr)) + copyto!(rowval, eltype(rowval).(S.rowval)) + copyto!(nzval, S.nzval) + return Dagger.DeviceSparseMatrixCSC(S.m, S.n, colptr, rowval, nzval) +end + +function SparseArrays.spzeros(p::Blocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) + d = Dagger.ArrayDomain(map(x->1:x, dims)) + N = length(dims) + # Route through `allocate_sparse_zeros` so a GPU compute scope yields + # device-resident sparse tiles (vendor sparse or DeviceSparseMatrixCSC). + a = Dagger.AllocateArray(T, (T, _dims) -> DSparseArray(Dagger.allocate_sparse_zeros(Dagger.task_processor(), T, _dims)), false, d, Dagger.partition(p, d), p, assignment; + return_type=DSparseArray{T,N}) + return Dagger._to_darray(a) +end +SparseArrays.spzeros(p::BlocksOrAuto, T::Type, dims::Integer...; assignment::AssignmentType = :arbitrary) = + SparseArrays.spzeros(p, T, dims; assignment) +SparseArrays.spzeros(p::BlocksOrAuto, dims::Integer...; assignment::AssignmentType = :arbitrary) = + SparseArrays.spzeros(p, Float64, dims; assignment) +SparseArrays.spzeros(p::BlocksOrAuto, dims::Dims; assignment::AssignmentType = :arbitrary) = + SparseArrays.spzeros(p, Float64, dims; assignment) +SparseArrays.spzeros(::AutoBlocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) = + SparseArrays.spzeros(Dagger.auto_blocks(dims), T, dims; assignment) + +function SparseArrays.sprand(p::Blocks, T::Type, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) + d = Dagger.ArrayDomain(map(x->1:x, dims)) + N = length(dims) + a = Dagger.AllocateArray(T, (T, _dims) -> DSparseArray(Dagger.allocate_sparse_rand(Dagger.task_processor(), T, _dims, sparsity)), false, d, Dagger.partition(p, d), p, assignment; + return_type=DSparseArray{T,N}) + return Dagger._to_darray(a) +end +SparseArrays.sprand(p::BlocksOrAuto, T::Type, dims_and_sparsity::Real...; assignment::AssignmentType = :arbitrary) = + SparseArrays.sprand(p, T, dims_and_sparsity[1:end-1], dims_and_sparsity[end]; assignment) +SparseArrays.sprand(p::BlocksOrAuto, dims_and_sparsity::Real...; assignment::AssignmentType = :arbitrary) = + SparseArrays.sprand(p, Float64, dims_and_sparsity[1:end-1], dims_and_sparsity[end]; assignment) +SparseArrays.sprand(p::BlocksOrAuto, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) = + SparseArrays.sprand(p, Float64, dims, sparsity; assignment) +SparseArrays.sprand(::AutoBlocks, T::Type, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) = + SparseArrays.sprand(Dagger.auto_blocks(dims), T, dims, sparsity; assignment) + +_apply_trans(X, t::Char) = + t == 'N' ? X : + t == 'T' ? transpose(X) : + t == 'C' ? adjoint(X) : + throw(ArgumentError("Invalid trans char: $t")) + +function _sparse_gemm_assign!(C::DSparseMatrix, prod, beta) + if iszero(beta) + C.mat = prod + elseif isone(beta) + C.mat = prod + C.mat + else + C.mat = prod + beta * C.mat + end + return C +end + +function Dagger.matmatmul!( + C::DSparseMatrix, + transA::Char, + transB::Char, + A::SparseMatrixCSC, + B::SparseMatrixCSC, + alpha, + beta +) + opA = _apply_trans(A, transA) + opB = _apply_trans(B, transB) + # Sparse*sparse yields a freshly-allocated sparse matrix, which we reassign + # into the wrapper (`DSparseMatrix` hides this reallocation from Datadeps). + # `SparseArrays` provides no efficient 5-arg `mul!` into a sparse `C` -- the + # output sparsity pattern is determined by the product -- so we form the + # product out-of-place and apply only the alpha/beta scaling that is actually + # needed. The transposed-operand products dispatch to specialized SparseArrays + # methods, so `opA`/`opB` are not materialized. + AB = opA * opB + prod = isone(alpha) ? AB : alpha * AB + return _sparse_gemm_assign!(C, prod, beta) +end + +# DeviceSparseMatrixCSC SpGEMM: gather to host, multiply, scatter back. +function Dagger.matmatmul!( + C::DSparseMatrix, + transA::Char, + transB::Char, + A::Dagger.DeviceSparseMatrixCSC, + B::Dagger.DeviceSparseMatrixCSC, + alpha, + beta +) + Ah = SparseMatrixCSC(A) + Bh = SparseMatrixCSC(B) + Ch = C.mat isa Dagger.DeviceSparseMatrixCSC ? SparseMatrixCSC(C.mat) : + C.mat isa SparseMatrixCSC ? C.mat : SparseMatrixCSC(C.mat) + opA = _apply_trans(Ah, transA) + opB = _apply_trans(Bh, transB) + AB = opA * opB + prod = isone(alpha) ? AB : alpha * AB + if iszero(beta) + result = prod + elseif isone(beta) + result = prod + Ch + else + result = prod + beta * Ch + end + C.mat = _to_device_sparse(A, SparseMatrixCSC(result)) + return C +end + +# Sparse matrix-vector multiply tile kernel: `C = alpha*op(A)*B + beta*C` with a +# `SparseMatrixCSC` `A` and dense vectors `B`/`C`. SparseArrays provides an +# efficient 5-arg `mul!` (SpMV) into a dense output, including for transposed and +# adjoint operands, so this updates `C` in place with no allocation. +function Dagger.matvecmul!(C::AbstractVector, transA::Char, A::SparseMatrixCSC, B::AbstractVector, alpha, beta) + LinearAlgebra.mul!(C, _apply_trans(A, transA), B, alpha, beta) + return C +end + +# DeviceSparseMatrixCSC SpMV: host fallback (works for any dense vector type +# that supports Array(::)/copyto!). +function Dagger.matvecmul!(C::AbstractVector, transA::Char, A::Dagger.DeviceSparseMatrixCSC, B::AbstractVector, alpha, beta) + Ah = SparseMatrixCSC(A) + Bh = Array(B) + Ch = Array(C) + LinearAlgebra.mul!(Ch, _apply_trans(Ah, transA), Bh, alpha, beta) + copyto!(C, Ch) + return C +end + +# Off-diagonal tile copy in `copytri!`: produce the (conjugate) transpose tile. +function Dagger.transpose_tile(B::SparseMatrixCSC) + return SparseArrays.sparse(B') +end +function Dagger.transpose_tile(B::Dagger.DeviceSparseMatrixCSC) + return _to_device_sparse(B, SparseArrays.sparse(SparseMatrixCSC(B)')) +end +# Diagonal tile symmetrization in `copytri!`: build the full Hermitian tile from +# its `uplo` triangle (matching the dense `copydiagtile!` semantics). +function Dagger.transpose_tile(B::SparseMatrixCSC, uplo::Char) + if uplo == 'U' + Bt = SparseArrays.triu(B) + elseif uplo == 'L' + Bt = SparseArrays.tril(B) + else + throw(ArgumentError("uplo must be 'U' or 'L', got $uplo")) + end + C = Bt + Bt' + # The shared diagonal was added twice; restore the original tile's diagonal. + for i in 1:LinearAlgebra.checksquare(B) + C[i, i] = B[i, i] + end + return C +end +function Dagger.transpose_tile(B::Dagger.DeviceSparseMatrixCSC, uplo::Char) + return _to_device_sparse(B, Dagger.transpose_tile(SparseMatrixCSC(B), uplo)) +end + +end # module SparseArraysExt diff --git a/src/Dagger.jl b/src/Dagger.jl index 81881716e..3fe2a54ba 100644 --- a/src/Dagger.jl +++ b/src/Dagger.jl @@ -2,7 +2,6 @@ module Dagger import Serialization import Serialization: AbstractSerializer, serialize, deserialize -import SparseArrays: sprand, SparseMatrixCSC import MemPool import MemPool: DRef, FileRef, poolget, poolset @@ -132,7 +131,7 @@ include("array/operators.jl") include("array/indexing.jl") include("array/setindex.jl") include("array/matrix.jl") -include("array/sparse_partition.jl") +include("array/sparse.jl") include("array/sort.jl") include("array/permute.jl") include("array/linalg.jl") @@ -142,6 +141,8 @@ include("array/trsm.jl") include("array/lu.jl") include("array/qr.jl") include("array/svd.jl") +include("array/iterativesolvers.jl") +include("array/sparsedirect.jl") # GPU include("gpu.jl") diff --git a/src/array/alloc.jl b/src/array/alloc.jl index 93862d545..6ffbc8f50 100644 --- a/src/array/alloc.jl +++ b/src/array/alloc.jl @@ -10,13 +10,20 @@ mutable struct AllocateArray{T,N} <: ArrayOp{T,N} domainchunks partitioning::AbstractBlocks{N} procgrid::Union{AbstractProcGrid{N}, Nothing} - - function AllocateArray(eltype::Type{T}, f, want_index::Bool, d::ArrayDomain{N}, domainchunks, p::AbstractBlocks{N}, assignment::Union{AssignmentType{N},Nothing} = nothing) where {T,N} + # Optional concrete tile type for staged `@spawn`s. When `nothing`, leave + # `return_type` unset so inference (and MPI `post_stage_array_chunks!`, + # which defaults dense tiles to `Array{T,N}`) decide. Sparse allocators + # pass e.g. `DSparseArray{T,N}` so MPI does not overwrite with `Array`. + return_type::Union{Type,Nothing} + + function AllocateArray(eltype::Type{T}, f, want_index::Bool, d::ArrayDomain{N}, + domainchunks, p::AbstractBlocks{N}, + assignment::Union{AssignmentType{N},Nothing} = nothing; + return_type::Union{Type,Nothing} = nothing) where {T,N} sizeA = map(length, d.indexes) procgrid = build_procgrid(something(assignment, :arbitrary), Tuple(sizeA), p.blocksize, current_acceleration()) - return new{T,N}(eltype, f, want_index, d, domainchunks, p, procgrid) + return new{T,N}(eltype, f, want_index, d, domainchunks, p, procgrid, return_type) end - end size(a::AllocateArray) = size(a.domain) @@ -45,12 +52,22 @@ function stage(ctx, A::AllocateArray) tasks = emit_chunk_tasks!(A.domainchunks, A.procgrid, A.eltype, (scope, I, i) -> begin x = A.domainchunks[I] - N = ndims(A.domainchunks) - ret_type = Array{A.eltype, N} - if A.want_index - Dagger.@spawn compute_scope=scope return_type=ret_type allocate_array(A.f, A.eltype, i, size(x)) + if A.f isa DArray + # `similar(::DArray)` carries the source tile type (e.g. sparse) forward. + chunk = A.f.chunks[I] + Dagger.@spawn compute_scope=scope similar(chunk, A.eltype, size(x)) + elseif A.want_index + if A.return_type !== nothing + Dagger.@spawn compute_scope=scope return_type=A.return_type allocate_array(A.f, A.eltype, i, size(x)) + else + Dagger.@spawn compute_scope=scope allocate_array(A.f, A.eltype, i, size(x)) + end else - Dagger.@spawn compute_scope=scope return_type=ret_type allocate_array(A.f, A.eltype, size(x)) + if A.return_type !== nothing + Dagger.@spawn compute_scope=scope return_type=A.return_type allocate_array(A.f, A.eltype, size(x)) + else + Dagger.@spawn compute_scope=scope allocate_array(A.f, A.eltype, size(x)) + end end end) return DArray(A.eltype, A.domain, A.domainchunks, tasks, A.partitioning) @@ -86,20 +103,6 @@ Base.randn(p::BlocksOrAuto, dims::Dims; assignment::AssignmentType = :arbitrary) Base.randn(::AutoBlocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) = randn(auto_blocks(dims), T, dims; assignment) -function sprand(p::Blocks, T::Type, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) - d = ArrayDomain(map(x->1:x, dims)) - a = AllocateArray(T, (T, _dims) -> sprand(T, _dims..., sparsity), false, d, partition(p, d), p, assignment) - return _to_darray(a) -end -sprand(p::BlocksOrAuto, T::Type, dims_and_sparsity::Real...; assignment::AssignmentType = :arbitrary) = - sprand(p, T, dims_and_sparsity[1:end-1], dims_and_sparsity[end]; assignment) -sprand(p::BlocksOrAuto, dims_and_sparsity::Real...; assignment::AssignmentType = :arbitrary) = - sprand(p, Float64, dims_and_sparsity[1:end-1], dims_and_sparsity[end]; assignment) -sprand(p::BlocksOrAuto, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) = - sprand(p, Float64, dims, sparsity; assignment) -sprand(::AutoBlocks, T::Type, dims::Dims, sparsity::AbstractFloat; assignment::AssignmentType = :arbitrary) = - sprand(auto_blocks(dims), T, dims, sparsity; assignment) - function Base.ones(p::Blocks, T::Type, dims::Dims; assignment::AssignmentType = :arbitrary) d = ArrayDomain(map(x->1:x, dims)) a = AllocateArray(T, ones, false, d, partition(p, d), p, assignment) @@ -167,3 +170,14 @@ function unsafe_free!(A::DArray) end end end + +# Initializers + +function Base.fill!(A::DArray, x) + spawn_datadeps() do + for chunk in A.chunks + Dagger.@spawn fill!(chunk, x) + end + end + return A +end diff --git a/src/array/copy.jl b/src/array/copy.jl index 7d92566e6..d703f49d1 100644 --- a/src/array/copy.jl +++ b/src/array/copy.jl @@ -36,6 +36,13 @@ function allocate_copy_buffer(part::Blocks{N}, A::DArray{T,N}) where {T,N} return DArray{T}(undef, part, size(A)) end +to_range(x::UnitRange) = x +to_range(x::Integer) = x:x +to_range(x::Base.OneTo{Int}) = UnitRange(x) +to_range(x::Base.Slice{Base.OneTo{Int}}) = UnitRange(x) +to_range(::StepRange) = throw(ArgumentError("Cannot convert StepRange to UnitRange")) +to_range(x) = throw(ArgumentError("Cannot convert $(typeof(x)) to UnitRange")) + function darray_copyto!(B::DArray{TB,NB}, A::DArray{TA,NA}, Binds=parentindices(B), Ainds=parentindices(A)) where {TB,NB,TA,NA} Nmax = max(NA, NB) @@ -45,13 +52,6 @@ function darray_copyto!(B::DArray{TB,NB}, A::DArray{TA,NA}, Binds=parentindices( padNmax(x) = ntuple(i->pad1range(x, i), Nmax) padNmax(x::ArrayDomain) = padNmax(x.indexes) - to_range(x::UnitRange) = x - to_range(x::Integer) = x:x - to_range(x::Base.OneTo{Int}) = UnitRange(x) - to_range(x::Base.Slice{Base.OneTo{Int}}) = UnitRange(x) - to_range(::StepRange) = throw(ArgumentError("Non-continuous ranges are not yet supported for DArray copy")) - to_range(x) = throw(ArgumentError("Unsupported range type for DArray copy: $(typeof(x))")) - if any(x->x isa Vector, Binds) || any(x->x isa Vector, Ainds) # Split the copy into multiple copies dims_with_vector = findall(x->x[1] isa Vector || x[2] isa Vector, collect(zip(Binds, Ainds))) @@ -154,3 +154,13 @@ StridedDArray{T,N} = Union{<:DArray{T,N}, SubArray{T,N,<:DArray{T,NP}} where NP} Base.copyto!(B::StridedDArray, A::StridedDArray) = darray_copyto!(parent(B), parent(A), parentindices(B), parentindices(A)) +function Base.copyto!(B::Array, A::StridedDArray) + DB = view(B, AutoBlocks()) + darray_copyto!(DB, parent(A), parentindices(DB), parentindices(A)) + return B +end +function Base.copyto!(B::SubArray, A::StridedDArray) + DB = view(parent(B), AutoBlocks()) + darray_copyto!(DB, parent(A), parentindices(B), parentindices(A)) + return B +end diff --git a/src/array/darray.jl b/src/array/darray.jl index bd0e5efc5..b774d913f 100644 --- a/src/array/darray.jl +++ b/src/array/darray.jl @@ -180,6 +180,11 @@ stage(ctx, c::DArray) = c # closures get non-uniform ArgumentWrapper hashes). _collect_cat(concat, dims::Int, xs...) = concat(xs...; dims) +# Finalize a gathered DArray to a dense `Array{T,N}`. Sparse tile cats may yield +# a `DSparseArray` / `SparseMatrixCSC`, and `Base.collect` does not densify those. +_collect_dense(::Type{T}, ::Val{N}, x::Array{T,N}) where {T,N} = x +_collect_dense(::Type{T}, ::Val{N}, x) where {T,N} = Array{T,N}(x) + function Base.collect(d::DArray{T,N}; tree=true, copyto=false) where {T,N} a = fetch(d) if isempty(d.chunks) @@ -208,13 +213,21 @@ function Base.collect(d::DArray{T,N}; tree=true, copyto=false) where {T,N} result = Dagger.spawn_datadeps() do treereduce_nd(spawn_catfuncs, a.chunks) end - return collect(fetch(result)) + return _collect_dense(T, Val(N), fetch(result)) end # Distributed: fetch chunks directly and concat in-process. This avoids # routing chunk data through datadeps aliasing, which requires an # aliasing-resolvable (e.g. isbits) element type. + # + # Sparse tiles (`DSparseArray`) are unwrapped to host storage before `cat`: + # GPU sparse types (CuSparse/ROCSparse/DeviceSparseMatrixCSC) disallow the + # scalar indexing that generic `cat` would use. dimcatfuncs = [(x...) -> concat(x..., dims=i) for i in 1:N] - return collect(treereduce_nd(dimcatfuncs, asyncmap(fetch, a.chunks))) + tiles = asyncmap(a.chunks) do c + x = fetch(c) + x isa DSparseArray ? _sparse_collect(x.mat) : x + end + return _collect_dense(T, Val(N), treereduce_nd(dimcatfuncs, tiles)) end Array{T,N}(A::DArray{S,N}) where {T,N,S} = convert(Array{T,N}, collect(A)) @@ -342,8 +355,16 @@ aliasing(x::DArray) = memory_space(x::DArray) = throw(ConcurrencyViolationError("DArray memory spaces may be mixed and unstable")) -Base.similar(D::DArray{T,N} where T, ::Type{S}, dims::Dims{N}) where {S,N} = - DArray{S,N}(undef, D.partitioning, dims) +function Base.similar(D::DArray{T,N} where T, ::Type{S}, dims::Dims{N}) where {S,N} + if dims == size(D) + domain = ArrayDomain(map(x->1:x, dims)) + subdomains = partition(D.partitioning, domain) + a = AllocateArray(S, D, false, domain, subdomains, D.partitioning, nothing) + return _to_darray(a) + else + return DArray{S,N}(undef, D.partitioning, dims) + end +end Base.similar(D::DArray{T,N1} where T, ::Type{S}, dims::Dims{N2}) where {S,N1,N2} = DArray{S,N2}(undef, auto_blocks(dims), dims) @@ -413,31 +434,21 @@ end """ Base.fetch(c::DArray) -If a `DArray` tree has a `Thunk` in it, make the whole thing a big thunk. +If a `DArray` tree has a `DTask` in it, `fetch` each one (as a `Chunk`, +without moving its data) and return a new `DArray` with all partitions +resolved to `Chunk`s. """ function Base.fetch(c::DArray{T}) where T - if any(istask, chunks(c)) - if uniform_execution() - # SPMD (MPI): every rank holds every task locally, and chunk - # records are rank-local views (placeholders for non-owned data). - # Resolve them locally instead of assembling on one rank and - # broadcasting (which would clobber the local records with the - # assembling rank's placeholders). - thunks = chunks(c) - new_chunks = Any[istask(t) ? fetch(t; raw=true) : t for t in thunks] - return DArray(T, domain(c), domainchunks(c), - reshape(new_chunks, size(thunks)), - c.partitioning, c.concat) - end - thunks = chunks(c) - sz = size(thunks) + thunks = chunks(c) + if any(istask, thunks) dmn = domain(c) dmnchunks = domainchunks(c) - return fetch(Dagger.spawn(Options(meta=true), thunks...) do results... - t = eltype(fetch(results[1])) - DArray(t, dmn, dmnchunks, reshape(Any[results...], sz), - c.partitioning, c.concat) - end) + new_chunks = similar(thunks, Any) + for idx in eachindex(thunks) + part = thunks[idx] + new_chunks[idx] = istask(part) ? fetch(part; raw=true) : part + end + return DArray(T, dmn, dmnchunks, new_chunks, c.partitioning, c.concat) else return c end @@ -503,7 +514,7 @@ function stage(ctx::Context, d::Distribute) cs = emit_chunk_tasks!(d.domainchunks, d.procgrid, T, (scope, I, i) -> begin c = d.domainchunks[I] - Dagger.@spawn compute_scope=scope identity(d.data[c]) + Dagger.@spawn compute_scope=scope maybe_wrap_tile(d.data[c]) end) end return DArray(eltype(d.data), diff --git a/src/array/iterativesolvers.jl b/src/array/iterativesolvers.jl new file mode 100644 index 000000000..f3d2f116c --- /dev/null +++ b/src/array/iterativesolvers.jl @@ -0,0 +1,358 @@ +# Distributed iterative (Krylov) linear solvers. +# +# The user-facing entry points (`Dagger.cg`, `Dagger.minres`, `Dagger.gmres`, +# `Dagger.bicgstab`, and the generic `Dagger.krylov_solve`) live here so the +# solver backend can evolve without changing user code. The actual solves are +# implemented in `ext/KrylovExt.jl`, which is loaded when `Krylov` is available +# (`using Krylov`). The matrix-free building blocks they rely on -- distributed +# `mul!`/SpMV and the BLAS-1 vector ops -- live in `mul.jl`/`linalg.jl`. + +""" + cg(A, b::DVector; M=I, atol, rtol, itmax, ...) -> (x::DVector, stats) + +Solve the symmetric positive-definite system `A x = b` with the conjugate +gradient method, distributed over `A`'s and `b`'s chunks. `A` may be a +`DMatrix` (dense or sparse-backed) or any object supporting `mul!(y, A, x)` over +`DVector`s (matrix-free). Requires `Krylov.jl` to be loaded. + +See also [`minres`](@ref), [`gmres`](@ref), [`bicgstab`](@ref), and +[`krylov_solve`](@ref). Keyword arguments are forwarded to `Krylov.cg!`. +""" +function cg end + +""" + minres(A, b::DVector; M=I, atol, rtol, itmax, ...) -> (x::DVector, stats) + +Solve the symmetric (possibly indefinite) system `A x = b` with MINRES. +Requires `Krylov.jl` to be loaded. See [`cg`](@ref). +""" +function minres end + +""" + gmres(A, b::DVector; M=I, N=I, restart=false, memory=20, ...) -> (x::DVector, stats) + +Solve the general (nonsymmetric) system `A x = b` with restarted GMRES. +`M`/`N` are left/right preconditioners. Requires `Krylov.jl` to be loaded. +See [`cg`](@ref). +""" +function gmres end + +""" + bicgstab(A, b::DVector; M=I, N=I, ...) -> (x::DVector, stats) + +Solve the general (nonsymmetric) system `A x = b` with BiCGStab (short +recurrence, low memory). Requires `Krylov.jl` to be loaded. See [`cg`](@ref). +""" +function bicgstab end + +""" + krylov_solve(method::Symbol, A, b::DVector; kwargs...) -> (x::DVector, stats) + +Generic entry point dispatching to the iterative `method` (`:cg`, `:minres`, +`:gmres`, `:bicgstab`). Requires `Krylov.jl` to be loaded. +""" +function krylov_solve end + +# Friendly fallbacks: these generic methods are shadowed by the more specific +# `(A, b::DVector)` methods added in `ext/KrylovExt.jl` once Krylov is loaded. +_krylov_required(name) = throw(ArgumentError( + "Dagger.$name requires Krylov.jl. Run `using Krylov` (or `import Krylov`) \ + to enable distributed iterative solvers.")) +cg(A, b; kwargs...) = _krylov_required(:cg) +minres(A, b; kwargs...) = _krylov_required(:minres) +gmres(A, b; kwargs...) = _krylov_required(:gmres) +bicgstab(A, b; kwargs...) = _krylov_required(:bicgstab) +krylov_solve(method::Symbol, A, b; kwargs...) = _krylov_required(:krylov_solve) + +# --- Preconditioners ------------------------------------------------------ +# A preconditioner object `P` represents the (approximate) *inverse* operator +# `M⁻¹`: applying it (`mul!(y, P, x)`) computes `y = M⁻¹ x`, which is exactly an +# ordinary matrix-vector product by the operator `P` stands for (no inversion +# happens at apply time -- any reciprocals/factorizations are precomputed once). +# This matches Krylov's `ldiv=false` convention, where the object passed as `M` +# is applied via `mul!(y, M, x)` to compute `y ← M⁻¹ x`, so these are passed +# straight through as `M=P`. +# +# These preconditioners are backend-agnostic: they operate on `DVector`s (and a +# `DMatrix`'s diagonal tiles) and need no solver backend, so they live in core. +# +# A note on the square-tile requirement below: the iterative solvers allocate +# *all* workspace vectors as `similar(b)` (one partitioning), and each must +# serve as both the length-`n` input and length-`n` output of `mul!(y, A, x)`. +# `gemv_dagger!` requires the input to match `A`'s column blocks and the output +# to match `A`'s row blocks, which is only simultaneously possible when those +# block sizes are equal. So the operator must use square tiles regardless of +# preconditioning; the preconditioners simply share that constraint. (Any +# uniform square tile size is fine, including a ragged final block.) + +""" + AbstractDaggerPreconditioner + +Supertype for Dagger's distributed preconditioners. A preconditioner `P` +represents an (approximate) inverse operator `M⁻¹` and applies it via +`mul!(y, P, x)` (`y = M⁻¹ x`) over `DVector`s, so it can be passed to the +solvers as `M=P` (with `ldiv=false`, the default). +""" +abstract type AbstractDaggerPreconditioner end + +""" + JacobiPreconditioner(A::DMatrix) + +Diagonal (Jacobi) preconditioner. The object represents the inverse-diagonal +operator `M⁻¹ = inv(Diagonal(A))`; it precomputes and stores the reciprocal +diagonal `dinv = 1 ./ diag(A)`, so applying it (`mul!(y, P, x)`) is a single +elementwise multiply `y = dinv .* x` per chunk -- an ordinary product by the +stored operator, not an inversion. Cheap to build (one diagonal extraction per +diagonal tile) and to apply. + +Requires square diagonal tiles (equal row/column block sizes); see the note +above on why the solver itself imposes this. Repartition `A` with equal block +sizes if needed. +""" +struct JacobiPreconditioner{V<:DVector} <: AbstractDaggerPreconditioner + dinv::V +end + +JacobiPreconditioner(A::DMatrix) = JacobiPreconditioner(_jacobi_dinv(A)) + +# Per-tile kernel: write the reciprocal diagonal of `tile` into `out`. +_set_inv_diag!(out, tile) = (out .= inv.(LinearAlgebra.diag(tile)); return nothing) + +# Validate that `A` has square diagonal tiles and return (n, Ac, mt, blocksize). +function _check_square_tiles(A::DMatrix, who) + n = LinearAlgebra.checksquare(A) + Ac = A.chunks + mt, nt = size(Ac) + mb, nb = A.partitioning.blocksize + # For a square matrix with square tiles, mt == nt follows automatically; we + # only need to reject non-square tiles (mb != nb). + mb == nb || throw(ArgumentError( + "$who requires square diagonal tiles (got block size $(mb)x$(nb)); \ + repartition A with equal block sizes (this matches the iterative \ + solver's own requirement on the operator)")) + return n, Ac, mt, mb +end + +function _jacobi_dinv(A::DMatrix{T}) where T + n, Ac, mt, mb = _check_square_tiles(A, "JacobiPreconditioner") + dinv = DVector{T}(undef, Blocks(mb), n) + dc = dinv.chunks + Dagger.spawn_datadeps() do + for i in 1:mt + Dagger.@spawn _set_inv_diag!(Out(dc[i]), In(Ac[i, i])) + end + end + return dinv +end + +function LinearAlgebra.mul!(y::DVector, P::JacobiPreconditioner, x::DVector) + part = P.dinv.partitioning + maybe_copy_buffered(P.dinv => part, x => part, y => part) do dinv, x, y + dc, xc, yc = dinv.chunks, x.chunks, y.chunks + Dagger.spawn_datadeps() do + for i in eachindex(yc) + Dagger.@spawn _jacobi_apply!(Out(yc[i]), In(dc[i]), In(xc[i])) + end + end + end + return y +end +_jacobi_apply!(y, dinv, x) = (y .= dinv .* x; return nothing) + +# --- Block-diagonal preconditioners --------------------------------------- +# A family of preconditioners of the form `M⁻¹ = blockdiag(op₁, …, op_k)`, where +# each `opⱼ` approximates the inverse of `A`'s `j`-th diagonal tile `Aⱼⱼ`. They +# share all machinery and differ only in how a per-tile operator is *built*: +# +# - BlockJacobiPreconditioner: exact tile solve via `lu` (dense or sparse). +# - BlockILUPreconditioner: incomplete LU per tile (needs IncompleteLU.jl). +# - AMGPreconditioner: an AMG hierarchy per tile (needs AlgebraicMultigrid.jl). +# +# Building per tile is embarrassingly parallel and a natural fit for the tiled +# layout. A single tile (`Blocks(n, n)`) makes any of these a *global* +# preconditioner over the whole matrix; many tiles make it a scalable +# block-Jacobi / additive-Schwarz variant that trades some convergence for +# parallelism. All require square diagonal tiles (see `JacobiPreconditioner`). +# +# A per-tile operator (`lu`/ILU factor, AMG hierarchy) generally cannot be moved +# between workers (dense `LU` has no `move!`; `UmfpackLU`/AMG hold process-bound +# resources). So each operator is built *once* and pinned (via +# `tochunk(..., ProcessScope)`) to the worker that owns its tile, and every apply +# for that block is scheduled there (`compute_scope`); datadeps then moves only +# the (movable) vector chunks to the operator, never the operator itself. + +""" + AbstractBlockPreconditioner <: AbstractDaggerPreconditioner + +Block-diagonal preconditioner: holds one per-diagonal-tile operator `opⱼ` +(`y ← opⱼ⁻¹ x`), pinned to its tile's worker, and applies them independently per +block. Concrete subtypes (`BlockJacobiPreconditioner`, `BlockILUPreconditioner`, +`AMGPreconditioner`) share these fields: `ops`, `scopes`, `part`, `n`. +""" +abstract type AbstractBlockPreconditioner <: AbstractDaggerPreconditioner end + +_tile_scope(c::Chunk) = ProcessScope(root_worker_id(c)) +_tile_scope(t::DTask) = ProcessScope(root_worker_id(fetch(t; raw=true))) + +# Build the per-tile operator on the current worker and pin the result there: the +# returned `Chunk` is process-scoped, so it can never be moved off this worker. +function _build_tile_pinned(build, tile) + op = build(tile) + proc = Dagger.task_processor() + return tochunk(op, proc, ProcessScope(root_worker_id(proc))) +end + +# Build a block preconditioner of type `Ctor` by applying `build` to each +# diagonal tile (pinned to the tile's worker). `build(tile) -> op` returns a +# per-tile operator supporting the block apply below. +function _build_block_preconditioner(Ctor, A::DMatrix, who, build) + n, Ac, mt, mb = _check_square_tiles(A, who) + ops = Vector{Any}(undef, mt) + scopes = Vector{ProcessScope}(undef, mt) + for i in 1:mt + tile = Ac[i, i] + scope = _tile_scope(tile) + scopes[i] = scope + ops[i] = Dagger.@spawn compute_scope=scope _build_tile_pinned(build, tile) + end + return Ctor(ops, scopes, Blocks(mb), n) +end + +function LinearAlgebra.mul!(y::DVector, P::AbstractBlockPreconditioner, x::DVector) + part = P.part + maybe_copy_buffered(x => part, y => part) do x, y + xc, yc = x.chunks, y.chunks + length(yc) == length(P.ops) || throw(DimensionMismatch( + "$(nameof(typeof(P))) has $(length(P.ops)) blocks but the vector has \ + $(length(yc)) chunks")) + Dagger.spawn_datadeps() do + for i in eachindex(yc) + # Pin the apply to the operator's worker; the operator is passed as + # an untracked arg (read-only, already-pinned) so datadeps never + # moves it -- only the vector chunks are moved to this worker. + Dagger.@spawn compute_scope=P.scopes[i] _block_apply!(Out(yc[i]), P.ops[i], In(xc[i])) + end + end + end + return y +end + +# Apply one block operator: `y = op⁻¹ x`. Default uses `\`; backends override for +# operator types where `\` is unavailable (e.g. ILU/AMG use `ldiv!`). +# Host factorizations (UMFPACK/`lu` of gathered CSC) need a host RHS; GPU vector +# chunks are gathered, solved, and written back. CPU `Array` short-circuits. +function _block_apply!(y, op, x) + if x isa Array + y .= op \ x + else + copyto!(y, op \ Adapt.adapt(Array, x)) + end + return nothing +end + +# Underlying matrix of a tile (overridden for `DSparseArray` in `sparse.jl`). +_tile_matrix(A) = A + +""" + BlockJacobiPreconditioner(A::DMatrix) + +Block-Jacobi preconditioner: `M⁻¹ = blockdiag(A₁₁, …, A_kk)⁻¹`. Each diagonal +tile is factorized *once* with `lu` (sparse or dense) and the apply solves +`Aᵢᵢ yᵢ = xᵢ`. Stronger than `JacobiPreconditioner` (captures intra-block +coupling); a single tile recovers an exact solve. See +[`AbstractBlockPreconditioner`](@ref). Requires square diagonal tiles. +""" +struct BlockJacobiPreconditioner{F,S} <: AbstractBlockPreconditioner + ops::F # cached per-tile factorizations (pinned to their workers) + scopes::S # the `ProcessScope` each operator/apply is pinned to + part::Blocks{1} # partitioning of the vectors it applies to + n::Int +end + +# Factorize a diagonal tile. Backends override for their storage (e.g. sparse +# tiles factorize the inner `SparseMatrixCSC`); the default is a dense LU factor. +_factorize_tile(A) = LinearAlgebra.lu(A) + +BlockJacobiPreconditioner(A::DMatrix) = + _build_block_preconditioner(BlockJacobiPreconditioner, A, "BlockJacobiPreconditioner", _factorize_tile) + +""" + BlockILUPreconditioner(A::DMatrix; τ=0.001, kwargs...) + +Block incomplete-LU preconditioner: an ILU factorization (with drop tolerance +`τ`) of each diagonal tile, applied per block. Cheaper setup than a full block +solve, good as a general-purpose preconditioner. Requires `IncompleteLU.jl` to +be loaded and sparse-backed tiles. See [`AbstractBlockPreconditioner`](@ref). +""" +struct BlockILUPreconditioner{F,S} <: AbstractBlockPreconditioner + ops::F + scopes::S + part::Blocks{1} + n::Int +end +# Friendly fallback (shadowed by the `::DMatrix` method added in `ext/IncompleteLUExt.jl`). +BlockILUPreconditioner(A; kwargs...) = throw(ArgumentError( + "Dagger.BlockILUPreconditioner requires IncompleteLU.jl. Run `using IncompleteLU` \ + to enable block incomplete-LU preconditioning.")) + +""" + AMGPreconditioner(A::DMatrix; method=:ruge_stuben, kwargs...) + +Algebraic-multigrid preconditioner: builds an AMG hierarchy (`method` is +`:ruge_stuben` or `:smoothed_aggregation`) for each diagonal tile and applies a +V-cycle per block. Near mesh-independent convergence for elliptic (Poisson-like) +operators. With one tile this is global AMG; with many tiles it is a scalable +block/additive-Schwarz AMG. Requires `AlgebraicMultigrid.jl` to be loaded and +sparse-backed tiles. See [`AbstractBlockPreconditioner`](@ref). +""" +struct AMGPreconditioner{F,S} <: AbstractBlockPreconditioner + ops::F + scopes::S + part::Blocks{1} + n::Int +end +# Friendly fallback (shadowed by the `::DMatrix` method added in `ext/AlgebraicMultigridExt.jl`). +AMGPreconditioner(A; kwargs...) = throw(ArgumentError( + "Dagger.AMGPreconditioner requires AlgebraicMultigrid.jl. Run \ + `using AlgebraicMultigrid` to enable algebraic-multigrid preconditioning.")) + +""" + BlockKLUPreconditioner(A::DMatrix; kwargs...) + +Block direct preconditioner using KLU: an exact (pure-Julia) KLU factorization of +each diagonal tile, applied per block. With a single tile this is an exact direct +solve; with many tiles it is an exact-block-Jacobi preconditioner. Requires +`PureKLU.jl` to be loaded and sparse-backed tiles. See also +[`AbstractBlockPreconditioner`](@ref) and the whole-matrix solver +[`Dagger.klu`](@ref). +""" +struct BlockKLUPreconditioner{F,S} <: AbstractBlockPreconditioner + ops::F + scopes::S + part::Blocks{1} + n::Int +end +# Friendly fallback (shadowed by the `::DMatrix` method added in `ext/PureKLUExt.jl`). +BlockKLUPreconditioner(A; kwargs...) = throw(ArgumentError( + "Dagger.BlockKLUPreconditioner requires PureKLU.jl. Run `using PureKLU` to \ + enable block KLU preconditioning.")) + +""" + BlockUMFPACKPreconditioner(A::DMatrix; kwargs...) + +Block direct preconditioner using a pure-Julia UMFPACK-style LU of each diagonal +tile, applied per block. With a single tile this is an exact direct solve; with +many tiles it is an exact-block-Jacobi preconditioner. Requires `PureUMFPACK.jl` +to be loaded and sparse-backed tiles. See also [`AbstractBlockPreconditioner`](@ref) +and the whole-matrix solver [`Dagger.splu`](@ref). +""" +struct BlockUMFPACKPreconditioner{F,S} <: AbstractBlockPreconditioner + ops::F + scopes::S + part::Blocks{1} + n::Int +end +# Friendly fallback (shadowed by the `::DMatrix` method added in `ext/PureUMFPACKExt.jl`). +BlockUMFPACKPreconditioner(A; kwargs...) = throw(ArgumentError( + "Dagger.BlockUMFPACKPreconditioner requires PureUMFPACK.jl. Run \ + `using PureUMFPACK` to enable block UMFPACK preconditioning.")) diff --git a/src/array/linalg.jl b/src/array/linalg.jl index 2fcaa42fa..89d0575e1 100644 --- a/src/array/linalg.jl +++ b/src/array/linalg.jl @@ -4,6 +4,71 @@ function LinearAlgebra.norm2(A::DArray{T,N}) where {T,N} zeroRT = zero(real(T)) return sqrt(sum(map(norm->fetch(norm)::real(T), norms); init=zeroRT)) end + +# --- BLAS-1 vector operations --------------------------------------------- +# Efficient, distributed level-1 operations needed by iterative (Krylov) +# solvers. These run one BLAS-1 kernel per chunk (avoiding the generic +# scalar-indexing fallbacks, which are slow and trigger scalar access on GPUs). +# Operands may be partitioned differently: we align them with +# `maybe_copy_buffered` (a no-op when they already share a layout, and which +# copies results back to the originals for in-place operations), so users keep +# full control over their own (possibly sub-optimal) partitioning. + +function LinearAlgebra.dot(x::DArray{Tx,N}, y::DArray{Ty,N}) where {Tx,Ty,N} + size(x) == size(y) || throw(DimensionMismatch("dot: x has size $(size(x)), y has size $(size(y))")) + R = typeof(LinearAlgebra.dot(zero(Tx), zero(Ty))) + return maybe_copy_buffered(x => x.partitioning, y => x.partitioning) do x, y + xc, yc = x.chunks, y.chunks + parts = [Dagger.@spawn LinearAlgebra.dot(xc[i], yc[i]) for i in eachindex(xc)] + sum(fetch, parts; init=zero(R)) + end +end + +function LinearAlgebra.axpy!(a::Number, x::DArray{Tx,N}, y::DArray{Ty,N}) where {Tx,Ty,N} + size(x) == size(y) || throw(DimensionMismatch("axpy!: x has size $(size(x)), y has size $(size(y))")) + # Align the read-only `x` to `y`'s layout so only `x` may be buffered; `y` is + # updated in place (and copied back by `maybe_copy_buffered` if it was buffered). + maybe_copy_buffered(x => y.partitioning, y => y.partitioning) do x, y + xc, yc = x.chunks, y.chunks + Dagger.spawn_datadeps() do + for i in eachindex(xc) + Dagger.@spawn LinearAlgebra.axpy!(a, In(xc[i]), InOut(yc[i])) + end + end + end + return y +end + +function LinearAlgebra.axpby!(a::Number, x::DArray{Tx,N}, b::Number, y::DArray{Ty,N}) where {Tx,Ty,N} + size(x) == size(y) || throw(DimensionMismatch("axpby!: x has size $(size(x)), y has size $(size(y))")) + maybe_copy_buffered(x => y.partitioning, y => y.partitioning) do x, y + xc, yc = x.chunks, y.chunks + Dagger.spawn_datadeps() do + for i in eachindex(xc) + Dagger.@spawn LinearAlgebra.axpby!(a, In(xc[i]), b, InOut(yc[i])) + end + end + end + return y +end + +function LinearAlgebra.rmul!(x::DArray, a::Number) + Dagger.spawn_datadeps() do + for c in x.chunks + Dagger.@spawn LinearAlgebra.rmul!(InOut(c), a) + end + end + return x +end + +function LinearAlgebra.lmul!(a::Number, x::DArray) + Dagger.spawn_datadeps() do + for c in x.chunks + Dagger.@spawn LinearAlgebra.lmul!(a, InOut(c)) + end + end + return x +end function LinearAlgebra.norm2(A::UpperTriangular{T,<:DArray{T,2}}) where T Ac = parent(A).chunks Ac_upper = [] diff --git a/src/array/mul.jl b/src/array/mul.jl index 14ed9f0ba..322b27760 100644 --- a/src/array/mul.jl +++ b/src/array/mul.jl @@ -1,3 +1,74 @@ +""" + matmatmul!(C, transA::Char, transB::Char, A, B, alpha, beta) + +A general-purpose matrix-matrix multiply, like `LinearAlgebra.generic_matmatmul!`, +but with extra functionality. May internally convert `A` and `B` to a type that +better matches `C` and provides optimal portability and, when possible, +better performance. The actual matrix multiply operation should happen in +`LinearAlgebra.generic_matmatmul!` or an equivalent call. + +The following automatic conversions are performed: +- If no `LinearAlgebra.generic_matmatmul!` method is available, convert `A` and `B` to dense Array-like +- If `C` is a `DSparseMatrix`, perform the operation out-of-place and then update `C` in-place +""" +function matmatmul!( + C, + transA::Char, + transB::Char, + A, + B, + alpha, + beta +) + EC = eltype(C) + EA = eltype(A) + EB = eltype(B) + + TC = typeof(C) + TA = typeof(A) + TB = typeof(B) + + mam = LinearAlgebra.MulAddMul(alpha, beta) + + # Check if C doesn't support in-place operations (e.g. DSparseMatrix) + # We'll get here if A and B don't have equivalent types + if isa(C, DSparseMatrix) + C.mat = alpha * A * B + beta * C.mat + return C + end + + # Check if the call will fail due to MethodError + sig = Tuple{TC, Char, Char, TA, TB, typeof(mam)} + if !hasmethod(LinearAlgebra.generic_matmatmul!, sig) + # Convert to Array-like + # FIXME: GPU support + C_new = C + A_new = collect(A) + B_new = collect(B) + alpha_new = alpha + beta_new = beta + # FIXME: Re-check hasmethod, and if no method, then convert and bounce C + @goto ready + end + + C_new = C + A_new = A + B_new = B + alpha_new = alpha + beta_new = beta + + @label ready + mam_new = LinearAlgebra.MulAddMul(alpha_new, beta_new) + return LinearAlgebra.generic_matmatmul!( + C_new, + transA, + transB, + A_new, + B_new, + mam_new + ) +end + function LinearAlgebra.generic_matmatmul!( C::DMatrix{T}, transA::Char, @@ -41,7 +112,8 @@ function LinearAlgebra.generic_matmatmul!( return gemm_dagger!(C, transA, transB, A, B, alpha, beta) end end -function _repartition_matmatmul(C, A, B, transA::Char, transB::Char)::Tuple{Blocks{2}, Blocks{2}, Blocks{2}} +# FIXME: Mixed-precision methods +function _repartition_matmatmul(C, A, B, transA::Char, transB::Char) partA = A.partitioning.blocksize partB = B.partitioning.blocksize istransA = transA == 'T' || transA == 'C' @@ -141,28 +213,28 @@ function gemm_dagger!( # A: NoTrans / B: NoTrans for k in range(1, Ant) mzone = k == 1 ? beta : T(1.0) - Dagger.@spawn BLAS.gemm!( + Dagger.@spawn matmatmul!( + InOut(Cc[m, n]), transA, transB, - alpha, In(Ac[m, k]), In(Bc[k, n]), + alpha, mzone, - InOut(Cc[m, n]), ) end else # A: NoTrans / B: [Conj]Trans for k in range(1, Ant) mzone = k == 1 ? beta : T(1.0) - Dagger.@spawn BLAS.gemm!( + Dagger.@spawn matmatmul!( + InOut(Cc[m, n]), transA, transB, - alpha, In(Ac[m, k]), In(Bc[n, k]), + alpha, mzone, - InOut(Cc[m, n]), ) end end @@ -171,28 +243,28 @@ function gemm_dagger!( # A: [Conj]Trans / B: NoTrans for k in range(1, Amt) mzone = k == 1 ? beta : T(1.0) - Dagger.@spawn BLAS.gemm!( + Dagger.@spawn matmatmul!( + InOut(Cc[m, n]), transA, transB, - alpha, In(Ac[k, m]), In(Bc[k, n]), + alpha, mzone, - InOut(Cc[m, n]), ) end else # A: [Conj]Trans / B: [Conj]Trans for k in range(1, Amt) mzone = k == 1 ? beta : T(1.0) - Dagger.@spawn BLAS.gemm!( + Dagger.@spawn matmatmul!( + InOut(Cc[m, n]), transA, transB, - alpha, In(Ac[k, m]), In(Bc[n, k]), + alpha, mzone, - InOut(Cc[m, n]), ) end end @@ -240,6 +312,7 @@ function syrk_dagger!( iscomplex = T <: Complex transs = iscomplex ? 'C' : 'T' + anti_transs = trans == 'N' ? transs : 'N' Dagger.spawn_datadeps() do for n in range(1, Cnt) @@ -247,116 +320,63 @@ function syrk_dagger!( # NoTrans for k in range(1, Ant) mzone = k == 1 ? real(beta) : one(real(T)) - if iscomplex - Dagger.@spawn BLAS.herk!( - uplo, + _alpha = iscomplex ? real(alpha) : alpha + Dagger.@spawn matmatmul!( + InOut(Cc[n, n]), + trans, + anti_transs, + In(Ac[n, k]), + In(Ac[n, k]), + _alpha, + mzone, + ) + end + # NoTrans / Upper + for m in range(n + 1, Cmt) + for k in range(1, Ant) + mzone = k == 1 ? beta : one(T) + Dagger.@spawn matmatmul!( + InOut(Cc[n, m]), trans, - real(alpha), + transs, In(Ac[n, k]), - mzone, - InOut(Cc[n, n]), - ) - else - Dagger.@spawn BLAS.syrk!( - uplo, - trans, + In(Ac[m, k]), alpha, - In(Ac[n, k]), mzone, - InOut(Cc[n, n]), ) end end - if uplo == 'L' - # NoTrans / Lower - for m in range(n + 1, Cmt) - for k in range(1, Ant) - mzone = k == 1 ? beta : one(T) - Dagger.@spawn BLAS.gemm!( - trans, - transs, - alpha, - In(Ac[m, k]), - In(Ac[n, k]), - mzone, - InOut(Cc[m, n]), - ) - end - end - else - # NoTrans / Upper - for m in range(n + 1, Cmt) - for k in range(1, Ant) - mzone = k == 1 ? beta : one(T) - Dagger.@spawn BLAS.gemm!( - trans, - transs, - alpha, - In(Ac[n, k]), - In(Ac[m, k]), - mzone, - InOut(Cc[n, m]), - ) - end - end - end else # [Conj]Trans for k in range(1, Amt) mzone = k == 1 ? real(beta) : one(real(T)) - if iscomplex - Dagger.@spawn BLAS.herk!( - uplo, + _alpha = iscomplex ? real(alpha) : alpha + _trans = iscomplex ? transs : trans + Dagger.@spawn matmatmul!( + InOut(Cc[n, n]), + _trans, + anti_transs, + In(Ac[k, n]), + In(Ac[k, n]), + _alpha, + mzone, + ) + end + # [Conj]Trans / Upper + for m in range(n + 1, Cmt) + for k in range(1, Amt) + mzone = k == 1 ? beta : one(T) + Dagger.@spawn matmatmul!( + InOut(Cc[n, m]), transs, - real(alpha), + 'N', In(Ac[k, n]), - mzone, - InOut(Cc[n, n]), - ) - else - Dagger.@spawn BLAS.syrk!( - uplo, - trans, + In(Ac[k, m]), alpha, - In(Ac[k, n]), mzone, - InOut(Cc[n, n]), ) end end - if uplo == 'L' - # [Conj]Trans / Lower - for m in range(n + 1, Cmt) - for k in range(1, Amt) - mzone = k == 1 ? beta : one(T) - Dagger.@spawn BLAS.gemm!( - transs, - 'N', - alpha, - In(Ac[k, m]), - In(Ac[k, n]), - mzone, - InOut(Cc[m, n]), - ) - end - end - else - # [Conj]Trans / Upper - for m in range(n + 1, Cmt) - for k in range(1, Amt) - mzone = k == 1 ? beta : one(T) - Dagger.@spawn BLAS.gemm!( - transs, - 'N', - alpha, - In(Ac[k, n]), - In(Ac[k, m]), - mzone, - InOut(Cc[n, m]), - ) - end - end - end end end end @@ -398,7 +418,7 @@ end return A end -@inline function copytile!(A::AbstractMatrix{T}, B::AbstractMatrix{T})::Nothing where {T} +function copytile!(A, B) m, n = size(A) C = B' @@ -408,15 +428,14 @@ end return nothing end -@inline function copydiagtile!(A::AbstractMatrix{T}, uplo::AbstractChar)::Nothing where {T} +function copydiagtile!(A, uplo) m, n = size(A) - Acpy = copy(A) if uplo == 'U' - C = UpperTriangular(Acpy)' + UpperTriangular(Acpy) + C = UpperTriangular(A)' + UpperTriangular(A) C[diagind(C)] .= A[diagind(A)] elseif uplo == 'L' - C = LowerTriangular(Acpy)' + Acpy - UpperTriangular(Acpy) + C = LowerTriangular(A)' + A - UpperTriangular(A) C[diagind(C)] .= A[diagind(A)] end @@ -425,6 +444,7 @@ end end return nothing end + function LinearAlgebra.generic_matvecmul!( C::DVector{T}, transA::Char, @@ -471,6 +491,21 @@ function _repartition_matvecmul(C, A, B, transA::Char)::Tuple{Blocks{1}, Blocks{ partC = (dimA,) return Blocks(partC...), Blocks(partA...), Blocks(partB...) end +""" + matvecmul!(C, transA::Char, A, B, alpha, beta) + +Tile-level matrix-vector multiply computing `C = alpha*op(A)*B + beta*C` in +place on the (dense) output vector `C`, where `op` is determined by `transA` +(`'N'`, `'T'`, `'C'`). Dispatches on the tile types: dense tiles use BLAS, while +sparse tiles (e.g. `DSparseArray`) provide their own method (in a package +extension) using a sparse matrix-vector product. This is the matvec analogue of +[`matmatmul!`](@ref). +""" +function matvecmul!(C, transA::Char, A, B, alpha, beta) + BLAS.gemv!(transA, alpha, A, B, beta, C) + return C +end + function gemv_dagger!( C::DVector{T}, transA::Char, @@ -513,26 +548,26 @@ function gemv_dagger!( # A: NoTrans for k in range(1, Ant) mzone = k == 1 ? beta : T(1.0) - Dagger.@spawn BLAS.gemv!( + Dagger.@spawn matvecmul!( + InOut(Cc[m]), transA, - alpha, In(Ac[m, k]), In(Bc[k]), + alpha, mzone, - InOut(Cc[m]), ) end else # A: [Conj]Trans — C's blocks index A's column-blocks for k in range(1, Amt) mzone = k == 1 ? beta : T(1.0) - Dagger.@spawn BLAS.gemv!( + Dagger.@spawn matvecmul!( + InOut(Cc[m]), transA, - alpha, In(Ac[k, m]), In(Bc[k]), + alpha, mzone, - InOut(Cc[m]), ) end end @@ -541,3 +576,53 @@ function gemv_dagger!( return C end + +wrap_as_darray(A::DArray) = A +wrap_as_darray(A::Array) = view(A, AutoBlocks()) +function wrap_as_darray(A::SubArray{T,Nv,Array{T,Na}}) where {T,Nv,Na} + Ap = parent(A) + part = auto_blocks(map(last, parentindices(Ap))) + partsize = part.blocksize + inds = parentindices(A) + inds_ranges_parent = ntuple(i->to_range(inds[i]), Val(Na)) + inds_ranges_view = ntuple(i->to_range(inds[i]), Val(Nv)) + subdomains = partition(part, ArrayDomain(inds_ranges_parent)) + nparts = size(subdomains) + chunks = Array{Any,Na}(undef, nparts...) + for idx in CartesianIndices(nparts) + subdomain_view = subdomains[idx] + subdomain_parent = ArrayDomain(ntuple(i->Nv >= i ? subdomain_view.indexes[i] : inds_ranges_parent[i], Val(Na))) + subinds = ntuple(i->subdomain_parent.indexes[i], Val(Na)) + subA = view(Ap, subinds...) + chunks[idx] = tochunk(subA) + end + return DArray(T, ArrayDomain(inds_ranges_parent), subdomains, chunks, part) +end + +# Generate generic_matvecmul! methods for all combinations of DArray, Array, and SubArray +for CT in (DVector, Vector, SubArray{<:Any,1,<:Array}), + AT in (DMatrix, Matrix, SubArray{<:Any,2,<:Array}), + BT in (DVector, Vector, SubArray{<:Any,1,<:Array}) + + # Don't commit type piracy + CT isa DArray || AT isa DArray || BT isa DArray || continue + + @eval function LinearAlgebra.generic_matvecmul!( + C::$(CT), + transA::Char, + A::$(AT), + B::$(BT), + _add::LinearAlgebra.MulAddMul, + ) + new_C = wrap_as_darray(C) + new_A = wrap_as_darray(A) + new_B = wrap_as_darray(B) + return LinearAlgebra.generic_matvecmul!( + new_C, + transA, + new_A, + new_B, + _add, + ) + end +end \ No newline at end of file diff --git a/src/array/sparse.jl b/src/array/sparse.jl new file mode 100644 index 000000000..c5ab69b91 --- /dev/null +++ b/src/array/sparse.jl @@ -0,0 +1,259 @@ +""" + DSparseArray{T,N} <: AbstractArray{T,N} + +A sparse array container, for which the contained array may be replaced with a +new one to support in-place operations. Designed to work well with Datadeps +algorithms: writes that reallocate (and grow/shrink) the inner sparse storage +are hidden behind the wrapper's stable identity, so Datadeps aliasing tracking +remains valid (see `aliases_as_whole`). + +`DSparseVector{T}` and `DSparseMatrix{T}` are aliases for the 1- and 2-dimensional +cases. The wrapper is general over `N` so it can hold sparse vectors, matrices, +and (eventually) higher-order sparse tensors (e.g. Finch tensors). +""" +mutable struct DSparseArray{T,N} <: AbstractArray{T,N} + mat +end +const DSparseVector{T} = DSparseArray{T,1} +const DSparseMatrix{T} = DSparseArray{T,2} + +# Convenience constructor inferring element type and dimensionality from `x`. +DSparseArray(x) = DSparseArray{eltype(x),ndims(x)}(x) + +_sparse_collect(M) = collect(M) +# Allocate a destination tile matching the inner storage's sparse backend. The +# default delegates to the storage type's own `similar`, which carries the +# concrete backend (e.g. `SparseMatrixCSC`) forward to the destination -- this is +# how `similar(::DArray)` propagates a tile's sparse type to a newly-allocated +# result DArray. Backends whose `similar` is unsuitable (e.g. Finch, whose generic +# `similar` yields tensor formats that destabilize later kernels) override this. +_sparse_similar(mat, ::Type{T}, dims::Dims) where {T} = similar(mat, T, dims) + +Base.eltype(M::DSparseArray{T}) where T = T +Base.size(M::DSparseArray) = size(M.mat) +Base.iterate(M::DSparseArray) = iterate(M.mat) +Base.iterate(M::DSparseArray, state) = iterate(M.mat, state) +Base.similar(M::DSparseArray, ::Type{T}, dims::Dims) where T = + DSparseArray{T,length(dims)}(_sparse_similar(M.mat, T, dims)) +Base.collect(M::DSparseArray) = _sparse_collect(M.mat) + +# N.B. hash and aliasing shouldn't change even if M.mat changes. +# This is the key property that makes `DSparseArray` safe for Datadeps: +# the aliasing identity is tied to the (stable) wrapper object, not to the +# inner storage, which may be reallocated (and grow/shrink) on writes. +Base.hash(M::DSparseArray, h::UInt) = hash(objectid(M), hash(DSparseArray, h)) + +# Sparse containers must alias as a single, indivisible unit. Their storage is +# reallocated (and may grow/shrink) on writes, so it is unsafe to alias any +# *part* of the container independently. We therefore force *all* access to a +# `DSparseArray` to resolve to the container's stable whole-object +# `ObjectAliasing` -- including access via views, transposes, adjoints, and +# reshapes -- rather than e.g. a `StridedAliasing` over storage that may have +# since moved. `aliases_as_whole` opts the type into this behavior, and Datadeps' +# `aliasing_root` resolves any wrapper of a `DSparseArray` back to the container +# before computing aliasing. Note: this must return a *bare* `ObjectAliasing`, as +# `aliases_as_whole(::ObjectAliasing)` relies on it to drive whole-object copies. +# +# The ObjectAliasing is stamped with the inner storage's memory space so GPU- +# resident sparse tiles (e.g. CuSparseMatrixCSC / ROCSparseMatrixCSC, or +# DeviceSparseMatrixCSC) are tracked in VRAM rather than host RAM. The pointer +# itself is still the host wrapper's `pointer_from_objref` (identity only); +# Datadeps never span-copies these objects (see `aliases_as_whole` / `FullCopy`). +function aliasing(M::DSparseArray, _=identity) + space = value_memory_space(M.mat) + ptr = RemotePtr{Cvoid}(UInt(pointer_from_objref(M)), space) + return ObjectAliasing(ptr, sizeof(typeof(M))) +end +aliases_as_whole(::DSparseArray) = true +# Chunk space follows the inner sparse storage (host CSC → CPURAM, GPU sparse → VRAM). +# `memory_space` (not only `value_memory_space`) must resolve correctly: Datadeps' +# `aliased_object!` uses `memory_space(x)` when deciding whether a same-space +# slot can share the original object. The generic fallback labels unknown values +# as CPURAM, which would force a discarded GPU copy and drop SpGEMM writes. +value_memory_space(M::DSparseArray) = value_memory_space(M.mat) +memory_space(M::DSparseArray) = value_memory_space(M) + +# A `DSparseArray` has no meaningful raw data pointer (its storage may be +# reallocated/resized). This trap ensures that if aliasing ever tries to treat a +# wrapper of a `DSparseArray` as strided memory -- e.g. a new array-wrapper type +# that `aliasing_root` failed to resolve via `Base.parent` -- it fails loudly +# instead of silently corrupting Datadeps' aliasing tracking. +function Base.pointer(::DSparseArray) + throw(ArgumentError("`pointer(::DSparseArray)` is intentionally unsupported: \ + a DSparseArray may reallocate its storage, so it must be aliased as a \ + whole object via `Dagger.aliasing`. If you reached here through Datadeps \ + aliasing of an array wrapper, ensure that wrapper implements `Base.parent` \ + so that `Dagger.aliasing_root` can resolve it to the DSparseArray.")) +end + +# Forward indexing to the inner array. Ranges/colons are supported so that +# views (used by copy-buffering between mismatched partitionings) work. +Base.getindex(M::DSparseArray, I...) = getindex(M.mat, I...) +Base.setindex!(M::DSparseArray, v, I...) = setindex!(M.mat, v, I...) + +# Whole-tile copy. This is how Datadeps moves a tile between workers via +# `move!`: we *reallocate* the inner storage rather than mutating in place, +# since sparse storage cannot generally be updated in place. The wrapper's +# identity (and thus its aliasing) is preserved. +function Base.copyto!(dst::DSparseArray, src::DSparseArray) + dst.mat = _sparse_copy(src.mat) + return dst +end +# Backend-overridable deep copy of the inner storage (Finch tensors don't define +# `Base.copy`). +_sparse_copy(mat) = copy(mat) + +# Adapt preserves the `DSparseArray` wrapper while adapting the inner storage +# (e.g. SparseMatrixCSC → CuSparseMatrixCSC under `adapt(CuArray, …)`). GPU +# package extensions override `move` for sparse types so OpenCL/Metal/oneAPI do +# not densify via the generic `adapt(DeviceArray, SparseMatrixCSC)` path. +Adapt.adapt_structure(to, M::DSparseArray) = DSparseArray(Adapt.adapt(to, M.mat)) + +# Whole-object in-place move for Datadeps `FullCopy`. Reassigns `to.mat` from a +# space-converted copy of `from.mat` rather than span-copying storage. +function move!(to_space::MemorySpace, from_space::MemorySpace, to::DSparseArray, from::DSparseArray) + if to_space == from_space + to.mat = _sparse_copy(from.mat) + else + to_proc = first(processors(to_space)) + from_proc = first(processors(from_space)) + moved = move(from_proc, to_proc, from) + to.mat = moved.mat + end + return +end + +# Wrapping hook used when materializing tiles (e.g. in `distribute`). Backends +# overload this (e.g. in package extensions) to wrap freshly-created tiles in a +# container that Datadeps can track (such as `DSparseArray` for sparse tiles); +# everything else is left untouched. +maybe_wrap_tile(x) = x + +# Partial-range tile copy used by copy-buffering (e.g. when matmul operands need +# to be repartitioned). Some backends (Finch) forbid `setindex!`, so we route +# through a backend hook that returns the (possibly reallocated) inner storage. +# `DSparseArray` hides the reallocation from Datadeps. +function copyto_view!(Bpart::DSparseArray, Brange, Apart, Arange) + Bpart.mat = _sparse_copyto_view!(Bpart.mat, Brange, view(Apart, Arange)) + return +end +# Default: in-place element-wise copy (works for `setindex!`-capable backends +# such as `SparseArrays`). May reallocate and return new storage for backends +# that can't update in place, so callers must use the returned value. +function _sparse_copyto_view!(mat, Brange, src) + copyto!(view(mat, Brange), src) + return mat +end + +#Base.BroadcastStyle(::Type{<:DSparseArray}) = FinchStyle() + +LinearAlgebra.norm2(M::DSparseArray) = LinearAlgebra.norm2(M.mat) + +# Matrix-specific operations dispatch on the 2-D alias `DSparseMatrix`. +function matmatmul!( + C::DSparseMatrix, + transA::Char, + transB::Char, + A::DSparseMatrix, + B::DSparseMatrix, + alpha, + beta +) + # Forward to sparse matmatmul!, but allow in-place update of C + return matmatmul!(C, transA, transB, A.mat, B.mat, alpha, beta) +end + +# Sparse matrix-vector multiply: `C = alpha*op(A)*B + beta*C` with a sparse +# matrix tile `A` and dense vectors `B`/`C`. Unwrap to the inner storage so each +# backend can provide an efficient (in-place on dense `C`) method. +function matvecmul!(C, transA::Char, A::DSparseMatrix, B, alpha, beta) + return matvecmul!(C, transA, A.mat, B, alpha, beta) +end + +# Factorize via host CSC so GPU sparse tiles (rocSPARSE / DeviceSparseMatrixCSC) +# can still feed block-Jacobi; CPU SparseMatrixCSC is a no-op collect. +_factorize_tile(M::DSparseArray) = LinearAlgebra.lu(_sparse_collect(M.mat)) + +# Expose the inner sparse storage so preconditioner backends (ILU, AMG) build on +# the actual `SparseMatrixCSC` rather than the `DSparseArray` wrapper. GPU mats +# are gathered to host CSC (those backends are host-only). +_tile_matrix(M::DSparseArray) = _sparse_collect(M.mat) + +# Jacobi diagonal extract without GPU scalar indexing. `diag(::SparseMatrixCSC)` +# returns a `SparseVector`; densify before `copyto!` into a device vector. +function _set_inv_diag!(out, tile::DSparseArray) + d = Vector(LinearAlgebra.diag(_sparse_collect(tile.mat))) + copyto!(out, inv.(d)) + return nothing +end + +function transpose_tile end +function copytile!(A::DSparseMatrix, B::DSparseMatrix) + A.mat = transpose_tile(B.mat) + return +end +function copydiagtile!(A::DSparseMatrix, uplo) + A.mat = transpose_tile(A.mat, uplo) + return +end + +#============================================================================== + Device-resident CSC for GPU backends without a vendor sparse library + (OpenCL / Metal / oneAPI). Storage vectors live in device memory; SpGEMM/SpMV + kernels fall back to host SparseArrays (see SparseArraysExt). CUDA/ROCm use + their native CuSparse/ROCSparse types instead. +==============================================================================# + +""" + DeviceSparseMatrixCSC{Tv,Ti,V,I} + +A CSC sparse matrix whose `colptr`/`rowval`/`nzval` live in device arrays +(`CLArray`, `MtlArray`, `oneArray`, …). Used as the inner storage of a +[`DSparseArray`](@ref) on GPU backends that lack a vendor sparse library. +""" +struct DeviceSparseMatrixCSC{Tv,Ti,V<:AbstractVector{Tv},I<:AbstractVector{Ti}} <: AbstractMatrix{Tv} + m::Int + n::Int + colptr::I + rowval::I + nzval::V +end +Base.size(A::DeviceSparseMatrixCSC) = (A.m, A.n) +Base.eltype(::DeviceSparseMatrixCSC{Tv}) where Tv = Tv +Base.IndexStyle(::Type{<:DeviceSparseMatrixCSC}) = IndexCartesian() +function Base.getindex(A::DeviceSparseMatrixCSC{Tv}, i::Integer, j::Integer) where Tv + @boundscheck checkbounds(A, i, j) + # Scalar device indexing; acceptable for rare host-side inspection. + col_start = Int(A.colptr[j]) + col_end = Int(A.colptr[j + 1]) - 1 + for p in col_start:col_end + Int(A.rowval[p]) == i && return A.nzval[p] + end + return zero(Tv) +end +# Memory space follows the nonzero values buffer. +value_memory_space(A::DeviceSparseMatrixCSC) = value_memory_space(A.nzval) +memory_space(A::DeviceSparseMatrixCSC) = value_memory_space(A) +# Whole-object aliasing (storage may be replaced on write); stamp with device space. +function aliasing(A::DeviceSparseMatrixCSC, _=identity) + space = value_memory_space(A) + ptr = RemotePtr{Cvoid}(UInt(pointer_from_objref(A)), space) + return ObjectAliasing(ptr, sizeof(typeof(A))) +end +aliases_as_whole(::DeviceSparseMatrixCSC) = true +function Base.pointer(::DeviceSparseMatrixCSC) + throw(ArgumentError("`pointer(::DeviceSparseMatrixCSC)` is unsupported; \ + alias as a whole object via `Dagger.aliasing`.")) +end + +# Allocation hooks used by `spzeros`/`sprand` under a GPU compute scope. +# GPU Exts override these for their processor type. +allocate_sparse_zeros(proc::Processor, ::Type{T}, dims::Dims) where T = + allocate_sparse_zeros_default(proc, T, dims) +allocate_sparse_rand(proc::Processor, ::Type{T}, dims::Dims, sparsity::AbstractFloat) where T = + allocate_sparse_rand_default(proc, T, dims, sparsity) +# Defaults are filled in by SparseArraysExt (host SparseMatrixCSC / SparseVector). +function allocate_sparse_zeros_default end +function allocate_sparse_rand_default end +# Upload host CSC onto device vectors of type `Arr` (filled in by SparseArraysExt). +function device_sparse_from_host end \ No newline at end of file diff --git a/src/array/sparse_partition.jl b/src/array/sparse_partition.jl deleted file mode 100644 index 26efe3e6a..000000000 --- a/src/array/sparse_partition.jl +++ /dev/null @@ -1,37 +0,0 @@ -export partition_sparse - -function partition_sparse(colptr, nnz, sz, nparts) - nnz_per_chunk = nnz/nparts - dom = ArrayDomain(map(x->1:x, sz)) - - nxt = nnz_per_chunk - colstart = 1 - splits = UnitRange[] - - for i=1:nparts-1 - idxs = searchsorted(colptr, nxt) - if length(idxs) == 0 - # the exact index was not found. latch on to the nearest column boundary - idx = first(idxs) - if abs(colptr[idx]-nxt) < abs(colptr[idx-1]-nxt) - nextidx = idx - else - nextidx = idx-1 - end - else - nextidx = last(idxs) - end - push!(splits, colstart:nextidx) - colstart=nextidx+1 - nxt=nxt + nnz_per_chunk - end - push!(splits, colstart:(length(colptr)-1)) - - - cumlength = cumsum(map(length, splits)) - subdomains = DomainBlocks((1,1), ([size(dom, 1)], cumlength)) - DomainSplit(dom, subdomains) -end - -partition_sparse(S::SparseMatrixCSC, nparts) = - partition_sparse(S.colptr, length(S.nzval), size(S), nparts) diff --git a/src/array/sparsedirect.jl b/src/array/sparsedirect.jl new file mode 100644 index 000000000..508b24ac4 --- /dev/null +++ b/src/array/sparsedirect.jl @@ -0,0 +1,518 @@ +# Sparse direct solvers. +# +# Pure-Julia sparse LU backends (KLU via `PureKLU.jl`, UMFPACK-style via +# `PureUMFPACK.jl`) provide *direct* solves for sparse systems. Unlike the +# C-bound `SparseArrays.UmfpackLU` (whose factorization holds process-local +# UMFPACK resources and can't be moved between workers), these factorizations are +# plain Julia data: serializable and movable, so Dagger can schedule them freely. +# +# Two entry points are exposed (implemented in `ext/PureKLUExt.jl` / +# `ext/PureUMFPACKExt.jl`): +# +# - A whole-matrix direct solve: `Dagger.klu(A)` / `Dagger.splu(A)` gather the +# (sparse) `DMatrix` onto a single worker (the one owning the most tiles), +# assemble one `SparseMatrixCSC`, factor it once, and return a +# `DaggerSparseLU` whose factor is pinned there. Solves move only the RHS / +# solution vectors. This is a correctness-first, single-worker factorization +# for systems that fit on one worker (e.g. coarse-grid solves). +# - An opt-in distributed solve path: `Dagger.splu(A; distributed=true)` +# (PureUMFPACK only) reuses that factorization, extracts the triangular +# factors into sparse `DMatrix`es, and solves with a tiled datadeps +# forward/backward substitution (`DistributedSparseLU`, `method=:trsv`). +# - A Schur-complement domain-decomposition path: +# `Dagger.splu(A; distributed=true, method=:schur)` factors interior blocks +# in parallel across workers and solves via a separator Schur system +# (`DistributedSchurLU`; requires Metis.jl). +# - Per-tile block direct solves (`BlockKLUPreconditioner` / +# `BlockUMFPACKPreconditioner`, in `iterativesolvers.jl`), which reuse the +# block-preconditioner machinery. + +""" + klu(A::DMatrix; kwargs...) -> DaggerSparseLU + +Direct sparse LU factorization of a sparse `DMatrix` using KLU. Tiles are +gathered onto one worker, assembled into a single `SparseMatrixCSC`, and factored +once; the factor stays pinned there. Solve with `F \\ b` or `ldiv!(x, F, b)` over +`DVector`s. Requires `PureKLU.jl` to be loaded. + +KLU is well suited to unsymmetric systems with near-triangular structure (e.g. +circuit simulation). See also [`splu`](@ref). +""" +function klu end + +""" + splu(A::DMatrix; distributed=false, method=:trsv, blocksize=nothing, nparts=nothing, kwargs...) + +Direct sparse LU factorization of a sparse `DMatrix` using a pure-Julia +UMFPACK-style multifrontal solver. + +- `distributed=false` → [`DaggerSparseLU`](@ref): gather+factor on one worker. +- `distributed=true, method=:trsv` (default) → [`DistributedSparseLU`](@ref): + Stage-4a factor plus tiled datadeps triangular solves (`blocksize` controls + the solve tiling). +- `distributed=true, method=:schur` → [`DistributedSchurLU`](@ref): single-level + Schur-complement domain decomposition (METIS separator; `nparts` defaults to + `max(2, nworkers())`). Requires `Metis.jl`. + +Requires `PureUMFPACK.jl`. See also [`klu`](@ref). +""" +function splu end + +_klu_required() = throw(ArgumentError( + "Dagger.klu requires PureKLU.jl. Run `using PureKLU` to enable sparse direct \ + (KLU) solves.")) +_splu_required() = throw(ArgumentError( + "Dagger.splu requires PureUMFPACK.jl. Run `using PureUMFPACK` to enable sparse \ + direct (UMFPACK) solves.")) +klu(A; kwargs...) = _klu_required() +splu(A; kwargs...) = _splu_required() + +""" + DaggerSparseLU + +A direct sparse factorization of a sparse `DMatrix`, produced by +[`Dagger.klu`](@ref) or [`Dagger.splu`](@ref). The underlying factorization is +pure-Julia and pinned to a single worker (`fact` is a `Chunk`/`DTask` with +`scope`). Solve `A x = b` with `F \\ b` (returns a `DVector` partitioned like +`b`) or `ldiv!(x, F, b)`. +""" +struct DaggerSparseLU{F,S,P} + fact::F # pinned factorization (`Chunk` / `DTask`) + scope::S # `ProcessScope` of the factor's worker + n::Int # system size (square) + part::P # default vector partitioning (column blocks of `A`) +end + +Base.size(F::DaggerSparseLU) = (F.n, F.n) +Base.size(F::DaggerSparseLU, i::Integer) = i <= 2 ? F.n : 1 + +""" + DistributedSparseLU + +A sparse LU factorization with triangular factors stored as sparse `DMatrix`es, +produced by `Dagger.splu(A; distributed=true)`. Solves use a tiled datadeps +forward/backward substitution over `L` and `U` (PureUMFPACK only). +""" +struct DistributedSparseLU{L,U,P,Q,R,Part} + L::L # unit-lower triangular factor (`DMatrix` of sparse tiles) + U::U # upper triangular factor (`DMatrix` of sparse tiles) + p::P # row permutation + q::Q # column permutation + Rs::R # row scaling + n::Int # system size (square) + part::Part # vector partitioning matching `L`/`U` tile rows +end + +Base.size(F::DistributedSparseLU) = (F.n, F.n) +Base.size(F::DistributedSparseLU, i::Integer) = i <= 2 ? F.n : 1 + +# Assemble local tiles into one `SparseMatrixCSC` (no densification). The real +# method is added by `SparseArraysExt`; without `SparseArrays` there are no +# sparse tiles to gather. +function _gather_sparse end +_gather_sparse(args...) = throw(ArgumentError( + "gathering a sparse DMatrix requires SparseArrays to be loaded")) + +# Backend-specific extraction of `(L, U, p, q, Rs)` from a factored object. +# Implemented for `PureUMFPACK.PureLU` in `PureUMFPACKExt`. +function _extract_lu_factors end +_extract_lu_factors(F) = throw(ArgumentError( + "distributed sparse LU requires a PureUMFPACK.PureLU factor (got $(typeof(F))). \ + Use `Dagger.splu(A; distributed=true)` with PureUMFPACK.jl loaded.")) + +# Choose the worker that already owns the most tiles of `A` (minimizes gather +# traffic). Ties break toward the lowest worker id. +function _select_factor_scope(A::DMatrix) + counts = Dict{Int,Int}() + for c in A.chunks + wid = _tile_scope(c).wid + counts[wid] = get(counts, wid, 0) + 1 + end + best_wid, best_count = 0, -1 + for (wid, cnt) in counts + if cnt > best_count || (cnt == best_count && wid < best_wid) + best_wid, best_count = wid, cnt + end + end + return ProcessScope(best_wid) +end + +# Runs on the chosen worker: assemble from tiles (already moved here by the +# scheduler), factorize, and pin the factor to this process. +function _assemble_and_factor_pinned(factorize, ::Type{T}, row_offsets, col_offsets, + m, n, tiles...) where T + S = _gather_sparse(T, tiles, row_offsets, col_offsets, m, n) + F = factorize(S) + proc = task_processor() + return tochunk(F, proc, ProcessScope(root_worker_id(proc))) +end + +# Shared path for `klu` / `splu`: pick a worker, spawn assemble+factor there. +function _spawn_direct_factorization(A::DMatrix{T}, factorize) where T + n = LinearAlgebra.checksquare(A) + scope = _select_factor_scope(A) + Ac = A.chunks + mt, nt = size(Ac) + ntiles = mt * nt + row_offsets = Vector{Int}(undef, ntiles) + col_offsets = Vector{Int}(undef, ntiles) + tiles = Vector{Any}(undef, ntiles) + idx = 1 + for i in 1:mt, j in 1:nt + dom = A.subdomains[i, j] + row_offsets[idx] = first(dom.indexes[1]) - 1 + col_offsets[idx] = first(dom.indexes[2]) - 1 + tiles[idx] = Ac[i, j] + idx += 1 + end + mA, nA = size(A) + part = Blocks(A.partitioning.blocksize[2]) + fact = spawn(_assemble_and_factor_pinned, Options(; compute_scope=scope), + factorize, T, row_offsets, col_offsets, mA, nA, tiles...) + return DaggerSparseLU(fact, scope, n, part) +end + +# Gather RHS chunks, solve against the pinned factor, return the dense solution. +function _direct_solve(fact, bparts...) + b = length(bparts) == 1 ? bparts[1] : reduce(vcat, bparts) + return fact \ b +end + +function Base.:\(F::DaggerSparseLU, b::DVector) + length(b) == F.n || throw(DimensionMismatch( + "factorization is $(F.n)×$(F.n) but b has length $(length(b))")) + # Solve on the factor's worker (factor stays pinned); only the O(n) solution + # returns here to be redistributed like `b`. + x = fetch(spawn(_direct_solve, Options(; compute_scope=F.scope), + F.fact, b.chunks...)) + return distribute(x, b.partitioning) +end + +function LinearAlgebra.ldiv!(x::DVector, F::DaggerSparseLU, b::DVector) + return copyto!(x, F \ b) +end + +# ---- Stage 4b: distributed triangular solves -------------------------------- + +_fetch_lu_factors(fact) = _extract_lu_factors(fact) + +function _resolve_solve_blocksize(A::DMatrix, blocksize) + if blocksize === nothing + return Int(min(A.partitioning.blocksize...)) + end + bs = Int(blocksize) + bs > 0 || throw(ArgumentError("blocksize must be positive, got $bs")) + return bs +end + +# Re-tile PureLU factors into sparse `DMatrix`es after a Stage-4a factorization. +function _make_distributed_sparse_lu(F::DaggerSparseLU, A::DMatrix; blocksize=nothing) + bs = _resolve_solve_blocksize(A, blocksize) + L, U, p, q, Rs = fetch(spawn(_fetch_lu_factors, Options(; compute_scope=F.scope), F.fact)) + part2 = Blocks(bs, bs) + part1 = Blocks(bs) + return DistributedSparseLU(distribute(L, part2), distribute(U, part2), + p, q, Rs, F.n, part1) +end + +# Tile kernels: unwrap `DSparseArray` → `SparseMatrixCSC`, never densify. +function _sparse_unit_ltrsv!(Ltile, x::AbstractVector) + LinearAlgebra.ldiv!(UnitLowerTriangular(_tile_matrix(Ltile)), x) + return +end + +function _sparse_utrsv!(Utile, x::AbstractVector) + LinearAlgebra.ldiv!(UpperTriangular(_tile_matrix(Utile)), x) + return +end + +function _sparse_gemv_sub!(Atile, x::AbstractVector, y::AbstractVector) + matvecmul!(y, 'N', Atile, x, -one(eltype(y)), one(eltype(y))) + return +end + +# Blocked forward substitution for unit-lower `L` (mirrors `trsv!` lower/'N'). +function _sparse_trsv_forward!(L::DMatrix, c::DVector) + Lc, Cc = L.chunks, c.chunks + nt = length(Cc) + size(Lc, 1) == nt && size(Lc, 2) == nt || throw(DimensionMismatch( + "L tile grid $(size(Lc)) incompatible with vector tiles $nt")) + Dagger.spawn_datadeps() do + for k in 1:nt + Dagger.@spawn _sparse_unit_ltrsv!(In(Lc[k, k]), InOut(Cc[k])) + for i in (k + 1):nt + Dagger.@spawn _sparse_gemv_sub!(In(Lc[i, k]), In(Cc[k]), InOut(Cc[i])) + end + end + end + return c +end + +# Blocked backward substitution for upper `U` (mirrors `trsv!` upper/'N'). +function _sparse_trsv_backward!(U::DMatrix, c::DVector) + Uc, Cc = U.chunks, c.chunks + nt = length(Cc) + size(Uc, 1) == nt && size(Uc, 2) == nt || throw(DimensionMismatch( + "U tile grid $(size(Uc)) incompatible with vector tiles $nt")) + Dagger.spawn_datadeps() do + for k in reverse(1:nt) + Dagger.@spawn _sparse_utrsv!(In(Uc[k, k]), InOut(Cc[k])) + for i in 1:(k - 1) + Dagger.@spawn _sparse_gemv_sub!(In(Uc[i, k]), In(Cc[k]), InOut(Cc[i])) + end + end + end + return c +end + +function Base.:\(F::DistributedSparseLU, b::DVector) + length(b) == F.n || throw(DimensionMismatch( + "factorization is $(F.n)×$(F.n) but b has length $(length(b))")) + # Permute + row-scale (driver O(n)); matches PureUMFPACK `_solve_factor`. + b_local = collect(b) + T = promote_type(eltype(b_local), eltype(F.Rs), eltype(F.U)) + c_local = Vector{T}(undef, F.n) + p, q, Rs = F.p, F.q, F.Rs + @inbounds for k in eachindex(p) + c_local[k] = Rs[p[k]] * b_local[p[k]] + end + c = distribute(c_local, F.part) + _sparse_trsv_forward!(F.L, c) + _sparse_trsv_backward!(F.U, c) + c_local = collect(c) + x_local = similar(c_local) + @inbounds for k in eachindex(q) + x_local[q[k]] = c_local[k] + end + return distribute(x_local, b.partitioning) +end + +function LinearAlgebra.ldiv!(x::DVector, F::DistributedSparseLU, b::DVector) + return copyto!(x, F \ b) +end + +# ---- Stage 4c: Schur-complement domain decomposition ------------------------- +# +# Single-level (non-recursive) k-way Schur factorization. Recursive / nested +# dissection of the separator is future work. + +""" + DistributedSchurLU + +Sparse LU via single-level Schur-complement domain decomposition, produced by +`Dagger.splu(A; distributed=true, method=:schur)`. Interior blocks are factored +in parallel (pinned round-robin across workers); the separator Schur system is +factored on one worker. Solve with `F \\ b` / `ldiv!(x, F, b)`. +""" +struct DistributedSchurLU{F,W,G,SF,S,P} + interior_facts::Vector{F} # pinned `Aᵢᵢ` factors + W::Vector{W} # pinned `Wᵢ = Aᵢᵢ⁻¹ AᵢΓ` (dense columns) + AΓi::Vector{G} # pinned `A_Γi` blocks + interior_scopes::Vector{S} + Sfact::SF # pinned Schur factor (`nothing` if `|Γ|==0`) + Sscope::S + interiors::Vector{Vector{Int}} + separator::Vector{Int} + perm::Vector{Int} + n::Int + part::P +end + +Base.size(F::DistributedSchurLU) = (F.n, F.n) +Base.size(F::DistributedSchurLU, i::Integer) = i <= 2 ? F.n : 1 + +# METIS partition + separator (implemented in `MetisExt`). +function _nested_dissection_partition end +_nested_dissection_partition(A, nparts) = throw(ArgumentError( + "Dagger.splu(...; method=:schur) requires Metis.jl. Run `using Metis` to \ + enable Schur-complement distributed sparse LU.")) + +# Gather a sparse `DMatrix` into one `SparseMatrixCSC` on the caller (pattern +# assembly for the partitioner; no densification). +function _collect_sparse_dmatrix(A::DMatrix{T}) where T + Ac = A.chunks + mt, nt = size(Ac) + ntiles = mt * nt + row_offsets = Vector{Int}(undef, ntiles) + col_offsets = Vector{Int}(undef, ntiles) + tiles = Vector{Any}(undef, ntiles) + idx = 1 + for i in 1:mt, j in 1:nt + dom = A.subdomains[i, j] + row_offsets[idx] = first(dom.indexes[1]) - 1 + col_offsets[idx] = first(dom.indexes[2]) - 1 + tiles[idx] = fetch(Ac[i, j]) + idx += 1 + end + mA, nA = size(A) + return _gather_sparse(T, tiles, row_offsets, col_offsets, mA, nA) +end + +_schur_worker_ids() = (w = workers(); isempty(w) ? Int[myid()] : Int[w...]) + +# Factor `Aᵢᵢ`, form dense `Wᵢ = Aᵢᵢ⁻¹ AᵢΓ` (column solves; densifying the thin +# RHS is simpler than sparse multi-RHS), and local Schur contribution `A_Γi Wᵢ`. +# Pins `(fact, W, A_Γi)` to the current worker. +function _schur_factor_interior_pinned(factorize, Aii, AiΓ, AΓi) + F = factorize(Aii) + ni, ng = size(AiΓ) + T = eltype(AiΓ) + W = zeros(T, ni, ng) + for j in 1:ng + W[:, j] = F \ Vector{T}(AiΓ[:, j]) + end + Si = AΓi * W + proc = task_processor() + scope = ProcessScope(root_worker_id(proc)) + return (tochunk(F, proc, scope), tochunk(W, proc, scope), + tochunk(AΓi, proc, scope), Si) +end + +function _schur_factor_interface_pinned(factorize, S) + F = factorize(S) + proc = task_processor() + scope = ProcessScope(root_worker_id(proc)) + return tochunk(F, proc, scope) +end + +# Forward: `y = Aᵢᵢ⁻¹ bᵢ`, `contrib = A_Γi y`. Pin `y` for the back-substitution. +function _schur_forward_pinned(fact, AΓi, bI) + y = fact \ bI + contrib = AΓi * y + proc = task_processor() + scope = ProcessScope(root_worker_id(proc)) + return tochunk(y, proc, scope), contrib +end + +_schur_back_pinned(y, W, xΓ) = y - W * xΓ + +_schur_interface_solve(Sfact, g) = Sfact \ g + +function _make_schur_sparse_lu(A::DMatrix, factorize; nparts=nothing) + n = LinearAlgebra.checksquare(A) + np = nparts === nothing ? max(2, nworkers()) : Int(nparts) + np >= 1 || throw(ArgumentError("nparts must be ≥ 1, got $np")) + + Asp = _collect_sparse_dmatrix(A) + nd = _nested_dissection_partition(Asp, np) + interiors = nd.interiors + separator = nd.separator + perm = nd.perm + + # Trivial partition → Stage-4a serial factorization. + if length(interiors) <= 1 && isempty(separator) + return _spawn_direct_factorization(A, factorize) + end + + wids = _schur_worker_ids() + k = length(interiors) + interior_facts = Vector{Any}(undef, k) + Ws = Vector{Any}(undef, k) + AΓis = Vector{Any}(undef, k) + scopes = Vector{ProcessScope}(undef, k) + Si_tasks = Vector{Any}(undef, k) + + Γ = separator + ng = length(Γ) + for i in 1:k + Ii = interiors[i] + scope = ProcessScope(wids[mod1(i, length(wids))]) + scopes[i] = scope + Aii = Asp[Ii, Ii] + AiΓ = Asp[Ii, Γ] + AΓi = Asp[Γ, Ii] + task = spawn(_schur_factor_interior_pinned, Options(; compute_scope=scope), + factorize, Aii, AiΓ, AΓi) + Si_tasks[i] = task + end + + # Fetch pinned chunks + local Schur contributions; reduce S on the driver. + T = eltype(Asp) + S = isempty(Γ) ? zeros(T, 0, 0) : Matrix{T}(Asp[Γ, Γ]) + for i in 1:k + Fi, Wi, AΓi_c, Si = fetch(Si_tasks[i]) + interior_facts[i] = Fi + Ws[i] = Wi + AΓis[i] = AΓi_c + isempty(Γ) || (S .-= Si) + end + + if isempty(Γ) + Sfact = nothing + Sscope = scopes[1] + else + Sscope = ProcessScope(wids[1]) + # Schur fill-in is expected; hand a sparse view to the backend factorizer. + Ssp = _sparse_copy_of(S) + Sfact = fetch(spawn(_schur_factor_interface_pinned, + Options(; compute_scope=Sscope), factorize, Ssp)) + end + + part = Blocks(A.partitioning.blocksize[2]) + return DistributedSchurLU(interior_facts, Ws, AΓis, scopes, Sfact, Sscope, + interiors, separator, perm, n, part) +end + +# Convert a dense Schur matrix to the sparse type the backend expects. Defined +# as a hook so the core package need not depend on SparseArrays; the real method +# lives in `SparseArraysExt`. +function _sparse_copy_of end +_sparse_copy_of(S) = throw(ArgumentError( + "Schur-complement sparse LU requires SparseArrays to be loaded")) + +function Base.:\(F::DistributedSchurLU, b::DVector) + length(b) == F.n || throw(DimensionMismatch( + "factorization is $(F.n)×$(F.n) but b has length $(length(b))")) + b_local = collect(b) + bp = b_local[F.perm] + k = length(F.interiors) + ng = length(F.separator) + + # Split permuted RHS into interior segments + separator. + ys = Vector{Any}(undef, k) + contrib_tasks = Vector{Any}(undef, k) + off = 0 + for i in 1:k + ni = length(F.interiors[i]) + bI = bp[off+1:off+ni] + off += ni + contrib_tasks[i] = spawn(_schur_forward_pinned, + Options(; compute_scope=F.interior_scopes[i]), + F.interior_facts[i], F.AΓi[i], bI) + end + + g = ng == 0 ? similar(bp, 0) : copy(bp[off+1:end]) + for i in 1:k + yi, contrib = fetch(contrib_tasks[i]) + ys[i] = yi + ng == 0 || (g .-= contrib) + end + + xΓ = if ng == 0 + similar(g) + else + fetch(spawn(_schur_interface_solve, Options(; compute_scope=F.Sscope), + F.Sfact, g)) + end + + x_parts = Vector{Any}(undef, k) + back_tasks = Vector{Any}(undef, k) + for i in 1:k + back_tasks[i] = spawn(_schur_back_pinned, + Options(; compute_scope=F.interior_scopes[i]), + ys[i], F.W[i], xΓ) + end + for i in 1:k + x_parts[i] = fetch(back_tasks[i]) + end + + xp = vcat(x_parts..., xΓ) + x_local = similar(b_local) + x_local[F.perm] = xp + return distribute(x_local, b.partitioning) +end + +function LinearAlgebra.ldiv!(x::DVector, F::DistributedSchurLU, b::DVector) + return copyto!(x, F \ b) +end + diff --git a/src/datadeps/aliasing.jl b/src/datadeps/aliasing.jl index dc6668486..22c7ac071 100644 --- a/src/datadeps/aliasing.jl +++ b/src/datadeps/aliasing.jl @@ -240,6 +240,11 @@ struct HistoryEntry ainfo::AliasingWrapper space::MemorySpace write_num::Int + # Producer of this write/copy. Remainder syncdeps wait on this task + # directly instead of re-resolving through live `ainfos_owner`, which may + # later name a different task for the same ainfo or miss the producer when + # only an overlapping ainfo is consulted. + task::DTask end struct AliasedObjectCacheStore @@ -610,7 +615,12 @@ function aliasing!(state::DataDepsState, target_space::MemorySpace, arg_w::Argum return state.ainfo_cache[remote_arg_w] end - # Calculate the ainfo + # Calculate the ainfo via the current acceleration so SPMD backends (MPI) + # can broadcast owner-computed aliasing and keep ranks uniform. The generic + # `aliasing(::Acceleration, ...)` fallback uses `aliasing_unwrapped`, which + # resolves wrappers of whole-object containers (e.g. a view of a + # `DSparseArray`) to the container itself. For `Chunk`s, unwrap+aliasing + # happen where the data lives (remotely / on the owning rank). ainfo = AliasingWrapper(aliasing(current_acceleration(), remote_arg, arg_w.dep_mod)) # Cache the result @@ -661,10 +671,8 @@ function merge_history!(state::DataDepsState, arg_w::ArgumentWrapper, other_arg_ history = state.arg_history[arg_w] @opcounter :merge_history @opcounter :merge_history_complexity length(history) - origin_space = state.arg_origin[other_arg_w.arg] for other_entry in state.arg_history[other_arg_w] - write_num_tuple = HistoryEntry(AliasingWrapper(NoAliasing()), origin_space, other_entry.write_num) - range = searchsorted(history, write_num_tuple; by=x->x.write_num) + range = searchsorted(history, other_entry; by=x->x.write_num) if !isempty(range) # Find and skip duplicates match = false @@ -672,7 +680,8 @@ function merge_history!(state::DataDepsState, arg_w::ArgumentWrapper, other_arg_ source_entry = history[source_idx] if source_entry.ainfo == other_entry.ainfo && source_entry.space == other_entry.space && - source_entry.write_num == other_entry.write_num + source_entry.write_num == other_entry.write_num && + source_entry.task == other_entry.task match = true break end @@ -839,12 +848,12 @@ function add_writer!(state::DataDepsState, arg_w::ArgumentWrapper, dest_space::M empty!(state.arg_history[arg_w]) # Add our own history - push!(state.arg_history[arg_w], HistoryEntry(ainfo, dest_space, write_num)) + push!(state.arg_history[arg_w], HistoryEntry(ainfo, dest_space, write_num, task)) # Find overlapping arguments and update their history for other_arg_w in state.arg_overlaps[arg_w] other_arg_w == arg_w && continue - push!(state.arg_history[other_arg_w], HistoryEntry(ainfo, dest_space, write_num)) + push!(state.arg_history[other_arg_w], HistoryEntry(ainfo, dest_space, write_num, task)) end # Track which spaces hold a fully-current replica of this region diff --git a/src/datadeps/chunkview.jl b/src/datadeps/chunkview.jl index 903785f93..86964e6a9 100644 --- a/src/datadeps/chunkview.jl +++ b/src/datadeps/chunkview.jl @@ -69,7 +69,10 @@ function aliasing(accel::Acceleration, x::ChunkView{N}, dep_mod) where N return remotecall_fetch(root_worker_id(x.chunk.processor), x.chunk, x.slices) do x, slices x = unwrap(x) v = view(x, slices...) - return aliasing(accel, v, dep_mod) + # A view of a whole-object container (e.g. `DSparseArray`) must alias the + # entire container; `aliasing_unwrapped` resolves that (otherwise it just + # aliases the view), and crucially does so here where `v` lives. + return aliasing_unwrapped(v) end end aliasing(x::ChunkView) = aliasing(current_acceleration(), x, identity) diff --git a/src/datadeps/queue.jl b/src/datadeps/queue.jl index 9441f3bbe..0ef6c6655 100644 --- a/src/datadeps/queue.jl +++ b/src/datadeps/queue.jl @@ -153,16 +153,13 @@ function distribute_tasks!(queue::DataDepsTaskQueue) check_uniform(arg_w) arg = arg_w.arg origin_space = state.arg_origin[arg] - # When the origin still holds a fully-current replica (the argument was - # only read, or copies merely propagated it), the write-back is elided. - # This is only safe here at region end: mid-region, the copy tasks also - # serialize readers against later writers, so they must not be skipped. - current = get(state.arg_current, arg_w, nothing) - if current !== nothing && origin_space in current - remainder = NoAliasing() - else - remainder, _ = compute_remainder_for_arg!(state, origin_space, arg_w, write_num) - end + # Always compute write-back from history. Eliding via `arg_current` + # (`origin in arg_current` ⇒ skip) is unsound under multi-worker + # Distributed: Krylov `similar` workspaces can leave origin marked + # current while a remote true write holds the only initialized bytes. + # Read-only copy-ins may therefore schedule a redundant write-back + # (copy history advances `arg_owner`); that is intentional and harmless. + remainder, _ = compute_remainder_for_arg!(state, origin_space, arg_w, write_num) if remainder isa MultiRemainderAliasing origin_scope = UnionScope(map(ExactScope, collect(processors(origin_space)))...) enqueue_remainder_copy_from!(state, origin_space, arg_w, remainder, origin_scope, write_num) diff --git a/src/datadeps/remainders.jl b/src/datadeps/remainders.jl index 1000471b2..cbd494fef 100644 --- a/src/datadeps/remainders.jl +++ b/src/datadeps/remainders.jl @@ -43,6 +43,7 @@ memory_spans(mra::MultiRemainderAliasing) = vcat(memory_spans.(mra.remainders).. Base.hash(mra::MultiRemainderAliasing, h::UInt) = hash(mra.remainders, hash(MultiRemainderAliasing, h)) Base.:(==)(mra1::MultiRemainderAliasing, mra2::MultiRemainderAliasing) = mra1.remainders == mra2.remainders +# A whole-object copy from the argument's recorded owner space. struct FullCopy end """ @@ -87,9 +88,34 @@ function compute_remainder_for_arg!(state::DataDepsState, target_space::MemorySpace, arg_w::ArgumentWrapper, write_num::Int; compute_syncdeps::Bool=true) + owner_space = state.arg_owner[arg_w] + + # Whole-object containers (e.g. `DSparseArray`) reallocate their entire + # storage on every write, so they can never be *partially* current in a + # memory space -- the whole object is current in exactly one space: the most + # recent writer (or the origin owner if never written). There is thus never a + # meaningful span "remainder" to compute, so we short-circuit to a whole-object + # `FullCopy` from the owner space (or `NoAliasing` if the target already holds + # it). This avoids the span-diff computation and the construction + transfer of + # a `RemainderAliasing` entirely. We detect this from the aliasing info, which + # is a bare `ObjectAliasing` for such containers (see `aliases_as_whole`). + # + # `arg_owner` is the most-recent *direct* writer of `arg_w`, while the merged + # `arg_history` records the latest write through *any* overlapping argument. + # For every path Datadeps currently exercises a whole-object container is only + # ever reached through a single argument, so these agree -- the assertion below + # guards that invariant. They could only diverge under cross-argument aliasing + # of a whole-object container (the same container written through one argument + # and read through a different, overlapping one); supporting that would require + # sourcing the copy from `history[end].space` instead of `arg_owner`. + if aliases_as_whole(aliasing!(state, target_space, arg_w)) + history = state.arg_history[arg_w] + @assert isempty(history) || history[end].space == owner_space "whole-object container reached through a different overlapping argument than its owner; cross-argument aliasing of whole-object containers is unsupported" + return owner_space == target_space ? (NoAliasing(), 0) : (FullCopy(), 0) + end + spaces_set = Set{MemorySpace}() push!(spaces_set, target_space) - owner_space = state.arg_owner[arg_w] push!(spaces_set, owner_space) @label restart @@ -152,6 +178,7 @@ function compute_remainder_for_arg!(state::DataDepsState, break end + other_entry = nothing if idx > 0 other_entry = state.arg_history[arg_w][idx] other_ainfo = other_entry.ainfo @@ -199,7 +226,17 @@ function compute_remainder_for_arg!(state::DataDepsState, has_overlap = schedule_remainder!(tracker_other_space[1], other_space_idx, target_space_idx, remainder, other_many_spans) if compute_syncdeps && has_overlap @assert haskey(state.ainfos_owner, other_ainfo) "[idx $idx] ainfo $(typeof(other_ainfo)) has no owner" - get_read_deps!(state, other_space, other_ainfo, write_num, tracker_other_space[3]) + if other_entry !== nothing + # Wait on the producer recorded in history. Re-resolving via + # live `ainfos_owner`/`get_read_deps!` is not equivalent: + # ownership of `other_ainfo` may have moved to a later copy, or + # the producer may only be registered on an overlapping ainfo + # missing from `ainfos_overlaps[other_ainfo]`. + push!(tracker_other_space[3], ThunkSyncdep(other_entry.task)) + else + # idx==0 owner fallback: no HistoryEntry, use live writer set. + get_read_deps!(state, other_space, other_ainfo, write_num, tracker_other_space[3]) + end push!(tracker_other_space[2], other_ainfo) end end @@ -437,6 +474,9 @@ function move!(dep_mod::RemainderAliasing{S}, to_space::MemorySpace, from_space: multi_span_copy!(to_s, from_s, dep_mod.spans) return end + # N.B. Whole-object containers (e.g. `DSparseArray`) never reach here: their + # storage reallocates on write, so they can't be span-copied. They are routed + # to a whole-object `FullCopy` in `compute_remainder_for_arg!` instead. # Gather source spans into a host buffer (device arrays pack via KA first). # Pin when the source is device memory so DtoH hits page-locked host RAM. @@ -548,17 +588,3 @@ function write_remainder!(copies::Vector{UInt8}, copies_offset::UInt64, to::Base write_remainder!(copies, copies_offset, to[], to_ptr, n) end end - -function find_object_holding_ptr(A::SparseMatrixCSC, ptr::UInt64) - span = LocalMemorySpan(pointer(A.nzval), length(A.nzval)*sizeof(eltype(A.nzval))) - if span_start(span) <= ptr <= span_end(span) - return A.nzval - end - span = LocalMemorySpan(pointer(A.colptr), length(A.colptr)*sizeof(eltype(A.colptr))) - if span_start(span) <= ptr <= span_end(span) - return A.colptr - end - span = LocalMemorySpan(pointer(A.rowval), length(A.rowval)*sizeof(eltype(A.rowval))) - @assert span_start(span) <= ptr <= span_end(span) "Pointer $ptr not found in SparseMatrixCSC" - return A.rowval -end diff --git a/src/file-io.jl b/src/file-io.jl index 1407fbced..15405d049 100644 --- a/src/file-io.jl +++ b/src/file-io.jl @@ -1,5 +1,4 @@ export save, load -using SparseArrays import MemPool: GenericFileDevice, FileRef struct File diff --git a/src/memory-spaces.jl b/src/memory-spaces.jl index 617f7ea8b..015ca502c 100644 --- a/src/memory-spaces.jl +++ b/src/memory-spaces.jl @@ -375,7 +375,11 @@ function memory_spans(oa::ObjectAliasing{S}) where S return [span] end -aliasing(accel::Acceleration, x, T) = aliasing(x, T) +# Acceleration entry point used by Datadeps. Unwrap whole-object containers +# (e.g. views of `DSparseArray`) before computing aliasing. SPMD backends +# (MPI) overload this for `Chunk`/`ChunkView` to broadcast owner-computed +# aliasing; those methods call `aliasing_unwrapped` on the owning rank. +aliasing(accel::Acceleration, x, T) = aliasing_unwrapped(x, T) function aliasing(x, dep_mod) if dep_mod isa Symbol return aliasing(getfield(x, dep_mod)) @@ -411,22 +415,95 @@ end aliasing(::String) = NoAliasing() # FIXME: Not necessarily true aliasing(::Symbol) = NoAliasing() aliasing(::Type) = NoAliasing() + +""" + aliases_as_whole(x) -> Bool + +Whether `x` -- either an object, or an already-computed aliasing -- must be +treated as a single, indivisible unit, i.e. it is never safe to alias (or to +consider current) only *part* of it. + +For an object: containers whose backing storage may be reallocated or resized on +write (such as `DSparseArray`) return `true`. This causes [`aliasing_root`](@ref) +to resolve any view/wrapper of them to the whole container, so all access funnels +through the container's own `aliasing`. By contract, such a container's `aliasing` +must be a bare [`ObjectAliasing`](@ref). + +For an aliasing: a bare `ObjectAliasing` (the aliasing of a whole-object +container) reports `true`. Datadeps uses this to copy such arguments as a whole +(a `FullCopy`) rather than computing per-span remainders, since they can never be +*partially* current in a memory space. +""" +aliases_as_whole(@nospecialize x) = false +aliases_as_whole(ainfo::AliasingWrapper) = aliases_as_whole(ainfo.inner) +aliases_as_whole(::ObjectAliasing) = true + +""" + aliasing_root(x) + +Resolve `x` down to the whole-object container it wraps: peel array wrappers +(views, transposes, adjoints, reshapes, permutations, ...) off of `x` until +reaching a container that must alias as a whole (see [`aliases_as_whole`](@ref)), +and return that container. If no such container is wrapped, return `x` unchanged. + +This relies solely on the standard `Base.parent` interface, so it transparently +handles *any* array wrapper without per-wrapper `aliasing` methods: a new wrapper +type needs no special-casing here, and a new whole-object container only needs to +define [`aliases_as_whole`](@ref) (plus its own `aliasing` method, which must +return a bare `ObjectAliasing`). A trapping `Base.pointer` on such containers +guards against a wrapper slipping through and being misinterpreted as strided +memory. + +!!! warning + The resolved object's *identity* defines its aliasing, so this is only + meaningful in the memory space where `x` physically resides. Never return its + result across a worker boundary and *then* compute `aliasing` -- the transfer + copies the object and changes its aliasing. Use [`aliasing_unwrapped`](@ref), + which fuses the two operations so they always run together. +""" +aliasing_root(@nospecialize x) = x +function aliasing_root(x::AbstractArray) + aliases_as_whole(x) && return x + p = parent(x) + # `Base.parent` returns the argument itself for non-wrapper arrays, which + # terminates the recursion. + p === x && return x + root = aliasing_root(p) + return aliases_as_whole(root) ? root : x +end + +""" + aliasing_unwrapped(x[, dep_mod]) + +Compute the `aliasing` of `x`, first resolving (via [`aliasing_root`](@ref)) any +whole-object container that `x` wraps. This is the entry point Datadeps uses to +compute the aliasing of a (possibly wrapped) argument. + +Unwrapping and `aliasing` are intentionally fused into a single call so they +always execute in the same place. It MUST be evaluated where `x` physically lives +-- e.g. *inside* a `Chunk`'s `remotecall_fetch` block -- because the unwrapped +object's identity determines its aliasing; returning the unwrapped object across +a worker boundary first would copy it and silently change the result. +""" +aliasing_unwrapped(x) = aliasing(aliasing_root(x)) +aliasing_unwrapped(x, dep_mod) = aliasing(aliasing_root(x), dep_mod) + function aliasing(x::Chunk, T) if root_worker_id(x.processor) == myid() - return aliasing(unwrap(x), T) + return aliasing_unwrapped(unwrap(x), T) end @assert x.handle isa DRef return remotecall_fetch(root_worker_id(x.processor), x, T) do x, T - aliasing(unwrap(x), T) + aliasing_unwrapped(unwrap(x), T) end end function aliasing(x::Chunk) if root_worker_id(x.processor) == myid() - return aliasing(unwrap(x)) + return aliasing_unwrapped(unwrap(x)) end @assert x.handle isa DRef return remotecall_fetch(root_worker_id(x.processor), x) do x - aliasing(unwrap(x)) + aliasing_unwrapped(unwrap(x)) end end aliasing(x::DTask, T) = aliasing(fetch(x; move_value=false, unwrap=false), T) diff --git a/test/Project.toml b/test/Project.toml index edc738961..47894c0bf 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,4 +1,5 @@ [deps] +AlgebraicMultigrid = "2169fc97-5a83-5252-b627-83903c6c433c" ArgParse = "c7e460c6-2fb9-53a9-8c5b-16f535851c63" Colors = "5ae59095-9a9b-59fe-a467-6f913c188581" CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab" @@ -9,15 +10,20 @@ FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341" GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" GraphViz = "f526b714-d49f-11e8-06ff-31ed36ee7ee0" Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6" +IncompleteLU = "40713840-3770-5561-ab4c-a76e7d0d7895" JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" +Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LinuxPerf = "b4c46c6c-4fb0-484d-a11a-41bc3392d094" MemPool = "f9f48841-c794-520a-933b-121f7ba6ed94" +Metis = "2679e427-3c69-5b7f-982b-ece356f1e94b" OnlineStats = "a15396b6-48d5-5d58-9928-6d29437db91e" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" Preferences = "21216c6a-2e73-6563-6e65-726566657250" +PureKLU = "0c0d3e7f-3a8b-4f7e-b6f1-9a4d2e7c1f01" +PureUMFPACK = "b7e1f0a2-3c4d-4e5f-9a0b-1c2d3e4f5a6b" PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b" diff --git a/test/array/allocation.jl b/test/array/allocation.jl index 83dba36d2..4548f3ed1 100644 --- a/test/array/allocation.jl +++ b/test/array/allocation.jl @@ -99,7 +99,8 @@ end @test size(Xsp) == dims AT = length(dims) == 2 ? SparseMatrixCSC : SparseVector Ach = fetch(Xsp.chunks[1]) - @test Ach isa AT{T} + @test Ach isa Dagger.DSparseArray{T,length(dims)} + @test Ach.mat isa AT{T} AXsp = collect(Xsp) @test AXsp isa Array{T,length(dims)} @test AXsp == collect(Xsp) diff --git a/test/array/linalg/core.jl b/test/array/linalg/core.jl index 4c66ac34a..5de6fefd1 100644 --- a/test/array/linalg/core.jl +++ b/test/array/linalg/core.jl @@ -23,3 +23,57 @@ end DA = DArray(A) @test isapprox(norm(A), norm(DA)) end + +@testset "BLAS-1 vector ops" begin + @testset "T=$T part=$(part.blocksize)" for T in (Float64, ComplexF64), part in (Blocks(16), Blocks(4)) + n = 16 + x = rand(T, n) + y = rand(T, n) + a = T(2) + b = T(3) + Dx = distribute(x, part) + + # dot (conjugates first arg for complex) + @test dot(Dx, distribute(y, part)) ≈ dot(x, y) + + # axpy!: y += a*x + Dy = distribute(copy(y), part) + axpy!(a, Dx, Dy) + @test collect(Dy) ≈ a .* x .+ y + + # axpby!: y = a*x + b*y + Dy = distribute(copy(y), part) + axpby!(a, Dx, b, Dy) + @test collect(Dy) ≈ a .* x .+ b .* y + + # rmul!/lmul!: scale in place + Dz = distribute(copy(x), part) + rmul!(Dz, a) + @test collect(Dz) ≈ a .* x + Dz = distribute(copy(x), part) + lmul!(a, Dz) + @test collect(Dz) ≈ a .* x + end + + # Operands with *different* partitionings must still work (aligned internally + # via maybe_copy_buffered), and the in-place result keeps its own layout. + @testset "mismatched partitionings T=$T" for T in (Float64, ComplexF64) + n = 16 + x = rand(T, n) + y = rand(T, n) + a, b = T(2), T(3) + Dx = distribute(x, Blocks(4)) + Dy = distribute(y, Blocks(8)) + + @test dot(Dx, Dy) ≈ dot(x, y) + + Dy2 = distribute(copy(y), Blocks(8)) + axpy!(a, Dx, Dy2) + @test collect(Dy2) ≈ a .* x .+ y + @test Dy2.partitioning == Blocks(8) # output layout preserved + + Dy3 = distribute(copy(y), Blocks(8)) + axpby!(a, Dx, b, Dy3) + @test collect(Dy3) ≈ a .* x .+ b .* y + end +end diff --git a/test/array/linalg/iterativesolvers.jl b/test/array/linalg/iterativesolvers.jl new file mode 100644 index 000000000..ca7865f29 --- /dev/null +++ b/test/array/linalg/iterativesolvers.jl @@ -0,0 +1,206 @@ +# Distributed iterative (Krylov) linear-solver tests. +# +# Exercises the matrix-free Krylov integration (`Dagger.cg`/`minres`/`gmres`/ +# `bicgstab` + the generic `krylov_solve`) over both dense and sparse-backed +# `DMatrix` operators, plus the Jacobi preconditioner. Reference solutions come +# from a dense direct solve. +# +# julia test/runtests.jl --test array/linalg/iterativesolvers + +using Krylov +using AlgebraicMultigrid +using IncompleteLU + +# Strongly diagonally-dominant tridiagonal SPD matrix. The large diagonal keeps +# the condition number small so the Krylov methods converge in a handful of +# iterations (keeping the distributed test fast). `inv(diag) == 1/4`. +const SPD_DIAG = 4.0 + +laplacian_1d(T, n) = SparseArrays.spdiagm( + -1 => fill(-one(T), n - 1), + 0 => fill(T(SPD_DIAG), n), + 1 => fill(-one(T), n - 1), +) + +# Add a first-order advection term -> nonsymmetric, still well-conditioned. +function advection_diffusion_1d(T, n) + return laplacian_1d(T, n) + SparseArrays.spdiagm( + -1 => fill(T(-3) / 10, n - 1), + 1 => fill(T(3) / 10, n - 1), + ) +end + +@testset "Iterative solvers (Krylov)" begin + n = 64 + k = 16 + Db_part = Blocks(k) + A_part = Blocks(k, k) + + @testset "SPD operator ($(backend))" for backend in (:dense, :sparse) + Asp = laplacian_1d(Float64, n) + Adense = Matrix(Asp) + b = rand(n) + xref = Adense \ b + + DA = backend === :dense ? distribute(Adense, A_part) : distribute(Asp, A_part) + Db = distribute(b, Db_part) + + @testset "$(nameof(solver))" for solver in (Dagger.cg, Dagger.minres, Dagger.gmres, Dagger.bicgstab) + x, stats = solver(DA, Db; atol = 1e-12, rtol = 1e-10, itmax = 500) + @test stats.solved + @test x isa Dagger.DVector + @test collect(x) ≈ xref rtol = 1e-6 + end + + # Generic entry point. + x, stats = Dagger.krylov_solve(:cg, DA, Db; atol = 1e-12, rtol = 1e-10) + @test stats.solved + @test collect(x) ≈ xref rtol = 1e-6 + end + + @testset "nonsymmetric operator ($(backend))" for backend in (:dense, :sparse) + Asp = advection_diffusion_1d(Float64, n) + b = rand(n) + xref = Matrix(Asp) \ b + + DA = backend === :dense ? distribute(Matrix(Asp), A_part) : distribute(Asp, A_part) + Db = distribute(b, Db_part) + + @testset "$(nameof(solver))" for solver in (Dagger.gmres, Dagger.bicgstab) + x, stats = solver(DA, Db; atol = 1e-12, rtol = 1e-10, itmax = 500) + @test stats.solved + @test collect(x) ≈ xref rtol = 1e-6 + end + end + + @testset "complex SPD (Hermitian) operator" begin + # Real SPD tridiagonal is Hermitian as a complex matrix. + Asp = SparseArrays.spdiagm( + -1 => fill(ComplexF64(-1), n - 1), + 0 => fill(ComplexF64(SPD_DIAG), n), + 1 => fill(ComplexF64(-1), n - 1), + ) + b = rand(ComplexF64, n) + xref = Matrix(Asp) \ b + DA = distribute(Asp, A_part) + Db = distribute(b, Db_part) + x, stats = Dagger.cg(DA, Db; atol = 1e-12, rtol = 1e-10, itmax = 500) + @test stats.solved + @test collect(x) ≈ xref rtol = 1e-6 + end + + @testset "Jacobi preconditioner" begin + Asp = laplacian_1d(Float64, n) + b = rand(n) + xref = Matrix(Asp) \ b + + @testset "build + apply ($(backend))" for backend in (:dense, :sparse) + DA = backend === :dense ? distribute(Matrix(Asp), A_part) : distribute(Asp, A_part) + Db = distribute(b, Db_part) + + P = Dagger.JacobiPreconditioner(DA) + @test collect(P.dinv) ≈ fill(1 / SPD_DIAG, n) # 1/diag + + # Apply: y = M⁻¹ x = dinv .* x. + y = similar(Db) + mul!(y, P, Db) + @test collect(y) ≈ (1 / SPD_DIAG) .* b + + x, stats = Dagger.cg(DA, Db; M = P, atol = 1e-12, rtol = 1e-10, itmax = 500) + @test stats.solved + @test collect(x) ≈ xref rtol = 1e-6 + end + + # Non-square block grid must be rejected with a helpful error. + DA_ragged = distribute(Matrix(Asp), Blocks(k, k ÷ 2)) + @test_throws ArgumentError Dagger.JacobiPreconditioner(DA_ragged) + end + + @testset "block-Jacobi preconditioner" begin + Asp = laplacian_1d(Float64, n) + Adense = Matrix(Asp) + b = rand(n) + xref = Adense \ b + + # Reference: apply the exact block-diagonal inverse. + yref = similar(b) + for s in 1:k:n + r = s:min(s + k - 1, n) + yref[r] = Adense[r, r] \ b[r] + end + + @testset "build + apply ($(backend))" for backend in (:dense, :sparse) + DA = backend === :dense ? distribute(Adense, A_part) : distribute(Asp, A_part) + Db = distribute(b, Db_part) + + P = Dagger.BlockJacobiPreconditioner(DA) + y = similar(Db) + mul!(y, P, Db) + @test collect(y) ≈ yref + + x, stats = Dagger.cg(DA, Db; M = P, atol = 1e-12, rtol = 1e-10, itmax = 500) + @test stats.solved + @test collect(x) ≈ xref rtol = 1e-6 + end + + # A single tile makes block-Jacobi an *exact* solve, so PCG converges + # essentially immediately. + DA1 = distribute(Adense, Blocks(n, n)) + Db1 = distribute(b, Blocks(n)) + P1 = Dagger.BlockJacobiPreconditioner(DA1) + x1, s1 = Dagger.cg(DA1, Db1; M = P1, atol = 1e-12, rtol = 1e-10, itmax = 500) + @test s1.solved + @test s1.niter <= 2 + @test collect(x1) ≈ xref rtol = 1e-8 + + @test_throws ArgumentError Dagger.BlockJacobiPreconditioner(distribute(Adense, Blocks(k, k ÷ 2))) + end + + @testset "AMG preconditioner ($(method))" for method in (:ruge_stuben, :smoothed_aggregation) + Asp = laplacian_1d(Float64, n) + Adense = Matrix(Asp) + b = rand(n) + xref = Adense \ b + + @testset "$(backend)" for backend in (:dense, :sparse) + DA = backend === :dense ? distribute(Adense, A_part) : distribute(Asp, A_part) + Db = distribute(b, Db_part) + + P = Dagger.AMGPreconditioner(DA; method = method) + # Apply is a V-cycle approximating M⁻¹ x; just check it runs + is finite + # (and repeatable, exercising the cached, pinned hierarchy). + y1 = similar(Db); mul!(y1, P, Db) + y2 = similar(Db); mul!(y2, P, Db) + @test all(isfinite, collect(y1)) + @test collect(y1) ≈ collect(y2) + + x, stats = Dagger.cg(DA, Db; M = P, atol = 1e-12, rtol = 1e-10, itmax = 500) + @test stats.solved + @test collect(x) ≈ xref rtol = 1e-6 + end + end + + @testset "block-ILU preconditioner" begin + Asp = laplacian_1d(Float64, n) + Adense = Matrix(Asp) + b = rand(n) + xref = Adense \ b + + @testset "build + apply ($(backend))" for backend in (:dense, :sparse) + DA = backend === :dense ? distribute(Adense, A_part) : distribute(Asp, A_part) + Db = distribute(b, Db_part) + + P = Dagger.BlockILUPreconditioner(DA; τ = 0.01) + y1 = similar(Db); mul!(y1, P, Db) + y2 = similar(Db); mul!(y2, P, Db) + @test all(isfinite, collect(y1)) + @test collect(y1) ≈ collect(y2) + + x, stats = Dagger.cg(DA, Db; M = P, atol = 1e-12, rtol = 1e-10, itmax = 500) + @test stats.solved + @test collect(x) ≈ xref rtol = 1e-6 + end + + @test_throws ArgumentError Dagger.BlockILUPreconditioner(distribute(Adense, Blocks(k, k ÷ 2)); τ = 0.01) + end +end diff --git a/test/array/linalg/matmul.jl b/test/array/linalg/matmul.jl index 26fd0afe3..f0ec73287 100644 --- a/test/array/linalg/matmul.jl +++ b/test/array/linalg/matmul.jl @@ -29,36 +29,71 @@ function test_gemm!(T, szA, szB, partA, partB) DA = distribute(A, partA) DB = distribute(B, partB) + SA = sprand(T, szA..., 0.1) + SB = sprand(T, szA..., 0.1) + + DSA = distribute(SA, partA) + DSB = distribute(SB, partB) + ## Out-of-place gemm # No transA, No transB + # Dense DC = DA * DB C = A * B @test collect(DC) ≈ C + # Sparse + DSC = DSA * DSB + SC = SA * SB + @test collect(DSC) ≈ SC if szA == szB # No transA, transB + # Dense DC = DA * DB' C = A * B' @test collect(DC) ≈ C + # Sparse + DSC = DSA * DSB' + SC = SA * SB' + @test collect(DSC) ≈ SC # transA, No transB + # Dense DC = DA' * DB C = A' * B @test collect(DC) ≈ C + # Sparse + DSC = DSA' * DSB + SC = SA' * SB + @test collect(DSC) ≈ SC end # transA, transB + # Dense DC = DA' * DB' C = A' * B' @test collect(DC) ≈ C + #= Sparse + DSC = DSA' * DSB' + SC = SA' * SB' + @test collect(DSC) ≈ SC + =# ## In-place gemm # No transA, No transB + # Dense C = zeros(T, szC...) DC = distribute(C, partC) mul!(C, A, B) mul!(DC, DA, DB) @test collect(DC) ≈ C + #= Sparse + SC = zeros(T, szC...) + DSC = distribute(SC, partC) + mul!(SC, SA, SB) + mul!(DSC, DSA, DSB) + @test collect(DSC) ≈ SC + =# if szA == szB # No transA, transB diff --git a/test/array/linalg/matmul_finch.jl b/test/array/linalg/matmul_finch.jl new file mode 100644 index 000000000..061874f1e --- /dev/null +++ b/test/array/linalg/matmul_finch.jl @@ -0,0 +1,105 @@ +# Fast, focused sparse matrix-multiply tests for the Finch backend. +# +# Mirrors `matmul_sparse.jl` (SparseArrays) but exercises Finch tensor tiles via +# `DSparseMatrix`. Finch is an optional (weak) dependency, so this test loads it +# on all workers and skips gracefully if it cannot be loaded. +# +# Run with: +# +# julia test/runtests.jl --test array/linalg/matmul_finch +# +# N.B. Finch precompilation is heavy (several minutes); this test is intended for +# local validation rather than fast CI. + +# Make the test-project environment (which provides Finch) resolvable on the +# workers, which are launched with Dagger's project. +const _FINCH_TEST_ENV = abspath(joinpath(@__DIR__, "..", "..")) +@everywhere pushfirst!(LOAD_PATH, $_FINCH_TEST_ENV) + +const FINCH_AVAILABLE = try + @eval using Finch + @everywhere using Finch + true +catch err + @warn "Finch unavailable; skipping Finch matmul tests" exception=err + false +end + +if !FINCH_AVAILABLE + @testset "Finch GEMM (quick)" begin + @test_skip false + end +else + const FINCH_DENSITY = 0.3 + + # Self-consistent: build distributed Finch operands, take the dense + # `collect` as the reference, and compare against the distributed result. + function test_finch_gemm!(T, sz, partA, partB) + partC = Blocks(partA.blocksize[1], partB.blocksize[2]) + + DSA = Finch.fsprand(partA, T, sz, FINCH_DENSITY) + DSB = Finch.fsprand(partB, T, sz, FINCH_DENSITY) + A = collect(DSA) + B = collect(DSB) + + ## Out-of-place + @test collect(DSA * DSB) ≈ A * B # N / N + @test collect(DSA * DSB') ≈ A * B' # N / C + @test collect(DSA' * DSB) ≈ A' * B # C / N + @test collect(DSA' * DSB') ≈ A' * B' # C / C + + ## Symmetric rank-k (syrk path: A === B with one operand transposed) + @test collect(DSA' * DSA) ≈ A' * A + @test collect(DSA * DSA') ≈ A * A' + + ## In-place + DSC = Finch.fspzeros(partC, T, sz) + mul!(DSC, DSA, DSB) + @test collect(DSC) ≈ A * B + end + + const FINCH_QUICK_CASES = [ + ((8, 8), Blocks(4, 4), Blocks(4, 4)), + ((8, 8), Blocks(2, 4), Blocks(4, 2)), + ] + + @testset "Finch GEMM (quick)" begin + @testset "size=$sz part=$(partA.blocksize)/$(partB.blocksize)" for (sz, partA, partB) in FINCH_QUICK_CASES + @testset "T=$T" for T in (Float64, ComplexF64) + test_finch_gemm!(T, sz, partA, partB) + end + end + end + + # Sparse matrix-vector multiply (SpMV): Finch sparse matrix tiles, dense vectors. + function test_finch_spmv!(T, n, part) + bs = part.blocksize[1] + DSA = Finch.fsprand(part, T, (n, n), FINCH_DENSITY) + A = collect(DSA) + x = rand(T, n) + Dx = distribute(x, Blocks(bs)) + + @test collect(DSA * Dx) ≈ A * x + @test collect(transpose(DSA) * Dx) ≈ transpose(A) * x + @test collect(DSA' * Dx) ≈ A' * x + + y = rand(T, n) + Dy = distribute(copy(y), Blocks(bs)) + alpha, beta = T(2), T(3) + mul!(Dy, DSA, Dx, alpha, beta) + @test collect(Dy) ≈ alpha * (A * x) + beta * y + end + + const FINCH_SPMV_CASES = [ + (8, Blocks(4, 4)), + (8, Blocks(2, 2)), + ] + + @testset "Finch SpMV (quick)" begin + @testset "n=$n part=$(part.blocksize)" for (n, part) in FINCH_SPMV_CASES + @testset "T=$T" for T in (Float64, ComplexF64) + test_finch_spmv!(T, n, part) + end + end + end +end diff --git a/test/array/linalg/matmul_sparse.jl b/test/array/linalg/matmul_sparse.jl new file mode 100644 index 000000000..cc2e8274e --- /dev/null +++ b/test/array/linalg/matmul_sparse.jl @@ -0,0 +1,139 @@ +# Fast, focused sparse matrix-multiply tests. +# +# This is intentionally a *small* subset of the full GEMM suite in `matmul.jl`, +# meant for quick iteration while developing sparse-array support (the full +# suite takes ~1hr). Run with: +# +# julia test/runtests.jl --test array/linalg/matmul_sparse +# +# Keep this fast: only a couple of sizes/partitionings, just enough to exercise +# the out-of-place and in-place paths (incl. transposes) for both real and +# complex element types. + +const SPARSE_DENSITY = 0.3 + +function test_sparse_gemm!(T, sz, partA, partB) + rows, cols = sz + @assert rows == cols "sparse quick-test uses square matrices so transposes line up" + partC = Blocks(partA.blocksize[1], partB.blocksize[2]) + + SA = sprand(T, sz..., SPARSE_DENSITY) + SB = sprand(T, sz..., SPARSE_DENSITY) + + DSA = distribute(SA, partA) + DSB = distribute(SB, partB) + + ## Out-of-place + @test collect(DSA * DSB) ≈ SA * SB # N / N + @test collect(DSA * DSB') ≈ SA * SB' # N / T + @test collect(DSA' * DSB) ≈ SA' * SB # T / N + @test collect(DSA' * DSB') ≈ SA' * SB' # T / T + + ## Symmetric rank-k (syrk path: A === B with one operand transposed) + @test collect(DSA' * DSA) ≈ Array(SA)' * Array(SA) + @test collect(DSA * DSA') ≈ Array(SA) * Array(SA)' + + ## In-place + SC = SA * SB + DSC = distribute(sparse(zeros(T, sz...)), partC) + mul!(DSC, DSA, DSB) + @test collect(DSC) ≈ SC +end + +const SPARSE_QUICK_CASES = [ + ((8, 8), Blocks(4, 4), Blocks(4, 4)), + ((8, 8), Blocks(2, 4), Blocks(4, 2)), + ((8, 8), Blocks(4, 2), Blocks(2, 4)), +] + +@testset "Sparse GEMM (quick)" begin + @testset "size=$sz part=$(partA.blocksize)/$(partB.blocksize)" for (sz, partA, partB) in SPARSE_QUICK_CASES + @testset "T=$T" for T in (Float64, ComplexF64) + test_sparse_gemm!(T, sz, partA, partB) + end + end +end + +# Sparse matrix-vector multiply (SpMV): sparse matrix tiles, dense vectors. This +# is the core kernel for distributed iterative (Krylov) solvers. +function test_sparse_spmv!(T, n, part) + bs = part.blocksize[1] + SA = sprand(T, n, n, SPARSE_DENSITY) + x = rand(T, n) + + DSA = distribute(SA, part) + Dx = distribute(x, Blocks(bs)) + + ## Out-of-place (N / T / C) + @test collect(DSA * Dx) ≈ SA * x + @test collect(transpose(DSA) * Dx) ≈ transpose(SA) * x + @test collect(DSA' * Dx) ≈ SA' * x + + ## In-place 5-arg mul!: y = alpha*A*x + beta*y + y = rand(T, n) + Dy = distribute(copy(y), Blocks(bs)) + alpha, beta = T(2), T(3) + mul!(Dy, DSA, Dx, alpha, beta) + @test collect(Dy) ≈ alpha * (SA * x) + beta * y +end + +const SPARSE_SPMV_CASES = [ + (8, Blocks(4, 4)), + (8, Blocks(2, 2)), +] + +@testset "Sparse SpMV (quick)" begin + @testset "n=$n part=$(part.blocksize)" for (n, part) in SPARSE_SPMV_CASES + @testset "T=$T" for T in (Float64, ComplexF64) + test_sparse_spmv!(T, n, part) + end + end +end + +# Any *partial* or *reinterpreted* access to a sparse container (views, +# transposes, adjoints, reshapes, and combinations thereof) must alias the +# *entire* container, so that Datadeps never tracks stale sub-spans of storage +# that may have been reallocated on write. This is enforced by `aliasing_root`, +# which resolves any such wrapper back to the container itself (via `Base.parent`) +# -- exposed to Datadeps through the fused `aliasing_unwrapped` -- plus a trapping +# `Base.pointer` that fires if a wrapper ever slips through to the strided path. +@testset "Sparse whole-container aliasing" begin + M = Dagger.DSparseMatrix{Float64}(sprand(Float64, 6, 6, 0.4)) + + @test Dagger.aliases_as_whole(M) + @test !Dagger.aliases_as_whole(rand(Float64, 6, 6)) + # The container's aliasing is a bare `ObjectAliasing`, which also reports + # whole-object (drives the Datadeps whole-object copy short-circuit). + @test Dagger.aliasing(M) isa Dagger.ObjectAliasing + @test Dagger.aliases_as_whole(Dagger.aliasing(M)) + @test !Dagger.aliases_as_whole(Dagger.aliasing(rand(Float64, 6, 6))) + + # `aliasing_root` peels any array wrapper back to the sparse container. + @test Dagger.aliasing_root(M) === M + @test Dagger.aliasing_root(view(M, 1:3, :)) === M + @test Dagger.aliasing_root(view(M, 4:6, 2:4)) === M + @test Dagger.aliasing_root(transpose(M)) === M + @test Dagger.aliasing_root(M') === M + @test Dagger.aliasing_root(reshape(M, 36)) === M + @test Dagger.aliasing_root(view(transpose(M), 1:2, :)) === M + + # After unwrapping, all wrappers alias the same whole container. + a_full = Dagger.aliasing(M) + @test Dagger.will_alias(a_full, Dagger.aliasing_unwrapped(view(M, 1:3, :))) + @test Dagger.will_alias(Dagger.aliasing_unwrapped(view(M, 1:3, :)), + Dagger.aliasing_unwrapped(view(M, 4:6, :))) + + # Distinct containers must not alias. + M2 = Dagger.DSparseMatrix{Float64}(sprand(Float64, 6, 6, 0.4)) + @test !Dagger.will_alias(a_full, Dagger.aliasing(M2)) + + # Dense arrays are untouched: not whole-object, views stay strided, and + # `aliasing_root` returns them unchanged. + A = rand(Float64, 6, 6) + Av = view(A, 1:3, :) + @test Dagger.aliasing_root(Av) === Av + @test Dagger.aliasing_unwrapped(Av) isa Dagger.StridedAliasing + + # The `pointer` trap fires if a sparse container is treated as raw memory. + @test_throws ArgumentError pointer(M) +end diff --git a/test/array/linalg/sparsedirect.jl b/test/array/linalg/sparsedirect.jl new file mode 100644 index 000000000..e2f94ad50 --- /dev/null +++ b/test/array/linalg/sparsedirect.jl @@ -0,0 +1,249 @@ +# Sparse direct-solver tests (KLU via PureKLU, UMFPACK-style via PureUMFPACK). +# +# Exercises the whole-matrix direct solves (`Dagger.klu` / `Dagger.splu`) and the +# per-tile block direct preconditioners (`BlockKLUPreconditioner` / +# `BlockUMFPACKPreconditioner`). Reference solutions come from a dense direct +# solve. +# +# julia test/runtests.jl --test array/linalg/sparsedirect + +using PureKLU +using PureUMFPACK +using Metis +using Krylov + +# A nonsymmetric, well-conditioned sparse system (advection-diffusion-like). +function nonsym_1d(T, n) + return SparseArrays.spdiagm( + -1 => fill(-one(T), n - 1), + 0 => fill(T(3), n), + 1 => fill(T(-7) / 10, n - 1), + ) +end + +# SPD, diagonally dominant tridiagonal (stable under pivoting). +function spd_1d(T, n) + return SparseArrays.spdiagm( + -1 => fill(-one(T), n - 1), + 0 => fill(T(4), n), + 1 => fill(-one(T), n - 1), + ) +end + +# 2D 5-point Laplacian on an N×N grid (SPD; clean METIS separators). +function lap2d(T, N) + n = N * N + Is = Int[]; Js = Int[]; Vs = T[] + idx(i, j) = (j - 1) * N + i + for j in 1:N, i in 1:N + k = idx(i, j) + push!(Is, k); push!(Js, k); push!(Vs, T(4)) + if i > 1; push!(Is, k); push!(Js, idx(i - 1, j)); push!(Vs, -one(T)); end + if i < N; push!(Is, k); push!(Js, idx(i + 1, j)); push!(Vs, -one(T)); end + if j > 1; push!(Is, k); push!(Js, idx(i, j - 1)); push!(Vs, -one(T)); end + if j < N; push!(Is, k); push!(Js, idx(i, j + 1)); push!(Vs, -one(T)); end + end + return SparseArrays.sparse(Is, Js, Vs, n, n) +end + +# Nonsymmetric, diagonally dominant 2D advection-diffusion stencil. +function advdiff2d(T, N) + n = N * N + Is = Int[]; Js = Int[]; Vs = T[] + idx(i, j) = (j - 1) * N + i + for j in 1:N, i in 1:N + k = idx(i, j) + push!(Is, k); push!(Js, k); push!(Vs, T(5)) + if i > 1; push!(Is, k); push!(Js, idx(i - 1, j)); push!(Vs, T(-12) / 10); end + if i < N; push!(Is, k); push!(Js, idx(i + 1, j)); push!(Vs, T(-8) / 10); end + if j > 1; push!(Is, k); push!(Js, idx(i, j - 1)); push!(Vs, T(-11) / 10); end + if j < N; push!(Is, k); push!(Js, idx(i, j + 1)); push!(Vs, T(-9) / 10); end + end + return SparseArrays.sparse(Is, Js, Vs, n, n) +end + +@testset "Sparse direct solvers" begin + n = 64 + k = 16 + Db_part = Blocks(k) + A_part = Blocks(k, k) + + Asp = nonsym_1d(Float64, n) + Adense = Matrix(Asp) + b = rand(n) + xref = Adense \ b + + @testset "whole-matrix direct ($(solver)) — $(backend)" for solver in (:klu, :splu), + backend in (:dense, :sparse, :singletile) + DA = if backend === :dense + distribute(Adense, A_part) + elseif backend === :sparse + distribute(Asp, A_part) + else + distribute(Asp, Blocks(n, n)) # single tile + end + Db = distribute(b, Db_part) + + F = solver === :klu ? Dagger.klu(DA) : Dagger.splu(DA) + @test size(F) == (n, n) + + # `F \ b` returns a DVector partitioned like `b`. + x = F \ Db + @test x isa Dagger.DVector + @test collect(x) ≈ xref + + # In-place `ldiv!` writes into a preallocated DVector. + y = similar(Db) + LinearAlgebra.ldiv!(y, F, Db) + @test collect(y) ≈ xref + + # Dimension mismatch is caught. + @test_throws DimensionMismatch F \ distribute(rand(n ÷ 2), Blocks(k)) + end + + @testset "block direct preconditioner ($(name))" for (name, build) in ( + ("KLU", Dagger.BlockKLUPreconditioner), + ("UMFPACK", Dagger.BlockUMFPACKPreconditioner), + ) + @testset "build + apply ($(backend))" for backend in (:dense, :sparse) + DA = backend === :dense ? distribute(Adense, A_part) : distribute(Asp, A_part) + Db = distribute(b, Db_part) + + # Reference: apply the exact block-diagonal inverse. + yref = similar(b) + for s in 1:k:n + r = s:min(s + k - 1, n) + yref[r] = Adense[r, r] \ b[r] + end + + P = build(DA) + y = similar(Db) + mul!(y, P, Db) + @test collect(y) ≈ yref + # Repeatable (exercises the cached, pinned factorization). + y2 = similar(Db) + mul!(y2, P, Db) + @test collect(y2) ≈ yref + + x, stats = Dagger.bicgstab(DA, Db; M = P, atol = 1e-12, rtol = 1e-10, itmax = 500) + @test stats.solved + @test collect(x) ≈ xref rtol = 1e-6 + end + + # A single tile makes the block solve an *exact* whole-matrix solve. + DA1 = distribute(Asp, Blocks(n, n)) + Db1 = distribute(b, Blocks(n)) + P1 = build(DA1) + y1 = similar(Db1) + mul!(y1, P1, Db1) + @test collect(y1) ≈ xref + + # Non-square block grid must be rejected with a helpful error. + @test_throws ArgumentError build(distribute(Adense, Blocks(k, k ÷ 2))) + end + + @testset "distributed triangular solve (splu)" begin + systems = ( + ("nonsym", nonsym_1d(Float64, n)), + ("spd", spd_1d(Float64, n)), + ) + # Include a case where A's tiling differs from the solve blocksize. + configs = ( + (Blocks(k, k), nothing), # default: match A's blocks + (Blocks(k, k), 8), # finer solve tiles than A + (Blocks(32, 16), 16), # non-square A tiling → square solve + (Blocks(n, n), 16), # single-tile A, multi-tile solve + ) + for (name, Asp_sys) in systems, (A_part_cfg, solve_bs) in configs + @testset "$name part=$(A_part_cfg.blocksize) bs=$(solve_bs)" begin + Adense_sys = Matrix(Asp_sys) + DA = distribute(Asp_sys, A_part_cfg) + Db = distribute(b, Db_part) + + F4a = Dagger.splu(DA) + F4b = if solve_bs === nothing + Dagger.splu(DA; distributed=true) + else + Dagger.splu(DA; distributed=true, blocksize=solve_bs) + end + @test F4b isa Dagger.DistributedSparseLU + @test size(F4b) == (n, n) + @test F4b.L isa Dagger.DMatrix + @test F4b.U isa Dagger.DMatrix + # Factors stay sparse (no densification). + @test fetch(F4b.L.chunks[1, 1]) isa Dagger.DSparseArray + @test fetch(F4b.U.chunks[1, 1]) isa Dagger.DSparseArray + + x4a = collect(F4a \ Db) + x = F4b \ Db + @test x isa Dagger.DVector + @test x.partitioning == Db.partitioning + xc = collect(x) + @test xc ≈ x4a rtol=1e-12 atol=1e-14 + @test norm(Adense_sys * xc - b) / norm(b) < 1e-10 + + # Multi-RHS / repeated solves against one factor. + for _ in 1:3 + bi = rand(n) + Dbi = distribute(bi, Db_part) + xi = collect(F4b \ Dbi) + @test xi ≈ collect(F4a \ Dbi) rtol=1e-12 atol=1e-14 + @test norm(Adense_sys * xi - bi) / norm(bi) < 1e-10 + end + + y = similar(Db) + LinearAlgebra.ldiv!(y, F4b, Db) + @test collect(y) ≈ x4a rtol=1e-12 atol=1e-14 + + @test_throws DimensionMismatch F4b \ distribute(rand(n ÷ 2), Blocks(k)) + end + end + end + + @testset "Schur distributed LU (4c)" begin + systems = ( + ("lap2d", N -> lap2d(Float64, N)), + ("advdiff2d", N -> advdiff2d(Float64, N)), + ) + for (name, buildA) in systems, N in (10, 20), nparts in (2, 4) + @testset "$name N=$N nparts=$nparts" begin + Asp_sys = buildA(N) + nsys = size(Asp_sys, 1) + bs = max(8, nsys ÷ 4) + DA = distribute(Asp_sys, Blocks(bs, bs)) + bsys = rand(nsys) + Db = distribute(bsys, Blocks(bs)) + + F4a = Dagger.splu(DA) + F4c = Dagger.splu(DA; distributed=true, method=:schur, nparts=nparts) + @test F4c isa Dagger.DistributedSchurLU + @test size(F4c) == (nsys, nsys) + + x = F4c \ Db + @test x isa Dagger.DVector + @test x.partitioning == Db.partitioning + xc = collect(x) + x4a = collect(F4a \ Db) + @test xc ≈ x4a rtol=1e-8 + @test norm(Asp_sys * xc - bsys) / norm(bsys) < 1e-9 + + y = similar(Db) + LinearAlgebra.ldiv!(y, F4c, Db) + @test collect(y) ≈ x4a rtol=1e-8 + @test norm(Asp_sys * collect(y) - bsys) / norm(bsys) < 1e-9 + + # Factor-once / solve-many. + for _ in 1:3 + bi = rand(nsys) + Dbi = distribute(bi, Blocks(bs)) + xi = collect(F4c \ Dbi) + @test xi ≈ collect(F4a \ Dbi) rtol=1e-8 + @test norm(Asp_sys * xi - bi) / norm(bi) < 1e-9 + end + + @test_throws DimensionMismatch F4c \ distribute(rand(nsys ÷ 2), Blocks(bs)) + end + end + end +end + diff --git a/test/cudaenv/Project.toml b/test/cudaenv/Project.toml index 6e245a12a..38ea2a004 100644 --- a/test/cudaenv/Project.toml +++ b/test/cudaenv/Project.toml @@ -4,6 +4,7 @@ Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [sources] diff --git a/test/gpu.jl b/test/gpu.jl index 474a0bd08..28b908eb6 100644 --- a/test/gpu.jl +++ b/test/gpu.jl @@ -1,11 +1,113 @@ using Random using LinearAlgebra +using SparseArrays +using Krylov @everywhere begin using Distributed, Dagger import Dagger: Kernel using KernelAbstractions end + +# Sparse DArray SpGEMM / SpMV under a GPU compute scope. `check_tile(chunk)` +# verifies a distributed tile is device-resident (backend-specific). +function test_gpu_sparse_darray(scope; check_tile) + SA = sprand(Float32, 8, 8, 0.35) + SB = sprand(Float32, 8, 8, 0.35) + x = rand(Float32, 8) + Dagger.with_options(;scope) do + DSA = distribute(SA, Blocks(4, 4)) + DSB = distribute(SB, Blocks(4, 4)) + for chunk in DSA.chunks + @test check_tile(fetch(chunk; raw=true)) + end + + @test collect(DSA * DSB) ≈ SA * SB + @test collect(DSA * DSB') ≈ SA * SB' + @test collect(DSA' * DSB) ≈ SA' * SB + + DSC = similar(DSA) + mul!(DSC, DSA, DSB) + @test collect(DSC) ≈ SA * SB + + Dx = distribute(x, Blocks(4)) + @test collect(DSA * Dx) ≈ SA * x + + Z = SparseArrays.spzeros(Blocks(4, 4), Float32, 8, 8) + # `collect` densifies; check emptiness on device tiles / dense gather. + @test iszero(sum(abs, collect(Z))) + for chunk in Z.chunks + @test check_tile(fetch(chunk; raw=true)) + end + end +end + +# Sparse Krylov solvers under a GPU scope (mirrors the sparse backend cases in +# `array/linalg/iterativesolvers.jl`, trimmed for GPU CI time). +function _gpu_sparse_laplacian(T, n) + return SparseArrays.spdiagm( + -1 => fill(-one(T), n - 1), + 0 => fill(T(4), n), + 1 => fill(-one(T), n - 1), + ) +end +function _gpu_sparse_advection(T, n) + return _gpu_sparse_laplacian(T, n) + SparseArrays.spdiagm( + -1 => fill(T(-3) / 10, n - 1), + 1 => fill(T(3) / 10, n - 1), + ) +end +function test_gpu_sparse_solvers(scope; check_tile=nothing) + n, k = 32, 8 + A_part, b_part = Blocks(k, k), Blocks(k) + T = Float32 + + Dagger.with_options(;scope) do + Asp = _gpu_sparse_laplacian(T, n) + b = rand(T, n) + xref = Matrix(Asp) \ b + DA = distribute(Asp, A_part) + Db = distribute(b, b_part) + if check_tile !== nothing + @test check_tile(fetch(DA.chunks[1]; raw=true)) + end + + @testset "$(nameof(solver))" for solver in (Dagger.cg, Dagger.minres, Dagger.gmres, Dagger.bicgstab) + x, stats = solver(DA, Db; atol = T(1e-6), rtol = T(1e-5), itmax = 500) + @test stats.solved + @test x isa Dagger.DVector + @test collect(x) ≈ xref rtol = 1e-3 + end + + x, stats = Dagger.krylov_solve(:cg, DA, Db; atol = T(1e-6), rtol = T(1e-5)) + @test stats.solved + @test collect(x) ≈ xref rtol = 1e-3 + + # Diagonal Jacobi is device-friendly. Block-Jacobi uses host UMFPACK + # factors that Datadeps cannot place under a GPU-only compute scope, so + # it is covered on CPU in `array/linalg/iterativesolvers.jl` instead. + P = Dagger.JacobiPreconditioner(DA) + @test collect(P.dinv) ≈ fill(T(1) / T(4), n) + y = similar(Db) + mul!(y, P, Db) + @test collect(y) ≈ (T(1) / T(4)) .* b + x, stats = Dagger.cg(DA, Db; M = P, atol = T(1e-6), rtol = T(1e-5), itmax = 500) + @test stats.solved + @test collect(x) ≈ xref rtol = 1e-3 + + Anonsym = _gpu_sparse_advection(T, n) + bn = rand(T, n) + xrefn = Matrix(Anonsym) \ bn + DAn = distribute(Anonsym, A_part) + Dbn = distribute(bn, b_part) + @testset "$(nameof(solver)) nonsym" for solver in (Dagger.gmres, Dagger.bicgstab) + x, stats = solver(DAn, Dbn; atol = T(1e-6), rtol = T(1e-5), itmax = 500) + @test stats.solved + @test collect(x) ≈ xrefn rtol = 1e-3 + end + end +end + @everywhere begin function isongpu(X) return !(X isa Array) @@ -227,6 +329,19 @@ end end @test collect(Db) ≈ b_ref rtol=1e-5 end + + @testset "Sparse DArray (GPU $gpu)" for gpu in single_gpu_configs + scope = Dagger.scope(worker=1, cuda_gpu=gpu) + CUDAExt = Base.get_extension(Dagger, :CUDAExt) + check_tile = chunk -> begin + v = Dagger.MemPool.poolget(chunk.handle) + return v isa Dagger.DSparseArray && + v.mat isa CUDA.CUSPARSE.CuSparseMatrixCSC && + chunk.space isa CUDAExt.CUDAVRAMMemorySpace + end + test_gpu_sparse_darray(scope; check_tile) + test_gpu_sparse_solvers(scope; check_tile) + end end end @@ -398,6 +513,19 @@ end end @test collect(Db) ≈ b_ref rtol=1e-5 end + + @testset "Sparse DArray (GPU $gpu)" for gpu in single_gpu_configs + scope = Dagger.scope(worker=1, rocm_gpu=gpu) + ROCExt = Base.get_extension(Dagger, :ROCExt) + check_tile = chunk -> begin + v = Dagger.MemPool.poolget(chunk.handle) + return v isa Dagger.DSparseArray && + v.mat isa AMDGPU.rocSPARSE.ROCSparseMatrixCSC && + chunk.space isa ROCExt.ROCVRAMMemorySpace + end + test_gpu_sparse_darray(scope; check_tile) + test_gpu_sparse_solvers(scope; check_tile) + end end end @@ -569,6 +697,19 @@ end end @test collect(Db) ≈ b_ref rtol=1e-5 end + + @testset "Sparse DArray (GPU $gpu)" for gpu in single_gpu_configs + scope = Dagger.scope(worker=1, intel_gpu=gpu) + IntelExt = Base.get_extension(Dagger, :IntelExt) + check_tile = chunk -> begin + v = Dagger.MemPool.poolget(chunk.handle) + return v isa Dagger.DSparseArray && + v.mat isa Dagger.DeviceSparseMatrixCSC && + chunk.space isa IntelExt.IntelVRAMMemorySpace + end + test_gpu_sparse_darray(scope; check_tile) + test_gpu_sparse_solvers(scope; check_tile) + end end end @@ -714,6 +855,19 @@ end @test_broken array[2, 1] == 4.0f0 @test_broken array[2, 2] == 5.0f0 end + + @testset "Sparse DArray" begin + scope = Dagger.scope(worker=1, metal_gpu=1) + MetalExt = Base.get_extension(Dagger, :MetalExt) + check_tile = chunk -> begin + v = Dagger.MemPool.poolget(chunk.handle) + return v isa Dagger.DSparseArray && + v.mat isa Dagger.DeviceSparseMatrixCSC && + chunk.space isa MetalExt.MetalVRAMMemorySpace + end + test_gpu_sparse_darray(scope; check_tile) + test_gpu_sparse_solvers(scope; check_tile) + end end end @@ -820,5 +974,18 @@ end end @test A ≈ ref .+ 1 end + + @testset "Sparse DArray (GPU $gpu)" for gpu in single_gpu_configs + scope = Dagger.scope(worker=1, cl_device=gpu) + OpenCLExt = Base.get_extension(Dagger, :OpenCLExt) + check_tile = chunk -> begin + v = Dagger.MemPool.poolget(chunk.handle) + return v isa Dagger.DSparseArray && + v.mat isa Dagger.DeviceSparseMatrixCSC && + chunk.space isa OpenCLExt.CLMemorySpace + end + test_gpu_sparse_darray(scope; check_tile) + test_gpu_sparse_solvers(scope; check_tile) + end end end diff --git a/test/metalenv/Project.toml b/test/metalenv/Project.toml index b3d5e9632..903bfcb7e 100644 --- a/test/metalenv/Project.toml +++ b/test/metalenv/Project.toml @@ -4,6 +4,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" Metal = "dde4c033-4e86-420c-a63e-0dd931031962" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [sources] diff --git a/test/mpi.jl b/test/mpi.jl index f5e7ffd6f..d8d25921f 100644 --- a/test/mpi.jl +++ b/test/mpi.jl @@ -543,14 +543,17 @@ read_only(X) = (sum(X); nothing) r1 = min(1, nranks-1) cross = r1 != 0 # single-rank runs degenerate to zero copies - # Read-only cross-rank arg: one copy-in, no writeback + # Read-only cross-rank arg: copy-in plus region-end write-back. The + # write-back is redundant (origin still holds the bytes) but is retained + # so write-back planning can stay history-based; `arg_current` elision is + # unsound for Krylov `similar` workspaces under Distributed. A = ones(4, 4) logs = with_logs() do Dagger.spawn_datadeps() do Dagger.@spawn scope=rank_scope(r1) read_only(In(A)) end end - @test count_move_tasks(logs) == (cross ? 1 : 0) + @test count_move_tasks(logs) == (cross ? 2 : 0) @test Dagger.check_uniform(count_move_tasks(logs)) # Same-rank InOut: the origin slot aliases the original, zero copies diff --git a/test/mpienv/Project.toml b/test/mpienv/Project.toml index 71e617954..581efcb0e 100644 --- a/test/mpienv/Project.toml +++ b/test/mpienv/Project.toml @@ -3,6 +3,7 @@ Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [sources] diff --git a/test/oneapienv/Project.toml b/test/oneapienv/Project.toml index 30e7631f3..1c5664279 100644 --- a/test/oneapienv/Project.toml +++ b/test/oneapienv/Project.toml @@ -3,6 +3,7 @@ Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" diff --git a/test/openclenv/Project.toml b/test/openclenv/Project.toml index b7718eca7..8c25145c4 100644 --- a/test/openclenv/Project.toml +++ b/test/openclenv/Project.toml @@ -1,10 +1,13 @@ [deps] Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54" +Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" OpenCL = "08131aa3-fb12-5dee-8b74-c09406e224a2" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +pocl_jll = "627d6b7a-bbe6-5189-83e7-98cc0a5aeadd" [sources] Dagger = {path = "../.."} diff --git a/test/rocmenv/Project.toml b/test/rocmenv/Project.toml index d350b4c4a..7b8ac67f3 100644 --- a/test/rocmenv/Project.toml +++ b/test/rocmenv/Project.toml @@ -1,9 +1,11 @@ [deps] AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54" +Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [sources] diff --git a/test/runtests.jl b/test/runtests.jl index bfc347ce0..d327ee50a 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -7,6 +7,11 @@ USE_METAL = parse(Bool, get(ENV, "CI_USE_METAL", "0")) USE_OPENCL = parse(Bool, get(ENV, "CI_USE_OPENCL", "0")) USE_GPU = USE_CUDA || USE_ROCM || USE_ONEAPI || USE_METAL || USE_OPENCL +# Finch is a heavy optional dependency (precompilation takes minutes), so it is +# not a default test dependency. The dedicated Finch CI job signals via this +# variable to `Pkg.add` Finch and run only the Finch-backed tests. +USE_FINCH = parse(Bool, get(ENV, "CI_TEST_FINCH", "0")) + tests = [ ("Thunk", "thunk.jl"), ("Scheduler", "scheduler.jl"), @@ -32,11 +37,15 @@ tests = [ ("Array - LinearAlgebra - Core", "array/linalg/core.jl"), ("Array - LinearAlgebra - Arithmetic", "array/linalg/arithmetic.jl"), ("Array - LinearAlgebra - Matmul", "array/linalg/matmul.jl"), + ("Array - LinearAlgebra - Matmul (Sparse, quick)", "array/linalg/matmul_sparse.jl"), + ("Array - LinearAlgebra - Matmul (Finch, quick)", "array/linalg/matmul_finch.jl"), ("Array - LinearAlgebra - Cholesky", "array/linalg/cholesky.jl"), ("Array - LinearAlgebra - LU", "array/linalg/lu.jl"), ("Array - LinearAlgebra - Solve", "array/linalg/solve.jl"), ("Array - LinearAlgebra - QR", "array/linalg/qr.jl"), ("Array - LinearAlgebra - SVD", "array/linalg/svd.jl"), + ("Array - LinearAlgebra - Iterative Solvers", "array/linalg/iterativesolvers.jl"), + ("Array - LinearAlgebra - Sparse Direct", "array/linalg/sparsedirect.jl"), ("Array - Permute", "array/permute.jl"), ("Array - Random", "array/random.jl"), ("Array - Stencils", "array/stencil.jl"), @@ -57,8 +66,22 @@ if USE_GPU ("Array - Stencils", "array/stencil.jl"), ] end +if USE_FINCH + # Only run the Finch-backed tests in the dedicated Finch CI job. + tests = [ + ("Array - LinearAlgebra - Matmul (Finch)", "array/linalg/matmul_finch.jl"), + ] +end all_test_names = map(test -> replace(last(test), ".jl"=>""), tests) +# Tests excluded from default runs; they only run when explicitly requested via +# `--test` or enabled by an environment signal. Finch's precompilation is very +# heavy, so its tests are opt-in (except in the dedicated Finch job, where we +# explicitly want them to run). +optin_test_names = USE_FINCH ? String[] : String[ + "array/linalg/matmul_finch", +] + additional_workers::Int = 3 worker_threads::Int = 1 @@ -71,6 +94,10 @@ if PROGRAM_FILE != "" && realpath(PROGRAM_FILE) == @__FILE__ Pkg.instantiate() catch end + if USE_FINCH + # Finch is not a default test dependency; add it on demand when signaled. + Pkg.add("Finch") + end using ArgParse s = ArgParseSettings(description = "Dagger Testsuite") @@ -109,7 +136,7 @@ if PROGRAM_FILE != "" && realpath(PROGRAM_FILE) == @__FILE__ parsed_args = parse_args(s) to_test = String[] if isempty(parsed_args["test"]) - to_test = copy(all_test_names) + to_test = filter(!in(optin_test_names), copy(all_test_names)) else for _test in parsed_args["test"] test = only(_test) @@ -166,7 +193,7 @@ if PROGRAM_FILE != "" && realpath(PROGRAM_FILE) == @__FILE__ Pkg.offline(true) end else - to_test = all_test_names + to_test = filter(!in(optin_test_names), all_test_names) @info "Running all tests" end