Implement sparse arrays for every back-end

Project.toml

@@ -1,6 +1,6 @@
 name = "GPUArrays"
 uuid = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
-version = "12.2.0"
+version = "13.0.0"
 
 [workspace]
 projects = ["lib/GPUArraysCore", "lib/JLArrays", "test", "docs"]

README.md

@@ -47,52 +47,13 @@ without further methods. See the [interface
 documentation](https://juliagpu.github.io/GPUArrays.jl/dev/interface/) and `JLArrays` for a
 complete example.
 
-### Sparse array support (optional)
-
-`GPUArrays.jl` provides **device-side** array types for `CSC`, `CSR`, `COO`, and `BSR` matrices, as well as sparse vectors.
-It also provides abstract types for these layouts that you can create concrete child types of in order to benefit from the
-backend-agnostic wrappers. In particular, `GPUArrays.jl` provides out-of-the-box support for broadcasting and `mapreduce` over
-GPU sparse arrays.
-
-For **host-side** types, your custom sparse types should implement:
-
-- `dense_array_type` - the corresponding dense array type. For example, for a `CuSparseVector` or `CuSparseMatrixCXX`, the `dense_array_type` is `CuArray`
-- `sparse_array_type` - the **untyped** sparse array type corresponding to a given parametrized type. A `CuSparseVector{Tv, Ti}` would have a `sparse_array_type` of `CuSparseVector` -- note the lack of type parameters!
-- `csc_type(::Type{T})` - the compressed sparse column type for your backend. A `CuSparseMatrixCSR` would have a `csc_type` of `CuSparseMatrixCSC`. 
-- `csr_type(::Type{T})` - the compressed sparse row type for your backend. A `CuSparseMatrixCSC` would have a `csr_type` of `CuSparseMatrixCSR`. 
-- `coo_type(::Type{T})` - the coordinate sparse matrix type for your backend. A `CuSparseMatrixCSC` would have a `coo_type` of `CuSparseMatrixCOO`.
-
-Additionally, you need to teach `GPUArrays.jl` how to translate your backend's specific types onto the device. `GPUArrays.jl` provides the device-side types:
-
-- `GPUSparseDeviceVector`
-- `GPUSparseDeviceMatrixCSC`
-- `GPUSparseDeviceMatrixCSR`
-- `GPUSparseDeviceMatrixBSR`
-- `GPUSparseDeviceMatrixCOO`
-
-You will need to create a method of `Adapt.adapt_structure` for each format your backend supports. **Note** that if your backend supports separate address spaces,
-as CUDA and ROCm do, you need to provide a parameter to these device-side arrays to indicate in which address space the underlying pointers live. An example of adapting
-an array to the device-side struct:
-
-```julia
-function GPUArrays.GPUSparseDeviceVector(iPtr::MyDeviceVector{Ti, A},
-                                         nzVal::MyDeviceVector{Tv, A},
-                                         len::Int,
-                                         nnz::Ti) where {Ti, Tv, A}
-    GPUArrays.GPUSparseDeviceVector{Tv, Ti, MyDeviceVector{Ti, A}, MyDeviceVector{Tv, A}, A}(iPtr, nzVal, len, nnz)
-end
-
-function Adapt.adapt_structure(to::MyAdaptor, x::MySparseVector)
-    return GPUArrays.GPUSparseDeviceVector(
-        adapt(to, x.iPtr),
-        adapt(to, x.nzVal),
-        length(x), x.nnz
-    )
-end
-```
-
-You'll also need to inform `GPUArrays.jl` and `GPUCompiler.jl` how to adapt your sparse arrays by extending `KernelAbstractions.jl`'s `get_backend()`:
-
-```julia
-KA.get_backend(::MySparseVector) = MyBackend()
-```  
+### Sparse array support
+
+`GPUArrays.jl` implements sparse vectors and CSR, CSC and COO matrices for every back-end
+(`GPUSparseVector`, `GPUSparseMatrixCSR`, `GPUSparseMatrixCSC`, `GPUSparseMatrixCOO`). The
+types are parametrized on the dense vectors that store their indices and values, so they
+work with any back-end's arrays, and the same structs are used inside kernels. A back-end
+needs no sparse-specific code; it can add aliases for its storage (e.g.
+`const MtlSparseMatrixCSR{Tv,Ti} = GPUSparseMatrixCSR{Tv,Ti,<:MtlVector{Ti},<:MtlVector{Tv}}`),
+a transfer policy for its adaptor, and vendor library overrides. See the
+[interface documentation](https://juliagpu.github.io/GPUArrays.jl/dev/interface/) for details.

docs/src/interface.md

@@ -42,6 +42,199 @@ precise:
 GPUArrays.memory_location
 ```
 
+## Sparse arrays
+
+GPUArrays implements the sparse array formats itself, for every back-end. There are four,
+all of them subtypes of `AbstractGPUSparseArray{Tv,Ti,N}` (which subtypes SparseArrays'
+`AbstractSparseArray`):
+
+| type | layout | fields |
+|:--|:--|:--|
+| `GPUSparseVector{Tv,Ti,Vi,Vv}` | like `SparseVector` | `nzInd`, `nzVal`, `len` |
+| `GPUSparseMatrixCSC{Tv,Ti,Vi,Vv}` | compressed sparse columns, like `SparseMatrixCSC` | `colPtr`, `rowVal`, `nzVal`, `dims` |
+| `GPUSparseMatrixCSR{Tv,Ti,Vi,Vv}` | compressed sparse rows | `rowPtr`, `colVal`, `nzVal`, `dims` |
+| `GPUSparseMatrixCOO{Tv,Ti,Vi,Vv}` | coordinates, sorted in row-major order | `rowInd`, `colInd`, `nzVal`, `dims` |
+
+`Tv` is the element type and `Ti` the index type. `Vi` and `Vv` are the types of the
+dense vectors that store the indices and values, typically the back-end's GPU vectors
+(`MtlVector{Int32,Metal.PrivateStorage}`, `CuVector{Float32}`, ...). They are part of the
+type so that generic code can allocate results with `similar` on the buffers it already
+has, and so that methods can dispatch on the storage. A matrix is constructed from its
+buffers, with the dimensions last:
+
+```julia
+A = GPUSparseMatrixCSR(rowPtr, colVal, nzVal, (m, n))
+```
+
+The constructors only check what is known on the host (buffer lengths, and whether the
+dimensions and number of stored entries fit the index type). Like SparseArrays, the
+formats are *canonical*: the indices within a row (CSR), column (CSC) or vector are sorted
+and unique, and COO entries are sorted by row and then column without duplicates.
+Explicitly stored zeros are allowed. Every operation relies on these invariants and
+preserves them; `GPUArrays.check_structure(A)` verifies them on the device.
+
+The structs are immutable, and the same struct is used inside kernels: `adapt` converts
+each buffer, so a kernel that receives a `GPUSparseMatrixCSR` can read `A.rowPtr`,
+`A.colVal` and `A.nzVal` directly, and index it with `A[i, j]`.
+
+The constructors above borrow their buffers, like SparseArrays' raw constructors; the
+allocating constructors (from other arrays), conversions and `copy` return arrays with
+their own storage. Operations that change the structure of an existing array (`copyto!`
+between sparse arrays, `dropzeros!`, `fkeep!`, in-place broadcasting) resize its buffers,
+which requires resizable dense vectors (not views) that no other object shares. Views do
+work for reading, and for updating the values through `nonzeros(A)` or a zero-preserving
+in-place broadcast.
+
+### What a back-end provides
+
+Nothing sparse-specific is required. The generic implementation relies on the back-end's
+dense support: its vector type, with `similar`, `copyto!` from and to the host, `resize!`,
+a KernelAbstractions back-end, and implementations of `sortperm` (at least for `Int64`
+keys), `accumulate!` and `findall` that run on the device.
+
+A back-end will usually add the following, all optional:
+
+1. **Aliases** for its storage, for naming and dispatch, with constructors on them:
+
+   ```julia
+   const MtlSparseMatrixCSR{Tv,Ti} = GPUSparseMatrixCSR{Tv,Ti,<:MtlVector{Ti},<:MtlVector{Tv}}
+   MtlSparseMatrixCSR(A::AbstractArray) = GPUSparseMatrixCSR(adapt(MtlArray, A))
+   ```
+
+   Constructors must be defined on the alias itself (including the `{Tv}` and `{Tv,Ti}`
+   forms): methods of the generic type do not apply to a constrained alias.
+2. **A transfer policy** for its opinionated adaptor (`mtl`, `cu`), one method per host
+   format on `adapt_structure`, e.g. producing CSR matrices with `Int32` indices and
+   narrowing the values as the dense adaptor does. Without it, those adaptors treat a host
+   sparse array like any other array.
+3. **A dense constructor** that densifies on the device, e.g.
+   `MtlArray{T,2}(A::MtlSparseMatrixCSR) = copyto!(similar(nonzeros(A), T, size(A)), A)`.
+   Otherwise the generic dense constructor copies through the host.
+4. **Vendor library overrides**, as methods of the public functions that are more specific
+   in storage, element and index type than the generic ones. When a vendor routine does
+   not apply, such a method can call the generic implementation, which never dispatches
+   back to the public function:
+
+   | operation | public function | generic implementation |
+   |:--|:--|:--|
+   | sparse × dense vector | `mul!(y, tA, A, x, α, β)` | `GPUArrays.generic_spmv!` |
+   | sparse × dense, dense × sparse | `mul!(C, tA, tB, A, B, α, β)` | `GPUArrays.generic_spmm!` |
+   | sparse × sparse | `*`, `mul!(C, tA, tB, A, B, α, β)` with a sparse `C` | `GPUArrays.generic_spgemm`, `GPUArrays.generic_spgemm!` |
+   | CSR ↔ CSC | the target format's constructor | `GPUArrays.generic_regroup` |
+   | CSR ↔ COO | the target format's constructor | `GPUArrays.generic_expand`, `GPUArrays.generic_compress` |
+   | assembly from coordinates | `sparse(I, J, V, m, n, combine)`, `sparsevec` | `GPUArrays.generic_assemble` |
+
+   The product methods are LinearAlgebra's storage-level `mul!`, which it calls on Julia
+   1.13 and later with a character per operand ('N', 'T', 'C', or 'S'/'s'/'H'/'h' for the
+   upper or lower triangle of a Symmetric or Hermitian matrix); on older versions GPUArrays
+   forwards LinearAlgebra's internal `generic_matvecmul!` and `generic_matmatmul!` to
+   them. A vendor method looks like:
+
+   ```julia
+   function LinearAlgebra.mul!(y::CuVector{T}, tA::AbstractChar, A::CuSparseMatrixCSR{T},
+                               x::DenseCuVector{T}, α::Number, β::Number) where {T<:BlasFloat}
+       tA in ('N', 'T', 'C') || return GPUArrays.generic_spmv!(y, tA, A, x, α, β)
+       # call the vendor library
+   end
+   ```
+
+   The generic implementations make no promise of bitwise reproducibility across versions
+   or back-ends.
+5. **Formats that GPUArrays does not implement** (block-sparse formats, batched matrices)
+   as back-end structs subtyping `AbstractGPUSparseArray`, with constructors to and from
+   the generic formats.
+
+### Allocation and conversion
+
+Generic code never names a back-end type; it allocates through the buffers it has:
+
+| need | idiom |
+|:--|:--|
+| a dense array on the same device and storage | `similar(nonzeros(A), T, dims)` |
+| a new index or value buffer | `similar(A.colVal, Ti, n)`, `similar(A.nzVal, Tv, n)` |
+| an empty matrix of the same format | `similar(A, Tv, (m, n))` |
+
+`similar` follows SparseArrays: `similar(A)`, `similar(A, Tv)` and `similar(A, Tv, Ti)`
+copy the structure and leave the values uninitialized, `similar(A, Tv, dims)` returns an
+array without stored entries (a matrix of the same format for two dimensions, a vector for
+one, a dense array for more), and `zero(A)` is an empty array of the same format.
+
+Conversions:
+
+- between formats (on the device), with the target format's constructor: `GPUSparseMatrixCSC(A)`, and
+  `GPUSparseMatrixCSC{Tv,Ti}(A)` to also change the element and index types;
+  `convert(GPUSparseMatrixCSC, A)` returns `A` itself if it already has that format;
+- to the host, with `SparseMatrixCSC(A)`, `SparseVector(x)` or `Array(A)`;
+- from a dense GPU array, dropping its zeros, with `sparse(A; fmt=:csc)` or the format's
+  constructor (`GPUSparseMatrixCSR(A)`); the index type defaults to `Int`;
+- to a dense array on the device, with `copyto!(similar(nonzeros(A), T, size(A)), A)`;
+- transposes: `transpose(A)` and `adjoint(A)` stay lazy, `copy(transpose(A))` and
+  `permutedims(A)` materialize them in the same format, and a format constructor
+  (`GPUSparseMatrixCSC(transpose(A))`) in another one;
+- between host and device, with `adapt`. `adapt(MtlArray, S)` and `adapt(Array, A)` move
+  the storage and keep the format, element and index type (a CSC matrix stays CSC, and
+  comes back as a `SparseMatrixCSC`); adapting never densifies. `adapt(T, S)` with a
+  sparse type `T` calls the constructor `T(S)`.
+
+The generic types cannot be constructed from host arrays, because they do not determine
+where to store the result: use a back-end alias, `adapt`, or the back-end's adaptor.
+
+These are the defaults of the generic implementations. A back-end can choose other index
+types for the constructors of its aliases, for `sparse` of its dense arrays and for its
+adaptor: CUDA.jl and AMDGPU.jl keep using `Cint`, as they did before GPUArrays 13.
+
+### Supported operations
+
+Beyond construction, transfer and conversion, the formats support products with dense
+and sparse arrays (including `Transpose`, `Adjoint`, `Symmetric` and `Hermitian`
+wrappers, `Diagonal` and `UniformScaling`), broadcasting and `map`, reductions, `norm`
+and `opnorm`, `dot`, `kron`, `exp`, indexing with scalars, ranges and logical masks,
+`hcat` and `vcat`, `==` (comparing values, so stored zeros do not matter), `fkeep!` and
+`dropzeros`, `triu`/`tril` and the structural predicates, `diag` and `tr`, `reshape`
+and `vec`, and assembly from coordinates with `sparse` and `sparsevec`.
+
+Some semantics differ from SparseArrays:
+
+- a broadcast that does not preserve zeros, or over dense arguments, returns a dense GPU
+  array rather than a sparse one, and cannot be stored into a sparse array in place;
+- a broadcast cannot combine sparse vectors with sparse matrices, or expand a sparse
+  operand along a dimension of length one;
+- `reshape` copies instead of sharing storage, and individual entries cannot be set;
+- `sparse(I, J, V, m, n, combine)` combines repeated coordinates in input order, but
+  without a defined parenthesization: `combine` must be associative, though it need not
+  commute. Floating-point results may differ from the host and between back-ends;
+- `hash` and `blockdiag` are not implemented, and neither are triangular solves
+  without a vendor library.
+
+### Migrating from GPUArrays 12
+
+GPUArrays 12 only defined abstract sparse types, which back-ends subtyped, and hooks
+(`dense_array_type`, `sparse_array_type`, `csc_type`, `csr_type`, `coo_type`,
+`_sptranspose`, `_spadjoint`) and device types (`GPUSparseDeviceVector`, ...) that they
+implemented. All of these are gone. `AbstractGPUSparseMatrixCSR`,
+`AbstractGPUSparseMatrixCSC` and `AbstractGPUSparseMatrixCOO` remain as deprecated names
+for the concrete formats, so that methods dispatching on them keep working, but they can
+no longer be subtyped. The `nnz` field is replaced by the `nnz` function, and the index
+field of sparse vectors is `nzInd`.
+
+### Reference
+
+```@docs
+GPUArrays.AbstractGPUSparseArray
+GPUArrays.GPUSparseMatrixCSR
+GPUArrays.GPUSparseMatrixCSC
+GPUArrays.GPUSparseMatrixCOO
+GPUArrays.GPUSparseVector
+GPUArrays.check_structure
+GPUArrays.generic_regroup
+GPUArrays.generic_expand
+GPUArrays.generic_assemble
+GPUArrays.generic_spmv!
+GPUArrays.generic_spmm!
+GPUArrays.generic_spgemm
+GPUArrays.generic_spgemm!
+```
+
 ## Caching Allocator
 
 ```@docs

lib/JLArrays/Project.toml

@@ -14,7 +14,7 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [compat]
 Adapt = "2.0, 3.0, 4.0"
-GPUArrays = "12.1"
+GPUArrays = "13"
 KernelAbstractions = "0.10.5"
 KernelInterface = "0.4.2"
 LinearAlgebra = "1"

lib/JLArrays/src/JLArrays.jl

@@ -6,15 +6,16 @@
 
 module JLArrays
 
-export JLArray, JLVector, JLMatrix, jl, JLBackend, JLSparseVector, JLSparseMatrixCSC, JLSparseMatrixCSR
+export JLArray, JLVector, JLMatrix, jl, JLBackend, JLSparseVector, JLSparseMatrixCSC, JLSparseMatrixCSR,
+       JLSparseMatrixCOO
 
 using GPUArrays
 
 using Adapt
 using SparseArrays, LinearAlgebra
 using Random
 
-import GPUArrays: dense_array_type
+using GPUArrays: GPUSparseVector, GPUSparseMatrixCSR, GPUSparseMatrixCSC, GPUSparseMatrixCOO
 
 import KernelAbstractions: KernelAbstractions, POCL
 

lib/JLArrays/src/JLKernels.jl

@@ -18,7 +18,6 @@ export JLBackend
 struct JLBackend <: KI.Backend end
 
 KI.get_backend(a::JLA) where JLA <: JLArray = JLBackend()
-KI.get_backend(a::JLA) where JLA <: Union{JLSparseMatrixCSC, JLSparseMatrixCSR, JLSparseVector} = JLBackend()
 
 KI.allocate(::JLBackend, ::Type{T}, dims::Tuple) where T = JLArray{T}(undef, dims)
 

lib/JLArrays/src/array.jl

@@ -52,10 +52,6 @@ GPUArrays.storage(a::JLArray) = a.data
 
 GPUArrays.memory_location(a::JLArray) = (UInt(pointer(a.data[])), a.offset)
 
-GPUArrays.dense_array_type(a::JLArray{T, N}) where {T, N} = JLArray{T, N}
-GPUArrays.dense_array_type(::Type{JLArray{T, N}}) where {T, N} = JLArray{T, N}
-GPUArrays.dense_vector_type(a::JLArray{T, N}) where {T, N} = JLArray{T, 1}
-GPUArrays.dense_vector_type(::Type{JLArray{T, N}}) where {T, N} = JLArray{T, 1}
 
 # the selector byte of element `i` of an array of isbits unions
 typetagdata(a::JLArray, i=1) = pointer(a.data[]) + a.tagoffset + i - 1