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