From 462e31cca684bdb664d5c4f0be29ad0d94b1be8b Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Fri, 25 Sep 2026 17:24:28 +0200 Subject: [PATCH 01/11] Replace the sparse back-end hooks with storage-parametric types GPUArrays now implements the sparse formats itself: GPUSparseVector and GPUSparseMatrixCSR/CSC/COO are immutable structs parametrized on the dense vectors that store their indices and values. Generic code allocates results with similar() on the buffers it already has, so the trait hooks (dense_array_type, sparse_array_type, csc_type, csr_type, coo_type, _sptranspose, _spadjoint) are gone, and a back-end needs no sparse-specific code: it can alias the types for its storage and add a transfer policy for its adaptor. The same struct is used on the device, replacing the GPUSparseDevice* types and the per-back-end adapt rules. The formats are canonical, as in SparseArrays (COO in row-major order). Constructors check host metadata only; check_structure verifies the invariants on the device. similar, copy, copyto!, zero and adapt follow SparseArrays and Adapt semantics: copies and conversions never share buffers, and adapting to a storage type keeps the format and types. Conversions between all formats run on the device (generic_regroup for CSR <-> CSC, generic_expand/generic_compress for CSR <-> COO). Broadcasting and reductions are ported to the new types. Broadcasting no longer reuses the input's pointer buffer for its output, and handles matrices without rows or columns. AbstractGPUSparseMatrixCSR/CSC/COO remain as deprecated aliases of the concrete types for downstream dispatch. JLArrays keeps only aliases and constructors on them. The testsuite builds its inputs from dense buffers and covers every format for every back-end. The raw constructors borrow their buffers, like SparseArrays' raw constructors. Operations that change the structure of an array resize its buffers, and check that they can be resized before modifying anything, so arrays backed by views remain usable for reading and for updating values. dataids covers every buffer, so that copyto! and broadcasts unalias their arguments. Scalar indexing reads the buffers, so it requires allowscalar exactly when the storage does. --- README.md | 59 +- docs/src/interface.md | 125 ++++ lib/JLArrays/src/JLArrays.jl | 5 +- lib/JLArrays/src/JLKernels.jl | 1 - lib/JLArrays/src/array.jl | 4 - lib/JLArrays/src/sparse.jl | 240 +------- src/GPUArrays.jl | 10 +- src/device/sparse.jl | 171 ------ src/host/sparse.jl | 1028 -------------------------------- src/host/sparse/broadcast.jl | 614 +++++++++++++++++++ src/host/sparse/conversions.jl | 304 ++++++++++ src/host/sparse/indexing.jl | 45 ++ src/host/sparse/primitives.jl | 106 ++++ src/host/sparse/reductions.jl | 136 +++++ src/host/sparse/show.jl | 37 ++ src/host/sparse/types.jl | 445 ++++++++++++++ test/runtests.jl | 19 +- test/testsuite.jl | 3 - test/testsuite/sparse.jl | 882 +++++++++++++++++---------- 19 files changed, 2437 insertions(+), 1797 deletions(-) delete mode 100644 src/device/sparse.jl delete mode 100644 src/host/sparse.jl create mode 100644 src/host/sparse/broadcast.jl create mode 100644 src/host/sparse/conversions.jl create mode 100644 src/host/sparse/indexing.jl create mode 100644 src/host/sparse/primitives.jl create mode 100644 src/host/sparse/reductions.jl create mode 100644 src/host/sparse/show.jl create mode 100644 src/host/sparse/types.jl diff --git a/README.md b/README.md index ebd08fbe3..0b3ee29c5 100644 --- a/README.md +++ b/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. diff --git a/docs/src/interface.md b/docs/src/interface.md index b45ae9767..9bb7edc96 100644 --- a/docs/src/interface.md +++ b/docs/src/interface.md @@ -42,6 +42,131 @@ 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 | + |:--|:--|:--| + | CSR ↔ CSC | the target format's constructor | `GPUArrays.generic_regroup` | + | CSR ↔ COO | the target format's constructor | `GPUArrays.generic_expand`, `GPUArrays.generic_compress` | + +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)`; +- to a dense array on the device, with `copyto!(similar(nonzeros(A), T, size(A)), A)`; +- 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. + +### Reference + +```@docs +GPUArrays.AbstractGPUSparseArray +GPUArrays.GPUSparseMatrixCSR +GPUArrays.GPUSparseMatrixCSC +GPUArrays.GPUSparseMatrixCOO +GPUArrays.GPUSparseVector +GPUArrays.check_structure +GPUArrays.generic_regroup +GPUArrays.generic_expand +``` + ## Caching Allocator ```@docs diff --git a/lib/JLArrays/src/JLArrays.jl b/lib/JLArrays/src/JLArrays.jl index 2d9fd9bde..6e55579ff 100644 --- a/lib/JLArrays/src/JLArrays.jl +++ b/lib/JLArrays/src/JLArrays.jl @@ -6,7 +6,8 @@ module JLArrays -export JLArray, JLVector, JLMatrix, jl, JLBackend, JLSparseVector, JLSparseMatrixCSC, JLSparseMatrixCSR +export JLArray, JLVector, JLMatrix, jl, JLBackend, JLSparseVector, JLSparseMatrixCSC, JLSparseMatrixCSR, + JLSparseMatrixCOO using GPUArrays @@ -14,7 +15,7 @@ using Adapt using SparseArrays, LinearAlgebra using Random -import GPUArrays: dense_array_type +using GPUArrays: GPUSparseVector, GPUSparseMatrixCSR, GPUSparseMatrixCSC, GPUSparseMatrixCOO import KernelAbstractions: KernelAbstractions, POCL diff --git a/lib/JLArrays/src/JLKernels.jl b/lib/JLArrays/src/JLKernels.jl index 44f68cbb7..8e37d93c7 100644 --- a/lib/JLArrays/src/JLKernels.jl +++ b/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) diff --git a/lib/JLArrays/src/array.jl b/lib/JLArrays/src/array.jl index d0cd68c1e..249a93d6a 100644 --- a/lib/JLArrays/src/array.jl +++ b/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 diff --git a/lib/JLArrays/src/sparse.jl b/lib/JLArrays/src/sparse.jl index 2acca591c..add207f5a 100644 --- a/lib/JLArrays/src/sparse.jl +++ b/lib/JLArrays/src/sparse.jl @@ -1,224 +1,22 @@ -# sparse array types - -mutable struct JLSparseVector{Tv, Ti} <: GPUArrays.AbstractGPUSparseVector{Tv, Ti} - iPtr::JLArray{Ti, 1} - nzVal::JLArray{Tv, 1} - len::Int - nnz::Ti - - function JLSparseVector{Tv, Ti}(iPtr::JLArray{<:Integer, 1}, nzVal::JLArray{Tv, 1}, - len::Integer) where {Tv, Ti <: Integer} - new{Tv, Ti}(iPtr, nzVal, len, length(nzVal)) - end -end -SparseArrays.nnz(x::JLSparseVector) = x.nnz -SparseArrays.nonzeroinds(x::JLSparseVector) = x.iPtr -SparseArrays.nonzeros(x::JLSparseVector) = x.nzVal - -mutable struct JLSparseMatrixCSC{Tv, Ti} <: GPUArrays.AbstractGPUSparseMatrixCSC{Tv, Ti} - colPtr::JLArray{Ti, 1} - rowVal::JLArray{Ti, 1} - nzVal::JLArray{Tv, 1} - dims::NTuple{2,Int} - nnz::Ti - - function JLSparseMatrixCSC{Tv, Ti}(colPtr::JLArray{<:Integer, 1}, rowVal::JLArray{<:Integer, 1}, - nzVal::JLArray{Tv, 1}, dims::NTuple{2,<:Integer}) where {Tv, Ti <: Integer} - new{Tv, Ti}(colPtr, rowVal, nzVal, dims, length(nzVal)) - end -end -function JLSparseMatrixCSC(colPtr::JLArray{Ti, 1}, rowVal::JLArray{Ti, 1}, nzVal::JLArray{Tv, 1}, dims::NTuple{2,<:Integer}) where {Tv, Ti <: Integer} - return JLSparseMatrixCSC{Tv, Ti}(colPtr, rowVal, nzVal, dims) -end -SparseArrays.SparseMatrixCSC(x::JLSparseMatrixCSC) = SparseMatrixCSC(size(x)..., Array(x.colPtr), Array(x.rowVal), Array(x.nzVal)) - -JLSparseMatrixCSC(A::JLSparseMatrixCSC) = A - -function Base.getindex(A::JLSparseMatrixCSC{Tv, Ti}, i::Integer, j::Integer) where {Tv, Ti} - @boundscheck checkbounds(A, i, j) - r1 = Int(@inbounds A.colPtr[j]) - r2 = Int(@inbounds A.colPtr[j+1]-1) - (r1 > r2) && return zero(Tv) - r1 = searchsortedfirst(view(A.rowVal, r1:r2), i) + r1 - 1 - ((r1 > r2) || (A.rowVal[r1] != i)) ? zero(Tv) : A.nzVal[r1] -end - -mutable struct JLSparseMatrixCSR{Tv, Ti} <: GPUArrays.AbstractGPUSparseMatrixCSR{Tv, Ti} - rowPtr::JLArray{Ti, 1} - colVal::JLArray{Ti, 1} - nzVal::JLArray{Tv, 1} - dims::NTuple{2,Int} - nnz::Ti - - function JLSparseMatrixCSR{Tv, Ti}(rowPtr::JLArray{<:Integer, 1}, colVal::JLArray{<:Integer, 1}, - nzVal::JLArray{Tv, 1}, dims::NTuple{2,<:Integer}) where {Tv, Ti<:Integer} - new{Tv, Ti}(rowPtr, colVal, nzVal, dims, length(nzVal)) - end -end -function JLSparseMatrixCSR(rowPtr::JLArray{Ti, 1}, colVal::JLArray{Ti, 1}, nzVal::JLArray{Tv, 1}, dims::NTuple{2,<:Integer}) where {Tv, Ti <: Integer} - return JLSparseMatrixCSR{Tv, Ti}(rowPtr, colVal, nzVal, dims) -end -function SparseArrays.SparseMatrixCSC(x::JLSparseMatrixCSR) - x_transpose = SparseMatrixCSC(size(x, 2), size(x, 1), Array(x.rowPtr), Array(x.colVal), Array(x.nzVal)) - return SparseMatrixCSC(transpose(x_transpose)) -end - -JLSparseMatrixCSC(Mat::Union{Transpose{Tv, <:SparseMatrixCSC}, Adjoint{Tv, <:SparseMatrixCSC}}) where {Tv} = JLSparseMatrixCSC(JLSparseMatrixCSR(Mat)) - -function Base.size(g::JLSparseMatrixCSR, d::Integer) - if 1 <= d <= 2 - return g.dims[d] - elseif d > 1 - return 1 - else - throw(ArgumentError("dimension must be ≥ 1, got $d")) - end -end - -JLSparseMatrixCSR(Mat::Transpose{Tv, <:SparseMatrixCSC}) where {Tv} = - JLSparseMatrixCSR(JLVector{Cint}(parent(Mat).colptr), JLVector{Cint}(parent(Mat).rowval), - JLVector(parent(Mat).nzval), size(Mat)) -JLSparseMatrixCSR(Mat::Adjoint{Tv, <:SparseMatrixCSC}) where {Tv} = - JLSparseMatrixCSR(JLVector{Cint}(parent(Mat).colptr), JLVector{Cint}(parent(Mat).rowval), - JLVector(conj.(parent(Mat).nzval)), size(Mat)) - -JLSparseMatrixCSR(A::JLSparseMatrixCSR) = A - -function Base.getindex(A::JLSparseMatrixCSR{Tv, Ti}, i0::Integer, i1::Integer) where {Tv, Ti} - @boundscheck checkbounds(A, i0, i1) - c1 = Int(A.rowPtr[i0]) - c2 = Int(A.rowPtr[i0+1]-1) - (c1 > c2) && return zero(Tv) - c1 = searchsortedfirst(A.colVal, i1, c1, c2, Base.Order.Forward) - (c1 > c2 || A.colVal[c1] != i1) && return zero(Tv) - nonzeros(A)[c1] -end - -GPUArrays.sparse_array_type(sa::JLSparseMatrixCSC) = JLSparseMatrixCSC -GPUArrays.sparse_array_type(::Type{<:JLSparseMatrixCSC}) = JLSparseMatrixCSC -GPUArrays.sparse_array_type(sa::JLSparseMatrixCSR) = JLSparseMatrixCSR -GPUArrays.sparse_array_type(::Type{<:JLSparseMatrixCSR}) = JLSparseMatrixCSR -GPUArrays.sparse_array_type(sa::JLSparseVector) = JLSparseVector -GPUArrays.sparse_array_type(::Type{<:JLSparseVector}) = JLSparseVector - -GPUArrays.dense_array_type(sa::JLSparseVector) = JLArray -GPUArrays.dense_array_type(::Type{<:JLSparseVector}) = JLArray -GPUArrays.dense_array_type(sa::JLSparseMatrixCSC) = JLArray -GPUArrays.dense_array_type(::Type{<:JLSparseMatrixCSC}) = JLArray -GPUArrays.dense_array_type(sa::JLSparseMatrixCSR) = JLArray -GPUArrays.dense_array_type(::Type{<:JLSparseMatrixCSR}) = JLArray - -GPUArrays.csc_type(::Type{<:JLSparseMatrixCSR}) = JLSparseMatrixCSC -GPUArrays.csr_type(::Type{<:JLSparseMatrixCSC}) = JLSparseMatrixCSR - -Base.similar(Mat::JLSparseMatrixCSR) = JLSparseMatrixCSR(copy(Mat.rowPtr), copy(Mat.colVal), similar(nonzeros(Mat)), size(Mat)) -Base.similar(Mat::JLSparseMatrixCSR, T::Type) = JLSparseMatrixCSR(copy(Mat.rowPtr), copy(Mat.colVal), similar(nonzeros(Mat), T), size(Mat)) - -Base.similar(Mat::JLSparseMatrixCSC, T::Type, N::Int, M::Int) = JLSparseMatrixCSC(JLVector([zero(Int32)]), JLVector{Int32}(undef, 0), JLVector{T}(undef, 0), (N, M)) -Base.similar(Mat::JLSparseMatrixCSR, T::Type, N::Int, M::Int) = JLSparseMatrixCSR(JLVector([zero(Int32)]), JLVector{Int32}(undef, 0), JLVector{T}(undef, 0), (N, M)) - -Base.similar(Mat::JLSparseMatrixCSC{Tv, Ti}, N::Int, M::Int) where {Tv, Ti} = similar(Mat, Tv, N, M) -Base.similar(Mat::JLSparseMatrixCSR{Tv, Ti}, N::Int, M::Int) where {Tv, Ti} = similar(Mat, Tv, N, M) - -Base.similar(Mat::JLSparseMatrixCSC, T::Type, dims::Tuple{Int, Int}) = similar(Mat, T, dims...) -Base.similar(Mat::JLSparseMatrixCSR, T::Type, dims::Tuple{Int, Int}) = similar(Mat, T, dims...) - -Base.similar(Mat::JLSparseMatrixCSC, dims::Tuple{Int, Int}) = similar(Mat, dims...) -Base.similar(Mat::JLSparseMatrixCSR, dims::Tuple{Int, Int}) = similar(Mat, dims...) - -JLArray(x::JLSparseVector) = JLArray(collect(SparseVector(x))) -JLArray(x::JLSparseMatrixCSC) = JLArray(collect(SparseMatrixCSC(x))) -JLArray(x::JLSparseMatrixCSR) = JLArray(collect(SparseMatrixCSC(x))) - - -## interop with SparseArrays - -function JLSparseVector(xs::SparseVector{Tv, Ti}) where {Ti, Tv} - iPtr = JLVector{Ti}(undef, length(xs.nzind)) - nzVal = JLVector{Tv}(undef, length(xs.nzval)) - copyto!(iPtr, convert(Vector{Ti}, xs.nzind)) - copyto!(nzVal, convert(Vector{Tv}, xs.nzval)) - return JLSparseVector{Tv, Ti}(iPtr, nzVal, length(xs),) -end -Base.length(x::JLSparseVector) = x.len -Base.size(x::JLSparseVector) = (x.len,) - -function JLSparseMatrixCSC(xs::SparseMatrixCSC{Tv, Ti}) where {Ti, Tv} - colPtr = JLVector{Ti}(undef, length(xs.colptr)) - rowVal = JLVector{Ti}(undef, length(xs.rowval)) - nzVal = JLVector{Tv}(undef, length(xs.nzval)) - copyto!(colPtr, convert(Vector{Ti}, xs.colptr)) - copyto!(rowVal, convert(Vector{Ti}, xs.rowval)) - copyto!(nzVal, convert(Vector{Tv}, xs.nzval)) - return JLSparseMatrixCSC{Tv, Ti}(colPtr, rowVal, nzVal, (xs.m, xs.n)) -end -JLSparseMatrixCSC(xs::SparseVector) = JLSparseMatrixCSC(SparseMatrixCSC(xs)) - -Base.length(x::JLSparseMatrixCSC) = prod(x.dims) -Base.size(x::JLSparseMatrixCSC) = x.dims - -function JLSparseMatrixCSR(xs::SparseMatrixCSC{Tv, Ti}) where {Ti, Tv} - csr_xs = SparseMatrixCSC(transpose(xs)) - rowPtr = JLVector{Ti}(undef, length(csr_xs.colptr)) - colVal = JLVector{Ti}(undef, length(csr_xs.rowval)) - nzVal = JLVector{Tv}(undef, length(csr_xs.nzval)) - copyto!(rowPtr, convert(Vector{Ti}, csr_xs.colptr)) - copyto!(colVal, convert(Vector{Ti}, csr_xs.rowval)) - copyto!(nzVal, convert(Vector{Tv}, csr_xs.nzval)) - return JLSparseMatrixCSR{Tv, Ti}(rowPtr, colVal, nzVal, (xs.m, xs.n)) -end -JLSparseMatrixCSR(xs::SparseVector{Tv, Ti}) where {Ti, Tv} = JLSparseMatrixCSR(SparseMatrixCSC(xs)) -function JLSparseMatrixCSR(xs::JLSparseMatrixCSC{Tv, Ti}) where {Ti, Tv} - return JLSparseMatrixCSR(SparseMatrixCSC(xs)) -end -function JLSparseMatrixCSC(xs::JLSparseMatrixCSR{Tv, Ti}) where {Ti, Tv} - return JLSparseMatrixCSC(SparseMatrixCSC(xs)) -end -function Base.copyto!(dst::JLSparseMatrixCSR, src::JLSparseMatrixCSR) - if size(dst) != size(src) - throw(ArgumentError("Inconsistent Sparse Matrix size")) +# sparse arrays with JLArray storage: GPUArrays implements the formats, JLArrays only names +# them and adds constructors + +const JLSparseVector{Tv,Ti} = GPUSparseVector{Tv,Ti,<:JLVector{Ti},<:JLVector{Tv}} +const JLSparseMatrixCSR{Tv,Ti} = GPUSparseMatrixCSR{Tv,Ti,<:JLVector{Ti},<:JLVector{Tv}} +const JLSparseMatrixCSC{Tv,Ti} = GPUSparseMatrixCSC{Tv,Ti,<:JLVector{Ti},<:JLVector{Tv}} +const JLSparseMatrixCOO{Tv,Ti} = GPUSparseMatrixCOO{Tv,Ti,<:JLVector{Ti},<:JLVector{Tv}} +const JLSparseMatrix = Union{JLSparseMatrixCSR,JLSparseMatrixCSC,JLSparseMatrixCOO} + +# move the input to JLArray storage (keeping its format), then convert on the device +for (alias, S) in ((:JLSparseVector, :GPUSparseVector), (:JLSparseMatrixCSR, :GPUSparseMatrixCSR), + (:JLSparseMatrixCSC, :GPUSparseMatrixCSC), (:JLSparseMatrixCOO, :GPUSparseMatrixCOO)) + @eval begin + $alias(A::AbstractArray) = $S(adapt(JLArray, A)) + $alias{Tv}(A::AbstractArray) where {Tv} = $S{Tv}(adapt(JLArray, A)) + $alias{Tv,Ti}(A::AbstractArray) where {Tv,Ti} = $S{Tv,Ti}(adapt(JLArray, A)) end - resize!(dst.rowPtr, length(src.rowPtr)) - resize!(dst.colVal, length(src.colVal)) - resize!(SparseArrays.nonzeros(dst), length(SparseArrays.nonzeros(src))) - copyto!(dst.rowPtr, src.rowPtr) - copyto!(dst.colVal, src.colVal) - copyto!(SparseArrays.nonzeros(dst), SparseArrays.nonzeros(src)) - dst.nnz = src.nnz - dst -end -Base.length(x::JLSparseMatrixCSR) = prod(x.dims) -Base.size(x::JLSparseMatrixCSR) = x.dims - -function GPUArrays._spadjoint(A::JLSparseMatrixCSR) - Aᴴ = JLSparseMatrixCSC(A.rowPtr, A.colVal, conj(A.nzVal), reverse(size(A))) - JLSparseMatrixCSR(Aᴴ) -end -function GPUArrays._sptranspose(A::JLSparseMatrixCSR) - Aᵀ = JLSparseMatrixCSC(A.rowPtr, A.colVal, A.nzVal, reverse(size(A))) - JLSparseMatrixCSR(Aᵀ) end -function GPUArrays._spadjoint(A::JLSparseMatrixCSC) - Aᴴ = JLSparseMatrixCSR(A.colPtr, A.rowVal, conj(A.nzVal), reverse(size(A))) - JLSparseMatrixCSC(Aᴴ) -end -function GPUArrays._sptranspose(A::JLSparseMatrixCSC) - Aᵀ = JLSparseMatrixCSR(A.colPtr, A.rowVal, A.nzVal, reverse(size(A))) - JLSparseMatrixCSC(Aᵀ) -end - - -## adapt for the device -function Adapt.adapt_structure(to::Adaptor, x::JLSparseMatrixCSC{Tv,Ti}) where {Tv,Ti} - colPtr, rowVal, nzVal = adapt(to, x.colPtr), adapt(to, x.rowVal), adapt(to, x.nzVal) - GPUSparseDeviceMatrixCSC{Tv,Ti,typeof(colPtr),typeof(nzVal),AS.Generic}(colPtr, rowVal, nzVal, x.dims, x.nnz) -end -function Adapt.adapt_structure(to::Adaptor, x::JLSparseMatrixCSR{Tv,Ti}) where {Tv,Ti} - rowPtr, colVal, nzVal = adapt(to, x.rowPtr), adapt(to, x.colVal), adapt(to, x.nzVal) - GPUSparseDeviceMatrixCSR{Tv,Ti,typeof(rowPtr),typeof(nzVal),AS.Generic}(rowPtr, colVal, nzVal, x.dims, x.nnz) -end -function Adapt.adapt_structure(to::Adaptor, x::JLSparseVector{Tv,Ti}) where {Tv,Ti} - iPtr, nzVal = adapt(to, x.iPtr), adapt(to, x.nzVal) - GPUSparseDeviceVector{Tv,Ti,typeof(iPtr),typeof(nzVal),AS.Generic}(iPtr, nzVal, x.len, x.nnz) -end +# densify on the device +JLArray{T,1}(x::JLSparseVector) where {T} = copyto!(similar(nonzeros(x), T, size(x)), x) +JLArray{T,2}(A::JLSparseMatrix) where {T} = copyto!(similar(nonzeros(A), T, size(A)), A) diff --git a/src/GPUArrays.jl b/src/GPUArrays.jl index 3db7a050e..a928c7e57 100644 --- a/src/GPUArrays.jl +++ b/src/GPUArrays.jl @@ -4,6 +4,7 @@ using KernelAbstractions using Serialization using Random using LinearAlgebra +using SparseArrays using Printf using LinearAlgebra.BLAS @@ -22,7 +23,6 @@ import AcceleratedKernels as AK # device functionality include("device/abstractarray.jl") -include("device/sparse.jl") # host abstractions include("host/abstractarray.jl") @@ -41,7 +41,13 @@ include("host/random.jl") include("host/quirks.jl") include("host/uniformscaling.jl") include("host/statistics.jl") -include("host/sparse.jl") +include("host/sparse/types.jl") +include("host/sparse/primitives.jl") +include("host/sparse/conversions.jl") +include("host/sparse/show.jl") +include("host/sparse/indexing.jl") +include("host/sparse/broadcast.jl") +include("host/sparse/reductions.jl") include("host/alloc_cache.jl") diff --git a/src/device/sparse.jl b/src/device/sparse.jl deleted file mode 100644 index b8346eafe..000000000 --- a/src/device/sparse.jl +++ /dev/null @@ -1,171 +0,0 @@ -# on-device sparse array types -# should be excluded from coverage counts -# COV_EXCL_START -using SparseArrays - -# NOTE: this functionality is currently very bare-bones, only defining the array types -# without any device-compatible sparse array functionality - - -# core types - -export GPUSparseDeviceVector, GPUSparseDeviceMatrixCSC, GPUSparseDeviceMatrixCSR, - GPUSparseDeviceMatrixBSR, GPUSparseDeviceMatrixCOO - -abstract type AbstractGPUSparseDeviceMatrix{Tv, Ti} <: AbstractSparseMatrix{Tv, Ti} end - - -struct GPUSparseDeviceVector{Tv,Ti,Vi,Vv, A} <: AbstractSparseVector{Tv,Ti} - iPtr::Vi - nzVal::Vv - len::Int - nnz::Ti -end - -Base.length(g::GPUSparseDeviceVector) = g.len -Base.size(g::GPUSparseDeviceVector) = (g.len,) -SparseArrays.nnz(g::GPUSparseDeviceVector) = g.nnz -SparseArrays.nonzeroinds(g::GPUSparseDeviceVector) = g.iPtr -SparseArrays.nonzeros(g::GPUSparseDeviceVector) = g.nzVal - -struct GPUSparseDeviceMatrixCSC{Tv,Ti,Vi,Vv,A} <: AbstractGPUSparseDeviceMatrix{Tv, Ti} - colPtr::Vi - rowVal::Vi - nzVal::Vv - dims::NTuple{2,Int} - nnz::Ti -end - -Base.length(g::GPUSparseDeviceMatrixCSC) = prod(g.dims) -Base.size(g::GPUSparseDeviceMatrixCSC) = g.dims -SparseArrays.nnz(g::GPUSparseDeviceMatrixCSC) = g.nnz -SparseArrays.rowvals(g::GPUSparseDeviceMatrixCSC) = g.rowVal -SparseArrays.getcolptr(g::GPUSparseDeviceMatrixCSC) = g.colPtr -SparseArrays.nzrange(g::GPUSparseDeviceMatrixCSC, col::Integer) = @inbounds SparseArrays.getcolptr(g)[col]:(SparseArrays.getcolptr(g)[col+1]-1) -SparseArrays.nonzeros(g::GPUSparseDeviceMatrixCSC) = g.nzVal - -const GPUSparseDeviceColumnView{Tv, Ti, Vi, Vv, A} = SubArray{Tv, 1, GPUSparseDeviceMatrixCSC{Tv, Ti, Vi, Vv, A}, Tuple{Base.Slice{Base.OneTo{Int}}, Int}} -function SparseArrays.nonzeros(x::GPUSparseDeviceColumnView) - rowidx, colidx = parentindices(x) - A = parent(x) - @inbounds y = view(SparseArrays.nonzeros(A), SparseArrays.nzrange(A, colidx)) - return y -end - -function SparseArrays.nonzeroinds(x::GPUSparseDeviceColumnView) - rowidx, colidx = parentindices(x) - A = parent(x) - @inbounds y = view(SparseArrays.rowvals(A), SparseArrays.nzrange(A, colidx)) - return y -end -SparseArrays.rowvals(x::GPUSparseDeviceColumnView) = SparseArrays.nonzeroinds(x) - -function SparseArrays.nnz(x::GPUSparseDeviceColumnView) - rowidx, colidx = parentindices(x) - A = parent(x) - return length(SparseArrays.nzrange(A, colidx)) -end - -struct GPUSparseDeviceMatrixCSR{Tv,Ti,Vi,Vv,A} <: AbstractGPUSparseDeviceMatrix{Tv,Ti} - rowPtr::Vi - colVal::Vi - nzVal::Vv - dims::NTuple{2, Int} - nnz::Ti -end - -@inline function _getindex(arg::Union{GPUSparseDeviceMatrixCSR{Tv}, - GPUSparseDeviceMatrixCSC{Tv}, - GPUSparseDeviceVector{Tv}}, I, ptr)::Tv where {Tv} - if ptr == 0 - return zero(Tv) - else - return @inbounds arg.nzVal[ptr]::Tv - end -end - -struct GPUSparseDeviceMatrixBSR{Tv,Ti,Vi,Vv,A} <: AbstractGPUSparseDeviceMatrix{Tv,Ti} - rowPtr::Vi - colVal::Vi - nzVal::Vv - dims::NTuple{2,Int} - blockDim::Ti - dir::Char - nnz::Ti -end - -struct GPUSparseDeviceMatrixCOO{Tv,Ti,Vi,Vv, A} <: AbstractGPUSparseDeviceMatrix{Tv,Ti} - rowInd::Vi - colInd::Vi - nzVal::Vv - dims::NTuple{2,Int} - nnz::Ti -end - -Base.length(g::AbstractGPUSparseDeviceMatrix) = prod(g.dims) -Base.size(g::AbstractGPUSparseDeviceMatrix) = g.dims -SparseArrays.nnz(g::AbstractGPUSparseDeviceMatrix) = g.nnz -SparseArrays.getnzval(g::AbstractGPUSparseDeviceMatrix) = g.nzVal - -struct GPUSparseDeviceArrayCSR{Tv, Ti, Vi, Vv, N, M, A} <: AbstractSparseArray{Tv, Ti, N} - rowPtr::Vi - colVal::Vi - nzVal::Vv - dims::NTuple{N, Int} - nnz::Ti -end - -function GPUSparseDeviceArrayCSR{Tv, Ti, Vi, Vv, N}(rowPtr::Vi, colVal::Vi, nzVal::Vv, dims::NTuple{N,<:Integer}) where {Tv, Ti<:Integer, M, Vi<:AbstractDeviceArray{<:Integer,M}, Vv<:AbstractDeviceArray{Tv, M}, N} - @assert M == N - 1 "GPUSparseDeviceArrayCSR requires ndims(rowPtr) == ndims(colVal) == ndims(nzVal) == length(dims) - 1" - GPUSparseDeviceArrayCSR{Tv, Ti, Vi, Vv, N, M}(rowPtr, colVal, nzVal, dims, length(nzVal)) -end - -Base.length(g::GPUSparseDeviceArrayCSR) = prod(g.dims) -Base.size(g::GPUSparseDeviceArrayCSR) = g.dims -SparseArrays.nnz(g::GPUSparseDeviceArrayCSR) = g.nnz -SparseArrays.getnzval(g::GPUSparseDeviceArrayCSR) = g.nzVal - -# input/output - -function Base.show(io::IO, ::MIME"text/plain", A::GPUSparseDeviceVector) - println(io, "$(length(A))-element device sparse vector at:") - println(io, " iPtr: $(A.iPtr)") - print(io, " nzVal: $(A.nzVal)") -end - -function Base.show(io::IO, ::MIME"text/plain", A::GPUSparseDeviceMatrixCSR) - println(io, "$(length(A))-element device sparse matrix CSR at:") - println(io, " rowPtr: $(A.rowPtr)") - println(io, " colVal: $(A.colVal)") - print(io, " nzVal: $(A.nzVal)") -end - -function Base.show(io::IO, ::MIME"text/plain", A::GPUSparseDeviceMatrixCSC) - println(io, "$(length(A))-element device sparse matrix CSC at:") - println(io, " colPtr: $(A.colPtr)") - println(io, " rowVal: $(A.rowVal)") - print(io, " nzVal: $(A.nzVal)") -end - -function Base.show(io::IO, ::MIME"text/plain", A::GPUSparseDeviceMatrixBSR) - println(io, "$(length(A))-element device sparse matrix BSR at:") - println(io, " rowPtr: $(A.rowPtr)") - println(io, " colVal: $(A.colVal)") - print(io, " nzVal: $(A.nzVal)") -end - -function Base.show(io::IO, ::MIME"text/plain", A::GPUSparseDeviceMatrixCOO) - println(io, "$(length(A))-element device sparse matrix COO at:") - println(io, " rowPtr: $(A.rowPtr)") - println(io, " colInd: $(A.colInd)") - print(io, " nzVal: $(A.nzVal)") -end - -function Base.show(io::IO, ::MIME"text/plain", A::GPUSparseDeviceArrayCSR) - println(io, "$(length(A))-element device sparse array CSR at:") - println(io, " rowPtr: $(A.rowPtr)") - println(io, " colVal: $(A.colVal)") - print(io, " nzVal: $(A.nzVal)") -end - -# COV_EXCL_STOP diff --git a/src/host/sparse.jl b/src/host/sparse.jl deleted file mode 100644 index 0c6357a35..000000000 --- a/src/host/sparse.jl +++ /dev/null @@ -1,1028 +0,0 @@ -using LinearAlgebra -using LinearAlgebra: BlasFloat - -abstract type AbstractGPUSparseArray{Tv, Ti, N} <: AbstractSparseArray{Tv, Ti, N} end -const AbstractGPUSparseVector{Tv, Ti} = AbstractGPUSparseArray{Tv, Ti, 1} -const AbstractGPUSparseMatrix{Tv, Ti} = AbstractGPUSparseArray{Tv, Ti, 2} - -abstract type AbstractGPUSparseMatrixCSC{Tv, Ti} <: AbstractGPUSparseArray{Tv, Ti, 2} end -abstract type AbstractGPUSparseMatrixCSR{Tv, Ti} <: AbstractGPUSparseArray{Tv, Ti, 2} end -abstract type AbstractGPUSparseMatrixCOO{Tv, Ti} <: AbstractGPUSparseArray{Tv, Ti, 2} end -abstract type AbstractGPUSparseMatrixBSR{Tv, Ti} <: AbstractGPUSparseArray{Tv, Ti, 2} end - -const AbstractGPUSparseVecOrMat = Union{AbstractGPUSparseVector,AbstractGPUSparseMatrix} - -SparseArrays.nnz(g::T) where {T<:AbstractGPUSparseArray} = g.nnz -SparseArrays.nonzeros(g::T) where {T<:AbstractGPUSparseArray} = g.nzVal - -SparseArrays.nonzeroinds(g::T) where {T<:AbstractGPUSparseVector} = g.iPtr -SparseArrays.rowvals(g::T) where {T<:AbstractGPUSparseVector} = SparseArrays.nonzeroinds(g) - -SparseArrays.rowvals(g::AbstractGPUSparseMatrixCSC) = g.rowVal -SparseArrays.getcolptr(S::AbstractGPUSparseMatrixCSC) = S.colPtr - -Base.convert(T::Type{<:AbstractGPUSparseArray}, m::AbstractArray) = m isa T ? m : T(m) - -# collect to Array -Base.collect(x::AbstractGPUSparseVector) = collect(SparseVector(x)) -Base.collect(x::AbstractGPUSparseMatrixCSC) = collect(SparseMatrixCSC(x)) -Base.collect(x::AbstractGPUSparseMatrixCSR) = collect(SparseMatrixCSC(x)) -Base.collect(x::AbstractGPUSparseMatrixBSR) = collect(SparseMatrixCSC(x)) -Base.collect(x::AbstractGPUSparseMatrixCOO) = collect(SparseMatrixCSC(x)) - -Base.Array(x::AbstractGPUSparseVector) = collect(SparseVector(x)) -Base.Array(x::AbstractGPUSparseMatrixCSC) = collect(SparseMatrixCSC(x)) -Base.Array(x::AbstractGPUSparseMatrixCSR) = collect(SparseMatrixCSC(x)) -Base.Array(x::AbstractGPUSparseMatrixBSR) = collect(SparseMatrixCSC(x)) -Base.Array(x::AbstractGPUSparseMatrixCOO) = collect(SparseMatrixCSC(x)) - -# iszero that avoids scalar indexing by using GPU-compatible reduction -# nnz(A) == 0 means no stored elements, so it's definitely zero -# all(iszero, nonzeros(A)) uses GPU reduction to check stored values -Base.iszero(A::AbstractGPUSparseArray) = SparseArrays.nnz(A) == 0 || all(iszero, SparseArrays.nonzeros(A)) - -SparseArrays.SparseVector(x::AbstractGPUSparseVector) = SparseVector(length(x), Array(SparseArrays.nonzeroinds(x)), Array(SparseArrays.nonzeros(x))) -SparseArrays.SparseMatrixCSC(x::AbstractGPUSparseMatrixCSC) = SparseMatrixCSC(size(x)..., Array(SparseArrays.getcolptr(x)), Array(SparseArrays.rowvals(x)), Array(SparseArrays.nonzeros(x))) - -# similar -Base.similar(Vec::V) where {V<:AbstractGPUSparseVector} = V(copy(SparseArrays.nonzeroinds(Vec)), similar(SparseArrays.nonzeros(Vec)), length(Vec)) -Base.similar(Mat::M) where {M<:AbstractGPUSparseMatrixCSC} = M(copy(SparseArrays.getcolptr(Mat)), copy(SparseArrays.rowvals(Mat)), similar(SparseArrays.nonzeros(Mat)), size(Mat)) - -Base.similar(Vec::V, T::Type) where {Tv, Ti, V<:AbstractGPUSparseVector{Tv, Ti}} = sparse_array_type(V){T, Ti}(copy(SparseArrays.nonzeroinds(Vec)), similar(SparseArrays.nonzeros(Vec), T), length(Vec)) -Base.similar(Mat::M, T::Type) where {M<:AbstractGPUSparseMatrixCSC} = sparse_array_type(M)(copy(SparseArrays.getcolptr(Mat)), copy(SparseArrays.rowvals(Mat)), similar(SparseArrays.nonzeros(Mat), T), size(Mat)) - -dense_array_type(a::AbstractArray) = dense_array_type(typeof(a)) -sparse_array_type(a::AbstractArray) = sparse_array_type(typeof(a)) -dense_vector_type(a::AbstractArray) = dense_vector_type(typeof(a)) - -dense_array_type(::Type{<:SparseVector}) = SparseVector -sparse_array_type(::Type{<:SparseVector}) = SparseVector -dense_vector_type(::Type{<:AbstractSparseArray}) = Vector -dense_vector_type(::Type{<:AbstractArray}) = Vector -dense_array_type(::Type{<:SparseMatrixCSC}) = SparseMatrixCSC -sparse_array_type(::Type{<:SparseMatrixCSC}) = SparseMatrixCSC - -coo_type(sa::AbstractArray) = coo_type(typeof(sa)) - -function _spadjoint end -function _sptranspose end - -function LinearAlgebra.opnorm(A::AbstractGPUSparseMatrixCSR, p::Real=2) - if p == Inf - return maximum(sum(abs, A; dims=2)) - elseif p == 1 - return maximum(sum(abs, A; dims=1)) - else - throw(ArgumentError("p=$p is not supported")) - end -end - -LinearAlgebra.opnorm(A::AbstractGPUSparseMatrixCSC, p::Real=2) = opnorm(csr_type(A)(A), p) - -function LinearAlgebra.norm(A::AbstractGPUSparseMatrix{T}, p::Real=2) where T - if p == Inf - return maximum(abs.(SparseArrays.nonzeros(A))) - elseif p == -Inf - return minimum(abs.(SparseArrays.nonzeros(A))) - elseif p == 0 - return Float64(SparseArrays.nnz(A)) - else - return sum(abs.(SparseArrays.nonzeros(A)).^p)^(1/p) - end -end - -function SparseArrays.findnz(S::MT) where {MT <: AbstractGPUSparseMatrix} - S2 = coo_type(MT)(S) - I = S2.rowInd - J = S2.colInd - V = S2.nzVal - - # To make it compatible with the SparseArrays.jl version - idxs = sortperm(J) - I = I[idxs] - J = J[idxs] - V = V[idxs] - - return (I, J, V) -end - -### WRAPPED ARRAYS -# Without these, `copy` of a lazy wrapper falls back to the generic `AbstractArray` method, -# which reads the parent one element at a time and so hits scalar indexing. -Base.copy(A::Transpose{<:Any, <:AbstractGPUSparseMatrix}) = _sptranspose(parent(A)) -Base.copy(A::Adjoint{<:Any, <:AbstractGPUSparseMatrix}) = _spadjoint(parent(A)) - -function Base.permutedims(A::AbstractGPUSparseMatrix, (a, b)) - (a, b) == (2, 1) && return _sptranspose(A) - (a, b) == (1, 2) && return copy(A) - throw(ArgumentError("no valid permutation of dimensions")) -end - -LinearAlgebra.issymmetric(M::Union{AbstractGPUSparseMatrixCSC,AbstractGPUSparseMatrixCSR}) = size(M, 1) == size(M, 2) ? norm(M - transpose(M), Inf) == 0 : false -LinearAlgebra.ishermitian(M::Union{AbstractGPUSparseMatrixCSC,AbstractGPUSparseMatrixCSR}) = size(M, 1) == size(M, 2) ? norm(M - adjoint(M), Inf) == 0 : false - -LinearAlgebra.istriu(M::UpperTriangular{T,S}) where {T<:BlasFloat, S<:Union{<:AbstractGPUSparseMatrix, Adjoint{<:Any, <:AbstractGPUSparseMatrix}, Transpose{<:Any, <:AbstractGPUSparseMatrix}}} = true -LinearAlgebra.istril(M::UpperTriangular{T,S}) where {T<:BlasFloat, S<:Union{<:AbstractGPUSparseMatrix, Adjoint{<:Any, <:AbstractGPUSparseMatrix}, Transpose{<:Any, <:AbstractGPUSparseMatrix}}} = false -LinearAlgebra.istriu(M::LowerTriangular{T,S}) where {T<:BlasFloat, S<:Union{<:AbstractGPUSparseMatrix, Adjoint{<:Any, <:AbstractGPUSparseMatrix}, Transpose{<:Any, <:AbstractGPUSparseMatrix}}} = false -LinearAlgebra.istril(M::LowerTriangular{T,S}) where {T<:BlasFloat, S<:Union{<:AbstractGPUSparseMatrix, Adjoint{<:Any, <:AbstractGPUSparseMatrix}, Transpose{<:Any, <:AbstractGPUSparseMatrix}}} = true - -Hermitian{T}(Mat::AbstractGPUSparseMatrix{T}) where {T} = Hermitian{eltype(Mat),typeof(Mat)}(Mat,'U') - -# work around upstream breakage from JuliaLang/julia#55547 -@static if VERSION >= v"1.11.2" - const GPUSparseUpperOrUnitUpperTriangular = LinearAlgebra.UpperOrUnitUpperTriangular{ - <:Any,<:Union{<:AbstractGPUSparseMatrix, Adjoint{<:Any, <:AbstractGPUSparseMatrix}, Transpose{<:Any, <:AbstractGPUSparseMatrix}}} - const GPUSparseLowerOrUnitLowerTriangular = LinearAlgebra.LowerOrUnitLowerTriangular{ - <:Any,<:Union{<:AbstractGPUSparseMatrix, Adjoint{<:Any, <:AbstractGPUSparseMatrix}, Transpose{<:Any, <:AbstractGPUSparseMatrix}}} - LinearAlgebra.istriu(::GPUSparseUpperOrUnitUpperTriangular) = true - LinearAlgebra.istril(::GPUSparseUpperOrUnitUpperTriangular) = false - LinearAlgebra.istriu(::GPUSparseLowerOrUnitLowerTriangular) = false - LinearAlgebra.istril(::GPUSparseLowerOrUnitLowerTriangular) = true -end - -for SparseMatrixType in [:AbstractGPUSparseMatrixCSC, :AbstractGPUSparseMatrixCSR] - @eval begin - LinearAlgebra.triu(A::ST, k::Integer) where {T, ST<:$SparseMatrixType{T}} = - ST( triu(coo_type(A)(A), k) ) - LinearAlgebra.triu(A::Transpose{T,<:ST}, k::Integer) where {T, ST<:$SparseMatrixType{T}} = - ST( triu(coo_type(A)(_sptranspose(parent(A))), k) ) - LinearAlgebra.triu(A::Adjoint{T,<:ST}, k::Integer) where {T, ST<:$SparseMatrixType{T}} = - ST( triu(coo_type(A)(_spadjoint(parent(A))), k) ) - - LinearAlgebra.tril(A::ST, k::Integer) where {T, ST<:$SparseMatrixType{T}} = - ST( tril(coo_type(A)(A), k) ) - LinearAlgebra.tril(A::Transpose{T,<:ST}, k::Integer) where {T, ST<:$SparseMatrixType{T}} = - ST( tril(coo_type(A)(_sptranspose(parent(A))), k) ) - LinearAlgebra.tril(A::Adjoint{T,<:ST}, k::Integer) where {T, ST<:$SparseMatrixType{T}} = - ST( tril(coo_type(A)(_spadjoint(parent(A))), k) ) - - LinearAlgebra.triu(A::Union{ST, Transpose{T,<:ST}, Adjoint{T,<:ST}}) where {T, ST<:$SparseMatrixType{T}} = - ST( triu(coo_type(A)(A), 0) ) - LinearAlgebra.tril(A::Union{ST,Transpose{T,<:ST}, Adjoint{T,<:ST}}) where {T, ST<:$SparseMatrixType{T}} = - ST( tril(coo_type(A)(A), 0) ) - - LinearAlgebra.kron(A::ST, B::ST) where {T, ST<:$SparseMatrixType{T}} = - ST( kron(coo_type(A)(A), coo_type(B)(B)) ) - LinearAlgebra.kron(A::ST, B::Diagonal) where {T, ST<:$SparseMatrixType{T}} = - ST( kron(coo_type(A)(A), B) ) - LinearAlgebra.kron(A::Diagonal, B::ST) where {T, ST<:$SparseMatrixType{T}} = - ST( kron(A, coo_type(B)(B)) ) - - LinearAlgebra.kron(A::Transpose{T,<:ST}, B::ST) where {T, ST<:$SparseMatrixType{T}} = - ST( kron(coo_type(A)(_sptranspose(parent(A))), coo_type(B)(B)) ) - LinearAlgebra.kron(A::ST, B::Transpose{T,<:ST}) where {T, ST<:$SparseMatrixType{T}} = - ST( kron(coo_type(A)(A), coo_type(parent(B))(_sptranspose(parent(B)))) ) - LinearAlgebra.kron(A::Transpose{T,<:ST}, B::Transpose{T,<:ST}) where {T, ST<:$SparseMatrixType{T}} = - ST( kron(coo_type(A)(_sptranspose(parent(A))), coo_type(parent(B))(_sptranspose(parent(B)))) ) - LinearAlgebra.kron(A::Transpose{T,<:ST}, B::Diagonal) where {T, ST<:$SparseMatrixType{T}} = - ST( kron(coo_type(A)(_sptranspose(parent(A))), B) ) - LinearAlgebra.kron(A::Diagonal, B::Transpose{T,<:ST}) where {T, ST<:$SparseMatrixType{T}} = - ST( kron(A, coo_type(B)(_sptranspose(parent(B)))) ) - - LinearAlgebra.kron(A::Adjoint{T,<:ST}, B::ST) where {T, ST<:$SparseMatrixType{T}} = - ST( kron(coo_type(A)(_spadjoint(parent(A))), coo_type(B)(B)) ) - LinearAlgebra.kron(A::ST, B::Adjoint{T,<:ST}) where {T, ST<:$SparseMatrixType{T}} = - ST( kron(coo_type(A)(A), coo_type(parent(B))(_spadjoint(parent(B)))) ) - LinearAlgebra.kron(A::Adjoint{T,<:ST}, B::Adjoint{T,<:ST}) where {T, ST<:$SparseMatrixType{T}} = - ST( kron(coo_type(parent(A))(_spadjoint(parent(A))), coo_type(parent(B))(_spadjoint(parent(B)))) ) - LinearAlgebra.kron(A::Adjoint{T,<:ST}, B::Diagonal) where {T, ST<:$SparseMatrixType{T}} = - ST( kron(coo_type(parent(A))(_spadjoint(parent(A))), B) ) - LinearAlgebra.kron(A::Diagonal, B::Adjoint{T,<:ST}) where {T, ST<:$SparseMatrixType{T}} = - ST( kron(A, coo_type(parent(B))(_spadjoint(parent(B)))) ) - - - function Base.reshape(A::ST, dims::Dims) where {ST<:$SparseMatrixType} - B = coo_type(A)(A) - ST(reshape(B, dims)) - end - - function SparseArrays.droptol!(A::ST, tol::Real) where {ST<:$SparseMatrixType} - B = coo_type(A)(A) - droptol!(B, tol) - copyto!(A, ST(B)) - end - - function LinearAlgebra.exp(A::ST; threshold = 1e-7, nonzero_tol = 1e-14) where {ST<:$SparseMatrixType} - rows = LinearAlgebra.checksquare(A) # Throws exception if not square - typeA = eltype(A) - - mat_norm = norm(A, Inf) - scaling_factor = nextpow(2, mat_norm) # Native routine, faster - A = A ./ scaling_factor - delta = 1 - - diag_vec = similar(SparseArrays.nonzeros(A), rows) - fill!(diag_vec, one(eltype(A))) - P = ST(spdiagm(0 => diag_vec)) - next_term = P - n = 1 - - while delta > threshold - next_term = typeA(1 / n) * A * next_term - droptol!(next_term, nonzero_tol) - delta = norm(next_term, Inf) - copyto!(P, P + next_term) - n = n + 1 - end - for n = 1:log2(scaling_factor) - P = P * P; - if nnz(P) / length(P) < 0.25 - droptol!(P, nonzero_tol) - end - end - P - end - end -end - - -### INDEXING -Base.getindex(A::AbstractGPUSparseVector, ::Colon) = copy(A) -Base.getindex(A::AbstractGPUSparseMatrix, ::Colon, ::Colon) = copy(A) -Base.getindex(A::AbstractGPUSparseMatrix, i, ::Colon) = getindex(A, i, 1:size(A, 2)) -Base.getindex(A::AbstractGPUSparseMatrix, ::Colon, i) = getindex(A, 1:size(A, 1), i) -Base.getindex(A::AbstractGPUSparseMatrix, I::Tuple{Integer,Integer}) = getindex(A, I[1], I[2]) - -function Base.getindex(A::AbstractGPUSparseVector{Tv, Ti}, i::Integer) where {Tv, Ti} - @boundscheck checkbounds(A, i) - ii = searchsortedfirst(SparseArrays.nonzeroinds(A), convert(Ti, i)) - (ii > SparseArrays.nnz(A) || SparseArrays.nonzeroinds(A)[ii] != i) && return zero(Tv) - SparseArrays.nonzeros(A)[ii] -end - -function Base.getindex(A::AbstractGPUSparseMatrixCSC{T}, i0::Integer, i1::Integer) where T - @boundscheck checkbounds(A, i0, i1) - r1 = Int(SparseArrays.getcolptr(A)[i1]) - r2 = Int(SparseArrays.getcolptr(A)[i1+1]-1) - (r1 > r2) && return zero(T) - r1 = searchsortedfirst(SparseArrays.rowvals(A), i0, r1, r2, Base.Order.Forward) - (r1 > r2 || SparseArrays.rowvals(A)[r1] != i0) && return zero(T) - SparseArrays.nonzeros(A)[r1] -end - -## copying between sparse GPU arrays -Base.copy(Vec::AbstractGPUSparseVector) = copyto!(similar(Vec), Vec) - -function Base.copyto!(dst::AbstractGPUSparseVector, src::AbstractGPUSparseVector) - if length(dst) != length(src) - throw(ArgumentError("Inconsistent Sparse Vector size")) - end - resize!(SparseArrays.nonzeroinds(dst), length(SparseArrays.nonzeroinds(src))) - resize!(SparseArrays.nonzeros(dst), length(SparseArrays.nonzeros(src))) - copyto!(SparseArrays.nonzeroinds(dst), SparseArrays.nonzeroinds(src)) - copyto!(SparseArrays.nonzeros(dst), SparseArrays.nonzeros(src)) - dst.nnz = src.nnz - dst -end - -function Base.copyto!(dst::AbstractGPUSparseMatrixCSC, src::AbstractGPUSparseMatrixCSC) - if size(dst) != size(src) - throw(ArgumentError("Inconsistent Sparse Matrix size")) - end - resize!(SparseArrays.getcolptr(dst), length(SparseArrays.getcolptr(src))) - resize!(SparseArrays.rowvals(dst), length(SparseArrays.rowvals(src))) - resize!(SparseArrays.nonzeros(dst), length(SparseArrays.nonzeros(src))) - copyto!(SparseArrays.getcolptr(dst), SparseArrays.getcolptr(src)) - copyto!(SparseArrays.rowvals(dst), SparseArrays.rowvals(src)) - copyto!(SparseArrays.nonzeros(dst), SparseArrays.nonzeros(src)) - dst.nnz = src.nnz - dst -end - -### BROADCAST - -# broadcast container type promotion for combinations of sparse arrays and other types -struct GPUSparseVecStyle <: Broadcast.AbstractArrayStyle{1} end -struct GPUSparseMatStyle <: Broadcast.AbstractArrayStyle{2} end -Broadcast.BroadcastStyle(::Type{<:AbstractGPUSparseVector}) = GPUSparseVecStyle() -Broadcast.BroadcastStyle(::Type{<:AbstractGPUSparseMatrix}) = GPUSparseMatStyle() -const SPVM = Union{GPUSparseVecStyle,GPUSparseMatStyle} - -# GPUSparseVecStyle handles 0-1 dimensions, GPUSparseMatStyle 0-2 dimensions. -# GPUSparseVecStyle promotes to GPUSparseMatStyle for 2 dimensions. -# Fall back to DefaultArrayStyle for higher dimensionality. -GPUSparseVecStyle(::Val{0}) = GPUSparseVecStyle() -GPUSparseVecStyle(::Val{1}) = GPUSparseVecStyle() -GPUSparseVecStyle(::Val{2}) = GPUSparseMatStyle() -GPUSparseVecStyle(::Val{N}) where N = Broadcast.DefaultArrayStyle{N}() -GPUSparseMatStyle(::Val{0}) = GPUSparseMatStyle() -GPUSparseMatStyle(::Val{1}) = GPUSparseMatStyle() -GPUSparseMatStyle(::Val{2}) = GPUSparseMatStyle() -GPUSparseMatStyle(::Val{N}) where N = Broadcast.DefaultArrayStyle{N}() - -Broadcast.BroadcastStyle(::GPUSparseVecStyle, ::AbstractGPUArrayStyle{1}) = GPUSparseVecStyle() -Broadcast.BroadcastStyle(::GPUSparseVecStyle, ::AbstractGPUArrayStyle{2}) = GPUSparseMatStyle() -Broadcast.BroadcastStyle(::GPUSparseMatStyle, ::AbstractGPUArrayStyle{2}) = GPUSparseMatStyle() - -# don't wrap sparse arrays with Extruded -Broadcast.extrude(x::AbstractGPUSparseVecOrMat) = x - -## detection of zero-preserving functions - -# modified from SparseArrays.jl - -# capturescalars takes a function (f) and a tuple of broadcast arguments, and returns a -# partially-evaluated function and a reduced argument tuple where all scalar operations have -# been applied already. -@inline function capturescalars(f, mixedargs) - let (passedsrcargstup, makeargs) = _capturescalars(mixedargs...) - parevalf = (passed...) -> f(makeargs(passed...)...) - return (parevalf, passedsrcargstup) - end -end - -## sparse broadcast style - -# Work around losing Type{T}s as DataTypes within the tuple that makeargs creates -@inline capturescalars(f, mixedargs::Tuple{Ref{Type{T}}, Vararg{Any}}) where {T} = - capturescalars((args...)->f(T, args...), Base.tail(mixedargs)) -@inline capturescalars(f, mixedargs::Tuple{Ref{Type{T}}, Ref{Type{S}}, Vararg{Any}}) where {T, S} = - # This definition is identical to the one above and necessary only for - # avoiding method ambiguity. - capturescalars((args...)->f(T, args...), Base.tail(mixedargs)) -@inline capturescalars(f, mixedargs::Tuple{AbstractGPUSparseVecOrMat, Ref{Type{T}}, Vararg{Any}}) where {T} = - capturescalars((a1, args...)->f(a1, T, args...), (mixedargs[1], Base.tail(Base.tail(mixedargs))...)) -@inline capturescalars(f, mixedargs::Tuple{Union{Ref,AbstractArray{<:Any,0}}, Ref{Type{T}}, Vararg{Any}}) where {T} = - capturescalars((args...)->f(mixedargs[1], T, args...), Base.tail(Base.tail(mixedargs))) - -scalararg(::Number) = true -scalararg(::Any) = false -scalarwrappedarg(::Union{AbstractArray{<:Any,0},Ref}) = true -scalarwrappedarg(::Any) = false - -@inline function _capturescalars() - return (), () -> () -end -@inline function _capturescalars(arg, mixedargs...) - let (rest, f) = _capturescalars(mixedargs...) - if scalararg(arg) - return rest, @inline function(tail...) - (arg, f(tail...)...) - end # add back scalararg after (in makeargs) - elseif scalarwrappedarg(arg) - return rest, @inline function(tail...) - (arg[], f(tail...)...) # TODO: This can put a Type{T} in a tuple - end # unwrap and add back scalararg after (in makeargs) - else - return (arg, rest...), @inline function(head, tail...) - (head, f(tail...)...) - end # pass-through to broadcast - end - end -end -@inline function _capturescalars(arg) # this definition is just an optimization (to bottom out the recursion slightly sooner) - if scalararg(arg) - return (), () -> (arg,) # add scalararg - elseif scalarwrappedarg(arg) - return (), () -> (arg[],) # unwrap - else - return (arg,), (head,) -> (head,) # pass-through - end -end - -@inline _iszero(x) = x == 0 -@inline _iszero(x::Number) = Base.iszero(x) -@inline _iszero(x::AbstractArray) = Base.iszero(x) -@inline _zeros_eltypes(A) = (zero(eltype(A)),) -@inline _zeros_eltypes(A, Bs...) = (zero(eltype(A)), _zeros_eltypes(Bs...)...) - -## COV_EXCL_START -## iteration helpers - -""" - CSRIterator{Ti}(row, args...) - -A GPU-compatible iterator for accessing the elements of a single row `row` of several CSR -matrices `args` in one go. The row should be in-bounds for every sparse argument. Each -iteration returns a 2-element tuple: The current column, and each arguments' pointer index -(or 0 if that input didn't have an element at that column). The pointers can then be used to -access the elements themselves. - -For convenience, this iterator can be passed non-sparse arguments as well, which will be -ignored (with the returned `col`/`ptr` values set to 0). -""" -struct CSRIterator{Ti,N,ATs} - row::Ti - col_ends::NTuple{N, Ti} - args::ATs -end - -function CSRIterator{Ti}(row, args::Vararg{Any, N}) where {Ti,N} - # check that `row` is valid for all arguments - @boundscheck begin - ntuple(Val(N)) do i - arg = @inbounds args[i] - arg isa GPUSparseDeviceMatrixCSR && checkbounds(arg, row, 1) - end - end - - col_ends = ntuple(Val(N)) do i - arg = @inbounds args[i] - if arg isa GPUSparseDeviceMatrixCSR - @inbounds(arg.rowPtr[row+1]) - else - zero(Ti) - end - end - - CSRIterator{Ti, N, typeof(args)}(row, col_ends, args) -end - -@inline function Base.iterate(iter::CSRIterator{Ti,N}, state=nothing) where {Ti,N} - # helper function to get the column of a sparse array at a specific pointer - @inline function get_col(i, ptr) - arg = @inbounds iter.args[i] - if arg isa GPUSparseDeviceMatrixCSR - col_end = @inbounds iter.col_ends[i] - if ptr < col_end - return @inbounds arg.colVal[ptr] % Ti - end - end - typemax(Ti) - end - - # initialize the state - # - ptr: the current index into the colVal/nzVal arrays - # - col: the current column index (cached so that we don't have to re-read each time) - state = something(state, - ntuple(Val(N)) do i - arg = @inbounds iter.args[i] - if arg isa GPUSparseDeviceMatrixCSR - ptr = @inbounds iter.args[i].rowPtr[iter.row] % Ti - col = @inbounds get_col(i, ptr) - else - ptr = typemax(Ti) - col = typemax(Ti) - end - (; ptr, col) - end - ) - - # determine the column we're currently processing - cols = ntuple(i -> @inbounds(state[i].col), Val(N)) - cur_col = min(cols...) - cur_col == typemax(Ti) && return - - # fetch the pointers (we don't look up the values, as the caller might want to index - # the sparse array directly, e.g., to mutate it). we don't return `ptrs` from the state - # directly, but first convert the `typemax(Ti)` to a more convenient zero value. - # NOTE: these values may end up unused by the caller (e.g. in the count_nnzs kernels), - # but LLVM appears smart enough to filter them away. - ptrs = ntuple(Val(N)) do i - ptr, col = @inbounds state[i] - col == cur_col ? ptr : zero(Ti) - end - - # advance the state - new_state = ntuple(Val(N)) do i - ptr, col = @inbounds state[i] - if col == cur_col - ptr += one(Ti) - col = get_col(i, ptr) - end - (; ptr, col) - end - - return (cur_col, ptrs), new_state -end - -struct CSCIterator{Ti,N,ATs} - col::Ti - row_ends::NTuple{N, Ti} - args::ATs -end - -function CSCIterator{Ti}(col, args::Vararg{Any, N}) where {Ti,N} - # check that `col` is valid for all arguments - @boundscheck begin - ntuple(Val(N)) do i - arg = @inbounds args[i] - arg isa GPUSparseDeviceMatrixCSR && checkbounds(arg, 1, col) - end - end - - row_ends = ntuple(Val(N)) do i - arg = @inbounds args[i] - x = if arg isa GPUSparseDeviceMatrixCSC - @inbounds(arg.colPtr[col+1]) - else - zero(Ti) - end - x - end - - CSCIterator{Ti, N, typeof(args)}(col, row_ends, args) -end - -@inline function Base.iterate(iter::CSCIterator{Ti,N}, state=nothing) where {Ti,N} - # helper function to get the column of a sparse array at a specific pointer - @inline function get_col(i, ptr) - arg = @inbounds iter.args[i] - if arg isa GPUSparseDeviceMatrixCSC - col_end = @inbounds iter.row_ends[i] - if ptr < col_end - return @inbounds arg.rowVal[ptr] % Ti - end - end - typemax(Ti) - end - - # initialize the state - # - ptr: the current index into the rowVal/nzVal arrays - # - row: the current row index (cached so that we don't have to re-read each time) - state = something(state, - ntuple(Val(N)) do i - arg = @inbounds iter.args[i] - if arg isa GPUSparseDeviceMatrixCSC - ptr = @inbounds iter.args[i].colPtr[iter.col] % Ti - row = @inbounds get_col(i, ptr) - else - ptr = typemax(Ti) - row = typemax(Ti) - end - (; ptr, row) - end - ) - - # determine the row we're currently processing - rows = ntuple(i -> @inbounds(state[i].row), Val(N)) - cur_row = min(rows...) - cur_row == typemax(Ti) && return - - # fetch the pointers (we don't look up the values, as the caller might want to index - # the sparse array directly, e.g., to mutate it). we don't return `ptrs` from the state - # directly, but first convert the `typemax(Ti)` to a more convenient zero value. - # NOTE: these values may end up unused by the caller (e.g. in the count_nnzs kernels), - # but LLVM appears smart enough to filter them away. - ptrs = ntuple(Val(N)) do i - ptr, row = @inbounds state[i] - row == cur_row ? ptr : zero(Ti) - end - - # advance the state - new_state = ntuple(Val(N)) do i - ptr, row = @inbounds state[i] - if row == cur_row - ptr += one(Ti) - row = get_col(i, ptr) - end - (; ptr, row) - end - - return (cur_row, ptrs), new_state -end - -# helpers to index a sparse or dense array -function _getindex(arg::Union{<:GPUSparseDeviceMatrixCSR,GPUSparseDeviceMatrixCSC}, I, ptr) - if ptr == 0 - zero(eltype(arg)) - else - @inbounds arg.nzVal[ptr] - end -end -@inline function _getindex(arg::DenseArray{Tv}, I, ptr)::Tv where {Tv} - return @inbounds arg[I]::Tv -end -@inline _getindex(arg, I, ptr) = Broadcast._broadcast_getindex(arg, I) - -## sparse broadcast implementation -iter_type(::Type{<:AbstractGPUSparseMatrixCSC}, ::Type{Ti}) where {Ti} = CSCIterator{Ti} -iter_type(::Type{<:AbstractGPUSparseMatrixCSR}, ::Type{Ti}) where {Ti} = CSRIterator{Ti} -iter_type(::Type{<:GPUSparseDeviceMatrixCSC}, ::Type{Ti}) where {Ti} = CSCIterator{Ti} -iter_type(::Type{<:GPUSparseDeviceMatrixCSR}, ::Type{Ti}) where {Ti} = CSRIterator{Ti} - -_has_row(A, offsets, row, fpreszeros::Bool) = fpreszeros ? 0 : row -_has_row(A::AbstractDeviceArray, offsets, row, ::Bool) = row -# the position of `row` among the stored indices of a sparse vector, or 0 if it isn't stored -function _has_row(A::GPUSparseDeviceVector, offsets, row, ::Bool) - ptr = searchsortedfirst(A.iPtr, row) - return (ptr <= length(A.iPtr) && @inbounds(A.iPtr[ptr]) == row) ? ptr : 0 -end - -# for every row, determine which arguments have an entry there (see `_has_row`), and whether -# the output has one (`key` is the row, or `typemax(Ti)` if not) -@kernel function compute_offsets_kernel(::Type{<:AbstractGPUSparseVector}, fpreszeros::Bool, - offsets::AbstractVector{Pair{Ti, NTuple{N, Ti}}}, - args...) where {Ti, N} - row = @index(Global, Linear) - if row ≤ length(offsets) - arg_row_is_nnz = ntuple(Val(N)) do i - arg = @inbounds args[i] - _has_row(arg, offsets, row, fpreszeros) % Ti - end - row_is_nnz = false - for i in 1:N - row_is_nnz |= @inbounds(arg_row_is_nnz[i]) != 0 - end - key = row_is_nnz ? row % Ti : typemax(Ti) - @inbounds offsets[row] = key => arg_row_is_nnz - end -end - -# kernel to count the number of non-zeros in a row, to determine the row offsets -@kernel function compute_offsets_kernel(T::Type{<:Union{AbstractGPUSparseMatrixCSR, AbstractGPUSparseMatrixCSC}}, - offsets::AbstractVector{Ti}, args...) where Ti - # every thread processes an entire row - leading_dim = @index(Global, Linear) - if leading_dim ≤ length(offsets)-1 - iter = @inbounds iter_type(T, Ti)(leading_dim, args...) - - # count the nonzero leading_dims of all inputs - accum = zero(Ti) - for (leading_dim, vals) in iter - accum += one(Ti) - end - - # the way we write the nnz counts is a bit strange, but done so that the result - # after accumulation can be directly used as the rowPtr/colPtr array of a CSR/CSC matrix. - @inbounds begin - if leading_dim == 1 - offsets[1] = 1 - end - offsets[leading_dim+1] = accum - end - end -end - -# `positions` is the inclusive prefix count of output entries, so an output entry for `row` -# goes to `positions[row]`. -@kernel function sparse_to_sparse_broadcast_kernel(f::F, output::GPUSparseDeviceVector{Tv,Ti}, - offsets::AbstractVector{Pair{Ti, NTuple{N, Ti}}}, - positions::AbstractVector{Ti}, - args...) where {Tv, Ti, N, F} - row = @index(Global, Linear) - if row ≤ length(offsets) - key, arg_ptrs = @inbounds offsets[row] - if key != typemax(Ti) - vals = ntuple(Val(N)) do i - @inline - arg = @inbounds args[i] - # ptr is 0 if the sparse vector doesn't have an element at this row - # ptr is 0 if the arg is a scalar AND f preserves zeros - ptr = @inbounds arg_ptrs[i] - _getindex(arg, row, ptr) - end - output_ix = @inbounds positions[row] - @inbounds output.iPtr[output_ix] = row - @inbounds output.nzVal[output_ix] = f(vals...) - end - end -end - -@kernel function sparse_to_sparse_broadcast_kernel(f, output::T, offsets::Union{<:AbstractArray,Nothing}, - args...) where {Ti, T<:Union{GPUSparseDeviceMatrixCSR{<:Any,Ti}, - GPUSparseDeviceMatrixCSC{<:Any,Ti}}} - # every thread processes an entire row - leading_dim = @index(Global, Linear) - leading_dim_size = output isa GPUSparseDeviceMatrixCSR ? size(output, 1) : size(output, 2) - if leading_dim ≤ leading_dim_size - iter = @inbounds iter_type(T, Ti)(leading_dim, args...) - - output_ptrs = output isa GPUSparseDeviceMatrixCSR ? output.rowPtr : output.colPtr - output_ivals = output isa GPUSparseDeviceMatrixCSR ? output.colVal : output.rowVal - # fetch the row offset, and write it to the output - @inbounds begin - output_ptr = output_ptrs[leading_dim] = offsets[leading_dim] - if leading_dim == leading_dim_size - output_ptrs[leading_dim+one(eltype(leading_dim))] = offsets[leading_dim+one(eltype(leading_dim))] - end - end - - # set the values for this row - for (sub_leading_dim, ptrs) in iter - index_first = output isa GPUSparseDeviceMatrixCSR ? leading_dim : sub_leading_dim - index_second = output isa GPUSparseDeviceMatrixCSR ? sub_leading_dim : leading_dim - I = CartesianIndex(index_first, index_second) - vals = ntuple(Val(length(args))) do i - arg = @inbounds args[i] - ptr = @inbounds ptrs[i] - _getindex(arg, I, ptr) - end - @inbounds output_ivals[output_ptr] = sub_leading_dim - @inbounds output.nzVal[output_ptr] = f(vals...) - output_ptr += one(Ti) - end - end -end -@kernel function sparse_to_dense_broadcast_kernel(T::Type{<:Union{AbstractGPUSparseMatrixCSR{Tv, Ti}, - AbstractGPUSparseMatrixCSC{Tv, Ti}}}, - f, output::AbstractArray, args...) where {Tv, Ti} - # every thread processes an entire row - leading_dim = @index(Global, Linear) - leading_dim_size = T <: AbstractGPUSparseMatrixCSR ? size(output, 1) : size(output, 2) - if leading_dim ≤ leading_dim_size - iter = @inbounds iter_type(T, Ti)(leading_dim, args...) - - # set the values for this row - for (sub_leading_dim, ptrs) in iter - index_first = T <: AbstractGPUSparseMatrixCSR ? leading_dim : sub_leading_dim - index_second = T <: AbstractGPUSparseMatrixCSR ? sub_leading_dim : leading_dim - I = CartesianIndex(index_first, index_second) - vals = ntuple(Val(length(args))) do i - arg = @inbounds args[i] - ptr = @inbounds ptrs[i] - _getindex(arg, I, ptr) - end - - @inbounds output[I] = f(vals...) - end - end -end - -@kernel function sparse_to_dense_broadcast_kernel(::Type{<:AbstractGPUSparseVector}, f::F, - output::AbstractArray{Tv}, - offsets::AbstractVector{Pair{Ti, NTuple{N, Ti}}}, - args...) where {Tv, F, N, Ti} - row = @index(Global, Linear) - if row ≤ length(output) - arg_ptrs = @inbounds offsets[row][2] - vals = ntuple(Val(length(args))) do i - @inline - arg = @inbounds args[i] - # ptr is 0 if the sparse vector doesn't have an element at this row - # ptr is row if the arg is dense OR a scalar with non-zero-preserving f - # ptr is 0 if the arg is a scalar AND f preserves zeros - ptr = @inbounds arg_ptrs[i] - _getindex(arg, row, ptr) - end - @inbounds output[row] = f(vals...) - end -end -## COV_EXCL_STOP - -function Broadcast.copy(bc::Broadcasted{<:Union{GPUSparseVecStyle,GPUSparseMatStyle}}) - # find the sparse inputs - bc = Broadcast.flatten(bc) - sparse_args = findall(bc.args) do arg - arg isa AbstractGPUSparseArray - end - sparse_types = unique(map(i->nameof(typeof(bc.args[i])), sparse_args)) - if length(sparse_types) > 1 - error("broadcast with multiple types of sparse arrays ($(join(sparse_types, ", "))) is not supported") - end - sparse_typ = typeof(bc.args[first(sparse_args)]) - sparse_typ <: Union{AbstractGPUSparseMatrixCSR,AbstractGPUSparseMatrixCSC,AbstractGPUSparseVector} || - error("broadcast with sparse arrays is currently only implemented for vectors and CSR and CSC matrices") - Ti = if sparse_typ <: AbstractGPUSparseMatrixCSR - reduce(promote_type, map(i->eltype(bc.args[i].rowPtr), sparse_args)) - elseif sparse_typ <: AbstractGPUSparseMatrixCSC - reduce(promote_type, map(i->eltype(bc.args[i].colPtr), sparse_args)) - elseif sparse_typ <: AbstractGPUSparseVector - reduce(promote_type, map(i->eltype(bc.args[i].iPtr), sparse_args)) - end - - # determine the output type - Tv = Broadcast.combine_eltypes(bc.f, eltype.(bc.args)) - if !Base.isconcretetype(Tv) - error("""GPU sparse broadcast resulted in non-concrete element type $Tv. - This probably means that the function you are broadcasting contains an error or type instability.""") - end - - # partially-evaluate the function, removing scalars. - parevalf, passedsrcargstup = capturescalars(bc.f, bc.args) - # check if the partially-evaluated function preserves zeros. if so, we'll only need to - # apply it to the sparse input arguments, preserving the sparse structure. - if all(arg->isa(arg, AbstractSparseArray), passedsrcargstup) - fofzeros = parevalf(_zeros_eltypes(passedsrcargstup...)...) - fpreszeros = _iszero(fofzeros) - else - fpreszeros = false - end - - # the kernels below parallelize across rows or cols, not elements, so it's unlikely - # we'll launch many threads. to maximize utilization, parallelize across blocks first. - rows, cols = get(size(bc), 1, 1), get(size(bc), 2, 1) - # `size(bc, ::Int)` is missing - offsets = nothing - positions = nothing - # allocate the output container - sparse_arg = bc.args[first(sparse_args)] - if !fpreszeros && sparse_typ <: Union{AbstractGPUSparseMatrixCSR, AbstractGPUSparseMatrixCSC} - # either we have dense inputs, or the function isn't preserving zeros, - # so use a dense output to broadcast into. - val_array = nonzeros(sparse_arg) - output = similar(val_array, Tv, size(bc)) - # since we'll be iterating the sparse inputs, we need to pre-fill the dense output - # with appropriate values (while setting the sparse inputs to zero). we do this by - # re-using the dense broadcast implementation. - nonsparse_args = map(bc.args) do arg - # NOTE: this assumes the broadcast is flattened, but not yet preprocessed - if arg isa AbstractGPUSparseArray - zero(eltype(arg)) - else - arg - end - end - broadcast!(bc.f, output, nonsparse_args...) - elseif length(sparse_args) == 1 && sparse_typ <: Union{AbstractGPUSparseMatrixCSR, AbstractGPUSparseMatrixCSC} - # we only have a single sparse input, so we can reuse its structure for the output. - # this avoids a kernel launch and costly synchronization. - if sparse_typ <: AbstractGPUSparseMatrixCSR - offsets = rowPtr = sparse_arg.rowPtr - colVal = similar(sparse_arg.colVal) - nzVal = similar(sparse_arg.nzVal, Tv) - output = sparse_array_type(sparse_typ)(rowPtr, colVal, nzVal, size(bc)) - elseif sparse_typ <: AbstractGPUSparseMatrixCSC - offsets = colPtr = sparse_arg.colPtr - rowVal = similar(sparse_arg.rowVal) - nzVal = similar(sparse_arg.nzVal, Tv) - output = sparse_array_type(sparse_typ)(colPtr, rowVal, nzVal, size(bc)) - end - else - # determine the number of non-zero elements per row so that we can create an - # appropriately-structured output container - offsets = if sparse_typ <: AbstractGPUSparseMatrixCSR - ptr_array = sparse_arg.rowPtr - similar(ptr_array, Ti, rows+1) - elseif sparse_typ <: AbstractGPUSparseMatrixCSC - ptr_array = sparse_arg.colPtr - similar(ptr_array, Ti, cols+1) - elseif sparse_typ <: AbstractGPUSparseVector - ptr_array = sparse_arg.iPtr - similar(ptr_array, Pair{Ti, NTuple{length(bc.args), Ti}}, rows) - end - let - args = if sparse_typ <: AbstractGPUSparseVector - (sparse_typ, fpreszeros, offsets, bc.args...) - else - (sparse_typ, offsets, bc.args...) - end - kernel = compute_offsets_kernel(get_backend(bc.args[first(sparse_args)])) - kernel(args...; ndrange=length(offsets)) - end - # accumulate these values so that we can use them directly as row pointer offsets, - # as well as to get the total nnz count to allocate the sparse output array. - # cusparseXcsrgeam2Nnz computes this in one go, but it doesn't seem worth the effort - if !(sparse_typ <: AbstractGPUSparseVector) - @allowscalar accumulate!(Base.add_sum, offsets, offsets) - total_nnz = @allowscalar last(offsets[end]) - 1 - elseif fpreszeros - # number the rows that have an output entry, which gives each its output position - positions = similar(offsets, Ti) - positions .= first.(offsets) .!= typemax(Ti) - @allowscalar accumulate!(Base.add_sum, positions, positions) - total_nnz = rows == 0 ? 0 : Int(@allowscalar positions[end]) - end - output = if sparse_typ <: Union{AbstractGPUSparseMatrixCSR,AbstractGPUSparseMatrixCSC} - ixVal = similar(offsets, Ti, total_nnz) - nzVal = similar(offsets, Tv, total_nnz) - output_sparse_typ = sparse_array_type(sparse_typ) - output_sparse_typ(offsets, ixVal, nzVal, size(bc)) - elseif sparse_typ <: AbstractGPUSparseVector && !fpreszeros - val_array = bc.args[first(sparse_args)].nzVal - similar(val_array, Tv, size(bc)) - elseif sparse_typ <: AbstractGPUSparseVector && fpreszeros - iPtr = similar(offsets, Ti, total_nnz) - nzVal = similar(offsets, Tv, total_nnz) - sparse_array_type(sparse_arg){Tv, Ti}(iPtr, nzVal, rows) - end - if sparse_typ <: AbstractGPUSparseVector && !fpreszeros - nonsparse_args = map(bc.args) do arg - # NOTE: this assumes the broadcst is flattened, but not yet preprocessed - if arg isa AbstractGPUSparseArray - zero(eltype(arg)) - else - arg - end - end - broadcast!(bc.f, output, nonsparse_args...) - end - end - # perform the actual broadcast - if output isa AbstractGPUSparseArray - args = sparse_typ <: AbstractGPUSparseVector ? (bc.f, output, offsets, positions, bc.args...) : - (bc.f, output, offsets, bc.args...) - kernel = sparse_to_sparse_broadcast_kernel(get_backend(bc.args[first(sparse_args)])) - ndrange = if sparse_typ <: AbstractGPUSparseVector - rows - elseif sparse_typ <: AbstractGPUSparseMatrixCSC - size(output, 2) - else - size(output, 1) - end - else - args = sparse_typ <: AbstractGPUSparseVector ? (sparse_typ, bc.f, output, offsets, bc.args...) : - (sparse_typ, bc.f, output, bc.args...) - kernel = sparse_to_dense_broadcast_kernel(get_backend(bc.args[first(sparse_args)])) - ndrange = sparse_typ <: AbstractGPUSparseMatrixCSC ? size(output, 2) : size(output, 1) - end - kernel(args...; ndrange) - return output -end -# `op` folded over `n ≥ 1` copies of `x`, as the implicit zeros of a sparse array contribute -# to a reduction: by repeated squaring for an associative `op` (e.g. `n * x` for `+`), and a -# single `x` for idempotent ones. -@inline function _fold_repeated(op, x, n::Integer) - _isidempotent(op) && return x - result = x - power = x - n -= one(n) - while n > 0 - if isodd(n) - result = op(result, power) - end - power = op(power, power) - n >>= 1 - end - return result -end -_isidempotent(op) = false -_isidempotent(::Union{typeof(max), typeof(min), typeof(|), typeof(&)}) = true - -## COV_EXCL_START -# reduce the stored values `nzVal[first_ptr:last_ptr]` of a slice that also contains -# `nzeros` implicit zeros, and has at least one element. As in SparseArrays, the stored -# values are reduced first and the implicit zeros folded in afterwards, so `op` is assumed -# to be commutative as well as associative. -@inline function _reduce_slice(f, op, ::Type{T}, nzVal, first_ptr, last_ptr, nzeros) where {T} - if first_ptr <= last_ptr - val = convert(T, f(@inbounds nzVal[first_ptr])) - for ptr in first_ptr+one(first_ptr):last_ptr - val = convert(T, op(val, f(@inbounds nzVal[ptr]))) - end - if nzeros > 0 - val = convert(T, op(val, _fold_repeated(op, convert(T, f(zero(eltype(nzVal)))), nzeros))) - end - val - else - _fold_repeated(op, convert(T, f(zero(eltype(nzVal)))), nzeros) - end -end - -@kernel function csr_reduce_kernel(f::F, op::OP, output::AbstractArray{T}, A) where {F, OP, T} - # every thread processes an entire row - row = @index(Global, Linear) - if row ≤ size(A, 1) - first_ptr = @inbounds A.rowPtr[row] - last_ptr = @inbounds A.rowPtr[row+1] - one(first_ptr) - nzeros = size(A, 2) - (last_ptr - first_ptr + 1) - @inbounds output[row] = _reduce_slice(f, op, T, A.nzVal, first_ptr, last_ptr, nzeros) - end -end - -@kernel function csc_reduce_kernel(f::F, op::OP, output::AbstractArray{T}, A) where {F, OP, T} - # every thread processes an entire column - col = @index(Global, Linear) - if col ≤ size(A, 2) - first_ptr = @inbounds A.colPtr[col] - last_ptr = @inbounds A.colPtr[col+1] - one(first_ptr) - nzeros = size(A, 1) - (last_ptr - first_ptr + 1) - @inbounds output[col] = _reduce_slice(f, op, T, A.nzVal, first_ptr, last_ptr, nzeros) - end -end -## COV_EXCL_STOP - -csc_type(A::AbstractGPUSparseMatrix) = csc_type(typeof(A)) -csr_type(A::AbstractGPUSparseMatrix) = csr_type(typeof(A)) - -# TODO: implement mapreducedim! -function Base.mapreduce(f, op, A::AbstractGPUSparseMatrix{Tv}; dims=:, init=nothing) where {Tv} - in(dims, [Colon(), 1, 2]) || error("only dims=:, dims=1 or dims=2 is supported") - - # every slice reduction `r` becomes `op(init, r)`; empty slices give `init`, or the - # reduction of an empty collection as in Base - empty_result() = init === nothing ? Base.mapreduce_empty(f, op, Tv) : init - - # figure out the destination element type by relying on inference to reason through - # the map and reduce functions, or by looking at the initializer - ET = Broadcast.combine_eltypes(f, (A,)) - ET = Base.promote_op(op, ET, ET) - if init !== nothing - ET = Base.promote_op(op, typeof(init), ET) - end - (ET === Union{} || ET === Any) && - error("mapreduce cannot figure the output element type, please pass an explicit init value") - - if dims == Colon() - # reduce the stored values, then fold in the implicit zeros - length(A) == 0 && return empty_result() - nzeros = length(A) - nnz(A) - zeros = nzeros > 0 ? _fold_repeated(op, convert(ET, f(zero(Tv))), nzeros) : nothing - if nnz(A) > 0 - stored = init === nothing ? mapreduce(f, op, nonzeros(A)) : - mapreduce(f, op, nonzeros(A); init) - return zeros === nothing ? stored : op(stored, zeros) - else - return init === nothing ? zeros : op(init, zeros) - end - end - - if A isa AbstractGPUSparseMatrixCSR && dims == 1 - A = csc_type(A)(A) - elseif A isa AbstractGPUSparseMatrixCSC && dims == 2 - A = csr_type(A)(A) - end - m, n = size(A) - output = similar(nonzeros(A), ET, dims == 1 ? (1, n) : (m, 1)) - isempty(output) && return output - if size(A, dims) == 0 - return fill!(output, empty_result()) - end - kernel = A isa AbstractGPUSparseMatrixCSR ? csr_reduce_kernel(get_backend(A)) : - csc_reduce_kernel(get_backend(A)) - kernel(f, op, output, A; ndrange=length(output)) - if init !== nothing - output .= op.(init, output) - end - return output -end diff --git a/src/host/sparse/broadcast.jl b/src/host/sparse/broadcast.jl new file mode 100644 index 000000000..c765d17f8 --- /dev/null +++ b/src/host/sparse/broadcast.jl @@ -0,0 +1,614 @@ +# broadcasting + +# broadcast container type promotion for combinations of sparse arrays and other types +struct GPUSparseVecStyle <: Broadcast.AbstractArrayStyle{1} end +struct GPUSparseMatStyle <: Broadcast.AbstractArrayStyle{2} end +Broadcast.BroadcastStyle(::Type{<:GPUSparseVector}) = GPUSparseVecStyle() +Broadcast.BroadcastStyle(::Type{<:GPUSparseMatrix}) = GPUSparseMatStyle() +const SPVM = Union{GPUSparseVecStyle,GPUSparseMatStyle} + +# GPUSparseVecStyle handles 0-1 dimensions, GPUSparseMatStyle 0-2 dimensions. Higher +# dimensions keep the sparse style, so that `copy` can reject them: Base's style would +# broadcast on the host, through scalar indexing. +GPUSparseVecStyle(::Val{0}) = GPUSparseVecStyle() +GPUSparseVecStyle(::Val{1}) = GPUSparseVecStyle() +GPUSparseVecStyle(::Val) = GPUSparseMatStyle() +GPUSparseMatStyle(::Val) = GPUSparseMatStyle() + +Broadcast.BroadcastStyle(::GPUSparseVecStyle, ::AbstractGPUArrayStyle{N}) where N = + GPUSparseVecStyle(Val(N)) +Broadcast.BroadcastStyle(::GPUSparseVecStyle, ::GPUSparseMatStyle) = GPUSparseMatStyle() +Broadcast.BroadcastStyle(::GPUSparseMatStyle, ::GPUSparseVecStyle) = GPUSparseMatStyle() +Broadcast.BroadcastStyle(::GPUSparseMatStyle, ::AbstractGPUArrayStyle) = GPUSparseMatStyle() + +# don't wrap sparse arrays with Extruded +Broadcast.extrude(x::GPUSparseArray) = x + +## detection of zero-preserving functions + +# modified from SparseArrays.jl + +# capturescalars takes a function (f) and a tuple of broadcast arguments, and returns a +# partially-evaluated function and a reduced argument tuple where all scalar operations have +# been applied already. +@inline function capturescalars(f, mixedargs) + let (passedsrcargstup, makeargs) = _capturescalars(mixedargs...) + parevalf = (passed...) -> f(makeargs(passed...)...) + return (parevalf, passedsrcargstup) + end +end + +## sparse broadcast style + +# Work around losing Type{T}s as DataTypes within the tuple that makeargs creates +@inline capturescalars(f, mixedargs::Tuple{Ref{Type{T}}, Vararg{Any}}) where {T} = + capturescalars((args...)->f(T, args...), Base.tail(mixedargs)) +@inline capturescalars(f, mixedargs::Tuple{Ref{Type{T}}, Ref{Type{S}}, Vararg{Any}}) where {T, S} = + # This definition is identical to the one above and necessary only for + # avoiding method ambiguity. + capturescalars((args...)->f(T, args...), Base.tail(mixedargs)) +@inline capturescalars(f, mixedargs::Tuple{GPUSparseArray, Ref{Type{T}}, Vararg{Any}}) where {T} = + capturescalars((a1, args...)->f(a1, T, args...), (mixedargs[1], Base.tail(Base.tail(mixedargs))...)) +@inline capturescalars(f, mixedargs::Tuple{Union{Ref,AbstractArray{<:Any,0}}, Ref{Type{T}}, Vararg{Any}}) where {T} = + capturescalars((args...)->f(mixedargs[1], T, args...), Base.tail(Base.tail(mixedargs))) + +scalararg(::Number) = true +scalararg(::Any) = false +scalarwrappedarg(::Union{AbstractArray{<:Any,0},Ref}) = true +scalarwrappedarg(::Any) = false + +@inline function _capturescalars() + return (), () -> () +end +@inline function _capturescalars(arg, mixedargs...) + let (rest, f) = _capturescalars(mixedargs...) + if scalararg(arg) + return rest, @inline function(tail...) + (arg, f(tail...)...) + end # add back scalararg after (in makeargs) + elseif scalarwrappedarg(arg) + return rest, @inline function(tail...) + (arg[], f(tail...)...) # TODO: This can put a Type{T} in a tuple + end # unwrap and add back scalararg after (in makeargs) + else + return (arg, rest...), @inline function(head, tail...) + (head, f(tail...)...) + end # pass-through to broadcast + end + end +end +@inline function _capturescalars(arg) # this definition is just an optimization (to bottom out the recursion slightly sooner) + if scalararg(arg) + return (), () -> (arg,) # add scalararg + elseif scalarwrappedarg(arg) + return (), () -> (arg[],) # unwrap + else + return (arg,), (head,) -> (head,) # pass-through + end +end + +@inline _iszero(x) = x == 0 +@inline _iszero(x::Number) = Base.iszero(x) +@inline _iszero(x::AbstractArray) = Base.iszero(x) +@inline _zeros_eltypes(A) = (zero(eltype(A)),) +@inline _zeros_eltypes(A, Bs...) = (zero(eltype(A)), _zeros_eltypes(Bs...)...) + +## COV_EXCL_START +## iteration helpers + +""" + CSRIterator{Ti}(row, args...) + +A GPU-compatible iterator for accessing the elements of a single row `row` of several CSR +matrices `args` in one go. The row should be in-bounds for every sparse argument. Each +iteration returns a 2-element tuple: The current column, and each arguments' pointer index +(or 0 if that input didn't have an element at that column). The pointers can then be used to +access the elements themselves. + +For convenience, this iterator can be passed non-sparse arguments as well, which will be +ignored (with the returned `col`/`ptr` values set to 0). +""" +struct CSRIterator{Ti,N,ATs} + row::Ti + col_ends::NTuple{N, Ti} + args::ATs +end + +function CSRIterator{Ti}(row, args::Vararg{Any, N}) where {Ti,N} + # check that `row` is valid for all arguments + @boundscheck begin + ntuple(Val(N)) do i + arg = @inbounds args[i] + arg isa GPUSparseMatrixCSR && checkbounds(axes(arg, 1), row) + end + end + + col_ends = ntuple(Val(N)) do i + arg = @inbounds args[i] + if arg isa GPUSparseMatrixCSR + @inbounds(arg.rowPtr[row+1]) + else + zero(Ti) + end + end + + CSRIterator{Ti, N, typeof(args)}(row, col_ends, args) +end + +@inline function Base.iterate(iter::CSRIterator{Ti,N}, state=nothing) where {Ti,N} + # helper function to get the column of a sparse array at a specific pointer + @inline function get_col(i, ptr) + arg = @inbounds iter.args[i] + if arg isa GPUSparseMatrixCSR + col_end = @inbounds iter.col_ends[i] + if ptr < col_end + return @inbounds arg.colVal[ptr] % Ti + end + end + zero(Ti) + end + + # initialize the state + # - ptr: the current index into the colVal/nzVal arrays + # - col: the current column index (cached so that we don't have to re-read each time) + state = something(state, + ntuple(Val(N)) do i + arg = @inbounds iter.args[i] + if arg isa GPUSparseMatrixCSR + ptr = @inbounds iter.args[i].rowPtr[iter.row] % Ti + col = @inbounds get_col(i, ptr) + else + ptr = zero(Ti) + col = zero(Ti) + end + (; ptr, col) + end + ) + + # determine the column we're currently processing + cols = ntuple(i -> @inbounds(state[i].col), Val(N)) + cur_col = reduce((a, b) -> iszero(a) ? b : iszero(b) ? a : min(a, b), cols) + iszero(cur_col) && return + + # fetch the pointers (we don't look up the values, as the caller might want to index + # the sparse array directly, e.g., to mutate it). we don't return `ptrs` from the state + # directly, but first convert the `typemax(Ti)` to a more convenient zero value. + # NOTE: these values may end up unused by the caller (e.g. in the count_nnzs kernels), + # but LLVM appears smart enough to filter them away. + ptrs = ntuple(Val(N)) do i + ptr, col = @inbounds state[i] + col == cur_col ? ptr : zero(Ti) + end + + # advance the state + new_state = ntuple(Val(N)) do i + ptr, col = @inbounds state[i] + if col == cur_col + ptr += one(Ti) + col = get_col(i, ptr) + end + (; ptr, col) + end + + return (cur_col, ptrs), new_state +end + +struct CSCIterator{Ti,N,ATs} + col::Ti + row_ends::NTuple{N, Ti} + args::ATs +end + +function CSCIterator{Ti}(col, args::Vararg{Any, N}) where {Ti,N} + # check that `col` is valid for all arguments + @boundscheck begin + ntuple(Val(N)) do i + arg = @inbounds args[i] + arg isa GPUSparseMatrixCSC && checkbounds(axes(arg, 2), col) + end + end + + row_ends = ntuple(Val(N)) do i + arg = @inbounds args[i] + x = if arg isa GPUSparseMatrixCSC + @inbounds(arg.colPtr[col+1]) + else + zero(Ti) + end + x + end + + CSCIterator{Ti, N, typeof(args)}(col, row_ends, args) +end + +@inline function Base.iterate(iter::CSCIterator{Ti,N}, state=nothing) where {Ti,N} + # helper function to get the column of a sparse array at a specific pointer + @inline function get_col(i, ptr) + arg = @inbounds iter.args[i] + if arg isa GPUSparseMatrixCSC + col_end = @inbounds iter.row_ends[i] + if ptr < col_end + return @inbounds arg.rowVal[ptr] % Ti + end + end + zero(Ti) + end + + # initialize the state + # - ptr: the current index into the rowVal/nzVal arrays + # - row: the current row index (cached so that we don't have to re-read each time) + state = something(state, + ntuple(Val(N)) do i + arg = @inbounds iter.args[i] + if arg isa GPUSparseMatrixCSC + ptr = @inbounds iter.args[i].colPtr[iter.col] % Ti + row = @inbounds get_col(i, ptr) + else + ptr = zero(Ti) + row = zero(Ti) + end + (; ptr, row) + end + ) + + # determine the row we're currently processing + rows = ntuple(i -> @inbounds(state[i].row), Val(N)) + cur_row = reduce((a, b) -> iszero(a) ? b : iszero(b) ? a : min(a, b), rows) + iszero(cur_row) && return + + # fetch the pointers (we don't look up the values, as the caller might want to index + # the sparse array directly, e.g., to mutate it). we don't return `ptrs` from the state + # directly, but first convert the `typemax(Ti)` to a more convenient zero value. + # NOTE: these values may end up unused by the caller (e.g. in the count_nnzs kernels), + # but LLVM appears smart enough to filter them away. + ptrs = ntuple(Val(N)) do i + ptr, row = @inbounds state[i] + row == cur_row ? ptr : zero(Ti) + end + + # advance the state + new_state = ntuple(Val(N)) do i + ptr, row = @inbounds state[i] + if row == cur_row + ptr += one(Ti) + row = get_col(i, ptr) + end + (; ptr, row) + end + + return (cur_row, ptrs), new_state +end + +# helpers to index a sparse or dense array +function _getindex(arg::Union{GPUSparseMatrixCSR,GPUSparseMatrixCSC,GPUSparseVector}, I, ptr) + if ptr == 0 + zero(eltype(arg)) + else + @inbounds arg.nzVal[ptr] + end +end +@inline function _getindex(arg::DenseArray{Tv}, I, ptr)::Tv where {Tv} + return @inbounds Broadcast._broadcast_getindex(arg, I)::Tv +end +@inline _getindex(arg, I, ptr) = Broadcast._broadcast_getindex(arg, I) + +## sparse broadcast implementation +iter_type(::Type{<:GPUSparseMatrixCSC}, ::Type{Ti}) where {Ti} = CSCIterator{Ti} +iter_type(::Type{<:GPUSparseMatrixCSR}, ::Type{Ti}) where {Ti} = CSRIterator{Ti} + +sparse_iterator(::Val{:csr}, row, args...) = CSRIterator{Int}(row, args...) +sparse_iterator(::Val{:csc}, col, args...) = CSCIterator{Int}(col, args...) + +_has_row(A, offsets, row, fpreszeros::Bool) = fpreszeros ? 0 : row +_has_row(A::AbstractDeviceArray, offsets, row, ::Bool) = row +# the position of `row` among the stored indices of a sparse vector, or 0 if it isn't stored +function _has_row(A::GPUSparseVector, offsets, row, ::Bool) + ptr = searchsortedfirst(A.nzInd, row) + return (ptr <= length(A.nzInd) && @inbounds(A.nzInd[ptr]) == row) ? ptr : 0 +end + +# for every row, determine which arguments have an entry there (see `_has_row`), and whether +# the output has one (`key` is the row, or `typemax(Ti)` if not) +@kernel function compute_offsets_kernel(::Type{<:GPUSparseVector}, fpreszeros::Bool, + offsets::AbstractVector{Pair{Ti, NTuple{N, Ti}}}, + args...) where {Ti, N} + row = @index(Global, Linear) + if row ≤ length(offsets) + arg_row_is_nnz = ntuple(Val(N)) do i + arg = @inbounds args[i] + _has_row(arg, offsets, row, fpreszeros) % Ti + end + row_is_nnz = false + for i in 1:N + row_is_nnz |= @inbounds(arg_row_is_nnz[i]) != 0 + end + key = row_is_nnz ? row % Ti : zero(Ti) + @inbounds offsets[row] = key => arg_row_is_nnz + end +end + +# kernel to count the number of non-zeros in a row, to determine the row offsets +@kernel function compute_offsets_kernel(layout::Val, + offsets::AbstractVector{Ti}, args...) where Ti + # every thread processes an entire row + leading_dim = @index(Global, Linear) + if leading_dim ≤ length(offsets)-1 + iter = @inbounds sparse_iterator(layout, leading_dim, args...) + + # count the nonzero leading_dims of all inputs + accum = zero(Ti) + for (leading_dim, vals) in iter + accum += one(Ti) + end + + # the counts are shifted by one (`offsets[1]` is set to 1 by the caller) so that + # after accumulation they can be used as the rowPtr/colPtr array of a CSR/CSC matrix + # Keep this store bounds checked: Metal miscompiles the count-only iterator + # loop with narrow input pointers and wide counts without this branch. + offsets[leading_dim+1] = accum + end +end + +# `positions` is the inclusive prefix count of output entries, so an output entry for `row` +# goes to `positions[row]`. +@kernel function sparse_to_sparse_broadcast_kernel(f::F, output::GPUSparseVector{Tv,Ti}, + offsets::AbstractVector{Pair{Int, NTuple{N, Int}}}, + positions::AbstractVector{Int}, + args...) where {Tv, Ti, N, F} + row = @index(Global, Linear) + if row ≤ length(offsets) + key, arg_ptrs = @inbounds offsets[row] + if key != 0 + vals = ntuple(Val(N)) do i + @inline + arg = @inbounds args[i] + # ptr is 0 if the sparse vector doesn't have an element at this row + # ptr is 0 if the arg is a scalar AND f preserves zeros + ptr = @inbounds arg_ptrs[i] + _getindex(arg, row, ptr) + end + output_ix = @inbounds positions[row] + @inbounds output.nzInd[output_ix] = row + @inbounds output.nzVal[output_ix] = f(vals...) + end + end +end + +@kernel function sparse_to_sparse_broadcast_kernel(f, output::T, offsets::Union{<:AbstractArray,Nothing}, + args...) where {Ti, T<:Union{GPUSparseMatrixCSR{<:Any,Ti}, + GPUSparseMatrixCSC{<:Any,Ti}}} + # every thread processes an entire row + leading_dim = @index(Global, Linear) + leading_dim_size = output isa GPUSparseMatrixCSR ? size(output, 1) : size(output, 2) + if leading_dim ≤ leading_dim_size + iter = @inbounds iter_type(T, Int)(leading_dim, args...) + + output_ptrs = output isa GPUSparseMatrixCSR ? output.rowPtr : output.colPtr + output_ivals = output isa GPUSparseMatrixCSR ? output.colVal : output.rowVal + # fetch the row offset, and write it to the output + @inbounds begin + output_ptr = output_ptrs[leading_dim] = offsets[leading_dim] + if leading_dim == leading_dim_size + output_ptrs[leading_dim+one(eltype(leading_dim))] = offsets[leading_dim+one(eltype(leading_dim))] + end + end + + # set the values for this row + for (sub_leading_dim, ptrs) in iter + index_first = output isa GPUSparseMatrixCSR ? leading_dim : sub_leading_dim + index_second = output isa GPUSparseMatrixCSR ? sub_leading_dim : leading_dim + I = CartesianIndex(index_first, index_second) + vals = ntuple(Val(length(args))) do i + arg = @inbounds args[i] + ptr = @inbounds ptrs[i] + _getindex(arg, I, ptr) + end + @inbounds output_ivals[output_ptr] = sub_leading_dim + @inbounds output.nzVal[output_ptr] = f(vals...) + output_ptr += one(Ti) + end + end +end +@kernel function sparse_to_dense_broadcast_kernel(T::Type{<:Union{GPUSparseMatrixCSR{Tv, Ti}, + GPUSparseMatrixCSC{Tv, Ti}}}, + f, output::AbstractArray, args...) where {Tv, Ti} + # every thread processes an entire row + leading_dim = @index(Global, Linear) + leading_dim_size = T <: GPUSparseMatrixCSR ? size(output, 1) : size(output, 2) + if leading_dim ≤ leading_dim_size + iter = @inbounds iter_type(T, Int)(leading_dim, args...) + + # set the values for this row + for (sub_leading_dim, ptrs) in iter + index_first = T <: GPUSparseMatrixCSR ? leading_dim : sub_leading_dim + index_second = T <: GPUSparseMatrixCSR ? sub_leading_dim : leading_dim + I = CartesianIndex(index_first, index_second) + vals = ntuple(Val(length(args))) do i + arg = @inbounds args[i] + ptr = @inbounds ptrs[i] + _getindex(arg, I, ptr) + end + + @inbounds output[I] = f(vals...) + end + end +end + +@kernel function sparse_to_dense_broadcast_kernel(::Type{<:GPUSparseVector}, f::F, + output::AbstractArray{Tv}, + offsets::AbstractVector{Pair{Ti, NTuple{N, Ti}}}, + args...) where {Tv, F, N, Ti} + row = @index(Global, Linear) + if row ≤ length(output) + arg_ptrs = @inbounds offsets[row][2] + vals = ntuple(Val(length(args))) do i + @inline + arg = @inbounds args[i] + # ptr is 0 if the sparse vector doesn't have an element at this row + # ptr is row if the arg is dense OR a scalar with non-zero-preserving f + # ptr is 0 if the arg is a scalar AND f preserves zeros + ptr = @inbounds arg_ptrs[i] + _getindex(arg, row, ptr) + end + @inbounds output[row] = f(vals...) + end +end +## COV_EXCL_STOP + +function Broadcast.copy(bc::Broadcasted{<:Union{GPUSparseVecStyle,GPUSparseMatStyle}}) + # find the sparse inputs + bc = Broadcast.flatten(bc) + sparse_args = findall(bc.args) do arg + arg isa GPUSparseArray + end + sparse_types = unique(map(i->nameof(typeof(bc.args[i])), sparse_args)) + if length(sparse_types) > 1 + error("broadcast with multiple types of sparse arrays ($(join(sparse_types, ", "))) is not supported") + end + sparse_typ = typeof(bc.args[first(sparse_args)]) + sparse_typ <: Union{GPUSparseMatrixCSR,GPUSparseMatrixCSC,GPUSparseVector} || + error("broadcast with sparse arrays is currently only implemented for vectors and CSR and CSC matrices") + Ti = reduce(promote_type, map(i -> indtype(bc.args[i]), sparse_args)) + + # determine the output type + Tv = Broadcast.combine_eltypes(bc.f, eltype.(bc.args)) + if !Base.isconcretetype(Tv) + error("""GPU sparse broadcast resulted in non-concrete element type $Tv. + This probably means that the function you are broadcasting contains an error or type instability.""") + end + + # partially-evaluate the function, removing scalars. + parevalf, passedsrcargstup = capturescalars(bc.f, bc.args) + # check if the partially-evaluated function preserves zeros. if so, we'll only need to + # apply it to the sparse input arguments, preserving the sparse structure. + if all(arg->isa(arg, AbstractSparseArray), passedsrcargstup) + fofzeros = parevalf(_zeros_eltypes(passedsrcargstup...)...) + fpreszeros = _iszero(fofzeros) + else + fpreszeros = false + end + + # the kernels below parallelize across rows or cols, not elements, so it's unlikely + # we'll launch many threads. to maximize utilization, parallelize across blocks first. + rows, cols = get(size(bc), 1, 1), get(size(bc), 2, 1) + # `size(bc, ::Int)` is missing + offsets = nothing + positions = nothing + # allocate the output container + sparse_arg = bc.args[first(sparse_args)] + if !fpreszeros && sparse_typ <: Union{GPUSparseMatrixCSR, GPUSparseMatrixCSC} + # either we have dense inputs, or the function isn't preserving zeros, + # so use a dense output to broadcast into. + val_array = nonzeros(sparse_arg) + output = similar(val_array, Tv, size(bc)) + # since we'll be iterating the sparse inputs, we need to pre-fill the dense output + # with appropriate values (while setting the sparse inputs to zero). we do this by + # re-using the dense broadcast implementation. + nonsparse_args = map(bc.args) do arg + # NOTE: this assumes the broadcast is flattened, but not yet preprocessed + if arg isa GPUSparseArray + zero(eltype(arg)) + else + arg + end + end + broadcast!(bc.f, output, nonsparse_args...) + elseif length(sparse_args) == 1 && sparse_typ <: Union{GPUSparseMatrixCSR, GPUSparseMatrixCSC} + # we only have a single sparse input, so we can reuse its structure for the output. + # this avoids a kernel launch and costly synchronization. + if sparse_typ <: GPUSparseMatrixCSR + offsets = rowPtr = copy(sparse_arg.rowPtr) + colVal = similar(sparse_arg.colVal) + nzVal = similar(sparse_arg.nzVal, Tv) + output = GPUSparseMatrixCSR(rowPtr, colVal, nzVal, size(bc)) + elseif sparse_typ <: GPUSparseMatrixCSC + offsets = colPtr = copy(sparse_arg.colPtr) + rowVal = similar(sparse_arg.rowVal) + nzVal = similar(sparse_arg.nzVal, Tv) + output = GPUSparseMatrixCSC(colPtr, rowVal, nzVal, size(bc)) + end + else + # determine the number of non-zero elements per row so that we can create an + # appropriately-structured output container. The counts are accumulated in `Int`: + # the union of the structures can have more entries than the index type holds. + offsets = if sparse_typ <: GPUSparseMatrixCSR + ptr_array = sparse_arg.rowPtr + fill!(similar(ptr_array, Int, rows + 1), 1) + elseif sparse_typ <: GPUSparseMatrixCSC + ptr_array = sparse_arg.colPtr + fill!(similar(ptr_array, Int, cols + 1), 1) + elseif sparse_typ <: GPUSparseVector + ptr_array = sparse_arg.nzInd + similar(ptr_array, Pair{Int, NTuple{length(bc.args), Int}}, rows) + end + let + args = if sparse_typ <: GPUSparseVector + (sparse_typ, fpreszeros, offsets, bc.args...) + else + (sparse_typ <: GPUSparseMatrixCSR ? Val(:csr) : Val(:csc), offsets, bc.args...) + end + kernel = compute_offsets_kernel(get_backend(bc.args[first(sparse_args)])) + # an empty launch fails on Metal (JuliaGPU/Metal.jl#980) + isempty(offsets) || kernel(args...; ndrange=length(offsets)) + end + # accumulate these values so that we can use them directly as row pointer offsets, + # as well as to get the total nnz count to allocate the sparse output array. + # cusparseXcsrgeam2Nnz computes this in one go, but it doesn't seem worth the effort + if !(sparse_typ <: GPUSparseVector) + accumulate!(Base.add_sum, offsets, offsets) + total_nnz = @allowscalar offsets[end] - 1 + elseif fpreszeros + # number the rows that have an output entry, which gives each its output position + positions = similar(offsets, Int) + positions .= first.(offsets) .!= 0 + rows == 0 || accumulate!(Base.add_sum, positions, positions) + total_nnz = rows == 0 ? 0 : Int(@allowscalar positions[end]) + end + output = if sparse_typ <: Union{GPUSparseMatrixCSR,GPUSparseMatrixCSC} + check_sparse_nnz(Ti, total_nnz) + offsets = convert_index_buffer(Ti, offsets) + ixVal = similar(offsets, Ti, total_nnz) + nzVal = similar(offsets, Tv, total_nnz) + sparse_format_type(sparse_arg)(offsets, ixVal, nzVal, size(bc)) + elseif sparse_typ <: GPUSparseVector && !fpreszeros + val_array = bc.args[first(sparse_args)].nzVal + similar(val_array, Tv, size(bc)) + elseif sparse_typ <: GPUSparseVector && fpreszeros + nzInd = similar(offsets, Ti, total_nnz) + nzVal = similar(offsets, Tv, total_nnz) + GPUSparseVector(nzInd, nzVal, rows) + end + if sparse_typ <: GPUSparseVector && !fpreszeros + nonsparse_args = map(bc.args) do arg + # NOTE: this assumes the broadcst is flattened, but not yet preprocessed + if arg isa GPUSparseArray + zero(eltype(arg)) + else + arg + end + end + broadcast!(bc.f, output, nonsparse_args...) + end + end + # perform the actual broadcast + if output isa GPUSparseArray + args = sparse_typ <: GPUSparseVector ? (bc.f, output, offsets, positions, bc.args...) : + (bc.f, output, offsets, bc.args...) + kernel = sparse_to_sparse_broadcast_kernel(get_backend(bc.args[first(sparse_args)])) + ndrange = if sparse_typ <: GPUSparseVector + rows + elseif sparse_typ <: GPUSparseMatrixCSC + size(output, 2) + else + size(output, 1) + end + else + args = sparse_typ <: GPUSparseVector ? (sparse_typ, bc.f, output, offsets, bc.args...) : + (sparse_typ, bc.f, output, bc.args...) + kernel = sparse_to_dense_broadcast_kernel(get_backend(bc.args[first(sparse_args)])) + ndrange = sparse_typ <: GPUSparseMatrixCSC ? size(output, 2) : size(output, 1) + end + # an empty launch fails on Metal (JuliaGPU/Metal.jl#980) + ndrange == 0 || kernel(args...; ndrange) + return output +end diff --git a/src/host/sparse/conversions.jl b/src/host/sparse/conversions.jl new file mode 100644 index 000000000..6ecf8d43f --- /dev/null +++ b/src/host/sparse/conversions.jl @@ -0,0 +1,304 @@ +# conversions between sparse formats, and to and from the host + +## to the host + +SparseArrays.SparseMatrixCSC(A::GPUSparseMatrixCSC) = + SparseMatrixCSC(size(A)..., Array(A.colPtr), Array(A.rowVal), Array(A.nzVal)) +# the CSR buffers are those of the transpose in CSC +SparseArrays.SparseMatrixCSC(A::GPUSparseMatrixCSR) = + copy(transpose(SparseMatrixCSC(reverse(size(A))..., Array(A.rowPtr), Array(A.colVal), + Array(A.nzVal)))) +function SparseArrays.SparseMatrixCSC(A::GPUSparseMatrixCOO) + rowInd = Array(A.rowInd) + rowPtr = zeros(eltype(rowInd), size(A, 1) + 1) + rowPtr[1] = 1 + for i in rowInd + rowPtr[i + 1] += one(eltype(rowPtr)) + end + cumsum!(rowPtr, rowPtr) + copy(transpose(SparseMatrixCSC(reverse(size(A))..., rowPtr, Array(A.colInd), + Array(A.nzVal)))) +end +SparseArrays.SparseMatrixCSC(x::GPUSparseVector) = SparseMatrixCSC(SparseVector(x)) +SparseArrays.SparseMatrixCSC{Tv}(A::GPUSparseArray) where {Tv} = + SparseMatrixCSC{Tv}(SparseMatrixCSC(A)) +SparseArrays.SparseMatrixCSC{Tv,Ti}(A::GPUSparseArray) where {Tv,Ti} = + SparseMatrixCSC{Tv,Ti}(SparseMatrixCSC(A)) + +SparseArrays.SparseVector(x::GPUSparseVector) = + SparseVector(length(x), Array(x.nzInd), Array(x.nzVal)) +SparseArrays.SparseVector{Tv}(x::GPUSparseVector) where {Tv} = SparseVector{Tv}(SparseVector(x)) +SparseArrays.SparseVector{Tv,Ti}(x::GPUSparseVector) where {Tv,Ti} = + SparseVector{Tv,Ti}(SparseVector(x)) + +host_sparse(A::GPUSparseMatrix) = SparseMatrixCSC(A) +host_sparse(x::GPUSparseVector) = SparseVector(x) + +Base.Array(A::GPUSparseArray) = Array(host_sparse(A)) +Base.Array{T}(A::GPUSparseArray) where {T} = Array{T}(host_sparse(A)) +Base.Array{T,1}(x::GPUSparseVector) where {T} = Array{T,1}(host_sparse(x)) +Base.Array{T,2}(A::GPUSparseMatrix) where {T} = Array{T,2}(host_sparse(A)) +Base.collect(A::GPUSparseArray) = Array(A) + + +## `adapt` + +# Kernel arguments and other adaptors: convert every buffer. With device vectors, the +# result is the isbits representation that kernels receive. Host buffers give SparseArrays' +# own types where there is one. +Adapt.adapt_structure(to, A::GPUSparseMatrixCSR) = + GPUSparseMatrixCSR(adapt(to, A.rowPtr), adapt(to, A.colVal), adapt(to, A.nzVal), size(A)) +Adapt.adapt_structure(to, A::GPUSparseMatrixCSC) = + sparse_csc(adapt(to, A.colPtr), adapt(to, A.rowVal), adapt(to, A.nzVal), size(A)) +Adapt.adapt_structure(to, A::GPUSparseMatrixCOO) = + GPUSparseMatrixCOO(adapt(to, A.rowInd), adapt(to, A.colInd), adapt(to, A.nzVal), size(A)) +Adapt.adapt_structure(to, x::GPUSparseVector) = + sparse_vector(adapt(to, x.nzInd), adapt(to, x.nzVal), length(x)) + +sparse_csc(colPtr::Vector, rowVal::Vector, nzVal::Vector, dims) = + SparseMatrixCSC(dims..., colPtr, rowVal, nzVal) +sparse_csc(colPtr, rowVal, nzVal, dims) = GPUSparseMatrixCSC(colPtr, rowVal, nzVal, dims) +sparse_vector(nzInd::Vector, nzVal::Vector, len) = SparseVector(len, nzInd, nzVal) +sparse_vector(nzInd, nzVal, len) = GPUSparseVector(nzInd, nzVal, len) + +# Storage types (`adapt(Array, A)`, `adapt(MtlArray, S)`): move the buffers, keeping the +# format, element and index types. The index buffers are allocated like the adapted +# values, so that a target with an element type (`Array{Float32}`, or the +# `CuArray{Float32,1,M}` that `cu` adapts to) only converts the values. +move_buffer(proto::AbstractVector, src::AbstractVector) = + copyto!(similar(proto, eltype(src), length(src)), src) + +function Adapt.adapt_structure(to::Type{<:DenseArray}, A::GPUSparseMatrixCSR) + nzVal = adapt(to, A.nzVal) + nzVal === A.nzVal && return A + GPUSparseMatrixCSR(move_buffer(nzVal, A.rowPtr), move_buffer(nzVal, A.colVal), nzVal, + size(A)) +end +function Adapt.adapt_structure(to::Type{<:DenseArray}, A::GPUSparseMatrixCSC) + nzVal = adapt(to, A.nzVal) + colPtr, rowVal = nzVal === A.nzVal ? (A.colPtr, A.rowVal) : + (move_buffer(nzVal, A.colPtr), move_buffer(nzVal, A.rowVal)) + sparse_csc(colPtr, rowVal, nzVal, size(A)) +end +function Adapt.adapt_structure(to::Type{<:DenseArray}, A::GPUSparseMatrixCOO) + nzVal = adapt(to, A.nzVal) + nzVal === A.nzVal && return A + GPUSparseMatrixCOO(move_buffer(nzVal, A.rowInd), move_buffer(nzVal, A.colInd), nzVal, + size(A)) +end +function Adapt.adapt_structure(to::Type{<:DenseArray}, x::GPUSparseVector) + nzVal = adapt(to, x.nzVal) + nzInd = nzVal === x.nzVal ? x.nzInd : move_buffer(nzVal, x.nzInd) + sparse_vector(nzInd, nzVal, length(x)) +end + +# host sparse arrays to GPU storage, likewise keeping format and types: a CSC matrix stays +# CSC. Adapting never densifies. +function Adapt.adapt_structure(to::Type{<:AbstractGPUArray}, S::SparseMatrixCSC) + nzVal = adapt(to, nonzeros(S)) + GPUSparseMatrixCSC(move_buffer(nzVal, getcolptr(S)), move_buffer(nzVal, rowvals(S)), + nzVal, size(S)) +end +function Adapt.adapt_structure(to::Type{<:AbstractGPUArray}, x::SparseVector) + nzVal = adapt(to, nonzeros(x)) + GPUSparseVector(move_buffer(nzVal, SparseArrays.nonzeroinds(x)), nzVal, length(x)) +end + +# explicit sparse types (`adapt(MtlSparseMatrixCSR{Float32,Int32}, S)`) are constructors +for S in (:SparseMatrixCSC, :SparseVector, :GPUSparseMatrixCSR, :GPUSparseMatrixCSC, + :GPUSparseMatrixCOO, :GPUSparseVector) + @eval Adapt.adapt_structure(::Type{T}, A::$S) where {T<:AbstractGPUSparseArray} = T(A) +end + + +## between formats + +# the generic types do not determine a storage, so they cannot be constructed from host +# arrays: that needs a back-end alias, or `adapt` to a storage type +for S in (:GPUSparseMatrixCSR, :GPUSparseMatrixCSC, :GPUSparseMatrixCOO, :GPUSparseVector) + @eval function $S(A::Union{SparseMatrixCSC,SparseVector}) + throw(ArgumentError("""$($S)(::$(nameof(typeof(A)))) does not know where to store the result. + Use a back-end alias (e.g. `MtlSparseMatrixCSR(S)`), `adapt(MtlArray, S)` to keep the format, or `mtl(S)`/`cu(S)`.""")) + end +end + +# the same format: a copy +GPUSparseMatrixCSR(A::GPUSparseMatrixCSR) = copy(A) +GPUSparseMatrixCSC(A::GPUSparseMatrixCSC) = copy(A) +GPUSparseMatrixCOO(A::GPUSparseMatrixCOO) = copy(A) +GPUSparseVector(x::GPUSparseVector) = copy(x) + +# between vectors and single-column matrices, as SparseArrays converts them +function GPUSparseVector(A::GPUSparseMatrix) + size(A, 2) == 1 || + throw(DimensionMismatch("a sparse vector can only be constructed from a single-column matrix")) + return vec(A) +end +function GPUSparseMatrixCSC(x::GPUSparseVector) + Ti = indtype(x) + colPtr = copyto!(similar(x.nzInd, 2), Ti[1, nnz(x) + 1]) + GPUSparseMatrixCSC(colPtr, copy(x.nzInd), copy(x.nzVal), (length(x), 1)) +end +GPUSparseMatrixCSR(x::GPUSparseVector) = GPUSparseMatrixCSR(GPUSparseMatrixCSC(x)) +GPUSparseMatrixCOO(x::GPUSparseVector) = GPUSparseMatrixCOO(GPUSparseMatrixCSR(x)) + +# typed forms: convert the format, then the element and index types (checking first that +# the result fits, as the conversion may allocate buffers sized by the dimensions) +for S in (:GPUSparseMatrixCSR, :GPUSparseMatrixCSC, :GPUSparseMatrixCOO, :GPUSparseVector) + @eval begin + $S{Tv}(A::AbstractGPUSparseArray) where {Tv} = $S{Tv,indtype(A)}(A) + function $S{Tv,Ti}(A::AbstractGPUSparseArray) where {Tv,Ti} + check_sparse_dims(Ti, size(A)) + check_sparse_nnz(Ti, nnz(A)) + return with_eltypes($S(A), Tv, Ti) + end + $S{Tv,Ti}(A::$S) where {Tv,Ti} = with_eltypes(A, Tv, Ti; copy=true) + end +end + +# `A` with the given element and index types, reusing its buffers where they already match +# unless `copy` is set +function with_eltypes(A::GPUSparseArray, ::Type{Tv}, ::Type{Ti}; copy::Bool=false) where {Tv,Ti} + Tv == eltype(A) && Ti == indtype(A) && return copy ? Base.copy(A) : A + check_sparse_dims(Ti, size(A)) + check_sparse_nnz(Ti, nnz(A)) + buffers = map(fieldnames(typeof(A))) do field + buf = getfield(A, field) + if field === :nzVal + convert_buffer(Tv, buf) + elseif buf isa AbstractVector + convert_index_buffer(Ti, buf) + else + buf + end + end + return sparse_format_type(A)(buffers...) +end +sparse_format_type(::GPUSparseMatrixCSR) = GPUSparseMatrixCSR +sparse_format_type(::GPUSparseMatrixCSC) = GPUSparseMatrixCSC +sparse_format_type(::GPUSparseMatrixCOO) = GPUSparseMatrixCOO +sparse_format_type(::GPUSparseVector) = GPUSparseVector + +""" + GPUArrays.generic_regroup(A::GPUSparseMatrixCSR)::GPUSparseMatrixCSC + GPUArrays.generic_regroup(A::GPUSparseMatrixCSC)::GPUSparseMatrixCSR + +Convert between CSR and CSC by regrouping the stored entries along the other dimension, on +the device. This is the generic implementation behind `GPUSparseMatrixCSC(::GPUSparseMatrixCSR)` +and vice versa, which back-ends can call when a vendor routine does not apply. +""" +function generic_regroup(A::GPUSparseMatrixCSR) + colPtr, rowVal, nzVal = regroup(A.rowPtr, A.colVal, A.nzVal, size(A)) + GPUSparseMatrixCSC(colPtr, rowVal, nzVal, size(A)) +end +function generic_regroup(A::GPUSparseMatrixCSC) + rowPtr, colVal, nzVal = regroup(A.colPtr, A.rowVal, A.nzVal, reverse(size(A))) + GPUSparseMatrixCSR(rowPtr, colVal, nzVal, size(A)) +end +GPUSparseMatrixCSC(A::GPUSparseMatrixCSR) = generic_regroup(A) +GPUSparseMatrixCSR(A::GPUSparseMatrixCSC) = generic_regroup(A) + +""" + GPUArrays.generic_expand(A::GPUSparseMatrixCSR)::GPUSparseMatrixCOO + GPUArrays.generic_compress(A::GPUSparseMatrixCOO)::GPUSparseMatrixCSR + +Convert between CSR and COO on the device, which have the same order of stored entries: +expand the row pointers into a row index per entry, or compress them back. +""" +generic_expand(A::GPUSparseMatrixCSR) = + GPUSparseMatrixCOO(expand_ptr(A.rowPtr, nnz(A)), copy(A.colVal), copy(A.nzVal), size(A)) +generic_compress(A::GPUSparseMatrixCOO) = + GPUSparseMatrixCSR(compress(A.rowInd, size(A, 1)), copy(A.colInd), copy(A.nzVal), size(A)) +GPUSparseMatrixCOO(A::GPUSparseMatrixCSR) = generic_expand(A) +GPUSparseMatrixCSR(A::GPUSparseMatrixCOO) = generic_compress(A) + +# CSC ↔ COO through CSR, borrowing the buffers of the intermediate +function GPUSparseMatrixCOO(A::GPUSparseMatrixCSC) + B = GPUSparseMatrixCSR(A) + GPUSparseMatrixCOO(expand_ptr(B.rowPtr, nnz(B)), B.colVal, B.nzVal, size(A)) +end +function GPUSparseMatrixCSC(A::GPUSparseMatrixCOO) + B = GPUSparseMatrixCSR(compress(A.rowInd, size(A, 1)), A.colInd, A.nzVal, size(A)) + generic_regroup(B) +end + +# The CSR buffers of `A` are the CSC buffers of `transpose(A)`, and vice versa. This +# borrowed view shares the buffers of `A`; it is only used internally, to share kernels +# between formats, and never outlives the operation that creates it. +transpose_view(A::GPUSparseMatrixCSR) = + GPUSparseMatrixCSC(A.rowPtr, A.colVal, A.nzVal, reverse(size(A))) +transpose_view(A::GPUSparseMatrixCSC) = + GPUSparseMatrixCSR(A.colPtr, A.rowVal, A.nzVal, reverse(size(A))) + +# materializing a transpose into the opposite format only copies the buffers +GPUSparseMatrixCSC(A::Transpose{<:Any,<:GPUSparseMatrixCSR}) = copy(transpose_view(parent(A))) +GPUSparseMatrixCSR(A::Transpose{<:Any,<:GPUSparseMatrixCSC}) = copy(transpose_view(parent(A))) +function GPUSparseMatrixCSC(A::Adjoint{<:Any,<:GPUSparseMatrixCSR}) + B = transpose_view(parent(A)) + GPUSparseMatrixCSC(copy(B.colPtr), copy(B.rowVal), conj.(B.nzVal), size(B)) +end +function GPUSparseMatrixCSR(A::Adjoint{<:Any,<:GPUSparseMatrixCSC}) + B = transpose_view(parent(A)) + GPUSparseMatrixCSR(copy(B.rowPtr), copy(B.colVal), conj.(B.nzVal), size(B)) +end + +Base.convert(::Type{T}, A::AbstractArray) where {T<:AbstractGPUSparseArray} = + A isa T ? A : T(A) + + +## to dense arrays + +## COV_EXCL_START +# every thread scatters one slice of a compressed matrix into `dst`, indexed linearly +# as an `m`-row matrix +@kernel function densify_compressed_kernel(dst, ptr, ind, nzVal, m, csr::Bool) + j = @index(Global, Linear) + if j < length(ptr) + for k in @inbounds(ptr[j]):@inbounds(ptr[j+1] - one(eltype(ptr))) + i = @inbounds ind[k] + row, col = csr ? (j, Int(i)) : (Int(i), j) + @inbounds dst[(col - 1) * m + row] = nzVal[k] + end + end +end + +@kernel function densify_coo_kernel(dst, rowInd, colInd, nzVal, m) + k = @index(Global, Linear) + if k <= length(nzVal) + @inbounds dst[(Int(colInd[k]) - 1) * m + rowInd[k]] = nzVal[k] + end +end + +@kernel function densify_vector_kernel(dst, nzInd, nzVal) + k = @index(Global, Linear) + if k <= length(nzVal) + @inbounds dst[nzInd[k]] = nzVal[k] + end +end +## COV_EXCL_STOP + +# densify on the device, into the first `length(src)` elements of `dst` (in column-major +# order, like `copyto!` between dense arrays). The formats are canonical, so every element +# is written at most once. +function Base.copyto!(dst::AnyGPUArray, src::GPUSparseArray) + length(dst) >= length(src) || throw(BoundsError(dst, length(src))) + # (`dst` is cleared before the entries are read) + src = Base.unalias(dst, src) + fill!(length(dst) == length(src) ? dst : view(dst, 1:length(src)), zero(eltype(dst))) + nnz(src) == 0 && return dst + backend = get_backend(dst) + if src isa GPUSparseMatrixCompressed + densify_compressed_kernel(backend)(dst, compressed_ptr(src), compressed_ind(src), + src.nzVal, size(src, 1), src isa GPUSparseMatrixCSR; + ndrange=major_dim(src)) + elseif src isa GPUSparseMatrixCOO + densify_coo_kernel(backend)(dst, src.rowInd, src.colInd, src.nzVal, size(src, 1); + ndrange=nnz(src)) + else + densify_vector_kernel(backend)(dst, src.nzInd, src.nzVal; ndrange=nnz(src)) + end + return dst +end + +Base.copyto!(dst::Array, src::GPUSparseArray) = copyto!(dst, Array(src)) + diff --git a/src/host/sparse/indexing.jl b/src/host/sparse/indexing.jl new file mode 100644 index 000000000..ddadaeaa1 --- /dev/null +++ b/src/host/sparse/indexing.jl @@ -0,0 +1,45 @@ +# indexing + +# Scalar indexing searches the stored indices of one slice. On the host this reads a few +# entries of the buffers, which the dense arrays only allow with `@allowscalar`; the same +# code looks up elements in kernels. + +# the value stored at `ind[lo:hi] == i` if any, else zero +function stored_value(A, ind, lo::Integer, hi::Integer, i::Integer) + lo > hi && return zero(eltype(A)) + k = searchsortedfirst(ind, i, lo, hi, Base.Order.Forward) + (k > hi || ind[k] != i) && return zero(eltype(A)) + return nonzeros(A)[k] +end + +function Base.getindex(A::GPUSparseMatrixCSR, i::Integer, j::Integer) + @boundscheck checkbounds(A, i, j) + stored_value(A, A.colVal, Int(A.rowPtr[i]), Int(A.rowPtr[i+1]) - 1, j) +end + +function Base.getindex(A::GPUSparseMatrixCSC, i::Integer, j::Integer) + @boundscheck checkbounds(A, i, j) + stored_value(A, A.rowVal, Int(A.colPtr[j]), Int(A.colPtr[j+1]) - 1, i) +end + +function Base.getindex(A::GPUSparseMatrixCOO, i::Integer, j::Integer) + @boundscheck checkbounds(A, i, j) + # the entries of row `i` are contiguous, and sorted by column + lo = searchsortedfirst(A.rowInd, i) + hi = searchsortedlast(A.rowInd, i) + stored_value(A, A.colInd, lo, hi, j) +end + +function Base.getindex(x::GPUSparseVector, i::Integer) + @boundscheck checkbounds(x, i) + stored_value(x, x.nzInd, 1, nnz(x), i) +end + +Base.getindex(A::GPUSparseMatrix, I::Tuple{Integer,Integer}) = getindex(A, I...) +Base.getindex(x::GPUSparseVector, ::Colon) = copy(x) +Base.getindex(A::GPUSparseMatrix, ::Colon, ::Colon) = copy(A) + +function Base.setindex!(::GPUSparseArray, v, I...) + throw(ArgumentError("""GPU sparse arrays do not support setting individual entries. + Assemble a new array with `sparse(I, J, V, m, n)` or `sparsevec(I, V, n)`, use broadcasting, or update stored values through `nonzeros(A)`.""")) +end diff --git a/src/host/sparse/primitives.jl b/src/host/sparse/primitives.jl new file mode 100644 index 000000000..fc953a228 --- /dev/null +++ b/src/host/sparse/primitives.jl @@ -0,0 +1,106 @@ +# sparse-specific primitives, as internal kernels +# +# These build on the Base verbs a back-end provides for its dense arrays (`sortperm`, +# `accumulate!`, `findall`, `mapreduce`, and gathers through `getindex`). They are +# candidates for AcceleratedKernels once proven. +# +# The sparse code never launches a kernel over an empty range, because not every back-end +# supports that (JuliaGPU/Metal.jl#980). + +# release a temporary buffer early, without waiting for the GC +free_buffer!(x::AbstractGPUArray) = unsafe_free!(x) +free_buffer!(x) = nothing + +## COV_EXCL_START +@kernel function expand_ptr_kernel(major, ptr) + k = @index(Global, Linear) + if k <= length(major) + @inbounds major[k] = searchsortedlast(ptr, k % eltype(ptr)) % eltype(major) + end +end + +@kernel function compress_kernel(ptr, major) + j = @index(Global, Linear) + if j <= length(ptr) + # `j` stays an `Int`: converting it to the index type would wrap when `n` is the + # largest index + @inbounds ptr[j] = searchsortedfirst(major, j) % eltype(ptr) + end +end + +@kernel function packed_key_kernel(keys, major, minor, nminor) + k = @index(Global, Linear) + if k <= length(keys) + @inbounds keys[k] = (Int64(major[k]) - 1) * nminor + Int64(minor[k]) + end +end +## COV_EXCL_STOP + +""" + expand_ptr(ptr, nnz) + +The slice index of each of the `nnz` stored entries of a compressed pointer buffer `ptr` +(e.g. the row of every entry of a CSR matrix), with the index type and storage of `ptr`. +Every stored entry is handled by its own thread, which finds its slice by binary search, +so that long slices do not serialize. +""" +function expand_ptr(ptr::AbstractVector{Ti}, nnz::Integer) where {Ti} + major = similar(ptr, Ti, nnz) + if nnz > 0 + expand_ptr_kernel(get_backend(ptr))(major, ptr; ndrange=nnz) + end + return major +end + +""" + compress(major, n) + +The compressed pointer buffer, of length `n+1`, for stored entries whose sorted slice +indices are `major`: the inverse of [`expand_ptr`](@ref). Every pointer is found by binary +search; an atomic histogram followed by a scan would avoid the searches. +""" +function compress(major::AbstractVector{Ti}, n::Integer) where {Ti} + ptr = similar(major, Ti, n + 1) + compress_kernel(get_backend(ptr))(ptr, major; ndrange=n + 1) + return ptr +end + +# Sort keys ordering entries by (major, minor) index. They are packed into a signed +# `Int64`, which fits for all practical sizes and which every back-end can sort on the +# device (Metal's `sortperm` does not support unsigned integers). +function packed_keys(major::AbstractVector, minor::AbstractVector, dims::Dims{2}) + nmajor, nminor = dims + widemul(nmajor, nminor) < typemax(Int64) || + throw(ArgumentError("sparse matrices with $nmajor × $nminor entries are too large to sort")) + keys = similar(major, Int64, length(major)) + if !isempty(keys) + packed_key_kernel(get_backend(keys))(keys, major, minor, Int64(nminor); + ndrange=length(keys)) + end + return keys +end + +# Regroup the stored entries of a compressed matrix along its other dimension: the CSC +# buffers of a CSR matrix, and vice versa. `dims` is (major, minor) of the input. +# +# The entries are sorted by (minor, major) and gathered. An atomic counting sort over the +# minor indices (histogram, scan, scatter) would be much faster than the general sort. +function regroup(ptr::AbstractVector{Ti}, ind::AbstractVector{Ti}, val::AbstractVector, + dims::Dims{2}) where {Ti} + nmajor, nminor = dims + n = length(val) + n == 0 && return empty_ptr(ptr, Ti, nminor), similar(ind, 0), similar(val, 0) + + major = expand_ptr(ptr, n) + keys = packed_keys(ind, major, (nminor, nmajor)) + perm = sortperm(keys) + free_buffer!(keys) + new_major = ind[perm] + new_ind = major[perm] + new_val = val[perm] + free_buffer!(major) + free_buffer!(perm) + new_ptr = compress(new_major, nminor) + free_buffer!(new_major) + return new_ptr, new_ind, new_val +end diff --git a/src/host/sparse/reductions.jl b/src/host/sparse/reductions.jl new file mode 100644 index 000000000..e28ee9c25 --- /dev/null +++ b/src/host/sparse/reductions.jl @@ -0,0 +1,136 @@ +# reductions + +# `op` folded over `n ≥ 1` copies of `x`, as the implicit zeros of a sparse array contribute +# to a reduction: by repeated squaring for an associative `op` (e.g. `n * x` for `+`), and a +# single `x` for idempotent ones. +@inline function _fold_repeated(op, x, n::Integer) + _isidempotent(op) && return x + result = x + power = x + n -= one(n) + while n > 0 + if isodd(n) + result = op(result, power) + end + power = op(power, power) + n >>= 1 + end + return result +end +_isidempotent(op) = false +_isidempotent(::Union{typeof(max), typeof(min), typeof(|), typeof(&)}) = true + +## COV_EXCL_START +# reduce the stored values `nzVal[first_ptr:last_ptr]` of a slice that also contains +# `nzeros` implicit zeros, and has at least one element. As in SparseArrays, the stored +# values are reduced first and the implicit zeros folded in afterwards, so `op` is assumed +# to be commutative as well as associative. +@inline function _reduce_slice(f, op, ::Type{T}, nzVal, first_ptr, last_ptr, nzeros) where {T} + if first_ptr <= last_ptr + val = convert(T, f(@inbounds nzVal[first_ptr])) + for ptr in first_ptr+one(first_ptr):last_ptr + val = convert(T, op(val, f(@inbounds nzVal[ptr]))) + end + if nzeros > 0 + val = convert(T, op(val, _fold_repeated(op, convert(T, f(zero(eltype(nzVal)))), nzeros))) + end + val + else + _fold_repeated(op, convert(T, f(zero(eltype(nzVal)))), nzeros) + end +end + +# every thread reduces one slice (a row of a CSR matrix, a column of a CSC one), which has +# `n` elements +@kernel function compressed_reduce_kernel(f::F, op::OP, output::AbstractArray{T}, ptr, nzVal, + n) where {F, OP, T} + j = @index(Global, Linear) + if j ≤ length(output) + first_ptr = @inbounds ptr[j] + last_ptr = @inbounds ptr[j+1] - one(first_ptr) + nzeros = n - (last_ptr - first_ptr + 1) + @inbounds output[j] = _reduce_slice(f, op, T, nzVal, first_ptr, last_ptr, nzeros) + end +end +## COV_EXCL_STOP + +# the element type of a reduction, relying on inference to reason through the map and +# reduce functions, or on the type of the initializer +function reduction_eltype(f, op, A, init) + ET = Broadcast.combine_eltypes(f, (A,)) + ET = Base.promote_op(op, ET, ET) + if init !== nothing + ET = Base.promote_op(op, typeof(init), ET) + end + (ET === Union{} || ET === Any) && + error("mapreduce cannot figure the output element type, please pass an explicit init value") + return ET +end + +# every slice reduction `r` becomes `op(init, r)`; empty slices give `init`, or the +# reduction of an empty collection as in Base +empty_reduction(f, op, ::Type{Tv}, init) where {Tv} = + init === nothing ? Base.mapreduce_empty(f, op, Tv) : init + +# reduce the stored values, then fold in the implicit zeros +function reduce_all(f, op, A::GPUSparseArray, init) + Tv = eltype(A) + length(A) == 0 && return empty_reduction(f, op, Tv, init) + ET = reduction_eltype(f, op, A, init) + nzeros = length(A) - nnz(A) + zeros = nzeros > 0 ? _fold_repeated(op, convert(ET, f(zero(Tv))), nzeros) : nothing + if nnz(A) > 0 + stored = init === nothing ? mapreduce(f, op, nonzeros(A)) : + mapreduce(f, op, nonzeros(A); init) + return zeros === nothing ? stored : op(stored, zeros) + else + return init === nothing ? zeros : op(init, zeros) + end +end + +Base.mapreduce(f, op, x::GPUSparseVector; dims=:, init=nothing) = + dims === Colon() ? reduce_all(f, op, x, init) : + error("only dims=: is supported for sparse vectors") + +function Base.mapreduce(f, op, A::GPUSparseMatrix; dims=:, init=nothing) + dims === Colon() && return reduce_all(f, op, A, init) + dims in (1, 2) || error("only dims=:, dims=1 or dims=2 is supported") + + # reduce along the compressed dimension, regrouping the entries if necessary. an atomic + # scatter could reduce along the other dimension directly, for `+` and friends. + B = dims == 1 ? GPUSparseMatrixCSC(A) : GPUSparseMatrixCSR(A) + ET = reduction_eltype(f, op, A, init) + output = similar(nonzeros(B), ET, dims == 1 ? (1, size(B, 2)) : (size(B, 1), 1)) + isempty(output) && return output + if size(B, dims) == 0 + return fill!(output, empty_reduction(f, op, eltype(A), init)) + end + compressed_reduce_kernel(get_backend(B))(f, op, output, compressed_ptr(B), nonzeros(B), + minor_dim(B); ndrange=length(output)) + if init !== nothing + output .= op.(init, output) + end + return output +end + +# `nnz(A) == 0` means nothing is stored, otherwise check the stored values +Base.iszero(A::GPUSparseArray) = nnz(A) == 0 || all(iszero, nonzeros(A)) + +# The norm of the stored values, as in SparseArrays (and using the dense implementation, +# which scales to avoid overflow). Unlike SparseArrays, an implicit zero does count towards +# norms of negative order, which are no larger than the smallest magnitude. +function LinearAlgebra.norm(A::GPUSparseArray, p::Real=2) + r = norm(nonzeros(A), p) + return p < 0 && nnz(A) < length(A) ? min(r, zero(r)) : r +end + +function LinearAlgebra.opnorm(A::GPUSparseMatrix, p::Real=2) + T = float(real(eltype(A))) + if p == Inf + return isempty(A) ? zero(T) : T(maximum(sum(abs, A; dims=2))) + elseif p == 1 + return isempty(A) ? zero(T) : T(maximum(sum(abs, A; dims=1))) + else + throw(ArgumentError("p=$p is not supported")) + end +end diff --git a/src/host/sparse/show.jl b/src/host/sparse/show.jl new file mode 100644 index 000000000..f4150715f --- /dev/null +++ b/src/host/sparse/show.jl @@ -0,0 +1,37 @@ +# printing, through a copy on the host + +# the name of the array type that stores the entries: "MtlArray", "CuArray", ... +storage_name(A) = nameof(typeof(A.nzVal)) + +# the concrete type is long and not what users write (they use the back-end aliases), so +# only show the format and the element and index types +function Base.summary(io::IO, A::GPUSparseArray) + print(io, Base.dims2string(size(A)), " ", nameof(typeof(A)), "{", eltype(A), ", ", + indtype(A), "} with ", nnz(A), " stored ", nnz(A) == 1 ? "entry" : "entries", + " in ", storage_name(A)) +end + +# the device-side representation, e.g. in an error message about a kernel argument, cannot +# be copied to the host +on_device(A) = A.nzVal isa AbstractDeviceArray + +Base.print_array(io::IO, A::GPUSparseMatrix) = + on_device(A) ? nothing : Base.print_array(io, SparseMatrixCSC(A)) + +Base.show(io::IO, A::GPUSparseMatrix) = + on_device(A) ? summary(io, A) : show(io, SparseMatrixCSC(A)) +Base.show(io::IO, x::GPUSparseVector) = + on_device(x) ? summary(io, x) : show(io, SparseVector(x)) +# (SparseArrays 1.10 has a method for this one) +Base.show(io::IOContext, x::GPUSparseVector) = + on_device(x) ? summary(io, x) : show(io, SparseVector(x)) + +# SparseArrays prints the stored entries of a vector as part of its `show` method, below +# a summary that we replace by our own +function Base.show(io::IO, mime::MIME"text/plain", x::GPUSparseVector) + summary(io, x) + (nnz(x) == 0 || on_device(x)) && return + println(io, ":") + entries = sprint(show, mime, SparseVector(x); context=io) + print(io, split(entries, '\n'; limit=2)[2]) +end diff --git a/src/host/sparse/types.jl b/src/host/sparse/types.jl new file mode 100644 index 000000000..4ea960fdc --- /dev/null +++ b/src/host/sparse/types.jl @@ -0,0 +1,445 @@ +# GPU sparse array types + +using SparseArrays: indtype, getcolptr, nonzeroinds + +""" + AbstractGPUSparseArray{Tv,Ti,N} + +Supertype of the sparse arrays whose storage lives on a GPU. GPUArrays provides the formats +[`GPUSparseVector`](@ref), [`GPUSparseMatrixCSR`](@ref), [`GPUSparseMatrixCSC`](@ref) and +[`GPUSparseMatrixCOO`](@ref); back-ends only subtype this directly for formats that +GPUArrays does not implement (e.g. block-sparse formats of a vendor library). +""" +abstract type AbstractGPUSparseArray{Tv,Ti,N} <: AbstractSparseArray{Tv,Ti,N} end +const AbstractGPUSparseVector{Tv,Ti} = AbstractGPUSparseArray{Tv,Ti,1} +const AbstractGPUSparseMatrix{Tv,Ti} = AbstractGPUSparseArray{Tv,Ti,2} +const AbstractGPUSparseVecOrMat = Union{AbstractGPUSparseVector,AbstractGPUSparseMatrix} + +# the checks below only look at host metadata; structural invariants (sorted, unique and +# in-bounds indices, consistent pointers) are the caller's responsibility, as for +# SparseArrays' raw constructors. `check_structure` verifies them on the device. +function check_sparse_dims(::Type{Ti}, dims::Dims) where {Ti} + all(>=(0), dims) || + throw(ArgumentError("dimensions must be non-negative, got $dims")) + all(d -> d <= typemax(Ti), dims) || + throw(ArgumentError("dimensions $dims do not fit index type $Ti")) + return +end +function check_sparse_nnz(::Type{Ti}, nnz::Integer) where {Ti} + # pointers go up to nnz+1 + nnz < typemax(Ti) || + throw(ArgumentError("$nnz stored entries do not fit index type $Ti")) + return +end +function check_sparse_entries(::Type{Ti}, inds::AbstractVector, nzVal::AbstractVector) where {Ti} + length(inds) == length(nzVal) || + throw(ArgumentError("index and value buffers must have the same length, got $(length(inds)) and $(length(nzVal))")) + check_sparse_nnz(Ti, length(nzVal)) +end +function check_sparse_ptr(ptr::AbstractVector, n::Integer) + length(ptr) == n + 1 || + throw(ArgumentError("the pointer buffer must have length $(n + 1), got $(length(ptr))")) + return +end + +""" + GPUSparseMatrixCSR(rowPtr, colVal, nzVal, dims) + +A sparse matrix in compressed sparse row format, stored in the dense vectors `rowPtr` +(length `m+1`), `colVal` and `nzVal` (length `nnz`). The stored entries of row `i` are at +positions `rowPtr[i]:rowPtr[i+1]-1`, with their columns in `colVal` and their values in +`nzVal`. The columns in a row must be sorted and unique; explicitly stored zeros are +allowed. The constructor borrows the buffers rather than copying them. + +The index and value buffers are typically GPU vectors, whose type is part of the matrix +type (`GPUSparseMatrixCSR{Tv,Ti,Vi,Vv}`), so that generic code allocates new buffers with +`similar` on the existing ones. Back-ends provide aliases for their storage, e.g. +`const MtlSparseMatrixCSR{Tv,Ti} = GPUSparseMatrixCSR{Tv,Ti,<:MtlVector{Ti},<:MtlVector{Tv}}`. + +The same struct is used on the device: `adapt` converts the buffers, so kernels can read +the fields of a sparse matrix passed as an argument. +""" +struct GPUSparseMatrixCSR{Tv,Ti<:Integer,Vi<:AbstractVector{Ti},Vv<:AbstractVector{Tv}} <: + AbstractGPUSparseMatrix{Tv,Ti} + rowPtr::Vi + colVal::Vi + nzVal::Vv + dims::Dims{2} + + function GPUSparseMatrixCSR{Tv,Ti,Vi,Vv}(rowPtr::Vi, colVal::Vi, nzVal::Vv, + dims::NTuple{2,Integer}) where {Tv,Ti,Vi,Vv} + dims = Dims{2}(dims) + check_sparse_dims(Ti, dims) + check_sparse_ptr(rowPtr, dims[1]) + check_sparse_entries(Ti, colVal, nzVal) + new{Tv,Ti,Vi,Vv}(rowPtr, colVal, nzVal, dims) + end +end + +""" + GPUSparseMatrixCSC(colPtr, rowVal, nzVal, dims) + +A sparse matrix in compressed sparse column format, the layout of `SparseMatrixCSC`: the +stored entries of column `j` are at positions `colPtr[j]:colPtr[j+1]-1`, with their rows +(sorted and unique) in `rowVal` and their values in `nzVal`. See +[`GPUSparseMatrixCSR`](@ref) for the storage type parameters. +""" +struct GPUSparseMatrixCSC{Tv,Ti<:Integer,Vi<:AbstractVector{Ti},Vv<:AbstractVector{Tv}} <: + AbstractGPUSparseMatrix{Tv,Ti} + colPtr::Vi + rowVal::Vi + nzVal::Vv + dims::Dims{2} + + function GPUSparseMatrixCSC{Tv,Ti,Vi,Vv}(colPtr::Vi, rowVal::Vi, nzVal::Vv, + dims::NTuple{2,Integer}) where {Tv,Ti,Vi,Vv} + dims = Dims{2}(dims) + check_sparse_dims(Ti, dims) + check_sparse_ptr(colPtr, dims[2]) + check_sparse_entries(Ti, rowVal, nzVal) + new{Tv,Ti,Vi,Vv}(colPtr, rowVal, nzVal, dims) + end +end + +""" + GPUSparseMatrixCOO(rowInd, colInd, nzVal, dims) + +A sparse matrix in coordinate format: stored entry `k` is at `(rowInd[k], colInd[k])` and +has value `nzVal[k]`. The entries must be sorted in row-major order without duplicate +coordinates, which makes the format a CSR matrix with explicit row indices. Unordered +coordinates with duplicates are the input of `sparse(I, J, V, m, n, combine)` instead. +See [`GPUSparseMatrixCSR`](@ref) for the storage type parameters. +""" +struct GPUSparseMatrixCOO{Tv,Ti<:Integer,Vi<:AbstractVector{Ti},Vv<:AbstractVector{Tv}} <: + AbstractGPUSparseMatrix{Tv,Ti} + rowInd::Vi + colInd::Vi + nzVal::Vv + dims::Dims{2} + + function GPUSparseMatrixCOO{Tv,Ti,Vi,Vv}(rowInd::Vi, colInd::Vi, nzVal::Vv, + dims::NTuple{2,Integer}) where {Tv,Ti,Vi,Vv} + dims = Dims{2}(dims) + check_sparse_dims(Ti, dims) + check_sparse_entries(Ti, rowInd, nzVal) + check_sparse_entries(Ti, colInd, nzVal) + new{Tv,Ti,Vi,Vv}(rowInd, colInd, nzVal, dims) + end +end + +""" + GPUSparseVector(nzInd, nzVal, len) + +A sparse vector of length `len` with stored entries at the sorted, unique indices `nzInd` +and values `nzVal`, the layout of `SparseVector`. See [`GPUSparseMatrixCSR`](@ref) for the +storage type parameters. +""" +struct GPUSparseVector{Tv,Ti<:Integer,Vi<:AbstractVector{Ti},Vv<:AbstractVector{Tv}} <: + AbstractGPUSparseVector{Tv,Ti} + nzInd::Vi + nzVal::Vv + len::Int + + function GPUSparseVector{Tv,Ti,Vi,Vv}(nzInd::Vi, nzVal::Vv, len::Integer) where {Tv,Ti,Vi,Vv} + check_sparse_dims(Ti, (Int(len),)) + check_sparse_entries(Ti, nzInd, nzVal) + new{Tv,Ti,Vi,Vv}(nzInd, nzVal, len) + end +end + +for (S, args) in ((:GPUSparseMatrixCSR, (:rowPtr, :colVal)), + (:GPUSparseMatrixCSC, (:colPtr, :rowVal)), + (:GPUSparseMatrixCOO, (:rowInd, :colInd))) + @eval $S($(args[1])::Vi, $(args[2])::Vi, nzVal::Vv, dims::NTuple{2,Integer}) where + {Ti<:Integer,Vi<:AbstractVector{Ti},Tv,Vv<:AbstractVector{Tv}} = + $S{Tv,Ti,Vi,Vv}($(args[1]), $(args[2]), nzVal, dims) +end +GPUSparseVector(nzInd::Vi, nzVal::Vv, len::Integer) where + {Ti<:Integer,Vi<:AbstractVector{Ti},Tv,Vv<:AbstractVector{Tv}} = + GPUSparseVector{Tv,Ti,Vi,Vv}(nzInd, nzVal, len) + +const GPUSparseMatrixCompressed = Union{GPUSparseMatrixCSR,GPUSparseMatrixCSC} +const GPUSparseMatrix = Union{GPUSparseMatrixCSR,GPUSparseMatrixCSC,GPUSparseMatrixCOO} +const GPUSparseArray = Union{GPUSparseVector,GPUSparseMatrix} + +# the abstract types of GPUArrays 12 remain available for dispatch, but the formats are +# concrete now, so they cannot be subtyped anymore +Base.@deprecate_binding AbstractGPUSparseMatrixCSR GPUSparseMatrixCSR false +Base.@deprecate_binding AbstractGPUSparseMatrixCSC GPUSparseMatrixCSC false +Base.@deprecate_binding AbstractGPUSparseMatrixCOO GPUSparseMatrixCOO false + + +## accessors + +Base.size(A::GPUSparseMatrix) = A.dims +Base.size(x::GPUSparseVector) = (x.len,) + +SparseArrays.nnz(A::GPUSparseArray) = length(A.nzVal) +SparseArrays.nonzeros(A::GPUSparseArray) = A.nzVal + +SparseArrays.getcolptr(A::GPUSparseMatrixCSC) = A.colPtr +SparseArrays.rowvals(A::GPUSparseMatrixCSC) = A.rowVal +SparseArrays.nzrange(A::GPUSparseMatrixCSC, col::Integer) = + @inbounds A.colPtr[col]:(A.colPtr[col+1] - one(eltype(A.colPtr))) + +# a view of a column of a CSC matrix exposes its stored entries, as in SparseArrays, which +# also works in kernels +const GPUSparseColumnView = SubArray{<:Any,1,<:GPUSparseMatrixCSC,<:Tuple{Base.Slice,Integer}} +SparseArrays.nzrange(v::GPUSparseColumnView) = nzrange(parent(v), last(parentindices(v))) +SparseArrays.nnz(v::GPUSparseColumnView) = length(nzrange(v)) +SparseArrays.nonzeros(v::GPUSparseColumnView) = view(nonzeros(parent(v)), nzrange(v)) +SparseArrays.rowvals(v::GPUSparseColumnView) = view(rowvals(parent(v)), nzrange(v)) +SparseArrays.nonzeroinds(v::GPUSparseColumnView) = rowvals(v) + +SparseArrays.nonzeroinds(x::GPUSparseVector) = x.nzInd +SparseArrays.rowvals(x::GPUSparseVector) = x.nzInd + +KernelAbstractions.get_backend(A::GPUSparseArray) = get_backend(A.nzVal) + +# the pointer and index buffers of a compressed matrix, whatever its orientation +compressed_ptr(A::GPUSparseMatrixCSR) = A.rowPtr +compressed_ptr(A::GPUSparseMatrixCSC) = A.colPtr +compressed_ind(A::GPUSparseMatrixCSR) = A.colVal +compressed_ind(A::GPUSparseMatrixCSC) = A.rowVal + +# the size of the dimension the pointers compress, and of the other one +major_dim(A::GPUSparseMatrixCSR) = size(A, 1) +major_dim(A::GPUSparseMatrixCSC) = size(A, 2) +minor_dim(A::GPUSparseMatrixCSR) = size(A, 2) +minor_dim(A::GPUSparseMatrixCSC) = size(A, 1) + +# the buffers of a sparse array, in the order of its fields +sparse_buffers(A::GPUSparseMatrixCSR) = (A.rowPtr, A.colVal, A.nzVal) +sparse_buffers(A::GPUSparseMatrixCSC) = (A.colPtr, A.rowVal, A.nzVal) +sparse_buffers(A::GPUSparseMatrixCOO) = (A.rowInd, A.colInd, A.nzVal) +sparse_buffers(A::GPUSparseVector) = (A.nzInd, A.nzVal) + +# two sparse arrays alias when any of their buffers do +Base.dataids(A::GPUSparseArray) = mapreduce(Base.dataids, (a, b) -> (a..., b...), sparse_buffers(A)) +Base.unaliascopy(A::GPUSparseArray) = copy(A) + + +## allocation + +# a copy of `src` converted to eltype `T` +function convert_buffer(::Type{T}, src::AbstractVector) where {T} + dst = similar(src, T) + dst .= src + return dst +end + +# the same for indices, which are converted without checks because the callers have +# verified that the dimensions and the number of entries fit the new type +function convert_index_buffer(::Type{T}, src::AbstractVector) where {T<:Integer} + T == eltype(src) && return copy(src) + dst = similar(src, T) + dst .= src .% T + return dst +end + +# an empty compressed pointer buffer for `n` slices +empty_ptr(proto::AbstractVector, ::Type{Ti}, n::Integer) where {Ti} = + fill!(similar(proto, Ti, n + 1), one(Ti)) + +# SparseArrays semantics: `similar(A[, Tv[, Ti]])` copies the structure and leaves the +# values uninitialized, while `similar(A, Tv, dims)` returns an array without stored +# entries (even when `dims == size(A)`): a sparse matrix of the same format for two +# dimensions, a sparse vector for one, and a dense array for more. +Base.similar(A::GPUSparseArray) = similar(A, eltype(A)) +Base.similar(A::GPUSparseArray, ::Type{Tv}) where {Tv} = similar(A, Tv, indtype(A)) +function Base.similar(A::GPUSparseMatrixCSR, ::Type{Tv}, ::Type{Ti}) where {Tv,Ti} + check_sparse_dims(Ti, size(A)) + check_sparse_nnz(Ti, nnz(A)) + GPUSparseMatrixCSR(convert_index_buffer(Ti, A.rowPtr), convert_index_buffer(Ti, A.colVal), + similar(A.nzVal, Tv), size(A)) +end +function Base.similar(A::GPUSparseMatrixCSC, ::Type{Tv}, ::Type{Ti}) where {Tv,Ti} + check_sparse_dims(Ti, size(A)) + check_sparse_nnz(Ti, nnz(A)) + GPUSparseMatrixCSC(convert_index_buffer(Ti, A.colPtr), convert_index_buffer(Ti, A.rowVal), + similar(A.nzVal, Tv), size(A)) +end +function Base.similar(A::GPUSparseMatrixCOO, ::Type{Tv}, ::Type{Ti}) where {Tv,Ti} + check_sparse_dims(Ti, size(A)) + check_sparse_nnz(Ti, nnz(A)) + GPUSparseMatrixCOO(convert_index_buffer(Ti, A.rowInd), convert_index_buffer(Ti, A.colInd), + similar(A.nzVal, Tv), size(A)) +end +function Base.similar(x::GPUSparseVector, ::Type{Tv}, ::Type{Ti}) where {Tv,Ti} + check_sparse_dims(Ti, size(x)) + check_sparse_nnz(Ti, nnz(x)) + GPUSparseVector(convert_index_buffer(Ti, x.nzInd), similar(x.nzVal, Tv), length(x)) +end + +Base.similar(A::GPUSparseArray, dims::Dims) = similar(A, eltype(A), dims) +Base.similar(A::GPUSparseArray, ::Type{Tv}, dims::Dims) where {Tv} = + similar(A, Tv, indtype(A), dims) +Base.similar(A::GPUSparseArray, ::Type{Tv}, ::Type{Ti}, dims::Dims) where {Tv,Ti} = + similar(A.nzVal, Tv, dims) +Base.similar(A::GPUSparseArray, ::Type{Tv}, ::Type{Ti}, m::Integer, n::Integer...) where {Tv,Ti} = + similar(A, Tv, Ti, Dims((m, n...))) +function Base.similar(A::GPUSparseArray, ::Type{Tv}, ::Type{Ti}, dims::Dims{1}) where {Tv,Ti} + GPUSparseVector(similar(A.nzVal, Ti, 0), similar(A.nzVal, Tv, 0), dims[1]) +end +function Base.similar(A::GPUSparseArray, ::Type{Tv}, ::Type{Ti}, dims::Dims{2}) where {Tv,Ti} + empty_sparse(matrix_format(A), A.nzVal, Tv, Ti, dims) +end + +# the matrix format that allocations derived from `A` use; vectors become CSC, as +# `similar(::SparseVector, T, (m, n))` does in SparseArrays +matrix_format(::GPUSparseMatrixCSR) = GPUSparseMatrixCSR +matrix_format(::GPUSparseMatrixCSC) = GPUSparseMatrixCSC +matrix_format(::GPUSparseMatrixCOO) = GPUSparseMatrixCOO +matrix_format(::GPUSparseVector) = GPUSparseMatrixCSC + +# a matrix without stored entries, in format `S` and with the storage of `proto` +function empty_sparse(::Type{GPUSparseMatrixCSR}, proto, ::Type{Tv}, ::Type{Ti}, + dims::Dims{2}) where {Tv,Ti} + GPUSparseMatrixCSR(empty_ptr(proto, Ti, dims[1]), similar(proto, Ti, 0), + similar(proto, Tv, 0), dims) +end +function empty_sparse(::Type{GPUSparseMatrixCSC}, proto, ::Type{Tv}, ::Type{Ti}, + dims::Dims{2}) where {Tv,Ti} + GPUSparseMatrixCSC(empty_ptr(proto, Ti, dims[2]), similar(proto, Ti, 0), + similar(proto, Tv, 0), dims) +end +function empty_sparse(::Type{GPUSparseMatrixCOO}, proto, ::Type{Tv}, ::Type{Ti}, + dims::Dims{2}) where {Tv,Ti} + GPUSparseMatrixCOO(similar(proto, Ti, 0), similar(proto, Ti, 0), + similar(proto, Tv, 0), dims) +end + +Base.zero(A::GPUSparseArray) = similar(A, eltype(A), size(A)) + +# allocating constructors and `copy` return independent storage; only the raw constructors +# above borrow their buffers +Base.copy(A::GPUSparseMatrixCSR) = + GPUSparseMatrixCSR(copy(A.rowPtr), copy(A.colVal), copy(A.nzVal), size(A)) +Base.copy(A::GPUSparseMatrixCSC) = + GPUSparseMatrixCSC(copy(A.colPtr), copy(A.rowVal), copy(A.nzVal), size(A)) +Base.copy(A::GPUSparseMatrixCOO) = + GPUSparseMatrixCOO(copy(A.rowInd), copy(A.colInd), copy(A.nzVal), size(A)) +Base.copy(x::GPUSparseVector) = + GPUSparseVector(copy(x.nzInd), copy(x.nzVal), length(x)) + +# Replace the contents of the buffers `dsts` by `srcs`, which may have other lengths. +# Buffers that cannot be resized (e.g. views) fail the operation before anything is +# modified, so that they can still be used for reads and for updating values. Growing a +# buffer can fail at run time (e.g. out of memory), so the buffers are grown first and +# shrunk back if that fails, which leaves `dsts` as they were; shrinking is assumed not +# to fail. +function replace_buffers!(dsts, srcs) + foreach(dsts, srcs) do dst, src + length(dst) == length(src) || hasmethod(resize!, Tuple{typeof(dst),Int}) || + throw(ArgumentError("cannot change the structure of a sparse array with $(nameof(typeof(dst))) buffers")) + end + lengths = map(length, dsts) + try + foreach(dsts, srcs) do dst, src + length(dst) < length(src) && resize!(dst, length(src)) + end + catch + foreach(dsts, lengths) do dst, n + length(dst) > n && resize!(dst, n) + end + rethrow() + end + foreach(dsts, srcs) do dst, src + length(dst) > length(src) && resize!(dst, length(src)) + dst .= src + end + return +end + +# `copyto!` between sparse arrays of the same format replaces the structure of `dst` by +# that of `src`, resizing its buffers, so `dst` must not share them with another array +function copy_sparse!(dst::GPUSparseArray, src::GPUSparseArray) + size(dst) == size(src) || + throw(DimensionMismatch("destination has size $(size(dst)), source has size $(size(src))")) + check_sparse_nnz(indtype(dst), nnz(src)) + dst === src && return dst + src = Base.mightalias(dst, src) ? copy(src) : src + replace_buffers!(sparse_buffers(dst), sparse_buffers(src)) + return dst +end +for S in (:GPUSparseMatrixCSR, :GPUSparseMatrixCSC, :GPUSparseMatrixCOO, :GPUSparseVector) + @eval Base.copyto!(dst::$S, src::$S) = copy_sparse!(dst, src) +end + + +## structural checks + +## COV_EXCL_START +# whether the indices in `ind[first:last]` are sorted, unique and within `1:n` +@inline function valid_slice(ind, first, last, n) + prev = zero(eltype(ind)) + for k in first:last + i = @inbounds ind[k] + (i > prev && i <= n) || return false + prev = i + end + return true +end + +@kernel function check_compressed_kernel(valid, ptr, ind, n) + j = @index(Global, Linear) + if j < length(ptr) + first = @inbounds ptr[j] + last = @inbounds ptr[j+1] - one(first) + @inbounds valid[j] = first >= 1 && last + 1 >= first && last <= length(ind) && + valid_slice(ind, first, last, n) + end +end + +@kernel function check_coo_kernel(valid, rowInd, colInd, dims) + k = @index(Global, Linear) + if k <= length(rowInd) + i = @inbounds rowInd[k] + j = @inbounds colInd[k] + ok = 1 <= i <= dims[1] && 1 <= j <= dims[2] + if k > 1 + # row-major order without duplicates + i′ = @inbounds rowInd[k-1] + j′ = @inbounds colInd[k-1] + ok &= i′ < i || (i′ == i && j′ < j) + end + @inbounds valid[k] = ok + end +end +## COV_EXCL_STOP + +""" + GPUArrays.check_structure(A) + +Verify on the device that the sparse array `A` satisfies the invariants of its format +(consistent pointers, sorted and unique indices within bounds), throwing an +`ArgumentError` otherwise. The constructors only check host metadata such as buffer +lengths, so this is useful for testing code that assembles sparse arrays from raw buffers. +""" +function check_structure(A::GPUSparseMatrixCompressed) + ptr, ind = compressed_ptr(A), compressed_ind(A) + n = major_dim(A) + valid = similar(A.nzVal, Bool, n) + if n > 0 + check_compressed_kernel(get_backend(A))(valid, ptr, ind, minor_dim(A); ndrange=n) + end + first, last = Array(ptr[[1, end]]) + (first == 1 && last == nnz(A) + 1 && all(valid)) || + throw(ArgumentError("invalid $(nameof(typeof(A))) structure")) + return A +end +function check_structure(A::GPUSparseMatrixCOO) + valid = similar(A.nzVal, Bool, nnz(A)) + if nnz(A) > 0 + check_coo_kernel(get_backend(A))(valid, A.rowInd, A.colInd, size(A); ndrange=nnz(A)) + end + all(valid) || throw(ArgumentError("invalid GPUSparseMatrixCOO structure")) + return A +end +function check_structure(x::GPUSparseVector) + valid = similar(x.nzVal, Bool, 1) + ptr = copyto!(similar(x.nzInd, 2), indtype(x)[1, nnz(x) + 1]) + check_compressed_kernel(get_backend(x))(valid, ptr, x.nzInd, length(x); ndrange=1) + all(valid) || throw(ArgumentError("invalid GPUSparseVector structure")) + return x +end diff --git a/test/runtests.jl b/test/runtests.jl index b43820e8d..fcf22bfbb 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -9,9 +9,6 @@ const init_worker_code = quote include("testsuite.jl") - TestSuite.sparse_types(::Type{<:JLArray}) = (JLSparseVector, JLSparseMatrixCSC, JLSparseMatrixCSR) - TestSuite.sparse_types(::Type{<:Array}) = (SparseVector, SparseMatrixCSC) - # JLArray kernels run on POCL, whose SPIR-V back-end lacks 128-bit integers TestSuite.supported_eltypes(::Type{<:JLArray}) = filter(!=(Int128), TestSuite.supported_eltypes()) @@ -56,6 +53,16 @@ testsuite["JLArray/kernelinterface"] = quote Testsuite.testsuite(JLBackend(), JLArray) end +# the sparse implementation should not add method ambiguities (the dense code has a few) +testsuite["sparse/ambiguities"] = quote + import GPUArrays + sparse_sources = (joinpath("src", "host", "sparse"), joinpath("JLArrays", "src", "sparse.jl")) + ambiguities = filter(Test.detect_ambiguities(GPUArrays, JLArrays; recursive=true)) do (a, b) + any(m -> any(src -> occursin(src, string(m.file)), sparse_sources), (a, b)) + end + @test isempty(ambiguities) +end + # captured variables that get boxed (`Core.Box`) lose their type if isdefined(Test, :detect_closure_boxes) testsuite["closure boxes"] = quote @@ -64,10 +71,4 @@ if isdefined(Test, :detect_closure_boxes) end end -# a type without a COO counterpart must hit a MethodError, not recurse (#780) -testsuite["sparse/coo_type"] = quote - import GPUArrays - @test_throws MethodError GPUArrays.coo_type(sparse([1], [1], [1.0], 2, 2)) -end - runtests(GPUArrays, args; init_code, init_worker_code, testsuite) diff --git a/test/testsuite.jl b/test/testsuite.jl index 93beae865..0a1a0fd3d 100644 --- a/test/testsuite.jl +++ b/test/testsuite.jl @@ -96,9 +96,6 @@ supported_eltypes() = (Int16, Int32, Int64, ComplexF16, ComplexF32, ComplexF64, Complex{Int16}, Complex{Int32}, Complex{Int64}) -# derived sparse container types that are supported by the array type -sparse_types(::Type{AT}) where {AT} = () - # some convenience predicates for filtering test eltypes isrealtype(T) = T <: Real iscomplextype(T) = T <: Complex diff --git a/test/testsuite/sparse.jl b/test/testsuite/sparse.jl index b4112aa7e..d66484366 100644 --- a/test/testsuite/sparse.jl +++ b/test/testsuite/sparse.jl @@ -1,179 +1,473 @@ +using SparseArrays +using SparseArrays: nonzeroinds, nonzeros, rowvals, getcolptr +using GPUArrays: GPUSparseMatrixCSR, GPUSparseMatrixCSC, GPUSparseMatrixCOO, GPUSparseVector, + AbstractGPUSparseArray, check_structure + @testsuite "sparse" (AT, eltypes)->begin - sparse_ATs = sparse_types(AT) - @testset "sparse_AT = $sparse_AT" for sparse_AT in sparse_ATs - if sparse_AT <: AbstractSparseVector - vector(sparse_AT, eltypes) - vector_construction(sparse_AT, eltypes) - broadcasting_vector(sparse_AT, eltypes) - iszero_vector(sparse_AT, eltypes) - elseif sparse_AT <: AbstractSparseMatrix - matrix(sparse_AT, eltypes) - matrix_construction(sparse_AT, eltypes) - broadcasting_matrix(sparse_AT, eltypes) - mapreduce_matrix(sparse_AT, eltypes) - linalg(sparse_AT, eltypes) - transpose_matrix(sparse_AT, eltypes) - iszero_matrix(sparse_AT, eltypes) - end - end + # the generic sparse formats use `AT`'s vectors as storage; plain arrays have no device + AT <: AbstractGPUArray || return + # nothing may fall back to scalar indexing unless a test asks for it + GPUArrays.allowscalar(false) + + sparse_construction(AT, eltypes) + sparse_allocation(AT, eltypes) + sparse_transfer(AT, eltypes) + sparse_kernels(AT, eltypes) + sparse_display(AT, eltypes) + sparse_indexing(AT, eltypes) + sparse_conversions(AT, eltypes) + broadcasting_vector(AT, eltypes) + broadcasting_matrix(AT, eltypes) + mapreduce_matrix(AT, eltypes) + sparse_linalg(AT, eltypes) + iszero_sparse(AT, eltypes) + sparse_bool(AT) end -using SparseArrays -using SparseArrays: nonzeroinds, nonzeros, rowvals +const sparse_matrix_formats = (GPUSparseMatrixCSR, GPUSparseMatrixCSC, GPUSparseMatrixCOO) + +# A GPU sparse matrix of format `S` with the entries of `A`, built directly from buffers +# computed on the host, so that the tests do not depend on the conversions they test. +function gpu_sparse(AT, S::Type, A::SparseMatrixCSC) + if S <: GPUSparseMatrixCSC + GPUSparseMatrixCSC(AT(getcolptr(A)), AT(rowvals(A)), AT(nonzeros(A)), size(A)) + else + # the CSC buffers of the transpose are the CSR buffers of `A` + At = copy(transpose(A)) + if S <: GPUSparseMatrixCSR + GPUSparseMatrixCSR(AT(getcolptr(At)), AT(rowvals(At)), AT(nonzeros(At)), size(A)) + else + rows = similar(rowvals(At)) + for i in 1:size(A, 1), k in nzrange(At, i) + rows[k] = i + end + GPUSparseMatrixCOO(AT(rows), AT(rowvals(At)), AT(nonzeros(At)), size(A)) + end + end +end +gpu_sparse(AT, x::SparseVector) = GPUSparseVector(AT(nonzeroinds(x)), AT(nonzeros(x)), length(x)) # `sprand`/`sprandn` can produce *stored* zeros: with a coarse element type like Float16, # `rand` returns an exact zero about once every 2048 draws, so roughly one in 60 vectors # here contains one. That's a valid sparse array, but CPU and GPU implementations # legitimately disagree on what to do with it -- `x .* 1` drops stored zeros on the CPU # while our kernels preserve the structure -- so the structural comparisons below fail at -# random. Keep the inputs canonical instead. +# random. Keep the inputs canonical instead, and test stored zeros explicitly. sprand_nozeros(args...) = dropzeros!(sprand(args...)) sprandn_nozeros(args...) = dropzeros!(sprandn(args...)) -function vector(AT, eltypes) - dense_AT = GPUArrays.dense_array_type(AT) - @testset "Sparse vector properties" begin - @testset "$ET" for ET in eltypes - m = 25 - n = 35 - k = 10 - p = 5 - blockdim = 5 - x = sprand_nozeros(ET,m,0.2) - d_x = AT(x) - @test length(d_x) == m - @test size(d_x) == (m,) - @test size(d_x,1) == m - @test size(d_x,2) == 1 - @test ndims(d_x) == 1 - dense_d_x = dense_AT(x) - dense_d_x2 = dense_AT(d_x) - @allowscalar begin - @test Array(d_x[:]) == x[:] - @test d_x[firstindex(d_x)] == x[firstindex(x)] - @test d_x[div(end, 2)] == x[div(end, 2)] - @test d_x[end] == x[end] - @test Array(d_x[firstindex(d_x):end]) == x[firstindex(x):end] - @test Array(dense_d_x[firstindex(d_x):end]) == x[firstindex(x):end] - @test Array(dense_d_x2[firstindex(d_x):end]) == x[firstindex(x):end] +# a random matrix with an empty row and column, a dense row, and a stored zero +function sprand_awkward(::Type{T}, m, n; Ti=Int) where {T} + A = Matrix(sprand_nozeros(T, m, n, 0.3)) + A[1, :] .= rand(T, n) .+ one(T) + A[:, 2] .= zero(T) + A[3, :] .= zero(T) + S = SparseMatrixCSC{T,Ti}(sparse(A)) + nonzeros(S)[end] = zero(T) + return S +end + +# the same entries, and the same explicitly stored ones +same_sparse(A::SparseMatrixCSC, B::SparseMatrixCSC) = + size(A) == size(B) && getcolptr(A) == getcolptr(B) && rowvals(A) == rowvals(B) && + nonzeros(A) == nonzeros(B) +same_sparse(x::SparseVector, y::SparseVector) = + length(x) == length(y) && nonzeroinds(x) == nonzeroinds(y) && nonzeros(x) == nonzeros(y) +same_sparse(A::AbstractGPUSparseArray, B) = same_sparse(host(A), B) +host(A::AbstractGPUSparseArray) = ndims(A) == 1 ? SparseVector(A) : SparseMatrixCSC(A) + +# the result of a reduction on the host, whether it is a scalar or an array +reduced(x::Number) = x +reduced(x::AbstractArray) = Array(x) + +function sparse_construction(AT, eltypes) + @testset "construction and properties" begin + @testset "$S{$ET, $Ti}" for S in sparse_matrix_formats, ET in eltypes, Ti in (Int32, Int64) + A = sprand_awkward(ET, 10, 8; Ti) + dA = gpu_sparse(AT, S, A) + @test dA isa S{ET,Ti} + @test dA isa AbstractGPUSparseArray{ET,Ti,2} + @test size(dA) == (10, 8) + @test length(dA) == 80 + @test ndims(dA) == 2 + @test nnz(dA) == nnz(A) + @test issparse(dA) + @test SparseArrays.indtype(dA) == Ti + @test nonzeros(dA) isa AT{ET} + check_structure(dA) + @test same_sparse(SparseMatrixCSC(dA), A) + @test SparseMatrixCSC(dA) isa SparseMatrixCSC{ET,Ti} + @test Array(dA) == Array(A) + @test collect(dA) == Array(A) + @test get_backend(dA) == get_backend(AT(ET[])) + if S <: GPUSparseMatrixCSC + @test Array(getcolptr(dA)) == getcolptr(A) + @test Array(rowvals(dA)) == rowvals(A) + end + end + + @testset "vector $ET, $Ti" for ET in eltypes, Ti in (Int32, Int64) + x = SparseVector{ET,Ti}(sprand_nozeros(ET, 20, 0.3)) + dx = gpu_sparse(AT, x) + @test dx isa GPUSparseVector{ET,Ti} + @test size(dx) == (20,) + @test nnz(dx) == nnz(x) + @test Array(nonzeroinds(dx)) == nonzeroinds(x) + @test Array(rowvals(dx)) == nonzeroinds(x) + @test Array(nonzeros(dx)) == nonzeros(x) + check_structure(dx) + @test same_sparse(SparseVector(dx), x) + @test Array(dx) == Array(x) + end + + @testset "shapes" begin + ET = first(eltypes) + for (m, n) in ((0, 0), (0, 4), (4, 0), (5, 3)), S in sparse_matrix_formats + A = spzeros(ET, m, n) + dA = gpu_sparse(AT, S, A) + check_structure(dA) + @test same_sparse(SparseMatrixCSC(dA), A) + @test Array(dA) == zeros(ET, m, n) + end + for n in (0, 5) + x = spzeros(ET, n) + @test same_sparse(SparseVector(gpu_sparse(AT, x)), x) + end + end + + @testset "checks" begin + ET = first(eltypes) + ptr = AT(Int32[1, 2, 3]) + ind = AT(Int32[1, 2]) + val = AT(ET[1, 2]) + @test GPUSparseMatrixCSR(ptr, ind, val, (2, 2)) isa GPUSparseMatrixCSR + # pointer length + @test_throws ArgumentError GPUSparseMatrixCSR(ptr, ind, val, (3, 2)) + @test_throws ArgumentError GPUSparseMatrixCSC(ptr, ind, val, (2, 3)) + # buffer lengths + @test_throws ArgumentError GPUSparseMatrixCSR(ptr, ind, AT(ET[1]), (2, 2)) + @test_throws ArgumentError GPUSparseMatrixCOO(ind, AT(Int32[1]), val, (2, 2)) + @test_throws ArgumentError GPUSparseVector(ind, AT(ET[1]), 2) + # dimensions + @test_throws ArgumentError GPUSparseMatrixCOO(ind, ind, val, (-1, 2)) + @test_throws ArgumentError GPUSparseVector(AT(Int8[1]), AT(ET[1]), 200) + # the structure is only checked on request + @test check_structure(GPUSparseMatrixCSR(AT(Int32[1, 3, 3]), AT(Int32[1, 2]), val, (2, 2))) isa GPUSparseMatrixCSR + @test_throws ArgumentError check_structure(GPUSparseMatrixCSR(AT(Int32[1, 3, 3]), AT(Int32[2, 1]), val, (2, 2))) + @test_throws ArgumentError check_structure(GPUSparseMatrixCSC(AT(Int32[1, 2, 3]), AT(Int32[1, 3]), val, (2, 2))) + @test_throws ArgumentError check_structure(GPUSparseMatrixCOO(AT(Int32[2, 1]), AT(Int32[1, 1]), val, (2, 2))) + @test_throws ArgumentError check_structure(GPUSparseMatrixCOO(AT(Int32[1, 1]), AT(Int32[1, 1]), val, (2, 2))) + @test_throws ArgumentError check_structure(GPUSparseVector(AT(Int32[2, 2]), val, 3)) + end + end +end + +function sparse_allocation(AT, eltypes) + @testset "allocation" begin + @testset "$S" for S in sparse_matrix_formats + ET = first(eltypes) + A = sprand_awkward(ET, 6, 7; Ti=Int32) + dA = gpu_sparse(AT, S, A) + + # structure-preserving + B = similar(dA) + @test B isa S{ET,Int32} + @test nnz(B) == nnz(dA) + copyto!(nonzeros(B), nonzeros(dA)) + @test same_sparse(B, A) + B = similar(dA, Float32) + @test B isa S{Float32,Int32} + @test nnz(B) == nnz(dA) + B = similar(dA, Float32, Int64) + @test B isa S{Float32,Int64} + @test nnz(B) == nnz(dA) + + # without stored entries + for dims in ((6, 7), (3, 4)) + B = similar(dA, Float32, dims) + @test B isa S{Float32,Int32} + @test size(B) == dims + @test nnz(B) == 0 + check_structure(B) end - @test_throws BoundsError d_x[firstindex(d_x) - 1] - @test_throws BoundsError d_x[end + 1] - @test nnz(d_x) == nnz(x) - @test Array(nonzeros(d_x)) == nonzeros(x) - @test Array(nonzeroinds(d_x)) == nonzeroinds(x) - @test Array(rowvals(d_x)) == nonzeroinds(x) - @test nnz(d_x) == length(nonzeros(d_x)) + B = similar(dA, (3, 4)) + @test B isa S{ET,Int32} + B = similar(dA, Float32, (5,)) + @test B isa GPUSparseVector{Float32,Int32} + @test length(B) == 5 && nnz(B) == 0 + B = similar(dA, Float32, (2, 3, 4)) + @test B isa AT{Float32,3} + B = zero(dA) + @test B isa S{ET,Int32} + @test size(B) == size(dA) && nnz(B) == 0 + + # copies share nothing with the original + B = copy(dA) + @test B isa S{ET,Int32} + @test same_sparse(B, A) + fill!(nonzeros(B), zero(ET)) + @test same_sparse(dA, A) + B = S(dA) + fill!(nonzeros(B), zero(ET)) + @test same_sparse(dA, A) + + # copyto! replaces the structure + A2 = sprand_awkward(ET, 6, 7; Ti=Int32) + dA2 = gpu_sparse(AT, S, A2) + B = copy(dA) + @test copyto!(B, dA2) === B + @test same_sparse(B, A2) + @test_throws DimensionMismatch copyto!(B, gpu_sparse(AT, S, spzeros(ET, 2, 2))) + + # densify + D = AT(zeros(ET, 6, 7)) + @test copyto!(D, dA) === D + @test Array(D) == Array(A) + D = AT(ones(ET, 50)) + copyto!(D, dA) + @test Array(D) == [vec(Array(A)); ones(ET, 8)] + end + + @testset "vector" begin + ET = first(eltypes) + x = SparseVector{ET,Int32}(sprand_nozeros(ET, 20, 0.3)) + dx = gpu_sparse(AT, x) + @test similar(dx) isa GPUSparseVector{ET,Int32} + @test similar(dx, Float32) isa GPUSparseVector{Float32,Int32} + @test similar(dx, Float32, Int64) isa GPUSparseVector{Float32,Int64} + @test nnz(similar(dx, Float32, (20,))) == 0 + B = similar(dx, Float32, (4, 5)) + @test B isa GPUSparseMatrixCSC{Float32,Int32} + @test nnz(B) == 0 + y = copy(dx) + fill!(nonzeros(y), zero(ET)) + @test same_sparse(dx, x) + y = zero(dx) + @test nnz(y) == 0 + y = gpu_sparse(AT, SparseVector{ET,Int32}(sprand_nozeros(ET, 20, 0.5))) + copyto!(y, dx) + @test same_sparse(y, x) + D = AT(zeros(ET, 20)) + copyto!(D, dx) + @test Array(D) == Array(x) end end end -function vector_construction(AT, eltypes) - dense_AT = GPUArrays.dense_array_type(AT) - for ET in eltypes - m = 25 - n = 35 - k = 10 - p = 5 - blockdim = 5 - x = sprand_nozeros(ET, m, 0.2) - d_x = AT(x) - @test collect(d_x) == collect(x) - @test similar(d_x) isa AT{ET} - @test similar(d_x, Float32) isa AT{Float32} +function sparse_transfer(AT, eltypes) + @testset "adapt" begin + ET = first(eltypes) + A = sprand_awkward(ET, 6, 7; Ti=Int32) + x = SparseVector{ET,Int32}(sprand_nozeros(ET, 20, 0.3)) + + # structural: keep format, element and index types + dA = adapt(AT, A) + @test dA isa GPUSparseMatrixCSC{ET,Int32} + @test nonzeros(dA) isa AT + @test same_sparse(dA, A) + dx = adapt(AT, x) + @test dx isa GPUSparseVector{ET,Int32} + @test same_sparse(dx, x) + @test adapt(AT, dA) === dA + + # to the host + B = adapt(Array, dA) + @test B isa SparseMatrixCSC{ET,Int32} + @test same_sparse(B, A) + @test adapt(Array, dx) isa SparseVector{ET,Int32} + for S in (GPUSparseMatrixCSR, GPUSparseMatrixCOO) + B = adapt(Array, gpu_sparse(AT, S, A)) + @test B isa S{ET,Int32,Vector{Int32},Vector{ET}} + @test same_sparse(SparseMatrixCSC(B), A) + end + + # a target with an element type only converts the values + T = ET <: Complex ? ComplexF32 : Float32 + if T in eltypes + dB = adapt(AT{T}, A) + @test dB isa GPUSparseMatrixCSC{T,Int32} + @test Array(dB) ≈ Array(A) + @test adapt(Array{T}, dA) isa SparseMatrixCSC{T,Int32} + end + + # to an explicit sparse type + dB = adapt(GPUSparseMatrixCSR, dA) + @test dB isa GPUSparseMatrixCSR{ET,Int32} + @test nonzeros(dB) isa AT + @test same_sparse(dB, A) + @test_throws ArgumentError adapt(GPUSparseMatrixCSR, A) end end -function matrix(AT, eltypes) - dense_AT = GPUArrays.dense_array_type(AT) - @testset "Sparse matrix properties" begin - @testset "$ET" for ET in eltypes - m = 25 - n = 35 - k = 10 - p = 5 - blockdim = 5 - x = sprand_nozeros(ET, m, 0.2) - d_x = AT(x) - @test size(d_x) == (m, 1) - x = sprand_nozeros(ET, m, n, 0.2) - d_x = AT(x) - d_tx = AT(transpose(x)) - d_ax = AT(adjoint(x)) - @test size(d_tx) == (n, m) - @test size(d_ax) == (n, m) - @test length(d_x) == m*n - @test size(d_x) == (m, n) - @test size(d_x, 1) == m - @test size(d_x, 2) == n - @test size(d_x, 3) == 1 - @test ndims(d_x) == 2 +@kernel function sparse_rowsum_kernel(out, A) + row = @index(Global, Linear) + acc = zero(eltype(out)) + for k in A.rowPtr[row]:(A.rowPtr[row+1] - 1) + acc += A.nzVal[k] + end + out[row] = acc +end + +@kernel function sparse_colsum_kernel(out, A) + col = @index(Global, Linear) + acc = zero(eltype(out)) + for k in nzrange(A, col) + acc += nonzeros(A)[k] + end + out[col] = acc +end + +function sparse_kernels(AT, eltypes) + @testset "user kernels" begin + ET = first(filter(isrealfloattype, eltypes)) + A = sprand_awkward(ET, 8, 6) + dA = gpu_sparse(AT, GPUSparseMatrixCSR, A) + out = AT(zeros(ET, 8)) + sparse_rowsum_kernel(get_backend(dA))(out, dA; ndrange=8) + @test Array(out) ≈ vec(sum(A; dims=2)) + dA = gpu_sparse(AT, GPUSparseMatrixCSC, A) + out = AT(zeros(ET, 6)) + sparse_colsum_kernel(get_backend(dA))(out, dA; ndrange=6) + @test Array(out) ≈ vec(sum(A; dims=1)) + end +end + +function sparse_display(AT, eltypes) + @testset "display" begin + ET = first(eltypes) + A = sprand_awkward(ET, 6, 7; Ti=Int32) + for S in sparse_matrix_formats + dA = gpu_sparse(AT, S, A) + str = summary(dA) + @test str == "6×7 $(nameof(S)){$ET, Int32} with $(nnz(A)) stored entries in $(nameof(AT))" + # the entries are rendered as for the host matrix (on some Julia versions, + # SparseArrays' MIME show draws large matrices as braille patterns instead) + host_str = sprint(Base.print_array, A; context=(:compact=>true, :typeinfo=>ET)) + dev_str = sprint(show, MIME"text/plain"(), dA) + @test split(dev_str, '\n'; limit=2)[2] == host_str + @test sprint(show, dA) == sprint(show, A) + end + x = SparseVector{ET,Int32}(sprand_nozeros(ET, 20, 0.3)) + dx = gpu_sparse(AT, x) + @test startswith(summary(dx), "20-element GPUSparseVector{$ET, Int32} with $(nnz(x)) stored") + host_str = sprint(show, MIME"text/plain"(), x) + dev_str = sprint(show, MIME"text/plain"(), dx) + @test split(dev_str, '\n')[2:end] == split(host_str, '\n')[2:end] + @test sprint(show, dx) == sprint(show, x) + @test sprint(show, MIME"text/plain"(), gpu_sparse(AT, spzeros(ET, 3))) == + summary(gpu_sparse(AT, spzeros(ET, 3))) + end +end + +function sparse_indexing(AT, eltypes) + @testset "indexing" begin + ET = first(eltypes) + A = sprand_awkward(ET, 6, 7) + # scalar indexing is allowed if the dense arrays allow it (e.g. shared memory) + scalar_allowed = try + AT(ET[1])[1] + true + catch err + err isa ErrorException || rethrow() + false + end + @testset "$S" for S in sparse_matrix_formats + dA = gpu_sparse(AT, S, A) + if scalar_allowed + @test dA[1, 1] == A[1, 1] + else + @test_throws ErrorException dA[1, 1] + end @allowscalar begin - @test d_x[:, :] == x[:, :] - @test d_tx[:, :] == transpose(x)[:, :] - @test d_ax[:, :] == adjoint(x)[:, :] - @test d_x[(1, 1)] == x[1, 1] - @test d_x[firstindex(d_x)] == x[firstindex(x)] - @test d_x[div(end, 2)] == x[div(end, 2)] - @test d_x[end] == x[end] - @test d_x[firstindex(d_x), firstindex(d_x)] == x[firstindex(x), firstindex(x)] - @test d_x[div(end, 2), div(end, 2)] == x[div(end, 2), div(end, 2)] - @test d_x[end, end] == x[end, end] - for i in 1:size(x, 2) - @test Array(d_x[:, i]) == x[:, i] - end + @test all(dA[i, j] == A[i, j] for i in 1:6, j in 1:7) + @test dA[(2, 3)] == A[2, 3] + @test dA[end] == A[end] + @test dA[5] == A[5] end - @test_throws BoundsError d_x[firstindex(d_x) - 1] - @test_throws BoundsError d_x[end + 1] - @test_throws BoundsError d_x[firstindex(d_x) - 1, firstindex(d_x) - 1] - @test_throws BoundsError d_x[end + 1, end + 1] - @test_throws BoundsError d_x[firstindex(d_x) - 1:end + 1, :] - @test_throws BoundsError d_x[firstindex(d_x) - 1, :] - @test_throws BoundsError d_x[end + 1, :] - @test_throws BoundsError d_x[:, firstindex(d_x) - 1:end + 1] - @test_throws BoundsError d_x[:, firstindex(d_x) - 1] - @test_throws BoundsError d_x[:, end + 1] - @test nnz(d_x) == nnz(x) - @test nnz(d_x) == length(nonzeros(d_x)) - @test !issymmetric(d_x) - @test !ishermitian(d_x) + @test_throws BoundsError dA[0, 1] + @test_throws BoundsError dA[7, 1] + @test_throws BoundsError dA[1, 8] + @test same_sparse(dA[:, :], A) + @test_throws ArgumentError dA[1, 1] = one(ET) + end + x = sprand_nozeros(ET, 20, 0.3) + dx = gpu_sparse(AT, x) + @allowscalar begin + @test all(dx[i] == x[i] for i in 1:20) + @test dx[end] == x[end] end + @test_throws BoundsError dx[21] + @test same_sparse(dx[:], x) + @test_throws ArgumentError dx[1] = one(ET) end end -function matrix_construction(AT, eltypes) - dense_AT = GPUArrays.dense_array_type(AT) - for ET in eltypes - m = 25 - n = 35 - k = 10 - p = 5 - blockdim = 5 - x = sprand_nozeros(ET, m, n, 0.2) - d_x = AT(x) - @test collect(d_x) == collect(x) - @test similar(d_x) isa AT{ET} - @test similar(d_x, Float32) isa AT{Float32} - @test similar(d_x, Float32, n, m) isa AT{Float32} - @test similar(d_x, Float32, (n, m)) isa AT{Float32} - @test similar(d_x, (3, 4)) isa AT{ET} - @test size(similar(d_x, (3, 4))) == (3, 4) +function sparse_conversions(AT, eltypes) + @testset "conversions" begin + @testset "$ET" for ET in eltypes + A = sprand_awkward(ET, 9, 7; Ti=Int32) + dcsr = gpu_sparse(AT, GPUSparseMatrixCSR, A) + dcsc = gpu_sparse(AT, GPUSparseMatrixCSC, A) + + # between CSR and CSC + B = GPUSparseMatrixCSC(dcsr) + @test B isa GPUSparseMatrixCSC{ET,Int32} + check_structure(B) + @test same_sparse(B, A) + B = GPUSparseMatrixCSR(dcsc) + @test B isa GPUSparseMatrixCSR{ET,Int32} + check_structure(B) + @test same_sparse(B, A) + @test convert(GPUSparseMatrixCSR, dcsr) === dcsr + @test same_sparse(convert(GPUSparseMatrixCSR, dcsc), A) + + # element and index types + B = GPUSparseMatrixCSR{ET,Int64}(dcsc) + @test B isa GPUSparseMatrixCSR{ET,Int64} + @test same_sparse(B, SparseMatrixCSC{ET,Int64}(A)) + B = GPUSparseMatrixCSC{ET,Int64}(dcsc) + @test B isa GPUSparseMatrixCSC{ET,Int64} + @test same_sparse(B, A) + if ET <: Real && Float32 in eltypes + B = GPUSparseMatrixCSC{Float32}(dcsr) + @test B isa GPUSparseMatrixCSC{Float32,Int32} + @test Array(B) ≈ Array(A) + end + @test_throws ArgumentError GPUSparseMatrixCSR{ET,Int8}(gpu_sparse(AT, GPUSparseMatrixCSR, spzeros(ET, 200, 2))) + + # materializing transposes into the opposite format + B = GPUSparseMatrixCSC(transpose(dcsr)) + @test B isa GPUSparseMatrixCSC{ET,Int32} + @test same_sparse(B, copy(transpose(A))) + B = GPUSparseMatrixCSR(adjoint(dcsc)) + @test B isa GPUSparseMatrixCSR{ET,Int32} + @test same_sparse(B, copy(adjoint(A))) + + # shapes + for (m, n) in ((0, 0), (0, 4), (4, 0)) + Z = spzeros(ET, Int32, m, n) + @test same_sparse(GPUSparseMatrixCSC(gpu_sparse(AT, GPUSparseMatrixCSR, Z)), Z) + @test same_sparse(GPUSparseMatrixCSR(gpu_sparse(AT, GPUSparseMatrixCSC, Z)), Z) + end + + # the host arrays have no storage to convert to + @test_throws ArgumentError GPUSparseMatrixCSR(A) + end end end function broadcasting_vector(AT, eltypes) - dense_AT = GPUArrays.dense_array_type(AT) @testset "SparseVector broadcasting" begin @testset "$ET" for ET in eltypes m = 64 p = 0.5 x = sprand_nozeros(ET, m, p) - dx = AT(x) + dx = gpu_sparse(AT, x) # zero-preserving y = x .* ET(1) dy = dx .* ET(1) - @test dy isa AT{ET} + @test dy isa GPUSparseVector{ET} @test collect(SparseArrays.nonzeroinds(dy)) == collect(SparseArrays.nonzeroinds(dx)) @test collect(SparseArrays.nonzeroinds(dy)) == SparseArrays.nonzeroinds(y) @test collect(SparseArrays.nonzeros(dy)) == SparseArrays.nonzeros(y) @@ -182,66 +476,66 @@ function broadcasting_vector(AT, eltypes) # not zero-preserving y = x .+ ET(1) dy = dx .+ ET(1) - @test dy isa dense_AT{ET} + @test dy isa AT{ET} hy = Array(dy) @test Array(y) == hy # involving something dense y = x .+ ones(ET, m) - dy = dx .+ dense_AT(ones(ET, m)) - @test dy isa dense_AT{ET} + dy = dx .+ AT(ones(ET, m)) + @test dy isa AT{ET} @test Array(y) == Array(dy) # sparse to sparse - dx = AT(x) + dx = gpu_sparse(AT, x) y = sprand_nozeros(ET, m, p) - dy = AT(y) + dy = gpu_sparse(AT, y) z = x .* y dz = dx .* dy - @test dz isa AT{ET} + @test dz isa GPUSparseVector{ET} @test z ≈ SparseVector(dz) # multiple inputs y = sprand_nozeros(ET, m, p) w = sprand_nozeros(ET, m, p) - dy = AT(y) - dx = AT(x) - dw = AT(w) + dy = gpu_sparse(AT, y) + dx = gpu_sparse(AT, x) + dw = gpu_sparse(AT, w) z = @. x * y * w dz = @. dx * dy * dw - @test dz isa AT{ET} + @test dz isa GPUSparseVector{ET} @test z ≈ SparseVector(dz) y = sprand_nozeros(ET, m, p) w = sprand_nozeros(ET, m, p) dense_arr = rand(ET, m) - d_dense_arr = dense_AT(dense_arr) - dy = AT(y) - dw = AT(w) + d_dense_arr = AT(dense_arr) + dy = gpu_sparse(AT, y) + dw = gpu_sparse(AT, w) z = @. x * y * w * dense_arr dz = @. dx * dy * dw * d_dense_arr - @test dz isa dense_AT{ET} + @test dz isa AT{ET} @test Array(z) ≈ Array(dz) y = sprand_nozeros(ET, m, p) - dy = AT(y) - dx = AT(x) + dy = gpu_sparse(AT, y) + dx = gpu_sparse(AT, x) z = x .* y .* ET(2) dz = dx .* dy .* ET(2) - @test dz isa AT{ET} + @test dz isa GPUSparseVector{ET} @test z ≈ SparseVector(dz) # type-mismatching ## non-zero-preserving - dx = AT(x) + dx = gpu_sparse(AT, x) dy = dx .+ 1 y = x .+ 1 - @test dy isa dense_AT{promote_type(ET, Int)} + @test dy isa AT{promote_type(ET, Int)} @test Array(y) == Array(dy) ## zero-preserving dy = dx .* 1 y = x .* 1 - @test dy isa AT{promote_type(ET, Int)} + @test dy isa GPUSparseVector{promote_type(ET, Int)} @test collect(SparseArrays.nonzeroinds(dy)) == collect(SparseArrays.nonzeroinds(dx)) @test collect(SparseArrays.nonzeroinds(dy)) == SparseArrays.nonzeroinds(y) @test collect(SparseArrays.nonzeros(dy)) == SparseArrays.nonzeros(y) @@ -249,13 +543,13 @@ function broadcasting_vector(AT, eltypes) # without stored entries x = spzeros(ET, m) - dx = AT(x) + dx = gpu_sparse(AT, x) y = sprand_nozeros(ET, m, p) - dy = AT(y) + dy = gpu_sparse(AT, y) @test SparseVector(dx .* ET(2)) == x .* ET(2) @test SparseVector(dx .* dy) == x .* y @test Array(dx .+ dy) == Array(x .+ y) - @test length(AT(spzeros(ET, 0)) .* AT(spzeros(ET, 0))) == 0 + @test length(gpu_sparse(AT, spzeros(ET, 0)) .* gpu_sparse(AT, spzeros(ET, 0))) == 0 end end @@ -266,129 +560,137 @@ function broadcasting_vector(AT, eltypes) m = 10^4 x = sprand_nozeros(ET, m, 0.1) y = sprand_nozeros(ET, m, 0.1) - dx = AT(x) - dy = AT(y) + dx = gpu_sparse(AT, x) + dy = gpu_sparse(AT, y) dz = dx .* dy - @test dz isa AT{ET} + @test dz isa GPUSparseVector{ET} @test SparseVector(dz) == x .* y @test Array(dx .+ dy) == Array(x .+ y) end end function broadcasting_matrix(AT, eltypes) - dense_AT = GPUArrays.dense_array_type(AT) @testset "SparseMatrix broadcasting" begin - @testset "$ET" for ET in eltypes + @testset "$S{$ET}" for S in (GPUSparseMatrixCSR, GPUSparseMatrixCSC), ET in eltypes m, n = 5, 6 p = 0.5 x = sprand_nozeros(ET, m, n, p) - dx = AT(x) + dx = gpu_sparse(AT, S, x) # zero-preserving y = x .* ET(1) dy = dx .* ET(1) - @test dy isa AT{ET} + @test dy isa S{ET} @test y == SparseMatrixCSC(dy) + # the output does not share buffers with the input + fill!(nonzeros(dy), zero(ET)) + @test same_sparse(dx, x) # not zero-preserving y = x .+ ET(1) dy = dx .+ ET(1) - @test dy isa dense_AT{ET} - hy = Array(dy) - dense_y = Array(y) + @test dy isa AT{ET} @test Array(y) == Array(dy) # involving something dense y = x .* ones(ET, m, n) - dy = dx .* dense_AT(ones(ET, m, n)) - @test dy isa dense_AT{ET} + dy = dx .* AT(ones(ET, m, n)) + @test dy isa AT{ET} @test Array(y) == Array(dy) # multiple inputs y = sprand_nozeros(ET, m, n, p) - dy = AT(y) + dy = gpu_sparse(AT, S, y) z = x .* y .* ET(2) dz = dx .* dy .* ET(2) - @test dz isa AT{ET} + @test dz isa S{ET} @test z ≈ SparseMatrixCSC(dz) # multiple inputs w = sprand_nozeros(ET, m, n, p) - dw = AT(w) + dw = gpu_sparse(AT, S, w) z = x .* y .* w dz = dx .* dy .* dw - @test dz isa AT{ET} + @test dz isa S{ET} @test z ≈ SparseMatrixCSC(dz) # create a matrix with nnz < leading_dim x = spdiagm(m, m, 2=>rand(ET, m - 2)) - dx = AT(x) + dx = gpu_sparse(AT, S, x) y = ET(3) * x dy = ET(3) * dx @test y == SparseMatrixCSC(dy) x = spdiagm(m, m, -2=>rand(ET, m - 2)) - dx = AT(x) + dx = gpu_sparse(AT, S, x) y = ET(3) * x dy = ET(3) * dx @test y == SparseMatrixCSC(dy) + + # without rows or columns + for dims in ((0, 4), (4, 0)) + x = spzeros(ET, dims...) + dx = gpu_sparse(AT, S, x) + @test same_sparse(SparseMatrixCSC(dx .* dx), x .* x) + @test same_sparse(SparseMatrixCSC(dx .* ET(2)), x .* ET(2)) + @test size(dx .+ dx) == dims + end + + # stored zeros are kept + x = sprand_awkward(ET, 6, 7) + dx = gpu_sparse(AT, S, x) + @test nnz(dx .* ET(2)) == nnz(x) end end end -# the result of a reduction on the host, whether it is a scalar or an array -reduced(x::Number) = x -reduced(x::AbstractArray) = Array(x) - function mapreduce_matrix(AT, eltypes) - dense_AT = GPUArrays.dense_array_type(AT) @testset "SparseMatrix mapreduce" begin - @testset "$ET" for ET in eltypes + @testset "$S{$ET}" for S in sparse_matrix_formats, ET in eltypes m,n = 5,6 p = 0.5 x = sprand_nozeros(ET, m, n, p) - dx = AT(x) + dx = gpu_sparse(AT, S, x) # dim=: y = sum(x) dy = sum(dx) @test y ≈ dy - y = mapreduce(abs, +, x) - dy = mapreduce(abs, +, dx) - @test y ≈ dy + # `abs` may give an element type the back-end does not support (Complex{Int16}) + if Base.promote_op(abs, ET) in eltypes + y = mapreduce(abs, +, x) + dy = mapreduce(abs, +, dx) + @test y ≈ dy + y = mapreduce(abs, +, x, dims=1) + dy = mapreduce(abs, +, dx, dims=1) + @test y ≈ Array(dy) + y = mapreduce(abs, +, x, dims=2) + dy = mapreduce(abs, +, dx, dims=2) + @test y ≈ Array(dy) + dy = mapreduce(abs, +, dx; init=zero(Base.promote_op(abs, ET))) + y = mapreduce(abs, +, x; init=zero(Base.promote_op(abs, ET))) + @test y ≈ dy + + # test with a matrix with fully empty rows + x = zeros(ET, m, n) + x[2, :] .= -one(ET) + x[2, end] = -ET(16) + dx = gpu_sparse(AT, S, sparse(x)) + y = mapreduce(abs, max, x) + dy = mapreduce(abs, max, dx) + @test y ≈ dy + end # dim=1 y = sum(x, dims=1) dy = sum(dx, dims=1) @test y ≈ Array(dy) - y = mapreduce(abs, +, x, dims=1) - dy = mapreduce(abs, +, dx, dims=1) - @test y ≈ Array(dy) - # dim=2 y = sum(x, dims=2) dy = sum(dx, dims=2) @test y ≈ Array(dy) - y = mapreduce(abs, +, x, dims=2) - dy = mapreduce(abs, +, dx, dims=2) - @test y ≈ Array(dy) - if ET in (Float32, Float64) - dy = mapreduce(abs, +, dx; init=zero(ET)) - y = mapreduce(abs, +, x; init=zero(ET)) - @test y ≈ dy - end - - # test with a matrix with fully empty rows - x = zeros(ET, m, n) - x[2, :] .= -one(ET) - x[2, end] = -ET(16) - dx = AT(sparse(x)) - y = mapreduce(abs, max, x) - dy = mapreduce(abs, max, dx) - @test y ≈ dy - # result shapes follow Base @test size(sum(dx; dims=1)) == (1, n) @test size(sum(dx; dims=2)) == (m, 1) @@ -396,7 +698,7 @@ function mapreduce_matrix(AT, eltypes) # functions that don't preserve zeros, with operators other than `+` o = one(ET) x = sprand_nozeros(ET, m, n, p) - dx = AT(x) + dx = gpu_sparse(AT, S, x) for dims in (:, 1, 2) @test mapreduce(v -> v + o, *, x; dims) ≈ reduced(mapreduce(v -> v + o, *, dx; dims)) if ET <: Real @@ -409,7 +711,7 @@ function mapreduce_matrix(AT, eltypes) # an operator without a known neutral element, relying on `init` @test mapreduce(identity, (a, b) -> a + b, dx; init=zero(ET)) ≈ mapreduce(identity, (a, b) -> a + b, x; init=zero(ET)) y = sparse([1, 2], [1, 3], ET[1, 2], 3, 3) - dy = AT(y) + dy = gpu_sparse(AT, S, y) @test mapreduce(v -> v + o, +, dy) == mapreduce(v -> v + o, +, y) @test mapreduce(v -> v + o, *, dy) == mapreduce(v -> v + o, *, y) if ET <: Real @@ -419,16 +721,16 @@ function mapreduce_matrix(AT, eltypes) # the stored values of a row or column don't start from zero if ET <: Real && !(ET <: Unsigned) x = sparse(-rand(ET(1):ET(10), m, n)) - dx = AT(x) + dx = gpu_sparse(AT, S, x) @test nnz(dx) == m * n for dims in (:, 1, 2) @test maximum(x; dims) == reduced(maximum(dx; dims)) - @test minimum(-x; dims) == reduced(minimum(-dx; dims)) + @test minimum(-x; dims) == reduced(minimum(gpu_sparse(AT, S, -x); dims)) end # the same with some implicit zeros x[1, :] .= 0 dropzeros!(x) - dx = AT(x) + dx = gpu_sparse(AT, S, x) for dims in (:, 1, 2) @test maximum(x; dims) == reduced(maximum(dx; dims)) end @@ -437,122 +739,88 @@ function mapreduce_matrix(AT, eltypes) # empty dimensions for (k, l) in ((0, n), (m, 0), (0, 0)) x = spzeros(ET, k, l) - dx = AT(x) + dx = gpu_sparse(AT, S, x) @test sum(dx) == sum(x) - @test collect(sum(dx; dims=1)) == sum(x; dims=1) - @test collect(sum(dx; dims=2)) == sum(x; dims=2) - @test collect(mapreduce(abs, +, dx; dims=1, init=zero(real(ET)))) == mapreduce(abs, +, x; dims=1, init=zero(real(ET))) + @test reduced(sum(dx; dims=1)) == sum(x; dims=1) + @test reduced(sum(dx; dims=2)) == sum(x; dims=2) + if Base.promote_op(abs, ET) in eltypes + T = Base.promote_op(abs, ET) + @test reduced(mapreduce(abs, +, dx; dims=1, init=zero(T))) == mapreduce(abs, +, x; dims=1, init=zero(T)) + end end end + + @testset "vector $ET" for ET in eltypes + o = one(ET) + x = sprand_nozeros(ET, 20, 0.3) + dx = gpu_sparse(AT, x) + @test sum(dx) ≈ sum(x) + @test mapreduce(v -> v + o, *, dx) ≈ mapreduce(v -> v + o, *, x) + @test sum(gpu_sparse(AT, spzeros(ET, 0))) == zero(ET) + end end end -function linalg(AT, eltypes) - dense_AT = GPUArrays.dense_array_type(AT) - @testset "Sparse matrix linear algebra" begin +function sparse_linalg(AT, eltypes) + @testset "linear algebra" begin # sprandn only works on real or complex float types - @testset "$ET" for ET in filter(isfloattype, eltypes) + @testset "$S{$ET}" for S in sparse_matrix_formats, ET in filter(isfloattype, eltypes) m = 10 A = sprandn_nozeros(ET, m, m, 0.2) - B = sprandn_nozeros(ET, m, m, 0.3) - ZA = spzeros(ET, m, m) - C = I(div(m, 2)) - dA = AT(A) - dB = AT(B) - dZA = AT(ZA) - @testset "opnorm and norm" begin - @test opnorm(A, Inf) ≈ opnorm(dA, Inf) - @test opnorm(A, 1) ≈ opnorm(dA, 1) - # SparseArrays.jl#859 implemented the 2-norm for SparseMatrixCSC - if dA isa SparseMatrixCSC && isdefined(SparseArrays, :opnorm2est) - @test opnorm(Matrix(A), 2) ≈ opnorm(dA, 2) - else - @test_throws ArgumentError opnorm(dA, 2) - end + dA = gpu_sparse(AT, S, A) + @test opnorm(A, Inf) ≈ opnorm(dA, Inf) + @test opnorm(A, 1) ≈ opnorm(dA, 1) + @test_throws ArgumentError opnorm(dA, 2) + for dims in ((0, 2), (2, 0)), p in (1, Inf) + @test opnorm(gpu_sparse(AT, S, spzeros(ET, dims...)), p) == 0 + end + for p in (1, 2, 3, Inf, -Inf, -2, -1, -0.5, 0) + # (SparseArrays ignores the implicit zeros for negative orders) + reference = p < 0 ? norm(Matrix(A), p) : norm(A, p) + @test reference ≈ norm(dA, p) end + @test norm(gpu_sparse(AT, S, spzeros(ET, 3, 3)), Inf) == 0 end end end -function transpose_matrix(AT, eltypes) - @testset "Sparse matrix transpose and permutedims" begin +function iszero_sparse(AT, eltypes) + @testset "iszero" begin @testset "$ET" for ET in eltypes m = 10 - n = 15 - x = sprand_nozeros(ET, m, n, 0.2) - d_x = AT(x) - - @test collect(copy(transpose(d_x))) == collect(transpose(x)) - @test collect(copy(adjoint(d_x))) == collect(adjoint(x)) - @test collect(permutedims(d_x, (2, 1))) == collect(permutedims(x, (2, 1))) - @test collect(permutedims(d_x, (1, 2))) == collect(x) - @test collect(permutedims(d_x, [2, 1])) == collect(permutedims(x, [2, 1])) - - # the result has to stay on the device, in the format of the input - @test copy(transpose(d_x)) isa AT - @test copy(adjoint(d_x)) isa AT - @test permutedims(d_x, (2, 1)) isa AT - - @test_throws ArgumentError permutedims(d_x, (2, 2)) - end - end -end - -function iszero_vector(AT, eltypes) - @testset "iszero SparseVector" begin - @testset "$ET" for ET in eltypes - m = 25 - - # Test non-zero sparse vector x = sprand_nozeros(ET, m, 0.5) while iszero(x) x = sprand_nozeros(ET, m, 0.5) end - d_x = AT(x) - @test iszero(d_x) == iszero(x) - @test iszero(d_x) == false - - # Test zero sparse vector (no stored elements) - z = spzeros(ET, m) - d_z = AT(z) - @test iszero(d_z) == iszero(z) - @test iszero(d_z) == true - - # Test sparse vector with stored zeros (e.g., after operations) - # Create a sparse vector then multiply by zero - x_zeros = x .* zero(ET) - d_x_zeros = d_x .* zero(ET) - @test iszero(d_x_zeros) == iszero(x_zeros) - @test iszero(d_x_zeros) == true + @test !iszero(gpu_sparse(AT, x)) + @test iszero(gpu_sparse(AT, spzeros(ET, m))) + @test iszero(gpu_sparse(AT, x) .* zero(ET)) + + A = sparse(Matrix(x * transpose(x))) + for S in sparse_matrix_formats + @test !iszero(gpu_sparse(AT, S, A)) + @test iszero(gpu_sparse(AT, S, spzeros(ET, m, m))) + # stored zeros + @test iszero(gpu_sparse(AT, S, A .* zero(ET))) + end end end end -function iszero_matrix(AT, eltypes) - @testset "iszero SparseMatrix" begin - @testset "$ET" for ET in eltypes - m, n = 10, 10 - - # Test non-zero sparse matrix - A = sprand_nozeros(ET, m, n, 0.5) - while iszero(A) - A = sprand_nozeros(ET, m, n, 0.5) - end - dA = AT(A) - @test iszero(dA) == iszero(A) - @test iszero(dA) == false - - # Test zero sparse matrix (no stored elements) - ZA = spzeros(ET, m, n) - dZA = AT(ZA) - @test iszero(dZA) == iszero(ZA) - @test iszero(dZA) == true - - # Test sparse matrix with stored zeros - A_zeros = A .* zero(ET) - dA_zeros = dA .* zero(ET) - @test iszero(dA_zeros) == iszero(A_zeros) - @test iszero(dA_zeros) == true +function sparse_bool(AT) + @testset "Bool" begin + A = sprand(Bool, 8, 9, 0.4) + B = sprand(Bool, 8, 9, 0.4) + for S in (GPUSparseMatrixCSR, GPUSparseMatrixCSC) + dA = gpu_sparse(AT, S, A) + dB = gpu_sparse(AT, S, B) + @test same_sparse(dA, A) + @test Array(dA .& dB) == Array(A .& B) + @test Array(dA .| dB) == Array(A .| B) + @test sum(dA) == sum(A) + @test reduced(sum(dA; dims=1)) == sum(A; dims=1) + @test reduced(mapreduce(identity, |, dA; dims=2)) == mapreduce(identity, |, A; dims=2) + @test same_sparse(GPUSparseMatrixCSC(gpu_sparse(AT, GPUSparseMatrixCSR, A)), A) end end end From 3a199466575327ad1337e90d6d6645939d17d228 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Fri, 25 Sep 2026 17:51:35 +0200 Subject: [PATCH 02/11] Convert between dense and sparse arrays and materialize transposes Dense GPU arrays convert to every sparse format on the device, through findall in column-major order (or on the dense transpose for row-major formats), and sparse(A; fmt) selects the format like CUDA.jl and AMDGPU.jl do. copy(transpose(A)), copy(adjoint(A)) and permutedims produce the same format by regrouping, and the format constructors accept transposed operands. findnz returns the coordinates in column-major order like SparseArrays, for every format. --- docs/src/interface.md | 5 ++ src/host/sparse/conversions.jl | 137 +++++++++++++++++++++++++++++++++ test/testsuite/sparse.jl | 93 ++++++++++++++++++++++ 3 files changed, 235 insertions(+) diff --git a/docs/src/interface.md b/docs/src/interface.md index 9bb7edc96..9933f10bb 100644 --- a/docs/src/interface.md +++ b/docs/src/interface.md @@ -145,7 +145,12 @@ Conversions: `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 diff --git a/src/host/sparse/conversions.jl b/src/host/sparse/conversions.jl index 6ecf8d43f..a07a8b110 100644 --- a/src/host/sparse/conversions.jl +++ b/src/host/sparse/conversions.jl @@ -301,4 +301,141 @@ function Base.copyto!(dst::AnyGPUArray, src::GPUSparseArray) end Base.copyto!(dst::Array, src::GPUSparseArray) = copyto!(dst, Array(src)) +# (more specific than SparseArrays' `copyto!(::AbstractVector, ::AbstractSparseVector)`) +Base.copyto!(dst::Vector, src::GPUSparseVector) = copyto!(dst, Array(src)) + +## from dense arrays + +## COV_EXCL_START +@kernel function split_linear_kernel(rows, cols, lin, m) + k = @index(Global, Linear) + if k <= length(lin) + l = @inbounds lin[k] - 1 + @inbounds rows[k] = (l % m + 1) % eltype(rows) + @inbounds cols[k] = (l ÷ m + 1) % eltype(cols) + end +end +## COV_EXCL_STOP + +# The nonzeros of a dense matrix, in column-major (`major = :col`) or row-major order: +# the sorted major and minor index of every entry, and its value. `findall` visits the +# elements in column-major order, so row-major order comes from a dense transpose. +function dense_entries(A::AbstractGPUMatrix, ::Type{Tv}, ::Type{Ti}, major::Symbol) where {Tv,Ti} + check_sparse_dims(Ti, size(A)) + B = major === :col ? A : permutedims(A) + lin = findall(!iszero, vec(B)) + n = length(lin) + check_sparse_nnz(Ti, n) + minor_ind = similar(lin, Ti, n) + major_ind = similar(lin, Ti, n) + if n > 0 + split_linear_kernel(get_backend(lin))(minor_ind, major_ind, lin, size(B, 1); ndrange=n) + end + vals = vec(B)[lin] + B === A || free_buffer!(B) + free_buffer!(lin) + return major_ind, minor_ind, eltype(vals) == Tv ? vals : convert_buffer(Tv, vals) +end + +function GPUSparseMatrixCSC{Tv,Ti}(A::AbstractGPUMatrix) where {Tv,Ti} + cols, rows, vals = dense_entries(A, Tv, Ti, :col) + GPUSparseMatrixCSC(compress(cols, size(A, 2)), rows, vals, size(A)) +end +function GPUSparseMatrixCSR{Tv,Ti}(A::AbstractGPUMatrix) where {Tv,Ti} + rows, cols, vals = dense_entries(A, Tv, Ti, :row) + GPUSparseMatrixCSR(compress(rows, size(A, 1)), cols, vals, size(A)) +end +function GPUSparseMatrixCOO{Tv,Ti}(A::AbstractGPUMatrix) where {Tv,Ti} + rows, cols, vals = dense_entries(A, Tv, Ti, :row) + GPUSparseMatrixCOO(rows, cols, vals, size(A)) +end +function GPUSparseVector{Tv,Ti}(x::AbstractGPUVector) where {Tv,Ti} + check_sparse_dims(Ti, size(x)) + lin = findall(!iszero, x) + check_sparse_nnz(Ti, length(lin)) + vals = x[lin] + GPUSparseVector(convert_index_buffer(Ti, lin), + eltype(vals) == Tv ? vals : convert_buffer(Tv, vals), length(x)) +end + +# like SparseArrays, the index type defaults to `Int` +for (S, D) in ((:GPUSparseMatrixCSR, :AbstractGPUMatrix), (:GPUSparseMatrixCSC, :AbstractGPUMatrix), + (:GPUSparseMatrixCOO, :AbstractGPUMatrix), (:GPUSparseVector, :AbstractGPUVector)) + @eval begin + $S(A::$D) = $S{eltype(A),Int}(A) + $S{Tv}(A::$D) where {Tv} = $S{Tv,Int}(A) + end +end + +# the format selected by a `fmt` keyword, as CUDA.jl and AMDGPU.jl accept it +sparse_format(fmt::Symbol) = + fmt === :csc ? GPUSparseMatrixCSC : + fmt === :csr ? GPUSparseMatrixCSR : + fmt === :coo ? GPUSparseMatrixCOO : + throw(ArgumentError("unknown sparse format :$fmt, expected :csc, :csr or :coo")) + +""" + sparse(A::AbstractGPUMatrix; fmt=:csc) + sparse(x::AbstractGPUVector) + +The nonzero entries of a dense GPU array as a sparse array on the same device, in CSC +format (like SparseArrays) or the format selected by `fmt` (`:csc`, `:csr` or `:coo`). +""" +SparseArrays.sparse(A::AbstractGPUMatrix; fmt::Symbol=:csc) = sparse_format(fmt)(A) +SparseArrays.sparse(x::AbstractGPUVector) = GPUSparseVector(x) +SparseArrays.sparsevec(x::AbstractGPUVector) = GPUSparseVector(x) + + +## transposes + +# The transpose with new storage, in the same format: the other compressed layout of `A`, +# read with reversed dimensions. +materialize_transpose(A::GPUSparseMatrixCSR) = transpose_view(GPUSparseMatrixCSC(A)) +materialize_transpose(A::GPUSparseMatrixCSC) = transpose_view(GPUSparseMatrixCSR(A)) +materialize_transpose(A::GPUSparseMatrixCOO) = + GPUSparseMatrixCOO(transpose_view(GPUSparseMatrixCSC(A))) + +Base.copy(A::Transpose{<:Any,<:GPUSparseMatrix}) = materialize_transpose(parent(A)) +function Base.copy(A::Adjoint{<:Any,<:GPUSparseMatrix}) + B = materialize_transpose(parent(A)) + eltype(B) <: Real || (B.nzVal .= conj.(B.nzVal)) + return B +end + +# into any format; the cheap cases (CSR ↔ CSC) are defined above +const GPUSparseMatrixAdjOrTrans = Union{Transpose{<:Any,<:GPUSparseMatrix}, + Adjoint{<:Any,<:GPUSparseMatrix}} +for S in (:GPUSparseMatrixCSR, :GPUSparseMatrixCSC, :GPUSparseMatrixCOO) + @eval begin + $S(A::GPUSparseMatrixAdjOrTrans) = convert($S, copy(A)) + $S{Tv}(A::GPUSparseMatrixAdjOrTrans) where {Tv} = $S{Tv,indtype(parent(A))}(A) + function $S{Tv,Ti}(A::GPUSparseMatrixAdjOrTrans) where {Tv,Ti} + check_sparse_dims(Ti, size(A)) + check_sparse_nnz(Ti, nnz(parent(A))) + return with_eltypes($S(A), Tv, Ti) + end + # (e.g. `Base.to_power_type` for `transpose(A)^2`; the storage stays that of the parent) + function Base.convert(::Type{T}, A::GPUSparseMatrixAdjOrTrans) where {T<:$S} + B = isconcretetype(T) ? $S{eltype(T),eltype(fieldtype(T, 1))}(A) : $S(A) + return convert(T, B) + end + end +end + +Base.permutedims(A::GPUSparseMatrix) = copy(transpose(A)) +function Base.permutedims(A::GPUSparseMatrix, perm) + Tuple(perm) == (2, 1) && return copy(transpose(A)) + Tuple(perm) == (1, 2) && return copy(A) + throw(ArgumentError("no valid permutation of dimensions")) +end + + +## coordinates + +# the coordinates and values of the stored entries, in column-major order like SparseArrays +function SparseArrays.findnz(A::GPUSparseMatrix) + C = A isa GPUSparseMatrixCSC ? copy(A) : GPUSparseMatrixCSC(A) + return C.rowVal, expand_ptr(C.colPtr, nnz(C)), C.nzVal +end +SparseArrays.findnz(x::GPUSparseVector) = (copy(x.nzInd), copy(x.nzVal)) diff --git a/test/testsuite/sparse.jl b/test/testsuite/sparse.jl index d66484366..b2a7364a8 100644 --- a/test/testsuite/sparse.jl +++ b/test/testsuite/sparse.jl @@ -16,6 +16,7 @@ using GPUArrays: GPUSparseMatrixCSR, GPUSparseMatrixCSC, GPUSparseMatrixCOO, GPU sparse_display(AT, eltypes) sparse_indexing(AT, eltypes) sparse_conversions(AT, eltypes) + sparse_dense_conversions(AT, eltypes) broadcasting_vector(AT, eltypes) broadcasting_matrix(AT, eltypes) mapreduce_matrix(AT, eltypes) @@ -824,3 +825,95 @@ function sparse_bool(AT) end end end + +function sparse_dense_conversions(AT, eltypes) + @testset "format and dense conversions" begin + @testset "$ET" for ET in eltypes + A = sprand_awkward(ET, 9, 7; Ti=Int32) + # every pair of formats, exactly, including the stored zero + for S in sparse_matrix_formats, S′ in sparse_matrix_formats + dA = gpu_sparse(AT, S, A) + B = S′(dA) + @test B isa S′{ET,Int32} + check_structure(B) + @test same_sparse(B, A) + # no hidden sharing with the source + fill!(nonzeros(B), zero(ET)) + @test same_sparse(dA, A) + end + + # transposes, into the same and the other formats + for S in sparse_matrix_formats + dA = gpu_sparse(AT, S, A) + B = copy(transpose(dA)) + @test B isa S{ET,Int32} + check_structure(B) + @test same_sparse(B, copy(transpose(A))) + B = copy(adjoint(dA)) + @test same_sparse(B, copy(adjoint(A))) + @test same_sparse(permutedims(dA), permutedims(A)) + @test same_sparse(permutedims(dA, (1, 2)), A) + for S′ in sparse_matrix_formats + @test same_sparse(S′(transpose(dA)), copy(transpose(A))) + @test same_sparse(S′(adjoint(dA)), copy(adjoint(A))) + end + end + + # dense to sparse drops the zeros, like `sparse` + D = Array(A) + dD = AT(D) + for S in sparse_matrix_formats + B = S(dD) + @test B isa S{ET,Int} + check_structure(B) + @test same_sparse(B, sparse(D)) + B = S{ET,Int32}(dD) + @test B isa S{ET,Int32} + @test same_sparse(B, SparseMatrixCSC{ET,Int32}(sparse(D))) + end + # (the index type of `sparse` is the back-end's choice: `Int` by default, but + # back-ends that defaulted to `Cint` before GPUArrays 13 keep doing so) + @test sparse(dD) isa GPUSparseMatrixCSC{ET} + @test same_sparse(sparse(dD), sparse(D)) + for (fmt, S) in ((:csc, GPUSparseMatrixCSC), (:csr, GPUSparseMatrixCSR), + (:coo, GPUSparseMatrixCOO)) + dS = sparse(dD; fmt) + @test dS isa S{ET} + check_structure(dS) + @test same_sparse(dS, sparse(D)) + end + @test_throws ArgumentError sparse(dD; fmt=:unknown) + for (m, n) in ((0, 0), (0, 3), (3, 0), (3, 4)) + Z = zeros(ET, m, n) + for S in sparse_matrix_formats + @test same_sparse(S(AT(Z)), sparse(Z)) + end + end + + x = Vector(sprand_nozeros(ET, 30, 0.3)) + dx = sparse(AT(x)) + @test dx isa GPUSparseVector{ET,Int} + check_structure(dx) + @test same_sparse(dx, sparse(x)) + @test GPUSparseVector{ET,Int32}(AT(x)) isa GPUSparseVector{ET,Int32} + @test same_sparse(sparsevec(AT(x)), sparse(x)) + @test nnz(sparse(AT(zeros(ET, 0)))) == 0 + end + end + + @testset "findnz" begin + ET = first(eltypes) + A = sprand_awkward(ET, 9, 7; Ti=Int32) + I, J, V = findnz(A) + # in particular for COO, whose conversion once recursed without end (CUDA.jl#3189) + for S in sparse_matrix_formats + dI, dJ, dV = findnz(gpu_sparse(AT, S, A)) + @test dI isa AT{Int32} && dJ isa AT{Int32} && dV isa AT{ET} + @test Array(dI) == I && Array(dJ) == J && Array(dV) == V + end + x = SparseVector{ET,Int32}(sprand_nozeros(ET, 20, 0.3)) + dI, dV = findnz(gpu_sparse(AT, x)) + @test (Array(dI), Array(dV)) == findnz(x) + @test all(isempty, findnz(gpu_sparse(AT, GPUSparseMatrixCOO, spzeros(ET, 3, 3)))) + end +end From 57495579b5fcd7add72726186e39e9c6c7457416 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Fri, 25 Sep 2026 17:52:22 +0200 Subject: [PATCH 03/11] Assemble sparse arrays from coordinates on the device sparse(I, J, V, m, n, combine; fmt) and sparsevec(I, V, n, combine) on GPU vectors sort packed coordinate keys, combine the values of repeated coordinates with a per-run fold (reduce_by_key), and compress. As in SparseArrays, repeated entries are folded in input order, with + by default (| for Bool) and explicit zeros kept, so that a call gives the same result on every back-end. Because the device sort need not be stable, a second sort on keys that are unique per entry restores the input order, unless the order cannot change the result (+ and * on integers, max, min, |, &). spdiagm builds on the same path. --- docs/src/interface.md | 2 + src/GPUArrays.jl | 1 + src/host/sparse/assembly.jl | 245 +++++++++++++++++++++++++++++++++ src/host/sparse/conversions.jl | 9 ++ src/host/sparse/primitives.jl | 51 +++++++ test/testsuite/sparse.jl | 92 +++++++++++++ 6 files changed, 400 insertions(+) create mode 100644 src/host/sparse/assembly.jl diff --git a/docs/src/interface.md b/docs/src/interface.md index 9933f10bb..0b0e4b5da 100644 --- a/docs/src/interface.md +++ b/docs/src/interface.md @@ -119,6 +119,7 @@ A back-end will usually add the following, all optional: |:--|:--|:--| | 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` | 5. **Formats that GPUArrays does not implement** (block-sparse formats, batched matrices) as back-end structs subtyping `AbstractGPUSparseArray`, with constructors to and from @@ -170,6 +171,7 @@ GPUArrays.GPUSparseVector GPUArrays.check_structure GPUArrays.generic_regroup GPUArrays.generic_expand +GPUArrays.generic_assemble ``` ## Caching Allocator diff --git a/src/GPUArrays.jl b/src/GPUArrays.jl index a928c7e57..1e68843ac 100644 --- a/src/GPUArrays.jl +++ b/src/GPUArrays.jl @@ -44,6 +44,7 @@ include("host/statistics.jl") include("host/sparse/types.jl") include("host/sparse/primitives.jl") include("host/sparse/conversions.jl") +include("host/sparse/assembly.jl") include("host/sparse/show.jl") include("host/sparse/indexing.jl") include("host/sparse/broadcast.jl") diff --git a/src/host/sparse/assembly.jl b/src/host/sparse/assembly.jl new file mode 100644 index 000000000..d566462c1 --- /dev/null +++ b/src/host/sparse/assembly.jl @@ -0,0 +1,245 @@ +# assembly of sparse arrays from coordinates + +## COV_EXCL_START +# a sort key that is unique per input entry, ordering the runs of equal coordinates and, +# within each run, the entries by their position in the input +@kernel function positional_key_kernel(keys, runs, perm, n) + k = @index(Global, Linear) + if k <= length(keys) + @inbounds keys[k] = (Int64(runs[k]) - 1) * n + (Int64(perm[k]) - 1) + end +end + +@kernel function unpack_key_kernel(major, minor, keys, nminor) + k = @index(Global, Linear) + if k <= length(keys) + key = @inbounds keys[k] - 1 + @inbounds major[k] = (key ÷ nminor + 1) % eltype(major) + @inbounds minor[k] = (key % nminor + 1) % eltype(minor) + end +end +## COV_EXCL_STOP + +# whether the order in which `combine` folds repeated entries of type `T` cannot change +# the result, so that the assembly does not need to restore the input order after an +# unstable sort. Floating-point `+` and `*` are commutative but not associative. +order_independent(combine, ::Type) = false +order_independent(::Union{typeof(max), typeof(min), typeof(|), typeof(&)}, ::Type) = true +order_independent(::Union{typeof(+), typeof(*)}, ::Type{<:Union{Integer,Complex{<:Integer}}}) = true + +# the order in which to visit the entries so that equal keys are contiguous and appear in +# their input order. `sortperm` need not be stable (Metal's is not), so the input order is +# restored with a second sort on keys that are unique per entry. +function stable_sortperm(keys::AbstractVector{Int64}; stable::Bool) + perm = sortperm(keys) + stable || return perm + n = length(keys) + widemul(n, n) < typemax(Int64) || + throw(ArgumentError("cannot assemble $n entries in input order")) + sorted = keys[perm] + runs = accumulate(Base.add_sum, run_heads(sorted)) + free_buffer!(sorted) + unique_keys = similar(keys) + if !isempty(keys) + positional_key_kernel(get_backend(keys))(unique_keys, runs, perm, n; ndrange=n) + end + free_buffer!(runs) + perm2 = sortperm(unique_keys) + free_buffer!(unique_keys) + return perm[perm2] +end + +function check_coordinates(I::AbstractVector, m::Integer, name) + isempty(I) || all(i -> 1 <= i <= m, I) || + throw(ArgumentError("all $name indices must be in 1:$m")) + return +end + +""" + GPUArrays.generic_assemble(I, J, V, m, n, combine; major=:col) + +Assemble the entries `V[k]` at coordinates `(I[k], J[k])` of an `m`×`n` sparse matrix on +the device, combining the values of repeated coordinates with `combine`, from left to +right in input order (as `sparse(I, J, V, m, n, combine)` does). Returns the sorted, unique +major and minor indices of the entries (rows and columns for `major=:row`, columns and +rows for `major=:col`) and their values, from which any format can be built. + +The coordinates are sorted on packed `Int64` keys; a second sort restores the input order +of repeated coordinates, unless the order cannot change the result (e.g. `+` on integers, +or `max`). `combine` must be associative. +""" +function generic_assemble(I::AbstractVector, J::AbstractVector, V::AbstractVector, + m::Integer, n::Integer, combine; major::Symbol=:col) + length(I) == length(J) == length(V) || + throw(ArgumentError("the coordinate and value vectors must have the same length")) + check_coordinates(I, m, "row") + check_coordinates(J, n, "column") + Ti = promote_type(eltype(I), eltype(J)) + check_sparse_dims(Ti, (Int(m), Int(n))) + majors, minors, dims = major === :col ? (J, I, (n, m)) : + major === :row ? (I, J, (m, n)) : + throw(ArgumentError("major must be :row or :col")) + + keys = packed_keys(majors, minors, dims) + perm = stable_sortperm(keys; stable=!order_independent(combine, eltype(V))) + sorted_keys = keys[perm] + sorted_vals = V[perm] + free_buffer!(keys) + free_buffer!(perm) + unique_keys, vals = reduce_by_key(combine, sorted_keys, sorted_vals) + free_buffer!(sorted_keys) + free_buffer!(sorted_vals) + + out_major = similar(I, Ti, length(unique_keys)) + out_minor = similar(I, Ti, length(unique_keys)) + if !isempty(unique_keys) + unpack_key_kernel(get_backend(unique_keys))(out_major, out_minor, unique_keys, + Int64(dims[2]); ndrange=length(unique_keys)) + end + free_buffer!(unique_keys) + return out_major, out_minor, vals +end + +default_combine(::Type{Bool}) = | +default_combine(::Type) = + + +const AnyGPUIndices = AnyGPUVector{<:Integer} + +# `combine` can be any callable, with separate methods for functions to be more specific +# than SparseArrays' own +function assemble_matrix(I, J, V, m::Integer, n::Integer, combine, fmt::Symbol) + m, n = Int(m), Int(n) + S = sparse_format(fmt) + major, minor, vals = generic_assemble(I, J, V, m, n, combine; + major=S <: GPUSparseMatrixCSC ? :col : :row) + if S <: GPUSparseMatrixCSC + GPUSparseMatrixCSC(compress(major, n), minor, vals, (m, n)) + elseif S <: GPUSparseMatrixCSR + GPUSparseMatrixCSR(compress(major, m), minor, vals, (m, n)) + else + GPUSparseMatrixCOO(major, minor, vals, (m, n)) + end +end +for C in (:Function, :Any) + @eval begin + SparseArrays.sparse(I::AnyGPUIndices, J::AnyGPUIndices, V::AnyGPUVector, m::Integer, + n::Integer, combine::$C; fmt::Symbol=:csc) = + assemble_matrix(I, J, V, m, n, combine, fmt) + SparseArrays.sparse(I::AnyGPUIndices, J::AnyGPUIndices, v::Number, m::Integer, + n::Integer, combine::$C; fmt::Symbol=:csc) = + assemble_matrix(I, J, fill!(similar(I, typeof(v)), v), m, n, combine, fmt) + end +end + +""" + sparse(I, J, V, [m, n, combine]; fmt=:csc) + +Assemble a GPU sparse matrix from GPU vectors of coordinates `I`, `J` and values `V` (or a +single value), like `SparseArrays.sparse`: repeated coordinates are combined with +`combine` (by default `+`, or `|` for `Bool` values) in input order, and explicit zeros are +kept. `fmt` selects the format: `:csc`, `:csr` or `:coo`. +""" +SparseArrays.sparse(::AnyGPUIndices, ::AnyGPUIndices, ::AnyGPUVector, ::Integer, ::Integer, ::Any) + +# the defaults: `+`, or `|` for Bool values, and the dimensions that hold the coordinates. +# (Separate methods for Bool values keep these more specific than SparseArrays' own.) +for V in (:AnyGPUVector, :(AnyGPUVector{Bool}), :Number, :Bool) + @eval begin + SparseArrays.sparse(I::AnyGPUIndices, J::AnyGPUIndices, V::$V, m::Integer, n::Integer; + fmt::Symbol=:csc) = + sparse(I, J, V, m, n, default_combine(eltype(V)); fmt) + SparseArrays.sparse(I::AnyGPUIndices, J::AnyGPUIndices, V::$V; fmt::Symbol=:csc) = + sparse(I, J, V, sparse_extent(I), sparse_extent(J); fmt) + end +end + +# the smallest dimension that holds the given indices +sparse_extent(I::AbstractVector) = isempty(I) ? 0 : Int(maximum(I)) + +function assemble_vector(I, V, n::Integer, combine) + n = Int(n) + J = fill!(similar(I), one(eltype(I))) + nzInd, _, vals = generic_assemble(I, J, V, n, 1, combine; major=:row) + GPUSparseVector(nzInd, vals, n) +end +for C in (:Function, :Any) + @eval begin + SparseArrays.sparsevec(I::AnyGPUIndices, V::AnyGPUVector, n::Integer, combine::$C) = + assemble_vector(I, V, n, combine) + SparseArrays.sparsevec(I::AnyGPUIndices, v::Number, n::Integer, combine::$C) = + assemble_vector(I, fill!(similar(I, typeof(v)), v), n, combine) + SparseArrays.sparsevec(I::AnyGPUIndices, V::Union{AnyGPUVector,Number}, combine::$C) = + sparsevec(I, V, sparse_extent(I), combine) + end +end +""" + sparsevec(I, V, [n, combine]) + +Assemble a GPU sparse vector from GPU vectors of indices `I` and values `V` (or a single +value), like `SparseArrays.sparsevec`. +""" +SparseArrays.sparsevec(::AnyGPUIndices, ::AnyGPUVector, ::Integer, ::Any) + +for V in (:AnyGPUVector, :(AnyGPUVector{Bool}), :Number, :Bool) + @eval begin + SparseArrays.sparsevec(I::AnyGPUIndices, V::$V, n::Integer) = + sparsevec(I, V, n, default_combine(eltype(V))) + SparseArrays.sparsevec(I::AnyGPUIndices, V::$V) = sparsevec(I, V, sparse_extent(I)) + end +end + +""" + spdiagm(kv::Pair{<:Integer,<:AbstractGPUVector}...) + spdiagm(m, n, kv...) + +A GPU sparse matrix in CSC format with the vectors `kv[i].second` on the diagonals +`kv[i].first`, like `SparseArrays.spdiagm`. Values on the same diagonal are added. +""" +const GPUDiagonal = Pair{<:Integer,<:AnyGPUVector} +SparseArrays.spdiagm(kv::GPUDiagonal, kvs::GPUDiagonal...) = spdiagm_gpu(nothing, kv, kvs...) +SparseArrays.spdiagm(m::Integer, n::Integer, kv::GPUDiagonal, kvs::GPUDiagonal...) = + spdiagm_gpu((Int(m), Int(n)), kv, kvs...) +SparseArrays.spdiagm(v::AnyGPUVector) = spdiagm(0 => v) +SparseArrays.spdiagm(m::Integer, n::Integer, v::AnyGPUVector) = spdiagm(m, n, 0 => v) + +## COV_EXCL_START +@kernel function diagonal_coordinates_kernel(I, J, k) + i = @index(Global, Linear) + if i <= length(I) + @inbounds I[i] = k >= 0 ? i : i - k + @inbounds J[i] = k >= 0 ? i + k : i + end +end +## COV_EXCL_STOP + +function spdiagm_gpu(dims, kv::Pair{<:Integer,<:AnyGPUVector}...) + # the extent that holds every diagonal, as in SparseArrays + m, n = if dims === nothing + extent = maximum(((k, v),) -> length(v) + abs(k), kv) + (extent, extent) + else + dims + end + Tv = mapreduce(((_, v),) -> eltype(v), promote_type, kv) + proto = last(first(kv)) + total = sum(((_, v),) -> length(v), kv) + I = similar(proto, Int, total) + J = similar(proto, Int, total) + V = similar(proto, Tv, total) + offset = 0 + for (k, v) in kv + len = length(v) + rows = k >= 0 ? len : len - k + cols = k >= 0 ? len + k : len + (rows <= m && cols <= n) || + throw(DimensionMismatch("diagonal $k of length $len does not fit in a $m×$n matrix")) + if len > 0 + range = offset+1:offset+len + diagonal_coordinates_kernel(get_backend(proto))(view(I, range), view(J, range), k; + ndrange=len) + view(V, range) .= v + end + offset += len + end + return sparse(I, J, V, m, n, +) +end diff --git a/src/host/sparse/conversions.jl b/src/host/sparse/conversions.jl index a07a8b110..96f1e2cb8 100644 --- a/src/host/sparse/conversions.jl +++ b/src/host/sparse/conversions.jl @@ -386,6 +386,15 @@ SparseArrays.sparse(A::AbstractGPUMatrix; fmt::Symbol=:csc) = sparse_format(fmt) SparseArrays.sparse(x::AbstractGPUVector) = GPUSparseVector(x) SparseArrays.sparsevec(x::AbstractGPUVector) = GPUSparseVector(x) +# a diagonal matrix, like `SparseMatrixCSC(::Diagonal)` +for S in (:GPUSparseMatrixCSR, :GPUSparseMatrixCSC, :GPUSparseMatrixCOO) + @eval begin + $S(D::Diagonal{<:Any,<:AnyGPUVector}) = convert($S, spdiagm(D.diag)) + $S{Tv}(D::Diagonal{<:Any,<:AnyGPUVector}) where {Tv} = $S{Tv,Int}(D) + $S{Tv,Ti}(D::Diagonal{<:Any,<:AnyGPUVector}) where {Tv,Ti} = with_eltypes($S(D), Tv, Ti) + end +end + ## transposes diff --git a/src/host/sparse/primitives.jl b/src/host/sparse/primitives.jl index fc953a228..39b374774 100644 --- a/src/host/sparse/primitives.jl +++ b/src/host/sparse/primitives.jl @@ -104,3 +104,54 @@ function regroup(ptr::AbstractVector{Ti}, ind::AbstractVector{Ti}, val::Abstract free_buffer!(new_major) return new_ptr, new_ind, new_val end + +## COV_EXCL_START +@kernel function run_heads_kernel(heads, keys) + k = @index(Global, Linear) + if k <= length(keys) + @inbounds heads[k] = k == 1 || keys[k] != keys[k-1] + end +end + +# every thread folds one run of equal keys, in order +@kernel function reduce_runs_kernel(out_keys, out_vals, op, starts, keys, vals) + r = @index(Global, Linear) + if r <= length(starts) + first = @inbounds starts[r] + last = r == length(starts) ? length(vals) : @inbounds(starts[r+1]) - 1 + acc = @inbounds vals[first] + for k in first+1:last + acc = convert(eltype(out_vals), op(acc, @inbounds vals[k])) + end + @inbounds out_keys[r] = keys[first] + @inbounds out_vals[r] = acc + end +end +## COV_EXCL_STOP + +# whether each of the sorted `keys` starts a run of equal keys +function run_heads(keys::AbstractVector) + heads = similar(keys, Bool) + isempty(keys) || run_heads_kernel(get_backend(keys))(heads, keys; ndrange=length(keys)) + return heads +end + +""" + reduce_by_key(op, keys, vals) + +Combine the values of every run of equal, sorted `keys` with `op`, folding each run from +left to right. Returns the unique keys and the combined values. Every run is folded by +one thread, so many repetitions of a key serialize; a tiled segmented scan would not. +""" +function reduce_by_key(op, keys::AbstractVector, vals::AbstractVector) + starts = findall(run_heads(keys)) + n = length(starts) + out_keys = similar(keys, n) + out_vals = similar(vals, n) + if n > 0 + reduce_runs_kernel(get_backend(vals))(out_keys, out_vals, op, starts, keys, vals; + ndrange=n) + end + free_buffer!(starts) + return out_keys, out_vals +end diff --git a/test/testsuite/sparse.jl b/test/testsuite/sparse.jl index b2a7364a8..50e8bd237 100644 --- a/test/testsuite/sparse.jl +++ b/test/testsuite/sparse.jl @@ -17,6 +17,7 @@ using GPUArrays: GPUSparseMatrixCSR, GPUSparseMatrixCSC, GPUSparseMatrixCOO, GPU sparse_indexing(AT, eltypes) sparse_conversions(AT, eltypes) sparse_dense_conversions(AT, eltypes) + sparse_assembly(AT, eltypes) broadcasting_vector(AT, eltypes) broadcasting_matrix(AT, eltypes) mapreduce_matrix(AT, eltypes) @@ -917,3 +918,94 @@ function sparse_dense_conversions(AT, eltypes) @test all(isempty, findnz(gpu_sparse(AT, GPUSparseMatrixCOO, spzeros(ET, 3, 3)))) end end + +# a callable that is not a `Function` +struct SparseSubtract end +(::SparseSubtract)(a, b) = a - b + +function sparse_assembly(AT, eltypes) + @testset "assembly" begin + @testset "$ET" for ET in eltypes + m, n = 7, 9 + N = 60 + # unordered, with many repeated coordinates + I = rand(1:m, N) + J = rand(1:n, N) + V = rand(ET, N) + dI, dJ, dV = AT(I), AT(J), AT(V) + + A = sparse(I, J, V, m, n) + for fmt in (:csc, :csr, :coo) + B = sparse(dI, dJ, dV, m, n; fmt) + @test B isa GPUArrays.sparse_format(fmt){ET,Int} + check_structure(B) + @test SparseMatrixCSC(B) ≈ A + end + if ET <: Real + @test same_sparse(sparse(dI, dJ, dV, m, n, max), sparse(I, J, V, m, n, max)) + end + # a combine function that depends on the order of the repeated entries + noncommutative(a, b) = a + a - b + B = sparse(dI, dJ, dV, m, n, noncommutative) + @test same_sparse(B, sparse(I, J, V, m, n, noncommutative)) + B = sparse(dI, dJ, dV, m, n, noncommutative; fmt=:csr) + @test same_sparse(B, sparse(I, J, V, m, n, noncommutative)) + + # repeated entries are folded in input order, also for `+`, which is not + # associative in floating point + if ET <: AbstractFloat + big = ET(2)^(precision(ET) + 1) + Is, Js, Vs = [2, 2, 2, 1], [3, 3, 3, 1], [big, -big, one(ET), one(ET)] + @test same_sparse(sparse(AT(Is), AT(Js), AT(Vs), 2, 3), sparse(Is, Js, Vs, 2, 3)) + @test same_sparse(sparsevec(AT(Is), AT(Vs), 2), sparsevec(Is, Vs, 2)) + end + # any callable + @test same_sparse(sparse(dI, dJ, dV, m, n, SparseSubtract()), sparse(I, J, V, m, n, -)) + @test same_sparse(sparsevec(dI, dV, m, SparseSubtract()), sparsevec(I, V, m, -)) + + # the dimensions follow from the coordinates + @test size(sparse(dI, dJ, dV)) == size(sparse(I, J, V)) + # a single value, and index types + @test same_sparse(sparse(dI, dJ, one(ET), m, n), sparse(I, J, one(ET), m, n)) + B = sparse(AT(Int32.(I)), AT(Int32.(J)), dV, m, n) + @test B isa GPUSparseMatrixCSC{ET,Int32} + @test SparseMatrixCSC(B) ≈ A + # explicit zeros are kept, as in SparseArrays + @test nnz(sparse(AT([1, 2]), AT([1, 2]), AT(zeros(ET, 2)), 2, 2)) == 2 + # nothing to assemble + @test nnz(sparse(AT(Int[]), AT(Int[]), AT(ET[]), 3, 4)) == 0 + @test_throws ArgumentError sparse(dI, dJ, dV, m - 1, n) + @test_throws ArgumentError sparse(dI, AT(J[1:end-1]), dV, m, n) + + # vectors + x = sparsevec(I, V, m) + dx = sparsevec(dI, dV, m) + @test dx isa GPUSparseVector{ET,Int} + check_structure(dx) + @test SparseVector(dx) ≈ x + @test same_sparse(sparsevec(dI, dV, m, noncommutative), sparsevec(I, V, m, noncommutative)) + @test length(sparsevec(dI, dV)) == length(sparsevec(I, V)) + @test same_sparse(sparsevec(dI, dV, noncommutative), sparsevec(I, V, noncommutative)) + + # diagonals + v = rand(ET, 5) + w = rand(ET, 4) + @test same_sparse(spdiagm(0 => AT(v), -1 => AT(w)), spdiagm(0 => v, -1 => w)) + @test same_sparse(spdiagm(7, 6, 1 => AT(w)), spdiagm(7, 6, 1 => w)) + @test same_sparse(spdiagm(AT(v)), spdiagm(v)) + @test_throws DimensionMismatch spdiagm(3, 3, 1 => AT(v)) + # without diagonals, SparseArrays applies + @test which(spdiagm, Tuple{Int,Int}).module === SparseArrays + end + + @testset "Bool" begin + I = rand(1:5, 40) + J = rand(1:6, 40) + V = rand(Bool, 40) + # repeated entries are combined with `|` by default + @test same_sparse(sparse(AT(I), AT(J), AT(V), 5, 6), sparse(I, J, V, 5, 6)) + @test same_sparse(sparse(AT(I), AT(J), true, 5, 6), sparse(I, J, true, 5, 6)) + @test same_sparse(sparsevec(AT(I), AT(V), 5), sparsevec(I, V, 5)) + end + end +end From 5c77cd8d2af2da2a7b579905af647acc358102e6 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Fri, 25 Sep 2026 17:52:22 +0200 Subject: [PATCH 04/11] Bump GPUArrays to 13.0.0 The sparse redesign removes the back-end hooks and changes the sparse types, which is breaking for back-ends. JLArrays requires the new version; its own breaking release (0.5, for KernelAbstractions 0.10) is not registered yet and covers the change of its sparse types. --- Project.toml | 2 +- lib/JLArrays/Project.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/Project.toml b/Project.toml index 8748f2195..8671e363b 100644 --- a/Project.toml +++ b/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"] diff --git a/lib/JLArrays/Project.toml b/lib/JLArrays/Project.toml index 961210ca2..82321f526 100644 --- a/lib/JLArrays/Project.toml +++ b/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" From baf58c6c90ee8c7c4a13b5aa3c21fd81233f027e Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Fri, 25 Sep 2026 18:18:13 +0200 Subject: [PATCH 05/11] =?UTF-8?q?Implement=20sparse=20=C3=97=20dense=20pro?= =?UTF-8?q?ducts?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Products hook into LinearAlgebra's storage-level mul!(C, tA, [tB,] A, B, α, β), which Julia 1.13 calls directly and which older versions reach through the forwarded generic_matvecmul!/generic_matmatmul!, like the dense products do. The generic implementations (generic_spmv!, generic_spmm!) gather over the CSR layout of op(A), or the CSC layout of op(B) for dense × sparse, so they are deterministic and work for every element type; Float16 accumulates in Float32. Transposed operands regroup the matrix first. Symmetric and Hermitian operands, which the old code treated as plain matrices, are computed as two passes over the used triangle and its transpose. LinearAlgebra allocates the destination of dense × sparse like the sparse operand, so * allocates a dense result for those. --- docs/src/interface.md | 20 +++ src/GPUArrays.jl | 1 + src/host/sparse/linalg.jl | 309 ++++++++++++++++++++++++++++++++++++++ test/testsuite/sparse.jl | 123 +++++++++++++++ 4 files changed, 453 insertions(+) create mode 100644 src/host/sparse/linalg.jl diff --git a/docs/src/interface.md b/docs/src/interface.md index 0b0e4b5da..9c3cbad61 100644 --- a/docs/src/interface.md +++ b/docs/src/interface.md @@ -117,10 +117,28 @@ A back-end will usually add the following, all optional: | 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!` | | 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. @@ -172,6 +190,8 @@ GPUArrays.check_structure GPUArrays.generic_regroup GPUArrays.generic_expand GPUArrays.generic_assemble +GPUArrays.generic_spmv! +GPUArrays.generic_spmm! ``` ## Caching Allocator diff --git a/src/GPUArrays.jl b/src/GPUArrays.jl index 1e68843ac..f58bf2150 100644 --- a/src/GPUArrays.jl +++ b/src/GPUArrays.jl @@ -49,6 +49,7 @@ include("host/sparse/show.jl") include("host/sparse/indexing.jl") include("host/sparse/broadcast.jl") include("host/sparse/reductions.jl") +include("host/sparse/linalg.jl") include("host/alloc_cache.jl") diff --git a/src/host/sparse/linalg.jl b/src/host/sparse/linalg.jl new file mode 100644 index 000000000..6c7e29bc8 --- /dev/null +++ b/src/host/sparse/linalg.jl @@ -0,0 +1,309 @@ +# products of sparse and dense arrays +# +# LinearAlgebra reduces products with transposed, adjoint, Symmetric and Hermitian operands +# to its storage-level `mul!(C, tA, [tB,] A, B, α, β)`, where a character describes each +# operand: 'N', 'T', 'C', and 'S'/'s' or 'H'/'h' for the upper or lower triangle of a +# Symmetric or Hermitian matrix. GPUArrays implements those methods for sparse operands +# (below), and back-ends override them for their vendor libraries, falling back to the +# generic implementations (`generic_spmv!`, `generic_spmm!`) where those do not apply. +# +# Every product is a gather over a compressed layout of the sparse operand, deterministic +# and valid for every element type. Transposed operands (Aᵀx for CSR, Ax for CSC) regroup +# the matrix first; an atomic scatter driven directly by the stored layout would avoid that. + +# the element type to accumulate products in: Float16 loses too much precision +product_accumulator(::Type{T}) where {T} = T +product_accumulator(::Type{Float16}) = Float32 +product_accumulator(::Type{ComplexF16}) = ComplexF32 +function product_accumulator(Ta::Type, Tb::Type) + T = typeof(zero(Ta) * zero(Tb) + zero(Ta) * zero(Tb)) + return product_accumulator(T) +end + +## the operands + +# The CSR buffers (pointers, indices, values) of `op(A)` for `t` in 'N', 'T' and 'C', with +# the conjugation of 'C' left to the kernel. Buffers of `A` itself are borrowed, others are +# temporary. +function csr_operand(A::GPUSparseMatrixCSR, t::AbstractChar) + B = t == 'N' ? A : transpose_view(generic_regroup(A)) + return B.rowPtr, B.colVal, B.nzVal +end +function csr_operand(A::GPUSparseMatrixCSC, t::AbstractChar) + B = t == 'N' ? generic_regroup(A) : transpose_view(A) + return B.rowPtr, B.colVal, B.nzVal +end +function csr_operand(A::GPUSparseMatrixCOO, t::AbstractChar) + t == 'N' || return csr_operand(generic_regroup(generic_compress(A)), t) + # a COO matrix is a CSR matrix with explicit row indices + return compress(A.rowInd, size(A, 1)), A.colInd, A.nzVal +end + +# The CSC buffers of `op(A)`, which are the CSR buffers of its transpose. The conjugation +# of 'C' is again left to the kernel. +csc_operand(A::GPUSparseMatrix, t::AbstractChar) = csr_operand(A, t == 'N' ? 'T' : 'N') + +# LinearAlgebra's characters for Symmetric and Hermitian operands combine two terms: the +# triangle of `A` that the wrapper uses, and the transpose (or adjoint) of that triangle +# without the diagonal. Each term is an operand character for the kernels, and a selection +# of the entries of `op(A)` by their position. +function operand_terms(t::AbstractChar) + t in ('N', 'T', 'C') && return ((t, all_entries),) + t == 'S' && return (('N', upper_entries), ('T', strict_lower_entries)) + t == 's' && return (('N', lower_entries), ('T', strict_upper_entries)) + t == 'H' && return (('N', upper_hermitian_entries), ('C', strict_lower_entries)) + t == 'h' && return (('N', lower_hermitian_entries), ('C', strict_upper_entries)) + throw(ArgumentError("unsupported operand character '$t'")) +end + +# the contribution of a stored entry `v` at `(i, j)`, as selected by a term +all_entries(i, j, v) = v +upper_entries(i, j, v) = j >= i ? v : zero(v) +lower_entries(i, j, v) = j <= i ? v : zero(v) +strict_upper_entries(i, j, v) = j > i ? v : zero(v) +strict_lower_entries(i, j, v) = j < i ? v : zero(v) +# the diagonal of a Hermitian matrix is real +upper_hermitian_entries(i, j, v) = j > i ? v : j == i ? oftype(v, real(v)) : zero(v) +lower_hermitian_entries(i, j, v) = j < i ? v : j == i ? oftype(v, real(v)) : zero(v) + +entry_selected(::typeof(all_entries), i, j) = true +entry_selected(::Union{typeof(upper_entries),typeof(upper_hermitian_entries)}, i, j) = j >= i +entry_selected(::Union{typeof(lower_entries),typeof(lower_hermitian_entries)}, i, j) = j <= i +entry_selected(::typeof(strict_upper_entries), i, j) = j > i +entry_selected(::typeof(strict_lower_entries), i, j) = j < i + +value_op(t::AbstractChar) = t == 'C' ? conj : identity + + +## kernels + +## COV_EXCL_START +# y = α op(A) x + β y, one thread per row of op(A), in CSR layout +@kernel function csr_matvec_kernel(y, ptr, ind, val, x, α, β, entry, valop, ::Type{T}) where {T} + i = @index(Global, Linear) + if i <= length(y) + acc = zero(T) + for k in @inbounds(ptr[i]):@inbounds(ptr[i+1] - one(eltype(ptr))) + j = @inbounds ind[k] + entry_selected(entry, i, j) || continue + v = entry(i, j, valop(@inbounds val[k])) + acc += T(v) * T(@inbounds x[j]) + end + # `y` is write-only when β is zero, so that it may hold NaN + @inbounds y[i] = iszero(β) ? α * acc : α * acc + β * y[i] + end +end + +# C = α op(A) op(B) + β C with a sparse op(A) in CSR layout, one thread per element of C +@kernel function csr_matmat_kernel(C, ptr, ind, val, B, α, β, entry, valop, tB::Char, + ::Type{T}) where {T} + i, j = @index(Global, NTuple) + if i <= size(C, 1) && j <= size(C, 2) + acc = zero(T) + for k in @inbounds(ptr[i]):@inbounds(ptr[i+1] - one(eltype(ptr))) + l = @inbounds ind[k] + entry_selected(entry, i, l) || continue + v = entry(i, l, valop(@inbounds val[k])) + b = tB == 'N' ? @inbounds(B[l, j]) : + tB == 'T' ? @inbounds(B[j, l]) : conj(@inbounds(B[j, l])) + acc += T(v) * T(b) + end + @inbounds C[i, j] = iszero(β) ? α * acc : α * acc + β * C[i, j] + end +end + +# C = α op(A) op(B) + β C with a sparse op(B) in CSC layout, one thread per element of C. +# Consecutive threads work down a column of C, reading a column of op(A) if it is not +# transposed, so that the dense accesses coalesce. +@kernel function csc_matmat_kernel(C, A, ptr, ind, val, α, β, entry, valop, tA::Char, + ::Type{T}) where {T} + i, j = @index(Global, NTuple) + if i <= size(C, 1) && j <= size(C, 2) + acc = zero(T) + for k in @inbounds(ptr[j]):@inbounds(ptr[j+1] - one(eltype(ptr))) + l = @inbounds ind[k] + entry_selected(entry, l, j) || continue + v = entry(l, j, valop(@inbounds val[k])) + a = tA == 'N' ? @inbounds(A[i, l]) : + tA == 'T' ? @inbounds(A[l, i]) : conj(@inbounds(A[l, i])) + acc += T(a) * T(v) + end + @inbounds C[i, j] = iszero(β) ? α * acc : α * acc + β * C[i, j] + end +end +## COV_EXCL_STOP + +## generic implementations + +# the dimensions of the result and of the operands of a product, as matrices +function check_product_dims(C::Dims, A::Dims, B::Dims) + mat(sz) = (sz[1], get(sz, 2, 1)) + C, A, B = mat(C), mat(A), mat(B) + A[2] == B[1] || + throw(DimensionMismatch("A has dimensions $A, B has dimensions $B")) + C == (A[1], B[2]) || + throw(DimensionMismatch("C has dimensions $C, needs $((A[1], B[2]))")) + return +end +op_size(A, t::AbstractChar) = t in ('T', 'C') ? reverse(size(A)) : size(A) + +# C = β C, treating C as write-only when β is zero +scale!(C, β) = iszero(β) ? fill!(C, zero(eltype(C))) : (C .*= β; C) + +## COV_EXCL_START +@kernel function materialize_symmetric_kernel(out, B, upper::Bool, hermitian::Bool) + i, j = @index(Global, NTuple) + if i <= size(out, 1) && j <= size(out, 2) + # the triangle of `B` that holds the values + stored = upper ? i <= j : i >= j + v = stored ? @inbounds(B[i, j]) : @inbounds(B[j, i]) + if hermitian + v = i == j ? oftype(v, real(v)) : stored ? v : conj(v) + end + @inbounds out[i, j] = v + end +end +## COV_EXCL_STOP + +# A dense operand wrapped in Symmetric or Hermitian is materialized; the kernels only index +# dense operands as they are, transposed, or adjoint. +function dense_operand(B, t::AbstractChar) + t in ('N', 'T', 'C') && return B, t + LinearAlgebra.checksquare(B) + out = similar(B, eltype(B), size(B)) + isempty(out) || + materialize_symmetric_kernel(get_backend(out))(out, B, t in ('S', 'H'), t in ('H', 'h'); + ndrange=size(out)) + return out, 'N' +end + +""" + GPUArrays.generic_spmv!(y, tA, A, x, α, β) + +Compute `y = α op(A) x + β y` for a sparse matrix `A` and dense vectors `x` and `y`, where +`op` is given by LinearAlgebra's character `tA` ('N', 'T', 'C', or 'S'/'s'/'H'/'h' for the +upper or lower triangle of a Symmetric or Hermitian `A`). This is the generic +implementation behind `mul!(y, tA, A, x, α, β)`, which back-ends can call when a vendor +routine does not apply. +""" +function generic_spmv!(y::AnyGPUVector, tA::AbstractChar, A::GPUSparseMatrix, + x::AnyGPUVector, α::Number, β::Number) + tA in ('S', 's', 'H', 'h') && LinearAlgebra.checksquare(A) + check_product_dims(size(y), op_size(A, tA), size(x)) + isempty(y) && return y + (isempty(x) || nnz(A) == 0) && return scale!(y, β) + T = product_accumulator(eltype(A), eltype(x)) + for (n, (t, entry)) in enumerate(operand_terms(tA)) + ptr, ind, val = csr_operand(A, t) + csr_matvec_kernel(get_backend(y))(y, ptr, ind, val, x, α, n == 1 ? β : true, + entry, value_op(t), T; ndrange=length(y)) + end + return y +end + +""" + GPUArrays.generic_spmm!(C, tA, tB, A, B, α, β) + +Compute `C = α op(A) op(B) + β C` where one of `A` and `B` is a sparse matrix and the other +one, like `C`, is dense; `tA` and `tB` are LinearAlgebra's characters as for +[`generic_spmv!`](@ref). This is the generic implementation behind +`mul!(C, tA, tB, A, B, α, β)`. +""" +function generic_spmm!(C::AnyGPUMatrix, tA::AbstractChar, tB::AbstractChar, + A::GPUSparseMatrix, B::AnyGPUMatrix, α::Number, β::Number) + tA in ('S', 's', 'H', 'h') && LinearAlgebra.checksquare(A) + B, tB = dense_operand(B, tB) + check_product_dims(size(C), op_size(A, tA), op_size(B, tB)) + isempty(C) && return C + (size(B, tB == 'N' ? 1 : 2) == 0 || nnz(A) == 0) && return scale!(C, β) + T = product_accumulator(eltype(A), eltype(B)) + for (n, (t, entry)) in enumerate(operand_terms(tA)) + ptr, ind, val = csr_operand(A, t) + csr_matmat_kernel(get_backend(C))(C, ptr, ind, val, B, α, n == 1 ? β : true, + entry, value_op(t), Char(tB), T; ndrange=size(C)) + end + return C +end +function generic_spmm!(C::AnyGPUMatrix, tA::AbstractChar, tB::AbstractChar, + A::AnyGPUMatrix, B::GPUSparseMatrix, α::Number, β::Number) + tB in ('S', 's', 'H', 'h') && LinearAlgebra.checksquare(B) + A, tA = dense_operand(A, tA) + check_product_dims(size(C), op_size(A, tA), op_size(B, tB)) + isempty(C) && return C + (size(A, tA == 'N' ? 2 : 1) == 0 || nnz(B) == 0) && return scale!(C, β) + T = product_accumulator(eltype(A), eltype(B)) + for (n, (t, entry)) in enumerate(operand_terms(tB)) + ptr, ind, val = csc_operand(B, t) + csc_matmat_kernel(get_backend(C))(C, A, ptr, ind, val, α, n == 1 ? β : true, + entry, value_op(t), Char(tA), T; ndrange=size(C)) + end + return C +end + + +## LinearAlgebra integration + +LinearAlgebra.mul!(C::AnyGPUVector, tA::AbstractChar, A::GPUSparseMatrix, B::AnyGPUVector, + α::Number, β::Number) = + generic_spmv!(C, tA, A, B, α, β) +LinearAlgebra.mul!(C::AnyGPUMatrix, tA::AbstractChar, tB::AbstractChar, A::GPUSparseMatrix, + B::AnyGPUMatrix, α::Number, β::Number) = + generic_spmm!(C, tA, tB, A, B, α, β) +LinearAlgebra.mul!(C::AnyGPUMatrix, tA::AbstractChar, tB::AbstractChar, A::AnyGPUMatrix, + B::GPUSparseMatrix, α::Number, β::Number) = + generic_spmm!(C, tA, tB, A, B, α, β) + +# Julia < 1.13 dispatches on the non-public `generic_matvecmul!` and `generic_matmatmul!` +# instead, as for dense arrays (see `linalg.jl`) +@static if VERSION < v"1.13.0-rc4" + for (TA, TB) in ((:GPUSparseMatrix, :AnyGPUMatrix), (:AnyGPUMatrix, :GPUSparseMatrix)) + @eval begin + LinearAlgebra.generic_matmatmul!(C::AnyGPUMatrix, tA::AbstractChar, tB::AbstractChar, + A::$TA, B::$TB, a::Number, b::Number) = + LinearAlgebra.mul!(C, tA, tB, A, B, a, b) + LinearAlgebra.generic_matmatmul!(C::AnyGPUMatrix, tA::AbstractChar, tB::AbstractChar, + A::$TA, B::$TB, _add::MulAddMul=MulAddMul()) = + LinearAlgebra.mul!(C, tA, tB, A, B, _add.alpha, _add.beta) + end + end + LinearAlgebra.generic_matvecmul!(C::AnyGPUVector, tA::AbstractChar, A::GPUSparseMatrix, + B::AnyGPUVector, a::Number, b::Number) = + LinearAlgebra.mul!(C, tA, A, B, a, b) + LinearAlgebra.generic_matvecmul!(C::AnyGPUVector, tA::AbstractChar, A::GPUSparseMatrix, + B::AnyGPUVector, _add::MulAddMul=MulAddMul()) = + LinearAlgebra.mul!(C, tA, A, B, _add.alpha, _add.beta) +end + +# LinearAlgebra allocates the result of a product like its second operand, which gives a +# sparse destination for dense × sparse +const GPUSparseMatrixOperand = Union{GPUSparseMatrix, Transpose{<:Any,<:GPUSparseMatrix}, + Adjoint{<:Any,<:GPUSparseMatrix}, + Symmetric{<:Any,<:GPUSparseMatrix}, + Hermitian{<:Any,<:GPUSparseMatrix}} +function Base.:(*)(A::AnyGPUMatrix, B::GPUSparseMatrixOperand) + T = Base.promote_op(LinearAlgebra.matprod, eltype(A), eltype(B)) + return mul!(similar(A, T, (size(A, 1), size(B, 2))), A, B) +end + +# A sparse vector operand is densified, as CUDA.jl does. (A product with a sparse vector +# could skip its implicit zeros, but that only pays for very sparse vectors.) +densify(x::GPUSparseVector) = copyto!(similar(x.nzVal, eltype(x), length(x)), x) + +const GPUMatrixOperand = Union{GPUSparseMatrix, AnyGPUMatrix} +LinearAlgebra.mul!(C::AnyGPUVector, tA::AbstractChar, A::GPUMatrixOperand, B::GPUSparseVector, + α::Number, β::Number) = + LinearAlgebra.mul!(C, tA, A, densify(B), α, β) +@static if VERSION < v"1.13.0-rc4" + LinearAlgebra.generic_matvecmul!(C::AnyGPUVector, tA::AbstractChar, A::GPUMatrixOperand, + B::GPUSparseVector, a::Number, b::Number) = + LinearAlgebra.mul!(C, tA, A, B, a, b) + LinearAlgebra.generic_matvecmul!(C::AnyGPUVector, tA::AbstractChar, A::GPUMatrixOperand, + B::GPUSparseVector, _add::MulAddMul=MulAddMul()) = + LinearAlgebra.mul!(C, tA, A, B, _add.alpha, _add.beta) +end + +# like dense × sparse, LinearAlgebra would allocate a sparse destination +function Base.:(*)(A::Union{AnyGPUMatrix,GPUSparseMatrixOperand}, x::GPUSparseVector) + T = Base.promote_op(LinearAlgebra.matprod, eltype(A), eltype(x)) + return mul!(similar(x.nzVal, T, size(A, 1)), A, x) +end diff --git a/test/testsuite/sparse.jl b/test/testsuite/sparse.jl index 50e8bd237..550a01268 100644 --- a/test/testsuite/sparse.jl +++ b/test/testsuite/sparse.jl @@ -18,6 +18,7 @@ using GPUArrays: GPUSparseMatrixCSR, GPUSparseMatrixCSC, GPUSparseMatrixCOO, GPU sparse_conversions(AT, eltypes) sparse_dense_conversions(AT, eltypes) sparse_assembly(AT, eltypes) + sparse_products(AT, eltypes) broadcasting_vector(AT, eltypes) broadcasting_matrix(AT, eltypes) mapreduce_matrix(AT, eltypes) @@ -1009,3 +1010,125 @@ function sparse_assembly(AT, eltypes) end end end + +# the operand wrappers LinearAlgebra encodes as characters in the storage-level `mul!` +function sparse_operands(A::AbstractMatrix) + ops = Any[identity, transpose, adjoint] + if size(A, 1) == size(A, 2) + append!(ops, [x -> Symmetric(x, :U), x -> Symmetric(x, :L)]) + eltype(A) <: Complex && append!(ops, [x -> Hermitian(x, :U), x -> Hermitian(x, :L)]) + end + return ops +end + +function sparse_products(AT, eltypes) + @testset "products" begin + @testset "$S{$ET}" for S in sparse_matrix_formats, ET in eltypes + # every operand character for one real and one complex element type, a few for + # the others + all_ops = ET == first(eltypes) || ET == first(filter(iscomplextype, eltypes)) + A = sprand_awkward(ET, 7, 7; Ti=Int32) + R = sprand_awkward(ET, 7, 5; Ti=Int32) # rectangular + dA = gpu_sparse(AT, S, A) + dR = gpu_sparse(AT, S, R) + α, β = ET(2), ET(3) + + # sparse × dense vector (with dense references: SparseArrays lacks some of + # these products for some element types) + for (M, dM) in ((A, dA), (R, dR)), op in (all_ops ? sparse_operands(M) : [identity, adjoint]) + x = rand(ET, size(op(M), 2)) + y = rand(ET, size(op(M), 1)) + @test Array(op(dM) * AT(x)) ≈ op(Matrix(M)) * x + dy = AT(copy(y)) + @test mul!(dy, op(dM), AT(x), α, β) === dy + @test Array(dy) ≈ α * (op(Matrix(M)) * x) + β * y + end + + # sparse × dense matrix, and dense × sparse + dense_ops = all_ops ? [identity, transpose, adjoint] : [identity] + for (M, dM) in ((A, dA), (R, dR)), op in (all_ops ? sparse_operands(M) : [identity, transpose]), dop in dense_ops + B = rand(ET, 4, size(op(M), 2)) + B = dop === identity ? permutedims(B) : B # op(M) * dop(B) + ref = op(Matrix(M)) * dop(B) + @test Array(op(dM) * dop(AT(B))) ≈ ref + C = rand(ET, size(ref)) + dC = AT(copy(C)) + mul!(dC, op(dM), dop(AT(B)), α, β) + @test Array(dC) ≈ α * ref + β * C + + D = rand(ET, size(op(M), 1), 4) + D = dop === identity ? permutedims(D) : D # dop(D) * op(M) + ref = dop(D) * op(Matrix(M)) + @test Array(dop(AT(D)) * op(dM)) ≈ ref + @test dop(AT(D)) * op(dM) isa AT + C = rand(ET, size(ref)) + dC = AT(copy(C)) + mul!(dC, dop(AT(D)), op(dM), α, β) + @test Array(dC) ≈ α * ref + β * C + end + if all_ops + # dense Symmetric and Hermitian operands + B = rand(ET, 7, 7) + for wrap in (Symmetric, Hermitian), uplo in (:U, :L) + @test Array(dA * wrap(AT(B), uplo)) ≈ A * wrap(B, uplo) + @test Array(wrap(AT(B), uplo) * dA) ≈ wrap(B, uplo) * A + end + end + + # with β = 0 the destination is write-only + if ET <: AbstractFloat + dy = AT(fill(ET(NaN), 7)) + mul!(dy, dA, AT(ones(ET, 7)), true, false) + @test Array(dy) ≈ A * ones(ET, 7) + dC = AT(fill(ET(NaN), 7, 3)) + mul!(dC, dA, AT(ones(ET, 7, 3)), true, false) + @test Array(dC) ≈ A * ones(ET, 7, 3) + dC = AT(fill(ET(NaN), 3, 7)) + mul!(dC, AT(ones(ET, 3, 7)), dA, true, false) + @test Array(dC) ≈ ones(ET, 3, 7) * A + end + + # sparse vector operands are densified + x = sprand_nozeros(ET, 7, 0.5) + dx = gpu_sparse(AT, x) + @test dA * dx isa AT{ET,1} + @test Array(dA * dx) ≈ A * Vector(x) + @test Array(AT(Matrix(A)) * dx) ≈ Matrix(A) * Vector(x) + dy = AT(ones(ET, 7)) + mul!(dy, dA, dx, α, β) + @test Array(dy) ≈ α * (A * Vector(x)) + β * ones(ET, 7) + + # empty operands + Z = spzeros(ET, 3, 0) + @test Array(gpu_sparse(AT, S, Z) * AT(zeros(ET, 0))) == zeros(ET, 3) + @test Array(gpu_sparse(AT, S, spzeros(ET, 0, 3)) * AT(ones(ET, 3))) == zeros(ET, 0) + @test Array(gpu_sparse(AT, S, spzeros(ET, 3, 4)) * AT(ones(ET, 4, 2))) == zeros(ET, 3, 2) + @test Array(AT(ones(ET, 2, 3)) * gpu_sparse(AT, S, spzeros(ET, 3, 4))) == zeros(ET, 2, 4) + @test_throws DimensionMismatch dR * AT(ones(ET, 7)) + end + + @testset "$T" for T in (Int, Bool) + A = sprand(T, 6, 5, 0.5) + x = rand(T, 5) + B = rand(T, 5, 3) + D = rand(T, 3, 6) + for S in sparse_matrix_formats + dA = gpu_sparse(AT, S, A) + @test Array(dA * AT(x)) == A * x + @test Array(transpose(dA) * AT(rand(T, 6))) isa Vector + @test Array(dA * AT(B)) == A * B + @test Array(AT(D) * dA) == D * A + end + end + + # Float16 products accumulate in Float32 + if Float16 in eltypes + A = sprand(Float16, 50, 200, 0.5) + x = rand(Float16, 200) + ref = Float16.(Float32.(A) * Float32.(x)) + for S in sparse_matrix_formats + @test Array(gpu_sparse(AT, S, A) * AT(x)) == ref + end + end + end +end From d964913e1a3fd48892efa93b22de67ff349efb60 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Fri, 25 Sep 2026 18:20:52 +0200 Subject: [PATCH 06/11] Broadcast over COO and mixed sparse formats The broadcast kernels iterate one compressed format, so arguments of other formats are converted to that of the first sparse matrix, COO matrices being read as CSR (which has the same order of entries) and a COO result converted back. map over sparse arrays is a broadcast, as in SparseArrays. SparseArrays implements +, - and map/broadcast of the basic operators for any AbstractSparseVector through scalar loops; GPUSparseVector gets its own methods. --- src/host/sparse/broadcast.jl | 64 ++++++++++++++++++++++++++++++++---- test/testsuite/sparse.jl | 51 ++++++++++++++++++++++++++++ 2 files changed, 108 insertions(+), 7 deletions(-) diff --git a/src/host/sparse/broadcast.jl b/src/host/sparse/broadcast.jl index c765d17f8..36a6d501c 100644 --- a/src/host/sparse/broadcast.jl +++ b/src/host/sparse/broadcast.jl @@ -455,19 +455,48 @@ end end ## COV_EXCL_STOP +# the CSR layout of a COO matrix, borrowing its buffers +csr_view(A::GPUSparseMatrixCOO) = + GPUSparseMatrixCSR(compress(A.rowInd, size(A, 1)), A.colInd, A.nzVal, size(A)) + function Broadcast.copy(bc::Broadcasted{<:Union{GPUSparseVecStyle,GPUSparseMatStyle}}) - # find the sparse inputs bc = Broadcast.flatten(bc) + sparse_inputs = filter(arg -> arg isa AbstractGPUSparseArray, bc.args) + all(arg -> arg isa GPUSparseArray, sparse_inputs) || + error("broadcast is only implemented for the sparse formats of GPUArrays") + (all(arg -> arg isa GPUSparseVector, sparse_inputs) || + all(arg -> arg isa GPUSparseMatrix, sparse_inputs)) || + error("broadcast with both sparse vectors and sparse matrices is not supported") + all(arg -> axes(arg) == axes(bc), sparse_inputs) || + throw(ArgumentError("broadcast expansion of sparse operands is not supported")) + # capturescalars evaluates wrapped scalars on the host to test zero preservation. + # GPU scalar arrays need device-side evaluation, which is not implemented here. + any(arg -> arg isa AnyGPUArray{<:Any,0}, bc.args) && + throw(ArgumentError("sparse broadcast with zero-dimensional GPU arrays is not supported")) + + # The kernels iterate a single compressed format: that of the first sparse matrix, with + # COO matrices read as CSR (which has the same order of entries). Other formats are + # converted, and a COO result is converted back. + lead = first(sparse_inputs) + if lead isa GPUSparseMatrix && !all(arg -> typeof(arg).name === typeof(lead).name && + !(arg isa GPUSparseMatrixCOO), sparse_inputs) + S = lead isa GPUSparseMatrixCSC ? GPUSparseMatrixCSC : GPUSparseMatrixCSR + args = map(bc.args) do arg + arg isa GPUSparseMatrixCOO && S === GPUSparseMatrixCSR ? csr_view(arg) : + arg isa GPUSparseMatrix ? convert(S, arg) : arg + end + out = copy(Broadcasted(bc.f, args, bc.axes)) + if lead isa GPUSparseMatrixCOO && out isa GPUSparseMatrixCSR + out = GPUSparseMatrixCOO(expand_ptr(out.rowPtr, nnz(out)), out.colVal, out.nzVal, + size(out)) + end + return out + end + sparse_args = findall(bc.args) do arg arg isa GPUSparseArray end - sparse_types = unique(map(i->nameof(typeof(bc.args[i])), sparse_args)) - if length(sparse_types) > 1 - error("broadcast with multiple types of sparse arrays ($(join(sparse_types, ", "))) is not supported") - end sparse_typ = typeof(bc.args[first(sparse_args)]) - sparse_typ <: Union{GPUSparseMatrixCSR,GPUSparseMatrixCSC,GPUSparseVector} || - error("broadcast with sparse arrays is currently only implemented for vectors and CSR and CSC matrices") Ti = reduce(promote_type, map(i -> indtype(bc.args[i]), sparse_args)) # determine the output type @@ -612,3 +641,24 @@ function Broadcast.copy(bc::Broadcasted{<:Union{GPUSparseVecStyle,GPUSparseMatSt ndrange == 0 || kernel(args...; ndrange) return output end + +# like SparseArrays, `map` over sparse arrays of the same size is a broadcast +function sparse_map(f, A, Bs...) + all(B -> size(B) == size(A), Bs) || + throw(DimensionMismatch("map requires arrays of the same size")) + return broadcast(f, A, Bs...) +end +Base.map(f, A::GPUSparseArray, Bs::Union{GPUSparseArray,AnyGPUArray}...) = sparse_map(f, A, Bs...) +Base.map(f, A::AnyGPUArray, B::GPUSparseArray, Cs::Union{GPUSparseArray,AnyGPUArray}...) = + sparse_map(f, A, B, Cs...) + +# SparseArrays implements these for any sparse vector, through scalar indexing +for f in (:+, :-, :*, :min, :max) + @eval begin + Base.map(::typeof($f), x::GPUSparseVector, y::GPUSparseVector) = sparse_map($f, x, y) + Base.broadcast(::typeof($f), x::GPUSparseVector, y::GPUSparseVector) = + Broadcast.materialize(Broadcast.broadcasted($f, x, y)) + end +end +Base.:(+)(x::GPUSparseVector, y::GPUSparseVector) = sparse_map(+, x, y) +Base.:(-)(x::GPUSparseVector, y::GPUSparseVector) = sparse_map(-, x, y) diff --git a/test/testsuite/sparse.jl b/test/testsuite/sparse.jl index 550a01268..74a3228b0 100644 --- a/test/testsuite/sparse.jl +++ b/test/testsuite/sparse.jl @@ -21,6 +21,7 @@ using GPUArrays: GPUSparseMatrixCSR, GPUSparseMatrixCSC, GPUSparseMatrixCOO, GPU sparse_products(AT, eltypes) broadcasting_vector(AT, eltypes) broadcasting_matrix(AT, eltypes) + broadcasting_mixed(AT, eltypes) mapreduce_matrix(AT, eltypes) sparse_linalg(AT, eltypes) iszero_sparse(AT, eltypes) @@ -646,6 +647,56 @@ function broadcasting_matrix(AT, eltypes) end end +function broadcasting_mixed(AT, eltypes) + @testset "mixed-format broadcasting" begin + @testset "$ET" for ET in eltypes + A = sprand_awkward(ET, 6, 7) + B = sprand_nozeros(ET, 6, 7, 0.4) + # the result has the format of the first sparse argument + for S in sparse_matrix_formats, S′ in sparse_matrix_formats + dA = gpu_sparse(AT, S, A) + dB = gpu_sparse(AT, S′, B) + dC = dA .* dB + @test dC isa S{ET} + check_structure(dC) + @test SparseMatrixCSC(dC) ≈ A .* B + dC = dA .+ dB + @test dC isa S{ET} + @test SparseMatrixCSC(dC) ≈ A .+ B + @test SparseMatrixCSC(dA + dB) ≈ A + B + @test SparseMatrixCSC(dA - dB) ≈ A - B + @test Array(dA .+ dB .+ ET(1)) ≈ Array(A .+ B .+ ET(1)) + end + for S in sparse_matrix_formats + dA = gpu_sparse(AT, S, A) + @test SparseMatrixCSC(ET(2) * dA) == ET(2) * A + @test SparseMatrixCSC(-dA) == -A + three = ET(3) + @test SparseMatrixCSC(map(v -> v * three, dA)) == map(v -> v * three, A) + @test Array(map(+, AT(Matrix(B)), dA)) ≈ Matrix(B) + A + @test SparseMatrixCSC(conj(dA)) == conj(A) + @test Array(dA .* AT(Matrix(B))) ≈ Array(A .* Matrix(B)) + end + @test_throws ErrorException gpu_sparse(AT, GPUSparseMatrixCSR, A) .* gpu_sparse(AT, sprand_nozeros(ET, 6, 0.5)) + + # SparseArrays' own methods for sparse vectors are replaced + x = sprand_nozeros(ET, 30, 0.3) + y = sprand_nozeros(ET, 30, 0.3) + dx, dy = gpu_sparse(AT, x), gpu_sparse(AT, y) + @test SparseVector(dx + dy) ≈ x + y + @test SparseVector(dx - dy) ≈ x - y + for f in (+, -, *) + @test SparseVector(map(f, dx, dy)) ≈ map(f, x, y) + @test SparseVector(broadcast(f, dx, dy)) ≈ broadcast(f, x, y) + end + if ET <: Real + @test SparseVector(map(max, dx, dy)) == map(max, x, y) + end + @test_throws DimensionMismatch dx + gpu_sparse(AT, sprand_nozeros(ET, 31, 0.3)) + end + end +end + function mapreduce_matrix(AT, eltypes) @testset "SparseMatrix mapreduce" begin @testset "$S{$ET}" for S in sparse_matrix_formats, ET in eltypes From 090e3b7783aed3336c19ecaef6f1802524d60a70 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Fri, 25 Sep 2026 18:20:52 +0200 Subject: [PATCH 07/11] Implement structural operations on sparse matrices fkeep!, dropzeros[!], droptol!, triu and tril filter the stored entries on the device and recompress; istriu, istril and isdiag check the nonzero values, and issymmetric/ishermitian compare with the transpose, which also accounts for stored zeros on one side. diag looks up each element with a binary search in its slice, reshape and vec reuse the column-major order of CSC, which reshaping preserves, and kron fills each column of the result from a pair of columns without sorting. The istriu/istril workaround for triangular wrappers of sparse matrices (JuliaLang/julia#55547) is carried over from the old code. --- src/GPUArrays.jl | 1 + src/host/sparse/structure.jl | 325 +++++++++++++++++++++++++++++++++++ test/testsuite/sparse.jl | 129 ++++++++++++++ 3 files changed, 455 insertions(+) create mode 100644 src/host/sparse/structure.jl diff --git a/src/GPUArrays.jl b/src/GPUArrays.jl index f58bf2150..36376c05b 100644 --- a/src/GPUArrays.jl +++ b/src/GPUArrays.jl @@ -50,6 +50,7 @@ include("host/sparse/indexing.jl") include("host/sparse/broadcast.jl") include("host/sparse/reductions.jl") include("host/sparse/linalg.jl") +include("host/sparse/structure.jl") include("host/alloc_cache.jl") diff --git a/src/host/sparse/structure.jl b/src/host/sparse/structure.jl new file mode 100644 index 000000000..7397242c7 --- /dev/null +++ b/src/host/sparse/structure.jl @@ -0,0 +1,325 @@ +# structural operations + +# the row and column of every stored entry, borrowing the index buffers of `A` +entry_coordinates(A::GPUSparseMatrixCSR) = expand_ptr(A.rowPtr, nnz(A)), A.colVal +entry_coordinates(A::GPUSparseMatrixCSC) = A.rowVal, expand_ptr(A.colPtr, nnz(A)) +entry_coordinates(A::GPUSparseMatrixCOO) = A.rowInd, A.colInd + +## COV_EXCL_START +@kernel function keep_matrix_entries_kernel(keep, f, rows, cols, vals) + k = @index(Global, Linear) + if k <= length(keep) + @inbounds keep[k] = f(rows[k], cols[k], vals[k]) + end +end + +@kernel function keep_vector_entries_kernel(keep, f, inds, vals) + k = @index(Global, Linear) + if k <= length(keep) + @inbounds keep[k] = f(inds[k], vals[k]) + end +end +## COV_EXCL_STOP + +""" + fkeep!(f, A::GPUSparseMatrix) + fkeep!(f, x::GPUSparseVector) + +Keep the stored entries of `A` for which `f(i, j, v)` (or `f(i, v)` for a vector) is true, +on the device, as in SparseArrays. `f` runs in a kernel. The buffers of `A` are resized, so +`A` must not share them with another array. +""" +function SparseArrays.fkeep!(f::F, A::GPUSparseMatrix) where {F} + nnz(A) == 0 && return A + rows, cols = entry_coordinates(A) + keep = similar(A.nzVal, Bool) + keep_matrix_entries_kernel(get_backend(A))(keep, f, rows, cols, A.nzVal; ndrange=nnz(A)) + kept = findall(keep) + free_buffer!(keep) + new_rows, new_cols, new_vals = rows[kept], cols[kept], A.nzVal[kept] + free_buffer!(kept) + new_buffers = if A isa GPUSparseMatrixCSR + (compress(new_rows, size(A, 1)), new_cols, new_vals) + elseif A isa GPUSparseMatrixCSC + (compress(new_cols, size(A, 2)), new_rows, new_vals) + else + (new_rows, new_cols, new_vals) + end + replace_buffers!(sparse_buffers(A), new_buffers) + return A +end +function SparseArrays.fkeep!(f::F, x::GPUSparseVector) where {F} + nnz(x) == 0 && return x + keep = similar(x.nzVal, Bool) + keep_vector_entries_kernel(get_backend(x))(keep, f, x.nzInd, x.nzVal; ndrange=nnz(x)) + kept = findall(keep) + free_buffer!(keep) + new_inds, new_vals = x.nzInd[kept], x.nzVal[kept] + free_buffer!(kept) + replace_buffers!(sparse_buffers(x), (new_inds, new_vals)) + return x +end + +# the tolerance is captured by value, so that the predicate can run in a kernel +struct AboveTolerance{T} + tol::T +end +(p::AboveTolerance)(i, j, x) = abs(x) > p.tol +(p::AboveTolerance)(i, x) = abs(x) > p.tol + +SparseArrays.dropzeros!(A::GPUSparseMatrix) = fkeep!((i, j, x) -> !iszero(x), A) +SparseArrays.dropzeros!(x::GPUSparseVector) = fkeep!((i, x) -> !iszero(x), x) +SparseArrays.dropzeros(A::GPUSparseArray) = dropzeros!(copy(A)) +# the tolerance is compared in the precision of the values, so that a Float64 tolerance +# (like the default of `exp`) does not end up in kernels on back-ends without Float64 +droptol_tolerance(::Type{T}, tol) where {T<:Union{AbstractFloat,Complex{<:AbstractFloat}}} = + convert(real(T), tol) +droptol_tolerance(::Type, tol) = tol +SparseArrays.droptol!(A::GPUSparseArray, tol::Real) = + fkeep!(AboveTolerance(droptol_tolerance(eltype(A), tol)), A) + + +## triangular parts + +struct BandSelection + lo::Int + hi::Int +end +# (in `Int`, so that unsigned indices do not wrap) +(b::BandSelection)(i, j, x) = b.lo <= (j % Int) - (i % Int) <= b.hi + +const GPUSparseMatrixOrTransposed = + Union{GPUSparseMatrix, Transpose{<:Any,<:GPUSparseMatrix}, Adjoint{<:Any,<:GPUSparseMatrix}} + +# a copy of the entries of `A` whose diagonal index `j - i` lies within `lo:hi` +band(A::GPUSparseMatrix, lo, hi) = fkeep!(BandSelection(lo, hi), copy(A)) +band(A::Union{Transpose,Adjoint}, lo, hi) = fkeep!(BandSelection(lo, hi), copy(A)) + +LinearAlgebra.triu(A::GPUSparseMatrixOrTransposed, k::Integer=0) = band(A, k, typemax(Int)) +LinearAlgebra.tril(A::GPUSparseMatrixOrTransposed, k::Integer=0) = band(A, typemin(Int), k) + +## COV_EXCL_START +@kernel function outside_band_kernel(outside, rows, cols, vals, lo, hi) + k = @index(Global, Linear) + if k <= length(outside) + d = @inbounds (cols[k] % Int) - (rows[k] % Int) + @inbounds outside[k] = !iszero(vals[k]) && !(lo <= d <= hi) + end +end +## COV_EXCL_STOP + +# whether all nonzero values lie within the diagonals `lo:hi` +function within_band(A::GPUSparseMatrix, lo, hi) + nnz(A) == 0 && return true + rows, cols = entry_coordinates(A) + outside = similar(A.nzVal, Bool) + outside_band_kernel(get_backend(A))(outside, rows, cols, A.nzVal, lo, hi; ndrange=nnz(A)) + return !any(outside) +end +LinearAlgebra.istriu(A::GPUSparseMatrix, k::Integer=0) = within_band(A, k, typemax(Int)) +LinearAlgebra.istril(A::GPUSparseMatrix, k::Integer=0) = within_band(A, typemin(Int), k) +LinearAlgebra.isdiag(A::GPUSparseMatrix) = within_band(A, 0, 0) + +# The wrapper determines where the nonzeros can be; LinearAlgebra's generic methods would +# index the parent (and are broken on some versions, JuliaLang/julia#55547). The answers +# are structural, not based on the values: LinearAlgebra's `ldiv!` picks the triangle +# to solve with from `istriu`, and a diagonal `LowerTriangular` reported as upper +# triangular would be rewrapped and miss the back-ends' triangular solves. +const GPUSparseUpperOrUnitUpperTriangular = + LinearAlgebra.UpperOrUnitUpperTriangular{<:Any,<:GPUSparseMatrixOrTransposed} +const GPUSparseLowerOrUnitLowerTriangular = + LinearAlgebra.LowerOrUnitLowerTriangular{<:Any,<:GPUSparseMatrixOrTransposed} +LinearAlgebra.istriu(::GPUSparseUpperOrUnitUpperTriangular) = true +LinearAlgebra.istril(::GPUSparseUpperOrUnitUpperTriangular) = false +LinearAlgebra.istril(::GPUSparseLowerOrUnitLowerTriangular) = true +LinearAlgebra.istriu(::GPUSparseLowerOrUnitLowerTriangular) = false + +# `A` equals its transpose (or adjoint) when no element differs, which broadcasting over +# the union of both structures also checks for stored zeros on one side only +function LinearAlgebra.issymmetric(A::GPUSparseMatrix) + size(A, 1) == size(A, 2) || return false + return !any(A .!= copy(transpose(A))) +end +function LinearAlgebra.ishermitian(A::GPUSparseMatrix) + size(A, 1) == size(A, 2) || return false + return !any(A .!= copy(adjoint(A))) +end + + +## diagonals + +## COV_EXCL_START +# every thread looks up one element of the diagonal `k` in its slice +@kernel function diagonal_kernel(d, ptr, ind, val, k, csr::Bool) + t = @index(Global, Linear) + if t <= length(d) + # the diagonal element in row/column `major`, at column/row `minor` + i, j = k >= 0 ? (t, t + k) : (t - k, t) + major, minor = csr ? (i, j) : (j, i) + lo = @inbounds ptr[major] + hi = @inbounds ptr[major+1] - one(lo) + p = searchsortedfirst(ind, minor % eltype(ind), lo, hi, Base.Order.Forward) + @inbounds d[t] = (p <= hi && ind[p] == minor) ? val[p] : zero(eltype(d)) + end +end +## COV_EXCL_STOP + +function LinearAlgebra.diag(A::GPUSparseMatrix, k::Integer=0) + # (as an `Int`: the offset is negated for transposes, and passed to a kernel) + k = Int(k) + m, n = size(A) + # like SparseArrays, diagonals outside of the matrix are empty + len = max(0, k >= 0 ? min(m, n - k) : min(m + k, n)) + d = similar(A.nzVal, len) + len == 0 && return d + B = A isa GPUSparseMatrixCOO ? csr_view(A) : A + diagonal_kernel(get_backend(B))(d, compressed_ptr(B), compressed_ind(B), B.nzVal, k, + B isa GPUSparseMatrixCSR; ndrange=len) + return d +end + +LinearAlgebra.diag(A::Transpose{<:Any,<:GPUSparseMatrix}, k::Integer=0) = + diag(parent(A), -Int(k)) +LinearAlgebra.diag(A::Adjoint{<:Any,<:GPUSparseMatrix}, k::Integer=0) = + conj.(diag(parent(A), -Int(k))) + +# (LinearAlgebra's constructor would index the elements) +LinearAlgebra.Diagonal(A::GPUSparseMatrix) = Diagonal(diag(A)) + +function LinearAlgebra.tr(A::GPUSparseMatrix) + LinearAlgebra.checksquare(A) + # `+` rather than `sum`, which would widen small integers + return reduce(+, diag(A); init=zero(eltype(A))) +end + + +## reshaping + +## COV_EXCL_START +@kernel function reshape_kernel(rows, cols, m, m′) + k = @index(Global, Linear) + if k <= length(rows) + l = (Int64(cols[k]) - 1) * m + Int64(rows[k]) - 1 + @inbounds rows[k] = (l % m′ + 1) % eltype(rows) + @inbounds cols[k] = (l ÷ m′ + 1) % eltype(cols) + end +end + +@kernel function linear_index_kernel(inds, rows, cols, m) + k = @index(Global, Linear) + if k <= length(inds) + @inbounds inds[k] = ((Int64(cols[k]) - 1) * m + rows[k]) % eltype(inds) + end +end +## COV_EXCL_STOP + +# Column-major order of the entries (that of CSC) is the order of their linear indices, so +# it is preserved by any reshape: new coordinates, in the same order, form the CSC buffers +# of the result. +function Base.reshape(A::GPUSparseMatrix, dims::Dims{2}) + prod(dims) == length(A) || + throw(DimensionMismatch("new dimensions $dims must be consistent with array size $(length(A))")) + dims == size(A) && return copy(A) + C = A isa GPUSparseMatrixCSC ? copy(A) : GPUSparseMatrixCSC(A) + rows = C.rowVal + cols = expand_ptr(C.colPtr, nnz(C)) + check_sparse_dims(indtype(A), dims) + if nnz(C) > 0 + reshape_kernel(get_backend(C))(rows, cols, size(A, 1), dims[1]; ndrange=nnz(C)) + end + B = GPUSparseMatrixCSC(compress(cols, dims[2]), rows, C.nzVal, dims) + free_buffer!(cols) + return convert(matrix_format(A), B) +end +Base.reshape(A::GPUSparseMatrix, dims::Tuple{Union{Int,Colon},Union{Int,Colon}}) = + reshape(A, Base._reshape_uncolon(A, dims)) + +function Base.vec(A::GPUSparseMatrix) + C = A isa GPUSparseMatrixCSC ? A : GPUSparseMatrixCSC(A) + check_sparse_dims(indtype(A), (length(A),)) + inds = similar(C.rowVal, nnz(C)) + if nnz(C) > 0 + cols = expand_ptr(C.colPtr, nnz(C)) + linear_index_kernel(get_backend(C))(inds, C.rowVal, cols, size(C, 1); ndrange=nnz(C)) + free_buffer!(cols) + end + return GPUSparseVector(inds, C === A ? copy(C.nzVal) : C.nzVal, length(A)) +end + + +## Kronecker products + +## COV_EXCL_START +# every thread fills a column of `kron(A, B)`: the products of column `j1` of `A` and +# column `j2` of `B`, in column-major order +@kernel function kron_kernel(colPtr, rowVal, nzVal, Aptr, Aind, Aval, Bptr, Bind, Bval, m2, n2) + c = @index(Global, Linear) + if c < length(colPtr) + j1 = (c - 1) ÷ n2 + 1 + j2 = (c - 1) % n2 + 1 + out = @inbounds colPtr[c] + for p in @inbounds(Aptr[j1]):@inbounds(Aptr[j1+1] - one(eltype(Aptr))) + for q in @inbounds(Bptr[j2]):@inbounds(Bptr[j2+1] - one(eltype(Bptr))) + @inbounds rowVal[out] = ((Aind[p] - 1) * m2 + Bind[q]) % eltype(rowVal) + @inbounds nzVal[out] = Aval[p] * Bval[q] + out += one(out) + end + end + end +end + +@kernel function kron_counts_kernel(counts, Aptr, Bptr, n2) + c = @index(Global, Linear) + if c <= length(counts) + j1 = (c - 1) ÷ n2 + 1 + j2 = (c - 1) % n2 + 1 + @inbounds counts[c] = (Aptr[j1+1] - Aptr[j1]) * (Bptr[j2+1] - Bptr[j2]) + end +end +## COV_EXCL_STOP + +function generic_kron(A::GPUSparseMatrixCSC, B::GPUSparseMatrixCSC) + m1, n1 = size(A) + m2, n2 = size(B) + Ti = promote_type(indtype(A), indtype(B)) + Tv = Base.promote_op(*, eltype(A), eltype(B)) + all(<=(typemax(Int)), (widemul(m1, m2), widemul(n1, n2))) || + throw(ArgumentError("the Kronecker product of $(size(A)) and $(size(B)) matrices is too large")) + dims = (m1 * m2, n1 * n2) + check_sparse_dims(Ti, dims) + check_sparse_nnz(Ti, widemul(nnz(A), nnz(B))) + n = nnz(A) * nnz(B) + colPtr = similar(A.colPtr, Ti, dims[2] + 1) + rowVal = similar(A.rowVal, Ti, n) + nzVal = similar(A.nzVal, Tv, n) + fill!(view(colPtr, 1:1), one(Ti)) + if dims[2] > 0 + kron_counts_kernel(get_backend(A))(view(colPtr, 2:dims[2]+1), A.colPtr, B.colPtr, n2; + ndrange=dims[2]) + accumulate!(Base.add_sum, colPtr, colPtr) + if n > 0 + kron_kernel(get_backend(A))(colPtr, rowVal, nzVal, A.colPtr, A.rowVal, A.nzVal, + B.colPtr, B.rowVal, B.nzVal, m2, n2; ndrange=dims[2]) + end + end + return GPUSparseMatrixCSC(colPtr, rowVal, nzVal, dims) +end + +# the result has the format of `A` +function LinearAlgebra.kron(A::GPUSparseMatrixOrTransposed, B::GPUSparseMatrixOrTransposed) + C = generic_kron(convert(GPUSparseMatrixCSC, materialize(A)), + convert(GPUSparseMatrixCSC, materialize(B))) + return convert(matrix_format(materialize_format(A)), C) +end +# (the diagonal can be on the host, as with `kron(A, I(n))`) +LinearAlgebra.kron(A::GPUSparseMatrixOrTransposed, D::Diagonal) = + kron(A, spdiagm(move_buffer(nonzeros(materialize_format(A)), D.diag))) +LinearAlgebra.kron(D::Diagonal, B::GPUSparseMatrixOrTransposed) = + convert(matrix_format(materialize_format(B)), + kron(spdiagm(move_buffer(nonzeros(materialize_format(B)), D.diag)), B)) + +# a sparse matrix without a lazy wrapper, and the format that has +materialize(A::GPUSparseMatrix) = A +materialize(A::Union{Transpose,Adjoint}) = copy(A) +materialize_format(A::GPUSparseMatrix) = A +materialize_format(A::Union{Transpose,Adjoint}) = parent(A) diff --git a/test/testsuite/sparse.jl b/test/testsuite/sparse.jl index 74a3228b0..dfa5f50e1 100644 --- a/test/testsuite/sparse.jl +++ b/test/testsuite/sparse.jl @@ -19,6 +19,7 @@ using GPUArrays: GPUSparseMatrixCSR, GPUSparseMatrixCSC, GPUSparseMatrixCOO, GPU sparse_dense_conversions(AT, eltypes) sparse_assembly(AT, eltypes) sparse_products(AT, eltypes) + sparse_structure(AT, eltypes) broadcasting_vector(AT, eltypes) broadcasting_matrix(AT, eltypes) broadcasting_mixed(AT, eltypes) @@ -78,6 +79,12 @@ same_sparse(A::SparseMatrixCSC, B::SparseMatrixCSC) = same_sparse(x::SparseVector, y::SparseVector) = length(x) == length(y) && nonzeroinds(x) == nonzeroinds(y) && nonzeros(x) == nonzeros(y) same_sparse(A::AbstractGPUSparseArray, B) = same_sparse(host(A), B) +# the same structure, with values that only need to be close (e.g. products of floats) +function same_sparse_approx(A::AbstractGPUSparseArray, B::SparseMatrixCSC) + C = host(A) + size(C) == size(B) && getcolptr(C) == getcolptr(B) && rowvals(C) == rowvals(B) && + nonzeros(C) ≈ nonzeros(B) +end host(A::AbstractGPUSparseArray) = ndims(A) == 1 ? SparseVector(A) : SparseMatrixCSC(A) # the result of a reduction on the host, whether it is a scalar or an array @@ -1183,3 +1190,125 @@ function sparse_products(AT, eltypes) end end end + +function sparse_structure(AT, eltypes) + @testset "structure" begin + @testset "$S{$ET}" for S in sparse_matrix_formats, ET in eltypes + A = sprand_awkward(ET, 7, 6; Ti=Int32) + dA = gpu_sparse(AT, S, A) + + # removing entries, in place + B = dropzeros!(copy(dA)) + @test B isa S{ET,Int32} + check_structure(B) + @test same_sparse(B, dropzeros(A)) + @test same_sparse(dropzeros(dA), dropzeros(A)) + @test nnz(dA) == nnz(A) + if ET <: Real + B = copy(dA) + z = zero(ET) + @test fkeep!((i, j, v) -> i != j && v > z, B) === B + check_structure(B) + @test same_sparse(B, fkeep!((i, j, v) -> i != j && v > z, copy(A))) + end + if !(ET <: Integer) && !(ET <: Complex{<:Integer}) + tol = real(ET)(0.5) + @test same_sparse(droptol!(copy(dA), tol), droptol!(copy(A), tol)) + end + + # triangular parts + for k in (-2, 0, 1) + @test same_sparse(triu(dA, k), triu(A, k)) + @test same_sparse(tril(dA, k), tril(A, k)) + @test triu(dA, k) isa S{ET,Int32} + @test same_sparse(triu(transpose(dA), k), triu(copy(transpose(A)), k)) + @test same_sparse(tril(adjoint(dA), k), tril(copy(adjoint(A)), k)) + end + @test same_sparse(triu(dA), triu(A)) + @test istriu(triu(dA)) && !istriu(dA) + @test istril(tril(dA, -1), -1) && !istril(triu(dA, 1)) + @test isdiag(gpu_sparse(AT, S, spdiagm(0 => rand(ET, 4)))) + @test !isdiag(dA) + Q = copy(A[1:6, 1:6]) + Q[1, 2] = Q[2, 1] = one(ET) + dQ = gpu_sparse(AT, S, Q) + @test istriu(UpperTriangular(dQ)) && !istril(UpperTriangular(dQ)) + @test istril(LowerTriangular(dQ)) && !istriu(LowerTriangular(dQ)) + dD = gpu_sparse(AT, S, spdiagm(0 => rand(ET, 4))) + # (structurally, as LinearAlgebra's triangular solves dispatch on it) + @test !istril(UpperTriangular(dD)) && !istriu(LowerTriangular(dD)) + + # symmetry, also with a stored zero on one side only + @test !issymmetric(dA) && !ishermitian(dA) + H = sparse(Matrix(A[1:6, 1:6]) + Matrix(A[1:6, 1:6])') + @test ishermitian(gpu_sparse(AT, S, H)) + @test issymmetric(gpu_sparse(AT, S, H)) == issymmetric(H) + Z = sparse([1, 2], [2, 1], ET[1, 0], 2, 2) + @test issymmetric(gpu_sparse(AT, S, Z)) == issymmetric(Z) + if ET <: AbstractFloat + Z = sparse([1, 2], [2, 1], ET[Inf, Inf], 2, 2) + @test issymmetric(gpu_sparse(AT, S, Z)) + end + + # diagonals + for k in (-3, 0, 2, 6, -7) + d = diag(dA, k) + @test d isa AT{ET} + @test Array(d) == diag(A, k) + end + @test isempty(diag(dA, 7)) + @test isempty(diag(dA, -7)) + @test Array(diag(transpose(dA), -1)) == diag(transpose(A), -1) + @test Array(diag(adjoint(dA), 1)) == diag(adjoint(A), 1) + # (offsets of other integer types) + @test Array(diag(dA, big(2))) == diag(A, 2) + @test isempty(diag(dA, UInt(7))) && isempty(diag(dA, UInt(8))) + @test_throws InexactError diag(dA, typemax(UInt)) + @test Array(diag(transpose(dA), UInt(1))) == diag(transpose(A), 1) + dB = gpu_sparse(AT, S, spdiagm(130, 130, 128 => ET[1, 2])) + @test Array(diag(transpose(dB), Int8(-128))) == ET[1, 2] + @test Diagonal(dA) isa Diagonal{ET,<:AT{ET}} + @test Array(Diagonal(dA).diag) == diag(A) + @test tr(gpu_sparse(AT, S, A[1:6, 1:6])) ≈ tr(A[1:6, 1:6]) + + # reshaping (SparseArrays gives a lazy wrapper), keeping the stored zero + B = reshape(dA, 14, 3) + @test B isa S{ET,Int32} + check_structure(B) + @test SparseMatrixCSC(B) == reshape(A, 14, 3) + @test nnz(B) == nnz(A) + @test SparseMatrixCSC(reshape(dA, (3, :))) == reshape(A, (3, :)) + @test_throws DimensionMismatch reshape(dA, 5, 5) + x = vec(dA) + @test x isa GPUSparseVector{ET,Int32} + check_structure(x) + @test SparseVector(x) == vec(A) + @test nnz(x) == nnz(A) + + # Kronecker products + B = SparseMatrixCSC{ET,Int32}(sprand_nozeros(ET, 3, 4, 0.5)) + dB = gpu_sparse(AT, S, B) + @test same_sparse_approx(kron(dA, dB), kron(A, B)) + @test kron(dA, dB) isa S{ET,Int32} + @test same_sparse_approx(kron(transpose(dA), dB), kron(copy(transpose(A)), B)) + v = rand(ET, 3) + @test same_sparse_approx(kron(dA, Diagonal(AT(v))), kron(A, Diagonal(v))) + @test same_sparse_approx(kron(Diagonal(AT(v)), dB), kron(Diagonal(v), B)) + # (with a diagonal on the host) + @test Array(kron(dA, I(3))) ≈ kron(Array(A), I(3)) + @test Array(kron(I(2), transpose(dB))) ≈ kron(I(2), transpose(Array(B))) + huge = gpu_sparse(AT, GPUSparseMatrixCSC, spzeros(ET, 2^32, 0)) + @test_throws ArgumentError kron(huge, huge) + end + + @testset "vector $ET" for ET in eltypes + x = SparseVector{ET,Int32}(sprand_nozeros(ET, 30, 0.4)) + nonzeros(x)[1] = zero(ET) + dx = gpu_sparse(AT, x) + @test same_sparse(dropzeros(dx), dropzeros(x)) + @test same_sparse(dropzeros!(copy(dx)), dropzeros(x)) + B = fkeep!((i, v) -> isodd(i), copy(dx)) + @test same_sparse(B, fkeep!((i, v) -> isodd(i), copy(x))) + end + end +end From 608629847322a18e13bcd856a433aedb38901ff3 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Fri, 25 Sep 2026 18:29:21 +0200 Subject: [PATCH 08/11] Index GPU sparse arrays with ranges and slices A[I, J] with ranges, colons and integers (giving a sparse vector for a row or column) and x[I] for vectors keep the stored entries in the selection and renumber them on the device. Ranges with a positive step preserve the order of the entries, so the result needs no sorting; other index vectors are rejected. The vector methods replace SparseArrays' own, which index through scalar loops. --- src/GPUArrays.jl | 2 +- src/host/sparse/indexing.jl | 142 ++++++++++++++++++++++++++++++++++++ test/testsuite/sparse.jl | 35 +++++++++ 3 files changed, 178 insertions(+), 1 deletion(-) diff --git a/src/GPUArrays.jl b/src/GPUArrays.jl index 36376c05b..1722e9eaf 100644 --- a/src/GPUArrays.jl +++ b/src/GPUArrays.jl @@ -46,11 +46,11 @@ include("host/sparse/primitives.jl") include("host/sparse/conversions.jl") include("host/sparse/assembly.jl") include("host/sparse/show.jl") -include("host/sparse/indexing.jl") include("host/sparse/broadcast.jl") include("host/sparse/reductions.jl") include("host/sparse/linalg.jl") include("host/sparse/structure.jl") +include("host/sparse/indexing.jl") include("host/alloc_cache.jl") diff --git a/src/host/sparse/indexing.jl b/src/host/sparse/indexing.jl index ddadaeaa1..ea7b94905 100644 --- a/src/host/sparse/indexing.jl +++ b/src/host/sparse/indexing.jl @@ -43,3 +43,145 @@ function Base.setindex!(::GPUSparseArray, v, I...) throw(ArgumentError("""GPU sparse arrays do not support setting individual entries. Assemble a new array with `sparse(I, J, V, m, n)` or `sparsevec(I, V, n)`, use broadcasting, or update stored values through `nonzeros(A)`.""")) end + + +## slices and ranges + +# Indexing with ranges (and colons) keeps the stored entries that fall in them and +# renumbers them. Ranges with a positive step preserve the order of the entries, so the +# result needs no sorting; other index vectors are not supported. +const SparseRangeIndex = Union{AbstractRange{<:Integer}, Colon} + +range_index(I::Colon, n) = Base.OneTo(n) +function range_index(I::AbstractRange{<:Integer}, n) + (length(I) <= 1 || step(I) > 0) || + throw(ArgumentError("indexing GPU sparse arrays requires ranges with a positive step")) + return I +end +range_index(i::Integer, n) = i:i + +# the position of `i` in the range `I` (from `first`, with `step`, of `len` elements), or 0 +@inline function range_position(i, first, step, len) + d = i - first + (d >= 0 && d % step == 0) || return 0 + p = d ÷ step + 1 + return 1 <= p <= len ? p : 0 +end + +## COV_EXCL_START +@kernel function select_entries_kernel(keep, new_rows, new_cols, rows, cols, I, J) + k = @index(Global, Linear) + if k <= length(keep) + i = range_position(Int(@inbounds rows[k]), I...) + j = range_position(Int(@inbounds cols[k]), J...) + @inbounds keep[k] = i != 0 && j != 0 + @inbounds new_rows[k] = i % eltype(new_rows) + @inbounds new_cols[k] = j % eltype(new_cols) + end +end + +@kernel function select_vector_entries_kernel(keep, new_inds, inds, I) + k = @index(Global, Linear) + if k <= length(keep) + i = range_position(Int(@inbounds inds[k]), I...) + @inbounds keep[k] = i != 0 + @inbounds new_inds[k] = i % eltype(new_inds) + end +end +## COV_EXCL_STOP + +# (a range of one element may have any step) +range_parameters(I::AbstractRange) = (Int(first(I)), length(I) > 1 ? Int(step(I)) : 1, length(I)) + +# the entries of `A[I, J]`, in the order they have in `A`: their new rows and columns, and +# their values +function select_entries(A::GPUSparseMatrix, I::AbstractRange, J::AbstractRange) + rows, cols = entry_coordinates(A) + n = nnz(A) + keep = similar(A.nzVal, Bool, n) + new_rows = similar(rows, n) + new_cols = similar(cols, n) + if n > 0 + select_entries_kernel(get_backend(A))(keep, new_rows, new_cols, rows, cols, + range_parameters(I), range_parameters(J); + ndrange=n) + end + kept = findall(keep) + return new_rows[kept], new_cols[kept], A.nzVal[kept] +end + +function sparse_getindex(A::GPUSparseMatrix, I, J) + @boundscheck checkbounds(A, I, J) + m, n = size(A) + rows, cols, vals = select_entries(A, range_index(I, m), range_index(J, n)) + if I isa Integer || J isa Integer + # entries of a single row or column are in order of the other index in every format + return GPUSparseVector(I isa Integer ? cols : rows, vals, + I isa Integer ? length(range_index(J, n)) : length(range_index(I, m))) + end + dims = (length(range_index(I, m)), length(range_index(J, n))) + if A isa GPUSparseMatrixCSR + GPUSparseMatrixCSR(compress(rows, dims[1]), cols, vals, dims) + elseif A isa GPUSparseMatrixCSC + GPUSparseMatrixCSC(compress(cols, dims[2]), rows, vals, dims) + else + GPUSparseMatrixCOO(rows, cols, vals, dims) + end +end + +Base.getindex(A::GPUSparseMatrix, I::SparseRangeIndex, J::SparseRangeIndex) = + sparse_getindex(A, I, J) +Base.getindex(A::GPUSparseMatrix, i::Integer, J::SparseRangeIndex) = sparse_getindex(A, i, J) +Base.getindex(A::GPUSparseMatrix, I::SparseRangeIndex, j::Integer) = sparse_getindex(A, I, j) + +# Logical masks also preserve the order of the entries: each one is renumbered by the +# number of selected rows and columns before it. +function Base.getindex(A::GPUSparseMatrix, I::AbstractVector{Bool}, J::AbstractVector{Bool}) + (length(I), length(J)) == size(A) || throw(BoundsError(A, (I, J))) + mask(M) = M isa AnyGPUVector ? M : copyto!(similar(A.nzVal, Bool, length(M)), Vector{Bool}(M)) + I, J = mask(I), mask(J) + rowrank = accumulate(+, Int.(I)) + colrank = accumulate(+, Int.(J)) + dims = (count(I), count(J)) + rows, cols = entry_coordinates(A) + kept = findall(I[rows] .& J[cols]) + Ti = indtype(A) + new_rows = convert_index_buffer(Ti, rowrank[rows[kept]]) + new_cols = convert_index_buffer(Ti, colrank[cols[kept]]) + vals = A.nzVal[kept] + if A isa GPUSparseMatrixCSR + GPUSparseMatrixCSR(compress(new_rows, dims[1]), new_cols, vals, dims) + elseif A isa GPUSparseMatrixCSC + GPUSparseMatrixCSC(compress(new_cols, dims[2]), new_rows, vals, dims) + else + GPUSparseMatrixCOO(new_rows, new_cols, vals, dims) + end +end + +# a range combined with a mask: both as masks +range_mask(I::AbstractVector{Bool}, n) = I +function range_mask(I::SparseRangeIndex, n) + mask = falses(n) + mask[range_index(I, n)] .= true + return mask +end +Base.getindex(A::GPUSparseMatrix, I::SparseRangeIndex, J::AbstractVector{Bool}) = + A[range_mask(I, size(A, 1)), J] +Base.getindex(A::GPUSparseMatrix, I::AbstractVector{Bool}, J::SparseRangeIndex) = + A[I, range_mask(J, size(A, 2))] + +function sparse_getindex(x::GPUSparseVector, I::AbstractRange{<:Integer}) + @boundscheck checkbounds(x, I) + I = range_index(I, length(x)) + keep = similar(x.nzVal, Bool, nnz(x)) + new_inds = similar(x.nzInd, nnz(x)) + if nnz(x) > 0 + select_vector_entries_kernel(get_backend(x))(keep, new_inds, x.nzInd, + range_parameters(I); ndrange=nnz(x)) + end + kept = findall(keep) + return GPUSparseVector(new_inds[kept], x.nzVal[kept], length(I)) +end +# (separate methods to be more specific than SparseArrays' for any sparse vector) +Base.getindex(x::GPUSparseVector, I::AbstractUnitRange{<:Integer}) = sparse_getindex(x, I) +Base.getindex(x::GPUSparseVector, I::StepRange{<:Integer}) = sparse_getindex(x, I) diff --git a/test/testsuite/sparse.jl b/test/testsuite/sparse.jl index dfa5f50e1..2219a7c16 100644 --- a/test/testsuite/sparse.jl +++ b/test/testsuite/sparse.jl @@ -20,6 +20,7 @@ using GPUArrays: GPUSparseMatrixCSR, GPUSparseMatrixCSC, GPUSparseMatrixCOO, GPU sparse_assembly(AT, eltypes) sparse_products(AT, eltypes) sparse_structure(AT, eltypes) + sparse_slicing(AT, eltypes) broadcasting_vector(AT, eltypes) broadcasting_matrix(AT, eltypes) broadcasting_mixed(AT, eltypes) @@ -1312,3 +1313,37 @@ function sparse_structure(AT, eltypes) end end end + +function sparse_slicing(AT, eltypes) + @testset "slicing" begin + @testset "$S{$ET}" for S in sparse_matrix_formats, ET in eltypes + A = sprand_awkward(ET, 9, 8; Ti=Int32) + dA = gpu_sparse(AT, S, A) + for (I, J) in ((2:5, 1:3:8), (:, 2:4), (3:3, :), (1:0, :), (9:-1:9, 2:2), (2:-1:2, 3:-5:3)) + B = dA[I, J] + @test B isa S{ET,Int32} + check_structure(B) + @test same_sparse(B, A[I, J]) + end + for (I, J) in ((:, 3), (4, :), (2:6, 2), (1, 2:2:8), (3, :)) + x = dA[I, J] + @test x isa GPUSparseVector{ET,Int32} + check_structure(x) + @test same_sparse(x, A[I, J]) + end + @test_throws BoundsError dA[1:10, :] + @test_throws BoundsError dA[:, 9] + @test_throws ArgumentError dA[5:-1:1, :] + end + @testset "vector $ET" for ET in eltypes + x = SparseVector{ET,Int32}(sprand_nozeros(ET, 30, 0.4)) + dx = gpu_sparse(AT, x) + for I in (3:17, 1:4:30, 31:30, 3:-1:3) + y = dx[I] + @test y isa GPUSparseVector{ET,Int32} + @test same_sparse(y, x[I]) + end + @test_throws BoundsError dx[0:3] + end + end +end From e4efdc965383a51f3833a964ad15f2c8645e1efc Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Fri, 25 Sep 2026 18:34:26 +0200 Subject: [PATCH 09/11] Add dot products, count/any/all, and element-type conversions dot between sparse vectors multiplies them with a broadcast, dot with a dense vector gathers its elements at the stored indices, and dot(x, A, y) is dot(x, A * y). count, any and all combine the stored values with the implicit zeros. float, complex, real and imag convert the stored values, sparse(A; fmt) and sparsevec(A) convert between sparse formats, and reductions of vectors support dims=1. Several of these replace SparseArrays' methods for any sparse vector, which loop over the elements on the host (Vector(x), imag). --- src/host/sparse/conversions.jl | 25 ++++++++++++++++ src/host/sparse/reductions.jl | 49 +++++++++++++++++++++++++++++-- test/testsuite/sparse.jl | 53 ++++++++++++++++++++++++++++++++++ 3 files changed, 124 insertions(+), 3 deletions(-) diff --git a/src/host/sparse/conversions.jl b/src/host/sparse/conversions.jl index 96f1e2cb8..142817970 100644 --- a/src/host/sparse/conversions.jl +++ b/src/host/sparse/conversions.jl @@ -39,6 +39,16 @@ Base.Array{T}(A::GPUSparseArray) where {T} = Array{T}(host_sparse(A)) Base.Array{T,1}(x::GPUSparseVector) where {T} = Array{T,1}(host_sparse(x)) Base.Array{T,2}(A::GPUSparseMatrix) where {T} = Array{T,2}(host_sparse(A)) Base.collect(A::GPUSparseArray) = Array(A) +# (SparseArrays has a scalar loop for any sparse vector) +Base.Vector(x::GPUSparseVector) = Vector(SparseVector(x)) + +# a dense GPU matrix becomes sparse on the device before its buffers are transferred +# (SparseArrays' constructor would index every element) +SparseArrays.SparseMatrixCSC(A::AbstractGPUMatrix) = SparseMatrixCSC(sparse(A)) +SparseArrays.SparseMatrixCSC{T}(A::AbstractGPUMatrix) where {T} = + SparseMatrixCSC(GPUSparseMatrixCSC{T}(A)) +SparseArrays.SparseMatrixCSC{T,Ti}(A::AbstractGPUMatrix) where {T,Ti} = + SparseMatrixCSC(GPUSparseMatrixCSC{T,Ti}(A)) ## `adapt` @@ -448,3 +458,18 @@ function SparseArrays.findnz(A::GPUSparseMatrix) return C.rowVal, expand_ptr(C.colPtr, nnz(C)), C.nzVal end SparseArrays.findnz(x::GPUSparseVector) = (copy(x.nzInd), copy(x.nzVal)) + + +## element types + +Base.float(A::GPUSparseArray) = + eltype(A) <: AbstractFloat ? A : with_eltypes(A, float(eltype(A)), indtype(A)) +Base.complex(A::GPUSparseArray) = + eltype(A) <: Complex ? A : with_eltypes(A, complex(eltype(A)), indtype(A)) +Base.real(A::GPUSparseArray) = eltype(A) <: Real ? A : real.(A) +Base.imag(A::GPUSparseArray) = eltype(A) <: Real ? zero(A) : imag.(A) + +# a sparse array is already sparse; the matrix format can be chosen as for dense arrays +SparseArrays.sparse(A::GPUSparseMatrix; fmt::Symbol=:csc) = sparse_format(fmt)(A) +SparseArrays.sparse(x::GPUSparseVector) = copy(x) +SparseArrays.sparsevec(A::GPUSparseMatrix) = vec(A) diff --git a/src/host/sparse/reductions.jl b/src/host/sparse/reductions.jl index e28ee9c25..451d628b8 100644 --- a/src/host/sparse/reductions.jl +++ b/src/host/sparse/reductions.jl @@ -88,9 +88,12 @@ function reduce_all(f, op, A::GPUSparseArray, init) end end -Base.mapreduce(f, op, x::GPUSparseVector; dims=:, init=nothing) = - dims === Colon() ? reduce_all(f, op, x, init) : - error("only dims=: is supported for sparse vectors") +function Base.mapreduce(f, op, x::GPUSparseVector; dims=:, init=nothing) + dims === Colon() && return reduce_all(f, op, x, init) + dims == 1 || error("only dims=: or dims=1 is supported for sparse vectors") + r = reduce_all(f, op, x, init) + return fill!(similar(x.nzVal, typeof(r), 1), r) +end function Base.mapreduce(f, op, A::GPUSparseMatrix; dims=:, init=nothing) dims === Colon() && return reduce_all(f, op, A, init) @@ -134,3 +137,43 @@ function LinearAlgebra.opnorm(A::GPUSparseMatrix, p::Real=2) throw(ArgumentError("p=$p is not supported")) end end + +# `f` of the stored values, and of the implicit zeros if there are any +function Base.count(f, A::GPUSparseArray) + n = nnz(A) == 0 ? 0 : count(f, nonzeros(A)) + return nnz(A) < length(A) && f(zero(eltype(A))) ? n + length(A) - nnz(A) : n +end +# (separate methods for functions, to be more specific than Base's) +for F in (:Function, :Any) + @eval begin + Base.any(f::$F, A::GPUSparseArray) = + (nnz(A) > 0 && any(f, nonzeros(A))) || (nnz(A) < length(A) && f(zero(eltype(A)))) + Base.all(f::$F, A::GPUSparseArray) = + (nnz(A) == 0 || all(f, nonzeros(A))) && (nnz(A) == length(A) || f(zero(eltype(A)))) + end +end +Base.any(A::GPUSparseArray) = any(identity, A) +Base.all(A::GPUSparseArray) = all(identity, A) + + +## dot products + +function LinearAlgebra.dot(x::GPUSparseVector, y::GPUSparseVector) + length(x) == length(y) || + throw(DimensionMismatch("x has length $(length(x)), y has length $(length(y))")) + # `+` rather than `sum`, which would widen small integers + return reduce(+, conj.(x) .* y) +end +function LinearAlgebra.dot(x::GPUSparseVector, y::AnyGPUVector) + length(x) == length(y) || + throw(DimensionMismatch("x has length $(length(x)), y has length $(length(y))")) + T = typeof(dot(zero(eltype(x)), zero(eltype(y)))) + nnz(x) == 0 && return zero(T) + return mapreduce((a, b) -> dot(a, b), +, x.nzVal, y[x.nzInd]) +end +LinearAlgebra.dot(y::AnyGPUVector, x::GPUSparseVector) = conj(dot(x, y)) +function LinearAlgebra.dot(x::AnyGPUVector, A::GPUSparseMatrix, y::AnyGPUVector) + (length(x), length(y)) == size(A) || + throw(DimensionMismatch("x has length $(length(x)), A has size $(size(A)), y has length $(length(y))")) + return dot(x, A * y) +end diff --git a/test/testsuite/sparse.jl b/test/testsuite/sparse.jl index 2219a7c16..e1d5027d7 100644 --- a/test/testsuite/sparse.jl +++ b/test/testsuite/sparse.jl @@ -21,6 +21,7 @@ using GPUArrays: GPUSparseMatrixCSR, GPUSparseMatrixCSC, GPUSparseMatrixCOO, GPU sparse_products(AT, eltypes) sparse_structure(AT, eltypes) sparse_slicing(AT, eltypes) + sparse_misc(AT, eltypes) broadcasting_vector(AT, eltypes) broadcasting_matrix(AT, eltypes) broadcasting_mixed(AT, eltypes) @@ -1347,3 +1348,55 @@ function sparse_slicing(AT, eltypes) end end end + +function sparse_misc(AT, eltypes) + @testset "dot, count and conversions" begin + @testset "$ET" for ET in eltypes + x = sprand_nozeros(ET, 30, 0.4) + y = sprand_nozeros(ET, 30, 0.4) + dx, dy = gpu_sparse(AT, x), gpu_sparse(AT, y) + v = rand(ET, 30) + @test dot(dx, dy) ≈ dot(x, y) + @test dot(dx, AT(v)) ≈ dot(x, v) + @test dot(AT(v), dx) ≈ dot(v, x) + @test dot(gpu_sparse(AT, spzeros(ET, 30)), AT(v)) == zero(ET) + @test_throws DimensionMismatch dot(dx, AT(rand(ET, 29))) + A = sprand_awkward(ET, 30, 30) + for S in sparse_matrix_formats + dA = gpu_sparse(AT, S, A) + w = rand(ET, 30) + @test dot(AT(w), dA, AT(v)) ≈ dot(w, A, v) + @test count(!iszero, dA) == count(!iszero, A) + @test count(iszero, dA) == count(iszero, A) + @test any(iszero, dA) == any(iszero, A) + @test all(iszero, dA) == all(iszero, A) + @test sparse(dA) isa GPUSparseMatrixCSC{ET} + @test same_sparse(sparse(dA), A) + @test sparse(dA; fmt=:coo) isa GPUSparseMatrixCOO{ET} + @test SparseVector(sparsevec(dA)) == vec(A) + end + @test Array(sum(dx; dims=1)) ≈ sum(x; dims=1) + @test Vector(dx) == Vector(x) + @test same_sparse(sparse(dx), x) + # (which may give an element type the back-end does not support) + if float(ET) in eltypes + @test SparseVector(float(dx)) == float(x) + end + if complex(ET) in eltypes + @test SparseVector(complex(dx)) == complex(x) + end + @test SparseVector(imag(dx)) == imag(x) + @test SparseVector(real(dx)) == real(x) + end + @testset "Bool" begin + A = sprand(Bool, 8, 9, 0.3) + for S in sparse_matrix_formats + dA = gpu_sparse(AT, S, A) + @test count(dA) == count(A) + @test any(dA) == any(A) + @test all(dA) == all(A) + @test !all(gpu_sparse(AT, S, sparse(trues(3, 3))) .& false) + end + end + end +end From a7da0088bda335e5d0f062a50c1e4fe6da52dc54 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Fri, 25 Sep 2026 18:37:18 +0200 Subject: [PATCH 10/11] Multiply sparse matrices MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit generic_spgemm computes A * B for any pair of formats by expansion, sorting and compression: every product of a stored entry of A with one of B is materialized, each thread finding its entry of A by binary search over the running count of products, and the products are summed per coordinate through the assembly path, which keeps the result deterministic. The memory needed grows with the number of products. The sparse-output mul!(C, A, B, α, β) gives C the union of the structures of α*A*B and β*C, resizing its buffers, rather than assuming that an equal number of stored entries means an equal structure. exp is ported from the old code on top of these, with a fix for matrices whose norm is below one, which it scaled up instead of down. --- docs/src/interface.md | 3 + src/host/sparse/linalg.jl | 156 ++++++++++++++++++++++++++++++++++++++ test/testsuite/sparse.jl | 126 ++++++++++++++++++++++++++++++ 3 files changed, 285 insertions(+) diff --git a/docs/src/interface.md b/docs/src/interface.md index 9c3cbad61..14ea554b1 100644 --- a/docs/src/interface.md +++ b/docs/src/interface.md @@ -119,6 +119,7 @@ A back-end will usually add the following, all optional: |:--|:--|:--| | 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` | @@ -192,6 +193,8 @@ GPUArrays.generic_expand GPUArrays.generic_assemble GPUArrays.generic_spmv! GPUArrays.generic_spmm! +GPUArrays.generic_spgemm +GPUArrays.generic_spgemm! ``` ## Caching Allocator diff --git a/src/host/sparse/linalg.jl b/src/host/sparse/linalg.jl index 6c7e29bc8..ce700012f 100644 --- a/src/host/sparse/linalg.jl +++ b/src/host/sparse/linalg.jl @@ -307,3 +307,159 @@ function Base.:(*)(A::Union{AnyGPUMatrix,GPUSparseMatrixOperand}, x::GPUSparseVe T = Base.promote_op(LinearAlgebra.matprod, eltype(A), eltype(x)) return mul!(similar(x.nzVal, T, size(A, 1)), A, x) end + + +## sparse × sparse + +## COV_EXCL_START +# the number of products every stored entry `A[i, l]` contributes: the length of row `l` +# of `B` +@kernel function spgemm_counts_kernel(counts, Aind, Bptr) + e = @index(Global, Linear) + if e <= length(counts) + l = @inbounds Aind[e] + @inbounds counts[e] = Int64(Bptr[l+1]) - Int64(Bptr[l]) + end +end + +# every thread computes one product `A[i, l] * B[l, j]`, finding its entry of `A` by binary +# search over the running total of the counts +@kernel function spgemm_expand_kernel(I, J, V, offsets, Arow, Aind, Aval, Bptr, Bind, Bval) + p = @index(Global, Linear) + if p <= length(I) + e = searchsortedfirst(offsets, Int64(p)) + q = p - (e == 1 ? 0 : @inbounds(offsets[e-1])) + l = @inbounds Aind[e] + k = @inbounds(Bptr[l]) + q - 1 + @inbounds I[p] = Arow[e] + @inbounds J[p] = Bind[k] + @inbounds V[p] = convert(eltype(V), Aval[e]) * convert(eltype(V), Bval[k]) + end +end +## COV_EXCL_STOP + +""" + GPUArrays.generic_spgemm(A, B) + +The product of two sparse matrices, in the format of `A`, computed on the device by +expansion, sorting and compression (ESC): every product of a stored entry of `A` with one +of `B` is materialized, then the products are summed per coordinate as in +`sparse(I, J, V, m, n, +)`. The memory needed grows with the number of products. The +generic implementation behind `A * B`; hash-based accumulation with atomic insertion would +avoid materializing all products. +""" +function generic_spgemm(A::GPUSparseMatrix, B::GPUSparseMatrix) + size(A, 2) == size(B, 1) || + throw(DimensionMismatch("A has dimensions $(size(A)), B has dimensions $(size(B))")) + m, n = size(A, 1), size(B, 2) + T = Base.promote_op(LinearAlgebra.matprod, eltype(A), eltype(B)) + Ti = promote_type(indtype(A), indtype(B)) + Aptr, Aind, Aval = csr_operand(A, 'N') + Bptr, Bind, Bval = csr_operand(B, 'N') + Arow = expand_ptr(Aptr, length(Aval)) + + counts = similar(Aval, Int64, length(Aval)) + if !isempty(counts) + spgemm_counts_kernel(get_backend(counts))(counts, Aind, Bptr; ndrange=length(counts)) + end + offsets = accumulate(+, counts) + free_buffer!(counts) + total = isempty(offsets) ? 0 : Int(@allowscalar offsets[end]) + + I = similar(Aind, Ti, total) + J = similar(Aind, Ti, total) + V = similar(Aval, product_accumulator(eltype(A), eltype(B)), total) + if total > 0 + spgemm_expand_kernel(get_backend(V))(I, J, V, offsets, Arow, Aind, Aval, Bptr, Bind, + Bval; ndrange=total) + end + free_buffer!(offsets) + free_buffer!(Arow) + fmt = A isa GPUSparseMatrixCSC ? :csc : A isa GPUSparseMatrixCSR ? :csr : :coo + C = assemble_matrix(I, J, V, m, n, +, fmt) + return eltype(C) == T ? C : with_eltypes(C, T, Ti) +end + +# a sparse operand without a lazy wrapper +function materialize(S::Symmetric{<:Any,<:GPUSparseMatrix}) + A = parent(S) + T = S.uplo == 'U' ? triu(A) : tril(A) + strict = S.uplo == 'U' ? triu(A, 1) : tril(A, -1) + return T .+ copy(transpose(strict)) +end +function materialize(H::Hermitian{<:Any,<:GPUSparseMatrix}) + A = parent(H) + strict = H.uplo == 'U' ? triu(A, 1) : tril(A, -1) + return strict .+ copy(adjoint(strict)) .+ real.(band(A, 0, 0)) +end + +Base.:(*)(A::GPUSparseMatrixOperand, B::GPUSparseMatrixOperand) = + generic_spgemm(materialize(A), materialize(B)) + +""" + GPUArrays.generic_spgemm!(C, tA, tB, A, B, α, β) + +Compute `C = α op(A) op(B) + β C` for sparse matrices, where `C` gets the union of the +structures of both terms (that of the product if `β` is zero). The buffers of `C` are +resized, so `C` must not share them with another array. The generic implementation behind +the sparse-output `mul!`. +""" +function generic_spgemm!(C::GPUSparseMatrix, tA::AbstractChar, tB::AbstractChar, + A::GPUSparseMatrix, B::GPUSparseMatrix, α::Number, β::Number) + P = generic_spgemm(materialize(LinearAlgebra.wrap(A, tA)), + materialize(LinearAlgebra.wrap(B, tB))) + size(P) == size(C) || + throw(DimensionMismatch("C has dimensions $(size(C)), needs $(size(P))")) + R = iszero(β) ? α .* P : α .* P .+ β .* C + return copyto!(C, convert(matrix_format(C), R)) +end + +LinearAlgebra.mul!(C::GPUSparseMatrix, tA::AbstractChar, tB::AbstractChar, + A::GPUSparseMatrix, B::GPUSparseMatrix, α::Number, β::Number) = + generic_spgemm!(C, tA, tB, A, B, α, β) +@static if VERSION < v"1.13.0-rc4" + LinearAlgebra.generic_matmatmul!(C::GPUSparseMatrix, tA::AbstractChar, tB::AbstractChar, + A::GPUSparseMatrix, B::GPUSparseMatrix, a::Number, b::Number) = + LinearAlgebra.mul!(C, tA, tB, A, B, a, b) + LinearAlgebra.generic_matmatmul!(C::GPUSparseMatrix, tA::AbstractChar, tB::AbstractChar, + A::GPUSparseMatrix, B::GPUSparseMatrix, + _add::MulAddMul=MulAddMul()) = + LinearAlgebra.mul!(C, tA, tB, A, B, _add.alpha, _add.beta) +end + + +## matrix exponential + +""" + exp(A::GPUSparseMatrix; threshold=1e-7, nonzero_tol=1e-14) + +The matrix exponential of a square sparse matrix, by scaling and squaring a truncated +Taylor series, dropping entries below `nonzero_tol` as it goes. The series is truncated +when a term's largest magnitude falls below `threshold`. +""" +function LinearAlgebra.exp(A::GPUSparseMatrix; threshold=1e-7, nonzero_tol=1e-14) + n = LinearAlgebra.checksquare(A) + A = float(A) + T = eltype(A) + # scale A to norm at most 1 by a power of two + scaling = max(one(real(T)), nextpow(2, max(norm(A, Inf), floatmin(real(T))))) + A = A ./ T(scaling) + Id = spdiagm(0 => fill!(similar(nonzeros(A), n), one(T))) + P = convert(matrix_format(A), Id) + term = P + k = 1 + while true + term = (T(1 / k) * A) * term + droptol!(term, nonzero_tol) + P = P + term + norm(term, Inf) > threshold || break + k += 1 + end + for _ in 1:Int(log2(scaling)) + P = P * P + if nnz(P) / length(P) < 0.25 + droptol!(P, nonzero_tol) + end + end + return P +end diff --git a/test/testsuite/sparse.jl b/test/testsuite/sparse.jl index e1d5027d7..96dcab5ee 100644 --- a/test/testsuite/sparse.jl +++ b/test/testsuite/sparse.jl @@ -22,6 +22,7 @@ using GPUArrays: GPUSparseMatrixCSR, GPUSparseMatrixCSC, GPUSparseMatrixCOO, GPU sparse_structure(AT, eltypes) sparse_slicing(AT, eltypes) sparse_misc(AT, eltypes) + sparse_spgemm(AT, eltypes) broadcasting_vector(AT, eltypes) broadcasting_matrix(AT, eltypes) broadcasting_mixed(AT, eltypes) @@ -29,6 +30,7 @@ using GPUArrays: GPUSparseMatrixCSR, GPUSparseMatrixCSC, GPUSparseMatrixCOO, GPU sparse_linalg(AT, eltypes) iszero_sparse(AT, eltypes) sparse_bool(AT) + sparse_generic(AT, eltypes) end const sparse_matrix_formats = (GPUSparseMatrixCSR, GPUSparseMatrixCSC, GPUSparseMatrixCOO) @@ -1400,3 +1402,127 @@ function sparse_misc(AT, eltypes) end end end + +function sparse_spgemm(AT, eltypes) + @testset "sparse × sparse" begin + @testset "$S{$ET}" for S in sparse_matrix_formats, ET in eltypes + A = sprand_awkward(ET, 8, 6; Ti=Int32) + B = SparseMatrixCSC{ET,Int32}(sprand_nozeros(ET, 6, 7, 0.4)) + dA = gpu_sparse(AT, S, A) + for S′ in sparse_matrix_formats + dB = gpu_sparse(AT, S′, B) + C = dA * dB + @test C isa S{Base.promote_op(LinearAlgebra.matprod, ET, ET),Int32} + check_structure(C) + @test SparseMatrixCSC(C) ≈ A * B + end + dB = gpu_sparse(AT, S, B) + @test SparseMatrixCSC(transpose(dB) * transpose(dA)) ≈ transpose(B) * transpose(A) + @test SparseMatrixCSC(dA' * dA) ≈ A' * A + Q = sprand_nozeros(ET, 6, 6, 0.4) + dQ = gpu_sparse(AT, S, Q) + @test SparseMatrixCSC(Symmetric(dQ) * dB) ≈ Symmetric(Matrix(Q)) * B + if ET <: Complex + @test SparseMatrixCSC(Hermitian(dQ, :L) * dB) ≈ Hermitian(Matrix(Q), :L) * B + end + @test nnz(dA * gpu_sparse(AT, S, spzeros(ET, Int32, 6, 3))) == 0 + @test size(gpu_sparse(AT, S, spzeros(ET, Int32, 0, 4)) * gpu_sparse(AT, S, spzeros(ET, Int32, 4, 2))) == (0, 2) + @test_throws DimensionMismatch dA * dA + + # the sparse-output mul! takes the union of both structures + if ET <: Union{AbstractFloat, Complex{<:AbstractFloat}} + α, β = ET(2), ET(3) + C = SparseMatrixCSC{ET,Int32}(sprand_nozeros(ET, 8, 7, 0.3)) + dC = gpu_sparse(AT, S, C) + @test mul!(dC, dA, dB, α, β) === dC + check_structure(dC) + @test SparseMatrixCSC(dC) ≈ α * (A * B) + β * C + # a destination with as many entries as the product, but elsewhere + P = A * B + D = copy(P) + D.rowval .= mod1.(D.rowval .+ 1, 8) + D = sparse(findnz(D)..., 8, 7) + dD = gpu_sparse(AT, S, D) + mul!(dD, dA, dB, true, true) + @test SparseMatrixCSC(dD) ≈ P + D + # with β = 0, the structure of the product + dD = gpu_sparse(AT, S, SparseMatrixCSC{ET,Int32}(D)) + mul!(dD, dA, dB) + @test SparseMatrixCSC(dD) ≈ P + @test nnz(dD) == nnz(dA * dB) + end + end + + @testset "$T" for T in (Int, Bool) + A = sprand(T, 6, 5, 0.5) + B = sprand(T, 5, 4, 0.5) + for S in sparse_matrix_formats + @test SparseMatrixCSC(gpu_sparse(AT, S, A) * gpu_sparse(AT, S, B)) == A * B + end + end + + @testset "exp" begin + # of an integer matrix (which gives Float64 values) + if Float64 in eltypes + A = sparse([1, 2], [1, 2], [1, 2], 2, 2) + @test Array(exp(gpu_sparse(AT, GPUSparseMatrixCSR, A))) ≈ exp(Matrix{Float64}(A)) + end + for ET in filter(T -> T <: Union{AbstractFloat, Complex{<:AbstractFloat}} && + real(T) != Float16, eltypes) + for scale in (0.25, 4) + A = sprand(ET, 7, 7, 0.4) * real(ET)(scale) + for S in sparse_matrix_formats + E = exp(gpu_sparse(AT, S, A)) + @test E isa S{ET} + @test Array(E) ≈ exp(Matrix(A)) rtol=1e-4 + end + end + end + end + end +end + +# the generic implementations, which back-ends with vendor libraries do not reach through +# the LinearAlgebra interface +function sparse_generic(AT, eltypes) + @testset "generic implementations" begin + ET = first(filter(isrealfloattype, eltypes)) + A = sprand_awkward(ET, 6, 7) + B = sprand_awkward(ET, 7, 5) + x = rand(ET, 7) + X = rand(ET, 7, 4) + Y = rand(ET, 4, 7) + @testset "$S" for S in sparse_matrix_formats + dA = gpu_sparse(AT, S, A) + dB = gpu_sparse(AT, S, B) + for t in ('N', 'T', 'C') + y = AT(rand(ET, t == 'N' ? 6 : 7)) + y0 = Array(y) + GPUArrays.generic_spmv!(y, t, dA, AT(t == 'N' ? x : x[1:6]), ET(2), ET(3)) + @test Array(y) ≈ 2 * LinearAlgebra.wrap(A, t) * (t == 'N' ? x : x[1:6]) + 3 * y0 + C = AT(rand(ET, t == 'N' ? 6 : 7, 4)) + D = Array(C) + GPUArrays.generic_spmm!(C, t, 'N', dA, AT(t == 'N' ? X : X[1:6, :]), ET(2), ET(3)) + @test Array(C) ≈ 2 * LinearAlgebra.wrap(A, t) * (t == 'N' ? X : X[1:6, :]) + 3 * D + C = AT(rand(ET, 4, t == 'N' ? 7 : 6)) + D = Array(C) + GPUArrays.generic_spmm!(C, 'N', t, AT(t == 'N' ? Y[:, 1:6] : Y), dA, ET(2), ET(3)) + @test Array(C) ≈ 2 * (t == 'N' ? Y[:, 1:6] : Y) * LinearAlgebra.wrap(A, t) + 3 * D + end + for t in ('S', 'H') + Q = A[1:6, 1:6] + dQ = gpu_sparse(AT, S, Q) + y = AT(zeros(ET, 6)) + GPUArrays.generic_spmv!(y, t, dQ, AT(x[1:6]), ET(1), ET(0)) + @test Array(y) ≈ LinearAlgebra.wrap(Q, t) * x[1:6] + end + dC = GPUArrays.generic_spgemm(dA, dB) + @test dC isa S + @test Array(check_structure(dC)) ≈ Array(A * B) + dC = gpu_sparse(AT, S, sprand_awkward(ET, 6, 5)) + C = SparseMatrixCSC(dC) + GPUArrays.generic_spgemm!(dC, 'N', 'N', dA, dB, ET(2), ET(3)) + @test Array(check_structure(dC)) ≈ 2 * Array(A * B) + 3 * Array(C) + end + end +end From 7f07bb1b79d0899264d6348a43240e445592a28b Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Fri, 6 Nov 2026 19:24:01 +0100 Subject: [PATCH 11/11] Add value equality, in-place broadcast, wrapper arithmetic and concatenation SparseArrays implements these for any AbstractSparseArray with scalar loops, which a GPU sparse array cannot run. == compares values, so stored zeros do not matter. Broadcasting into a sparse array materializes the result and copies it over, so the destination takes the structure of the result. Arithmetic with transposed, Symmetric and Hermitian wrappers, UniformScaling and Diagonal materializes the wrapper and combines the operands on the device, and hcat/vcat of sparse arrays offset the coordinates and reassemble them. blockdiag throws instead of falling back to scalar indexing. --- docs/src/interface.md | 38 ++++++ src/GPUArrays.jl | 1 + src/host/sparse/array.jl | 254 ++++++++++++++++++++++++++++++++++++ test/testsuite/sparse.jl | 270 +++++++++++++++++++++++++++++++++++++++ 4 files changed, 563 insertions(+) create mode 100644 src/host/sparse/array.jl diff --git a/docs/src/interface.md b/docs/src/interface.md index 14ea554b1..1ff89b6dc 100644 --- a/docs/src/interface.md +++ b/docs/src/interface.md @@ -179,6 +179,44 @@ Conversions: 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 diff --git a/src/GPUArrays.jl b/src/GPUArrays.jl index 1722e9eaf..b0bb07d19 100644 --- a/src/GPUArrays.jl +++ b/src/GPUArrays.jl @@ -51,6 +51,7 @@ include("host/sparse/reductions.jl") include("host/sparse/linalg.jl") include("host/sparse/structure.jl") include("host/sparse/indexing.jl") +include("host/sparse/array.jl") include("host/alloc_cache.jl") diff --git a/src/host/sparse/array.jl b/src/host/sparse/array.jl new file mode 100644 index 000000000..aad277287 --- /dev/null +++ b/src/host/sparse/array.jl @@ -0,0 +1,254 @@ +# array semantics that SparseArrays implements with scalar loops on the host + +## equality + +# compare the values, not the stored structure, so that stored zeros do not matter +Base.:(==)(A::GPUSparseArray, B::GPUSparseArray) = size(A) == size(B) && !any(A .!= B) +Base.:(==)(A::GPUSparseArray, B::AnyGPUArray) = size(A) == size(B) && !any(A .!= B) +Base.:(==)(A::AnyGPUArray, B::GPUSparseArray) = B == A + + +## in-place broadcast + +# A broadcast into a sparse destination is materialized first, which gives it the +# structure of the result (and copes with the destination being an argument), and then +# copied over. A result that is not sparse (dense arguments, or a function that does not +# preserve zeros) cannot be stored. +function sparse_broadcast!(dest::GPUSparseArray, bc::Broadcasted) + bc = Broadcast.instantiate(Broadcast.broadcasted(bc.f, bc.args...)) + if isempty(axes(bc)) + # `A .= x`: only a zero can fill a sparse array + fill!(dest, bc[]) + return dest + end + axes(dest) == axes(bc) || + throw(DimensionMismatch("destination has axes $(axes(dest)), result has axes $(axes(bc))")) + result = copy(bc) + result isa GPUSparseArray || + throw(ArgumentError("cannot broadcast a dense result into a sparse array")) + if dest isa GPUSparseMatrix + return copyto!(dest, convert(matrix_format(dest), result)) + else + return copyto!(dest, result) + end +end +Broadcast.materialize!(dest::GPUSparseArray, bc::Broadcasted) = sparse_broadcast!(dest, bc) +Base.copyto!(dest::GPUSparseArray, bc::Broadcasted) = sparse_broadcast!(dest, bc) +Base.copyto!(dest::GPUSparseArray, bc::Broadcasted{Nothing}) = sparse_broadcast!(dest, bc) +# (disambiguation with the dense GPU and the zero-dimensional methods) +for Style in (AbstractGPUArrayStyle, Broadcast.AbstractArrayStyle{0}, AbstractGPUArrayStyle{0}) + @eval Base.copyto!(dest::GPUSparseArray, bc::Broadcasted{<:$Style}) = + sparse_broadcast!(dest, bc) +end + +function Base.fill!(A::GPUSparseArray, x) + iszero(x) || throw(ArgumentError("a sparse array can only be filled with zeros")) + fill!(nonzeros(A), x) + return A +end + + +## linear indexing + +Base.getindex(A::GPUSparseMatrix, ::Colon) = vec(A) +Base.getindex(A::GPUSparseMatrix, I::AbstractRange{<:Integer}) = vec(A)[I] + + +## lazy wrappers + +# Operations with transposed, Symmetric and Hermitian sparse matrices materialize the +# wrapper, as the sparse kernels only index the formats themselves. (Products are the +# exception; see `linalg.jl`.) +const GPUSparseTransposed = Union{Transpose{<:Any,<:GPUSparseMatrix}, + Adjoint{<:Any,<:GPUSparseMatrix}} +const GPUSparseStructured = Union{Symmetric{<:Any,<:GPUSparseMatrix}, + Hermitian{<:Any,<:GPUSparseMatrix}} +const GPUSparseWrapped = Union{GPUSparseTransposed,GPUSparseStructured} + +for op in (:+, :-) + @eval begin + Base.$op(A::GPUSparseWrapped, B::GPUSparseWrapped) = + sparse_map($op, materialize(A), materialize(B)) + Base.$op(A::GPUSparseMatrix, B::GPUSparseWrapped) = sparse_map($op, A, materialize(B)) + Base.$op(A::GPUSparseWrapped, B::GPUSparseMatrix) = sparse_map($op, materialize(A), B) + + # (disambiguation with SparseArrays) + Base.$op(A::Symmetric{<:Real,<:GPUSparseMatrix}, B::Hermitian{<:Any,<:GPUSparseMatrix}) = + sparse_map($op, materialize(A), materialize(B)) + Base.$op(A::Hermitian{<:Any,<:GPUSparseMatrix}, B::Symmetric{<:Real,<:GPUSparseMatrix}) = + sparse_map($op, materialize(A), materialize(B)) + end +end + +# (LinearAlgebra keeps Symmetric and Hermitian wrappers when scaling by a real number, +# which works through the parent, but not a Hermitian matrix by a complex one) +for (S, T) in ((:GPUSparseTransposed, :Number), (:(Hermitian{<:Any,<:GPUSparseMatrix}), :Complex)) + @eval begin + Base.:(*)(A::$S, x::$T) = materialize(A) .* x + Base.:(*)(x::$T, A::$S) = x .* materialize(A) + Base.:(/)(A::$S, x::$T) = materialize(A) ./ x + Base.:(\)(x::$T, A::$S) = x .\ materialize(A) + end +end + +# (Julia 1.10's `power_by_squaring` would convert the wrapper to the inferred, non-concrete +# type of its product) +Base.:(^)(A::GPUSparseTransposed, p::Integer) = materialize(A)^p + +SparseArrays.nnz(A::GPUSparseTransposed) = nnz(parent(A)) +SparseArrays.nonzeros(A::Transpose{<:Any,<:GPUSparseMatrix}) = nonzeros(parent(A)) +SparseArrays.nonzeros(A::Adjoint{<:Any,<:GPUSparseMatrix}) = conj.(nonzeros(parent(A))) +Base.mapreduce(f, op, A::GPUSparseTransposed; kwargs...) = mapreduce(f, op, copy(A); kwargs...) +Base.getindex(A::GPUSparseTransposed, I::SparseRangeIndex, J::SparseRangeIndex) = copy(A)[I, J] + +SparseArrays.sparse(A::GPUSparseWrapped) = materialize(A) +SparseArrays.SparseMatrixCSC(A::GPUSparseWrapped) = SparseMatrixCSC(materialize(A)) +SparseArrays.SparseMatrixCSC{Tv}(A::GPUSparseWrapped) where {Tv} = + SparseMatrixCSC{Tv}(SparseMatrixCSC(A)) +SparseArrays.SparseMatrixCSC{Tv,Ti}(A::GPUSparseWrapped) where {Tv,Ti} = + SparseMatrixCSC{Tv,Ti}(SparseMatrixCSC(A)) +Base.Array(A::GPUSparseWrapped) = Array(materialize(A)) +Base.collect(A::GPUSparseWrapped) = Array(A) + +# (LinearAlgebra's constructors would convert the parent through scalar indexing. The +# element type is taken from the wrapper type because Julia 1.10 cannot bind it from +# `::Type{Symmetric{T}}`.) +(W::Type{Symmetric{T}} where {T})(A::GPUSparseMatrix, uplo::Symbol=:U) = + Symmetric(with_eltypes(A, eltype(W), indtype(A)), uplo) +(W::Type{Hermitian{T}} where {T})(A::GPUSparseMatrix, uplo::Symbol=:U) = + Hermitian(with_eltypes(A, eltype(W), indtype(A)), uplo) + + +## diagonal matrices and uniform scaling + +# a diagonal matrix with the format and storage of `A` +function sparse_diagonal(A::GPUSparseMatrix, d::AbstractVector, dims::Dims{2}=size(A)) + length(d) == min(dims...) || + throw(DimensionMismatch("a $(dims[1])×$(dims[2]) matrix has $(min(dims...)) diagonal entries, got $(length(d))")) + D = spdiagm(dims..., 0 => move_buffer(nonzeros(A), d)) + return convert(matrix_format(A), with_eltypes(D, eltype(D), indtype(A))) +end +sparse_identity(A::GPUSparseMatrix, λ) = + sparse_diagonal(A, fill!(similar(nonzeros(A), typeof(λ), min(size(A)...)), λ)) + +Base.one(A::GPUSparseMatrix) = (LinearAlgebra.checksquare(A); sparse_identity(A, one(eltype(A)))) +Base.oneunit(A::GPUSparseMatrix) = (LinearAlgebra.checksquare(A); sparse_identity(A, oneunit(eltype(A)))) + +const GPUSparseMatrixOrWrapped = Union{GPUSparseMatrix,GPUSparseWrapped} + +# (unlike with `UniformScaling`, which SparseArrays also adds to rectangular matrices) +check_diagonal_dims(A, D::Diagonal) = size(A) == size(D) || + throw(DimensionMismatch("cannot combine a $(join(size(A), '×')) matrix with a $(join(size(D), '×')) diagonal matrix")) + +# wrappers compare by value too +Base.:(==)(A::GPUSparseWrapped, B::GPUSparseMatrixOrWrapped) = materialize(A) == materialize(B) +Base.:(==)(A::GPUSparseMatrix, B::GPUSparseWrapped) = A == materialize(B) +Base.:(*)(A::GPUSparseMatrixOrWrapped, J::UniformScaling) = materialize(A) .* J.λ +Base.:(*)(J::UniformScaling, A::GPUSparseMatrixOrWrapped) = J.λ .* materialize(A) +for op in (:+, :-) + @eval begin + function Base.$op(A::GPUSparseMatrixOrWrapped, J::UniformScaling) + B = materialize(A) + sparse_map($op, B, sparse_identity(B, J.λ)) + end + function Base.$op(J::UniformScaling, A::GPUSparseMatrixOrWrapped) + B = materialize(A) + sparse_map($op, sparse_identity(B, J.λ), B) + end + function Base.$op(A::GPUSparseMatrixOrWrapped, D::Diagonal) + B = materialize(A) + check_diagonal_dims(B, D) + sparse_map($op, B, sparse_diagonal(B, D.diag)) + end + function Base.$op(D::Diagonal, A::GPUSparseMatrixOrWrapped) + B = materialize(A) + check_diagonal_dims(B, D) + sparse_map($op, sparse_diagonal(B, D.diag), B) + end + end +end +# (disambiguation with LinearAlgebra's Symmetric and Hermitian methods) +for op in (:+, :-), (Wrapper, Tλ, Td) in ((:Symmetric, :Number, :Number), + (:Hermitian, :Complex, :Real)) + @eval begin + Base.$op(A::$Wrapper{<:Any,<:GPUSparseMatrix}, J::UniformScaling{<:$Tλ}) = + invoke($op, Tuple{GPUSparseMatrixOrWrapped,UniformScaling}, A, J) + Base.$op(J::UniformScaling{<:$Tλ}, A::$Wrapper{<:Any,<:GPUSparseMatrix}) = + invoke($op, Tuple{UniformScaling,GPUSparseMatrixOrWrapped}, J, A) + Base.$op(A::$Wrapper{<:Any,<:GPUSparseMatrix}, D::Diagonal{<:$Td}) = + invoke($op, Tuple{GPUSparseMatrixOrWrapped,Diagonal}, A, D) + Base.$op(D::Diagonal{<:$Td}, A::$Wrapper{<:Any,<:GPUSparseMatrix}) = + invoke($op, Tuple{Diagonal,GPUSparseMatrixOrWrapped}, D, A) + end +end + +# scaling the rows or columns of a sparse matrix keeps its structure +function scale_sparse(A::GPUSparseMatrixOrWrapped, D::Diagonal, side::Symbol) + B = materialize(A) + n = side === :right ? size(B, 2) : size(B, 1) + size(D, 1) == n || + throw(DimensionMismatch("cannot multiply a $(join(size(B), '×')) matrix with a $(join(size(D), '×')) diagonal matrix on the $side")) + d = move_buffer(nonzeros(B), D.diag) + rows, cols = entry_coordinates(B) + vals = side === :right ? nonzeros(B) .* d[cols] : d[rows] .* nonzeros(B) + return sparse_format_type(B)(map(copy, Base.front(sparse_buffers(B)))..., vals, size(B)) +end +Base.:(*)(A::GPUSparseMatrixOrWrapped, D::Diagonal) = scale_sparse(A, D, :right) +Base.:(*)(D::Diagonal, A::GPUSparseMatrixOrWrapped) = scale_sparse(A, D, :left) + + +## concatenation + +# The coordinates of the entries are offset and concatenated, and the result assembled +# from them: there are no duplicates, and the entries of each array come out in order. +function sparse_cat(dim::Integer, As::GPUSparseMatrix...) + A = first(As) + other = 3 - dim + all(B -> size(B, other) == size(A, other), As) || + throw(DimensionMismatch("mismatch in dimension $other of the concatenated matrices")) + Ti = mapreduce(indtype, promote_type, As) + Tv = mapreduce(eltype, promote_type, As) + dims = ntuple(d -> d == dim ? sum(B -> size(B, dim), As) : size(A, d), 2) + check_sparse_dims(Ti, dims) + Is, Js, Vs = [], [], [] + offset = 0 + for B in As + rows, cols = entry_coordinates(B) + push!(Is, rows .+ Ti(dim == 1 ? offset : 0)) + push!(Js, cols .+ Ti(dim == 2 ? offset : 0)) + push!(Vs, convert_buffer(Tv, nonzeros(B))) + offset += size(B, dim) + end + return assemble_matrix(reduce(vcat, Is), reduce(vcat, Js), reduce(vcat, Vs), dims..., +, + matrix_format_symbol(A)) +end +matrix_format_symbol(::GPUSparseMatrixCSR) = :csr +matrix_format_symbol(::GPUSparseMatrixCSC) = :csc +matrix_format_symbol(::GPUSparseMatrixCOO) = :coo + +Base.hcat(A::GPUSparseMatrixOrWrapped, As::GPUSparseMatrixOrWrapped...) = + sparse_cat(2, materialize(A), map(materialize, As)...) +Base.vcat(A::GPUSparseMatrixOrWrapped, As::GPUSparseMatrixOrWrapped...) = + sparse_cat(1, materialize(A), map(materialize, As)...) + +# vectors concatenate to a vector, or side by side to a matrix, as in SparseArrays +function Base.vcat(x::GPUSparseVector, rest::GPUSparseVector...) + xs = (x, rest...) + Ti = mapreduce(indtype, promote_type, xs) + Tv = mapreduce(eltype, promote_type, xs) + len = sum(length, xs) + check_sparse_dims(Ti, (len,)) + inds, vals = [], [] + offset = 0 + for x in xs + push!(inds, x.nzInd .+ Ti(offset)) + push!(vals, convert_buffer(Tv, nonzeros(x))) + offset += length(x) + end + return GPUSparseVector(reduce(vcat, inds), reduce(vcat, vals), len) +end +Base.hcat(x::GPUSparseVector, xs::GPUSparseVector...) = + sparse_cat(2, GPUSparseMatrixCSC(x), map(GPUSparseMatrixCSC, xs)...) + +SparseArrays.blockdiag(A::GPUSparseMatrix, Bs::GPUSparseMatrix...) = + throw(ArgumentError("blockdiag of GPU sparse arrays is not supported")) diff --git a/test/testsuite/sparse.jl b/test/testsuite/sparse.jl index 96dcab5ee..47fca9711 100644 --- a/test/testsuite/sparse.jl +++ b/test/testsuite/sparse.jl @@ -30,6 +30,7 @@ using GPUArrays: GPUSparseMatrixCSR, GPUSparseMatrixCSC, GPUSparseMatrixCOO, GPU sparse_linalg(AT, eltypes) iszero_sparse(AT, eltypes) sparse_bool(AT) + sparse_semantics(AT, eltypes) sparse_generic(AT, eltypes) end @@ -1482,6 +1483,275 @@ function sparse_spgemm(AT, eltypes) end end +@kernel function sparse_lookup_kernel(out, A) + i = @index(Global, Linear) + out[i] = A[i, i] +end + +# the composition of affine maps modulo a prime: associative but not commutative, and +# represented by an isbits integer that every back-end can handle +function compose_affine(x::UInt32, y::UInt32) + a, b = x & 0xffff, x >> 16 + c, d = y & 0xffff, y >> 16 + ((a * c) % 251) | (((a * d + b) % 251) << 16) +end + +# semantics that SparseArrays implements with scalar loops, and corner cases +function sparse_semantics(AT, eltypes) + @testset "semantics" begin + @testset "$S" for S in sparse_matrix_formats + host = sparse(Float32[1 0 0; 2 0 0; 0 0 3]) + A = gpu_sparse(AT, S, host) + v = gpu_sparse(AT, sparsevec([1, 3], Float32[2, 4], 3)) + dense_v = AT(Float32[2, 3, 4]) + + # equality compares values, not stored zeros + @test A == copy(A) + @test A == AT(Array(host)) + @test AT(Array(host)) == A + @test A != gpu_sparse(AT, S, sparse(Float32[1 0 0; 2 0 0; 0 0 4])) + @test A == gpu_sparse(AT, S, SparseMatrixCSC(3, 3, [1, 3, 4, 5], [1, 2, 1, 3], Float32[1, 2, 0, 3])) + # (also through wrappers) + Dg = gpu_sparse(AT, S, sparse(Float32[1 0; 0 2])) + @test Dg == transpose(Dg) && transpose(Dg) == Dg && Dg == Symmetric(Dg) + @test transpose(Dg) == adjoint(Dg) + @test A != transpose(A) + + # in-place broadcast + B = copy(A) + B .*= 2 + @test Array(B) == 2 .* Array(host) + B .= B .+ A + @test Array(B) == 3 .* Array(host) + B .= 0 + @test iszero(B) && nnz(B) == nnz(A) + @test_throws ArgumentError fill!(B, 1) + before = copy(A) + @test_throws ArgumentError (A .= A .+ 1) + @test A == before + C = 2f0 .* gpu_sparse(AT, S, sparse(Float32[1 2 3; 0 0 0; 0 0 0])) .+ + 3f0 .* gpu_sparse(AT, S, sparse(Float32[0 0 0; 4 0 5; 0 6 0])) + @test Array(check_structure(C)) == Float32[2 4 6; 12 0 15; 0 18 0] + + # unsupported broadcasts fail instead of running on the host + @test_throws ErrorException A .* v + @test_throws ErrorException v .* A + @test_throws ErrorException (A .= A .* v) + @test A == before + scalar = AT(fill(0f0)) + @test_throws ArgumentError A .+ scalar + @test_throws ArgumentError v .+ scalar + for shape in ((3, 3, 1), (3, 3, 2)) + @test_throws ArgumentError A .+ AT(zeros(Float32, shape)) + @test_throws DimensionMismatch (A .= A .+ AT(zeros(Float32, shape))) + end + @test Array(A .* dense_v) == Array(host) .* Float32[2, 3, 4] + @test Array(dense_v .+ A) == Float32[2, 3, 4] .+ Array(host) + row = gpu_sparse(AT, S, sparse(reshape(Float32[1, 0, 2], 1, 3))) + @test_throws ArgumentError row .+ AT(Float32[2, 3]) + + # lookup in kernels, and from the host + out = AT(zeros(Float32, 3)) + sparse_lookup_kernel(get_backend(out))(out, A; ndrange=3) + @test Array(out) == diag(host) + @test Array(A[:]) == vec(Array(host)) + @test Array(A[2:5]) == vec(Array(host))[2:5] + mask = Bool[true, false, true] + for I in (mask, AT(mask)) + @test Array(A[I, I]) == Array(host)[mask, mask] + @test Array(A[I, 1:2]) == Array(host)[mask, 1:2] + @test Array(A[:, I]) == Array(host)[:, mask] + end + @test_throws BoundsError A[AT(trues(2)), AT(trues(3))] + + # lazy wrappers + @test nnz(transpose(A)) == nnz(A) + @test sum(transpose(A)) == sum(host) + @test Array(sum(adjoint(A); dims=1)) == sum(adjoint(host); dims=1) + @test Array(transpose(A)[1:2, 2:3]) == Array(transpose(host))[1:2, 2:3] + @test SparseMatrixCSC(transpose(A)) == transpose(host) + for wrap in (transpose, adjoint) + @test Array(2 * wrap(A)) == 2 * wrap(Array(host)) + @test Array(wrap(A) * 2) == wrap(Array(host)) * 2 + @test Array(wrap(A) / 2) == wrap(Array(host)) / 2 + @test Array(wrap(A)^2) ≈ wrap(Array(host))^2 + end + @test Array(Hermitian(A) * 2im) == Hermitian(Array(host)) * 2im + @test Array(2im \ Hermitian(A)) ≈ 2im \ Hermitian(Array(host)) + for wrap in (transpose, adjoint, Symmetric, Hermitian), op in (+, -) + @test Array(op(A, wrap(A))) == op(Array(host), wrap(Array(host))) + @test Array(op(wrap(A), A)) == op(wrap(Array(host)), Array(host)) + @test Array(op(wrap(A), wrap(A))) == op(wrap(Array(host)), wrap(Array(host))) + end + for wrap in (Symmetric, Hermitian) + @test sparse(wrap(A)) isa S + @test Array(sparse(wrap(A))) == Array(wrap(host)) + @test SparseMatrixCSC(wrap(A)) == sparse(wrap(host)) + @test Array(wrap(A)) == Array(wrap(host)) + @test wrap{Float16}(A) isa wrap{Float16,<:S{Float16}} + end + # (excluded entries are skipped, not multiplied by zero) + for wrap in (Symmetric, Hermitian), uplo in (:U, :L) + H = sparse(Float32[1 2; 4 3]) + @test Array(wrap(gpu_sparse(AT, S, H), uplo) * AT(Float32[Inf, 1])) == + wrap(H, uplo) * Float32[Inf, 1] + end + + # uniform scaling and diagonal matrices + @test Array(one(A)) == Matrix{Float32}(I, 3, 3) + @test Array(oneunit(A)) == Matrix{Float32}(I, 3, 3) + for J in (I, 2.5f0 * I), op in (+, -) + @test Array(op(A, J)) == op(Array(host), J) + @test Array(op(J, A)) == op(J, Array(host)) + @test Array(op(transpose(A), J)) == op(transpose(Array(host)), J) + end + @test Array(A * 2I) == Array(host) * 2 + @test Array(2I * A) == 2 * Array(host) + D = Diagonal(AT(Float32[2, 3, 4])) + Dh = Diagonal(Float32[2, 3, 4]) + for op in (+, -) + @test Array(op(A, D)) == op(Array(host), Dh) + @test Array(op(D, A)) == op(Dh, Array(host)) + end + @test Array(A * D) == Array(host) * Dh + @test Array(D * A) == Dh * Array(host) + @test (A * D) isa S + @test_throws DimensionMismatch A * Diagonal(AT(Float32[1, 2])) + R = gpu_sparse(AT, S, sparse(Float32[1 2 3; 4 5 6])) + # (SparseArrays adds `I` to rectangular matrices, but not a `Diagonal`) + @test Array(R + I) == Array(sparse(Float32[1 2 3; 4 5 6]) + I) + @test_throws DimensionMismatch R + Diagonal(AT(Float32[1, 2])) + @test_throws DimensionMismatch Diagonal(AT(Float32[1, 2])) - R + @test_throws DimensionMismatch one(R) + + # concatenation + host2 = sparse(Float32[0 1 0; 0 0 5; 6 0 0]) + A2 = gpu_sparse(AT, S, host2) + for f in (hcat, vcat) + C = f(A, A2, transpose(A)) + @test C isa S + @test Array(check_structure(C)) == f(Array(host), Array(host2), transpose(Array(host))) + @test Array(f(A, gpu_sparse(AT, GPUSparseMatrixCOO, host2))) == f(Array(host), Array(host2)) + end + @test_throws DimensionMismatch hcat(A, gpu_sparse(AT, S, sparse(ones(Float32, 2, 2)))) + x = gpu_sparse(AT, sparsevec([2], Float32[3], 3)) + @test Array(vcat(v, x)) == vcat(Array(v), Array(x)) + @test Array(hcat(v, x)) == hcat(Array(v), Array(x)) + @test_throws ArgumentError blockdiag(A, A) + + # norms are scaled, and negative orders see the implicit zeros + for val in (1f30, 1f-30) + @test norm(gpu_sparse(AT, S, sparse([1], [1], [val], 1, 1))) ≈ val + end + @test norm(A, -Inf) == 0 + @test norm(gpu_sparse(AT, S, sparse(Float32[1 2; 3 4])), -Inf) == 1 + N = gpu_sparse(AT, S, sparse(Float32[NaN 0; 0 1])) + @test isnan(norm(N, -1)) + + # dot products check their dimensions + x = AT(Float32[1, 2, 3]) + @test dot(x, A, x) ≈ dot(Float32[1, 2, 3], host, Float32[1, 2, 3]) + @test_throws DimensionMismatch dot(AT(Float32[1, 2]), A, x) + @test_throws DimensionMismatch dot(x, A, AT(Float32[1, 2])) + + # the trace is accumulated in the element type + T32 = gpu_sparse(AT, S, sparse(Int32[typemax(Int32) 0; 0 typemax(Int32)])) + @test tr(T32) === tr(sparse(Int32[typemax(Int32) 0; 0 typemax(Int32)])) + end + + @testset "mutation" begin + # a view cannot be resized, so structural mutation fails before changing + # anything + p = AT([1, 2, 3]); i = AT([1, 2]); v = AT(Float32[1, 7, 0, 9]) + A = GPUSparseMatrixCSC(p, i, view(v, 1:2:4), (2, 2)) + before = (Array(p), Array(i), Array(v)) + @test_throws ArgumentError dropzeros!(A) + @test (Array(p), Array(i), Array(v)) == before + src = GPUSparseMatrixCSC(AT([1, 2, 2]), AT([1]), AT(Float32[1]), (2, 2)) + @test_throws ArgumentError copyto!(A, src) + @test (Array(p), Array(i), Array(v)) == before + # values can still be updated through a view + A .*= 2 + @test Array(v) == [2., 7., 0., 9.] + + # narrowing values is checked (the error depends on the back-end) + x = GPUSparseVector(AT(Int32[1]), AT(Int64[300]), 2) + @test_throws Exception Array(nonzeros(GPUSparseVector{Int8,Int32}(x))) + dest = GPUSparseVector(AT(Int32[1]), AT(Int8[1]), 2) + @test_throws Exception Array(nonzeros(dest .= 300 .* dest)) + + # aliasing is detected through the buffers + A = gpu_sparse(AT, GPUSparseMatrixCSC, sparse(Float32[1 0; 2 3])) + B = GPUSparseMatrixCSC(A.colPtr, A.rowVal, copy(A.nzVal), size(A)) + @test Base.mightalias(A, B) + @test !Base.mightalias(A, copy(A)) + copyto!(B, A .* 2) + @test Array(B) == 2 .* Array(A) + # densifying into a buffer of the source + w = AT(Float32[2, 3]) + x = GPUSparseVector(AT([1, 2]), w, 2) + copyto!(w, x) + @test Array(w) == [2, 3] + # and into a view + dest = AT(zeros(Float32, 4)) + copyto!(view(dest, 1:2:4), GPUSparseVector(AT([2]), AT(Float32[5]), 2)) + @test Array(dest) == [0, 0, 5, 0] + end + + @testset "index types" begin + # unsigned indices (the diagonal offsets are negative below the diagonal) + U = GPUSparseMatrixCSR(AT(UInt32[1, 1, 2]), AT(UInt32[1]), AT(Float32[3]), (2, 2)) + @test nnz(triu(U)) == 0 && nnz(tril(U)) == 1 + @test !istriu(U) && istril(U) + # the largest index is a valid coordinate + x = GPUSparseVector(AT(Int8[127]), AT(Float32[2]), 127) + y = x .+ x + @test nnz(y) == 1 + @test Array(y)[end] == 4 + @test Array(GPUArrays.compress(AT(Int8[127]), 127))[end] == 2 + # dimensions and entries must fit the index type + coo = GPUSparseMatrixCOO(AT(Int64[]), AT(Int64[]), AT(Float32[]), (2^32, 1)) + @test_throws ArgumentError GPUSparseMatrixCSR{Float32,Int32}(coo) + @test_throws ArgumentError GPUSparseMatrixCSC{Float32,Int32}(transpose(coo)) + @test_throws ArgumentError GPUSparseMatrixCSR{Float32,Int8}(Diagonal(AT(ones(Float32, 127)))) + D = GPUSparseMatrixCSR{Float32,Int8}(Diagonal(AT(zeros(Float32, 126)))) + @test nnz(D) == 126 && size(D) == (126, 126) && iszero(D) + # mixed-width inputs are promoted + A = GPUSparseMatrixCSR(AT(Int32[1, 2, 3]), AT(Int32[1, 2]), AT(Float32[1, 2]), (2, 2)) + B = GPUSparseMatrixCSR(AT(Int64[1, 2, 3]), AT(Int64[2, 1]), AT(Float32[1, 2]), (2, 2)) + @test SparseArrays.indtype(A .+ B) == Int64 + @test SparseArrays.indtype(A * B) == Int64 + @test Array(A * B) == Array(A) * Array(B) + # Int32 dimensions + @test Array(sparse(AT(Int32[1]), AT(Int32[1]), AT(Float32[2]), Int32(1), Int32(1))) == reshape(Float32[2], 1, 1) + end + + @testset "column views" begin + A = GPUSparseMatrixCSC(AT(Int32[1, 3, 3, 4]), AT(Int32[1, 3, 2]), AT(Float32[2, 4, 6]), (3, 3)) + @allowscalar for j in 1:3 + column = view(A, :, j) + @test nnz(column) == (j == 1 ? 2 : j == 2 ? 0 : 1) + @test Array(nonzeros(column)) == Array(nonzeros(A))[nzrange(A, j)] + @test Array(rowvals(column)) == Array(rowvals(A))[nzrange(A, j)] + end + end + + @testset "assembly order" begin + I = AT(Int32[2, 1, 2, 1, 2, 1]) + V = AT(UInt32[2 | (3 << 16), 4 | (5 << 16), 6 | (7 << 16), + 8 | (9 << 16), 10 | (11 << 16), 12 | (13 << 16)]) + x = sparsevec(I, V, 3, compose_affine) + @test Array(x) == Array(sparsevec(Array(I), Array(V), 3, compose_affine)) + end + + @testset "deprecated names" begin + A = gpu_sparse(AT, GPUSparseMatrixCSR, sparse(Int32[1, 2], Int32[1, 2], Float32[1, 1])) + @test A isa GPUArrays.AbstractGPUSparseMatrixCSR + @test A isa GPUArrays.AbstractGPUSparseMatrixCSR{Float32,Int32} + end + end +end + # the generic implementations, which back-ends with vendor libraries do not reach through # the LinearAlgebra interface function sparse_generic(AT, eltypes)