From fc0b33ba84954486ae56f40e2857c0336b92fb22 Mon Sep 17 00:00:00 2001 From: Michel Schanen Date: Thu, 24 Sep 2026 10:34:09 -0500 Subject: [PATCH 1/2] Make oneMKL sparse matrices AbstractGPUSparseArrays Subtype GPUArrays' AbstractGPUSparseMatrixCSR/CSC/COO so that the generic sparse functionality of GPUArrays (broadcast, mapreduce, norms, findnz, triu/tril/kron, indexing, similar/copy) applies to oneSparseMatrixCSR/CSC/COO, and run the GPUArrays sparse testsuite on the CSR and CSC types. The oneMKL matrix handle is now created lazily on the first oneMKL operation and invalidated when the storage vectors are replaced (copyto!, unsafe_free!), since GPUArrays' generic code constructs sparse matrices freely and resizes their storage. As a consequence any element type can be stored; the oneMKL operations still require Float32/Float64/ComplexF32/ComplexF64 values with Int32/Int64 indices and error otherwise. Conversions between the three formats, transposition/adjoint, sparse addition and the COO structural operations (triu, tril, kron, reshape, droptol!) are implemented on the device with generic operations (sortperm, broadcast). adapt(oneArray, ::SparseMatrixCSC) now returns a oneSparseMatrixCSC, matching CUDA.jl and AMDGPU.jl. Fixes #627. --- docs/src/onemkl.md | 31 +++- lib/mkl/array.jl | 279 +++++++++++++++++++++++++++++++-- lib/mkl/oneMKL.jl | 1 + lib/mkl/sparse_conversions.jl | 206 +++++++++++++++++++++++++ lib/mkl/wrappers_sparse.jl | 282 ++++++++++++++++++---------------- src/oneAPIKernels.jl | 9 ++ test/onemkl.jl | 141 ++++++++++++++++- test/runtests.jl | 14 ++ 8 files changed, 812 insertions(+), 151 deletions(-) create mode 100644 lib/mkl/sparse_conversions.jl diff --git a/docs/src/onemkl.md b/docs/src/onemkl.md index 640f2a68..c98e826f 100644 --- a/docs/src/onemkl.md +++ b/docs/src/onemkl.md @@ -40,10 +40,33 @@ y = dA * x ``` Three storage formats are available: `oneSparseMatrixCSR`, `oneSparseMatrixCSC` and -`oneSparseMatrixCOO`. oneMKL's sparse back-end is CSR-based, and a `oneSparseMatrixCSC` is -therefore stored as the CSR representation of its transpose. As a consequence the triangular -operations (`sparse_trmv!`, `sparse_trsv!`, `sparse_trsm!`) cannot be expressed for CSC -matrices and throw an `ArgumentError`. Prefer CSR when you have the choice. +`oneSparseMatrixCOO`. They are subtypes of the corresponding GPUArrays.jl abstract types +(`AbstractGPUSparseMatrixCSR`, `AbstractGPUSparseMatrixCSC`, `AbstractGPUSparseMatrixCOO`), so the +generic sparse functionality of GPUArrays.jl is available: broadcasting (zero-preserving functions +return a sparse matrix, others a dense `oneArray`), `sum`/`mapreduce` (also along a dimension), +`norm`/`opnorm`, `findnz`, `triu`/`tril`/`kron`, `iszero`, and scalar indexing under +`GPUArrays.@allowscalar`. Matrices can be converted between the three formats, transposed and +added on the device, and `adapt(oneArray, A)` of a `SparseMatrixCSC` yields a `oneSparseMatrixCSC`. + +```julia +dA = oneSparseMatrixCSR(sprand(Float32, 100, 100, 0.1)) +dB = dA .* 2f0 # oneSparseMatrixCSR +dC = dA .+ 1f0 # dense oneMatrix +sum(dA; dims=1) # row vector +dAt = oneSparseMatrixCSC(dA) # format conversion, on the device +dS = dA + transpose(dA) # sparse addition +``` + +Any element type can be stored in these matrices, but the oneMKL operations (`*`, `mul!`, the +triangular solves, and the `sparse_*!` wrappers) require `Float32`, `Float64`, `ComplexF32` or +`ComplexF64` values with `Int32` or `Int64` indices. The oneMKL matrix handle is created lazily +when such an operation is first invoked; it refers to the storage vectors of the matrix, which +therefore must not be modified in place afterwards (use `copyto!` or create a new matrix instead). + +oneMKL's sparse back-end is CSR-based, and a `oneSparseMatrixCSC` is therefore handed to oneMKL as +the CSR representation of its transpose (this requires oneMKL 2025.3 or later). As a consequence +the triangular operations (`sparse_trmv!`, `sparse_trsv!`, `sparse_trsm!`) cannot be expressed +for CSC matrices and throw an `ArgumentError`. Prefer CSR when you have the choice. ## FFTs diff --git a/lib/mkl/array.jl b/lib/mkl/array.jl index 3f117bc3..7d85d718 100644 --- a/lib/mkl/array.jl +++ b/lib/mkl/array.jl @@ -1,49 +1,296 @@ export oneSparseMatrixCSR, oneSparseMatrixCSC, oneSparseMatrixCOO -abstract type oneAbstractSparseArray{Tv, Ti, N} <: AbstractSparseArray{Tv, Ti, N} end -const oneAbstractSparseVector{Tv, Ti} = oneAbstractSparseArray{Tv, Ti, 1} -const oneAbstractSparseMatrix{Tv, Ti} = oneAbstractSparseArray{Tv, Ti, 2} +using ..oneAPI: KernelAdaptor, oneDeviceVector, AS +using LinearAlgebra: Transpose, Adjoint +using SparseArrays: SparseVector, SparseMatrixCSC, nnz, nonzeros +import Adapt +import Adapt: adapt -mutable struct oneSparseMatrixCSR{Tv, Ti} <: oneAbstractSparseMatrix{Tv, Ti} +# The oneMKL sparse types are children of the GPUArrays sparse hierarchy, so all the generic +# functionality from GPUArrays (broadcast, mapreduce, norms, findnz, ...) applies to them. +# +# The `handle` field caches the oneMKL matrix handle. It is created lazily, on the first oneMKL +# operation (see `matrix_handle` in wrappers_sparse.jl), so that generic code can construct and +# convert sparse matrices freely without paying for oneMKL handle set-up, and so that the struct +# can hold element types oneMKL does not support. + +mutable struct oneSparseMatrixCSR{Tv, Ti} <: GPUArrays.AbstractGPUSparseMatrixCSR{Tv, Ti} handle::Union{Nothing, matrix_handle_t} rowPtr::oneVector{Ti} colVal::oneVector{Ti} nzVal::oneVector{Tv} - dims::NTuple{2,Int} + dims::NTuple{2, Int} nnz::Ti + + function oneSparseMatrixCSR{Tv, Ti}( + rowPtr::oneVector{Ti}, colVal::oneVector{Ti}, nzVal::oneVector{Tv}, + dims::NTuple{2, <:Integer} + ) where {Tv, Ti <: Integer} + A = new{Tv, Ti}(nothing, rowPtr, colVal, nzVal, Int.(dims), Ti(length(nzVal))) + return finalizer(sparse_release_matrix_handle, A) + end end -mutable struct oneSparseMatrixCSC{Tv, Ti} <: oneAbstractSparseMatrix{Tv, Ti} +mutable struct oneSparseMatrixCSC{Tv, Ti} <: GPUArrays.AbstractGPUSparseMatrixCSC{Tv, Ti} handle::Union{Nothing, matrix_handle_t} colPtr::oneVector{Ti} rowVal::oneVector{Ti} nzVal::oneVector{Tv} - dims::NTuple{2,Int} + dims::NTuple{2, Int} nnz::Ti + + function oneSparseMatrixCSC{Tv, Ti}( + colPtr::oneVector{Ti}, rowVal::oneVector{Ti}, nzVal::oneVector{Tv}, + dims::NTuple{2, <:Integer} + ) where {Tv, Ti <: Integer} + A = new{Tv, Ti}(nothing, colPtr, rowVal, nzVal, Int.(dims), Ti(length(nzVal))) + return finalizer(sparse_release_matrix_handle, A) + end end -mutable struct oneSparseMatrixCOO{Tv, Ti} <: oneAbstractSparseMatrix{Tv, Ti} +mutable struct oneSparseMatrixCOO{Tv, Ti} <: GPUArrays.AbstractGPUSparseMatrixCOO{Tv, Ti} handle::Union{Nothing, matrix_handle_t} rowInd::oneVector{Ti} colInd::oneVector{Ti} nzVal::oneVector{Tv} - dims::NTuple{2,Int} + dims::NTuple{2, Int} nnz::Ti + + function oneSparseMatrixCOO{Tv, Ti}( + rowInd::oneVector{Ti}, colInd::oneVector{Ti}, nzVal::oneVector{Tv}, + dims::NTuple{2, <:Integer} + ) where {Tv, Ti <: Integer} + A = new{Tv, Ti}(nothing, rowInd, colInd, nzVal, Int.(dims), Ti(length(nzVal))) + return finalizer(sparse_release_matrix_handle, A) + end end +const oneAbstractSparseMatrix{Tv, Ti} = Union{ + oneSparseMatrixCSR{Tv, Ti}, oneSparseMatrixCSC{Tv, Ti}, oneSparseMatrixCOO{Tv, Ti}, +} +const oneSparseMatrixAdjOrTrans = Union{ + Transpose{<:Any, <:oneAbstractSparseMatrix}, Adjoint{<:Any, <:oneAbstractSparseMatrix}, +} + +# untyped constructors from device vectors (GPUArrays' generic code relies on these) +oneSparseMatrixCSR( + rowPtr::oneVector{Ti}, colVal::oneVector{Ti}, nzVal::oneVector{Tv}, dims::NTuple{2, <:Integer} +) where {Tv, Ti <: Integer} = oneSparseMatrixCSR{Tv, Ti}(rowPtr, colVal, nzVal, dims) +oneSparseMatrixCSC( + colPtr::oneVector{Ti}, rowVal::oneVector{Ti}, nzVal::oneVector{Tv}, dims::NTuple{2, <:Integer} +) where {Tv, Ti <: Integer} = oneSparseMatrixCSC{Tv, Ti}(colPtr, rowVal, nzVal, dims) +oneSparseMatrixCOO( + rowInd::oneVector{Ti}, colInd::oneVector{Ti}, nzVal::oneVector{Tv}, dims::NTuple{2, <:Integer} +) where {Tv, Ti <: Integer} = oneSparseMatrixCOO{Tv, Ti}(rowInd, colInd, nzVal, dims) + +# the storage vectors of a matrix, in the order (pointer/first index, second index, values) +_storage(A::oneSparseMatrixCSR) = (A.rowPtr, A.colVal, A.nzVal) +_storage(A::oneSparseMatrixCSC) = (A.colPtr, A.rowVal, A.nzVal) +_storage(A::oneSparseMatrixCOO) = (A.rowInd, A.colInd, A.nzVal) + + +## GPUArrays interface + +GPUArrays.sparse_array_type(::Type{<:oneSparseMatrixCSR}) = oneSparseMatrixCSR +GPUArrays.sparse_array_type(::Type{<:oneSparseMatrixCSC}) = oneSparseMatrixCSC +GPUArrays.sparse_array_type(::Type{<:oneSparseMatrixCOO}) = oneSparseMatrixCOO + +GPUArrays.dense_array_type(::Type{<:oneAbstractSparseMatrix}) = oneArray + +GPUArrays.csr_type(::Type{<:Union{oneAbstractSparseMatrix, oneSparseMatrixAdjOrTrans}}) = oneSparseMatrixCSR +GPUArrays.csc_type(::Type{<:Union{oneAbstractSparseMatrix, oneSparseMatrixAdjOrTrans}}) = oneSparseMatrixCSC +GPUArrays.coo_type(::Type{<:Union{oneAbstractSparseMatrix, oneSparseMatrixAdjOrTrans}}) = oneSparseMatrixCOO + + +## array interface + Base.length(A::oneAbstractSparseMatrix) = prod(A.dims) Base.size(A::oneAbstractSparseMatrix) = A.dims -function Base.size(A::oneAbstractSparseMatrix, d::Integer) - if d == 1 || d == 2 - return A.dims[d] - else - throw(ArgumentError("dimension must be 1 or 2, got $d")) +# `similar` preserving the sparsity structure +Base.similar(A::oneSparseMatrixCSR{Tv, Ti}) where {Tv, Ti} = + oneSparseMatrixCSR{Tv, Ti}(copy(A.rowPtr), copy(A.colVal), similar(A.nzVal), size(A)) +Base.similar(A::oneSparseMatrixCSC{Tv, Ti}) where {Tv, Ti} = + oneSparseMatrixCSC{Tv, Ti}(copy(A.colPtr), copy(A.rowVal), similar(A.nzVal), size(A)) +Base.similar(A::oneSparseMatrixCOO{Tv, Ti}) where {Tv, Ti} = + oneSparseMatrixCOO{Tv, Ti}(copy(A.rowInd), copy(A.colInd), similar(A.nzVal), size(A)) + +Base.similar(A::oneSparseMatrixCSR{<:Any, Ti}, ::Type{T}) where {T, Ti} = + oneSparseMatrixCSR{T, Ti}(copy(A.rowPtr), copy(A.colVal), similar(A.nzVal, T), size(A)) +Base.similar(A::oneSparseMatrixCSC{<:Any, Ti}, ::Type{T}) where {T, Ti} = + oneSparseMatrixCSC{T, Ti}(copy(A.colPtr), copy(A.rowVal), similar(A.nzVal, T), size(A)) +Base.similar(A::oneSparseMatrixCOO{<:Any, Ti}, ::Type{T}) where {T, Ti} = + oneSparseMatrixCOO{T, Ti}(copy(A.rowInd), copy(A.colInd), similar(A.nzVal, T), size(A)) + +# `similar` with a different shape: an empty (all-zero) matrix of the same format +_empty_ptr(::Type{Ti}, len::Integer) where {Ti} = fill!(oneVector{Ti}(undef, len), one(Ti)) +Base.similar(::oneSparseMatrixCSR{<:Any, Ti}, ::Type{T}, m::Integer, n::Integer) where {T, Ti} = + oneSparseMatrixCSR{T, Ti}(_empty_ptr(Ti, m + 1), oneVector{Ti}(undef, 0), oneVector{T}(undef, 0), (m, n)) +Base.similar(::oneSparseMatrixCSC{<:Any, Ti}, ::Type{T}, m::Integer, n::Integer) where {T, Ti} = + oneSparseMatrixCSC{T, Ti}(_empty_ptr(Ti, n + 1), oneVector{Ti}(undef, 0), oneVector{T}(undef, 0), (m, n)) +Base.similar(::oneSparseMatrixCOO{<:Any, Ti}, ::Type{T}, m::Integer, n::Integer) where {T, Ti} = + oneSparseMatrixCOO{T, Ti}(oneVector{Ti}(undef, 0), oneVector{Ti}(undef, 0), oneVector{T}(undef, 0), (m, n)) + +Base.similar(A::oneAbstractSparseMatrix, ::Type{T}, dims::Dims{2}) where {T} = similar(A, T, dims...) +Base.similar(A::oneAbstractSparseMatrix{Tv}, m::Integer, n::Integer) where {Tv} = similar(A, Tv, m, n) +Base.similar(A::oneAbstractSparseMatrix{Tv}, dims::Dims{2}) where {Tv} = similar(A, Tv, dims...) + +# other dimensionalities (e.g. a column slice) are dense +Base.similar(::oneAbstractSparseMatrix, ::Type{T}, dims::Dims) where {T} = oneArray{T}(undef, dims) +Base.similar(A::oneAbstractSparseMatrix{Tv}, dims::Dims) where {Tv} = similar(A, Tv, dims) + +# scalar indexing (GPUArrays provides the CSC method) +function Base.getindex(A::oneSparseMatrixCSR{Tv}, i0::Integer, i1::Integer) where {Tv} + @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) + return A.nzVal[c1] +end + +function Base.getindex(A::oneSparseMatrixCOO{Tv}, i0::Integer, i1::Integer) where {Tv} + @boundscheck checkbounds(A, i0, i1) + # COO entries are not guaranteed to be sorted, so search for the entry + k = findfirst((A.rowInd .== i0) .& (A.colInd .== i1)) + k === nothing && return zero(Tv) + return A.nzVal[k] +end + +# non-scalar indexing goes through a dense copy: the generic fallback would launch a kernel that +# indexes the device-side sparse matrix, which GPUArrays does not implement (yet) +Base.getindex(A::oneAbstractSparseMatrix, ::Colon, ::Colon) = copy(A) +for I in (:Colon, :Integer, :AbstractVector), J in (:Colon, :Integer, :AbstractVector) + (I === :Integer && J === :Integer) && continue + (I === :Colon && J === :Colon) && continue + @eval Base.getindex(A::oneAbstractSparseMatrix, I::$I, J::$J) = oneArray(A)[I, J] +end + +# copying +Base.copy(A::oneSparseMatrixCSR{Tv, Ti}) where {Tv, Ti} = + oneSparseMatrixCSR{Tv, Ti}(copy(A.rowPtr), copy(A.colVal), copy(A.nzVal), size(A)) +Base.copy(A::oneSparseMatrixCSC{Tv, Ti}) where {Tv, Ti} = + oneSparseMatrixCSC{Tv, Ti}(copy(A.colPtr), copy(A.rowVal), copy(A.nzVal), size(A)) +Base.copy(A::oneSparseMatrixCOO{Tv, Ti}) where {Tv, Ti} = + oneSparseMatrixCOO{Tv, Ti}(copy(A.rowInd), copy(A.colInd), copy(A.nzVal), size(A)) + +# `copyto!` may change the sparsity structure, so it replaces the storage vectors (they may be +# shared with other matrices, e.g. the output of a single-input broadcast reuses the pointer +# array of its input) and drops the oneMKL handle, which refers to the old storage. +function Base.copyto!(dst::oneSparseMatrixCSR{Tv, Ti}, src::oneSparseMatrixCSR{Tv, Ti}) where {Tv, Ti} + size(dst) == size(src) || throw(ArgumentError("Inconsistent Sparse Matrix size")) + _invalidate_handle!(dst) + dst.rowPtr = copy(src.rowPtr) + dst.colVal = copy(src.colVal) + dst.nzVal = copy(src.nzVal) + dst.nnz = src.nnz + return dst +end +function Base.copyto!(dst::oneSparseMatrixCSC{Tv, Ti}, src::oneSparseMatrixCSC{Tv, Ti}) where {Tv, Ti} + size(dst) == size(src) || throw(ArgumentError("Inconsistent Sparse Matrix size")) + _invalidate_handle!(dst) + dst.colPtr = copy(src.colPtr) + dst.rowVal = copy(src.rowVal) + dst.nzVal = copy(src.nzVal) + dst.nnz = src.nnz + return dst +end +function Base.copyto!(dst::oneSparseMatrixCOO{Tv, Ti}, src::oneSparseMatrixCOO{Tv, Ti}) where {Tv, Ti} + size(dst) == size(src) || throw(ArgumentError("Inconsistent Sparse Matrix size")) + _invalidate_handle!(dst) + dst.rowInd = copy(src.rowInd) + dst.colInd = copy(src.colInd) + dst.nzVal = copy(src.nzVal) + dst.nnz = src.nnz + return dst +end + +# dense conversion, on the device (sparse .+ dense broadcasts to a dense array) +oneAPI.oneArray(A::Union{oneSparseMatrixCSR{Tv}, oneSparseMatrixCSC{Tv}}) where {Tv} = + A .+ fill!(similar(A.nzVal, Tv, size(A)), zero(Tv)) +oneAPI.oneArray(A::oneSparseMatrixCOO) = oneArray(collect(A)) + + +## interop with SparseArrays + +# CPU to GPU +function oneSparseMatrixCSR(A::SparseMatrixCSC{Tv, Ti}) where {Tv, Ti} + At = SparseMatrixCSC(transpose(A)) + return oneSparseMatrixCSR{Tv, Ti}( + oneVector{Ti}(At.colptr), oneVector{Ti}(At.rowval), oneVector{Tv}(At.nzval), size(A) + ) +end +oneSparseMatrixCSC(A::SparseMatrixCSC{Tv, Ti}) where {Tv, Ti} = + oneSparseMatrixCSC{Tv, Ti}( + oneVector{Ti}(A.colptr), oneVector{Ti}(A.rowval), oneVector{Tv}(A.nzval), size(A) +) +function oneSparseMatrixCOO(A::SparseMatrixCSC{Tv, Ti}) where {Tv, Ti} + row, col, val = findnz(A) + return oneSparseMatrixCOO{Tv, Ti}(oneVector{Ti}(row), oneVector{Ti}(col), oneVector{Tv}(val), size(A)) +end + +# transposes of CPU matrices: CSR(Aᵀ) is CSC(A) with the roles of the index arrays swapped +function oneSparseMatrixCSR(t::Transpose{Tv, <:SparseMatrixCSC{Tv, Ti}}) where {Tv, Ti} + A = parent(t) + return oneSparseMatrixCSR{Tv, Ti}( + oneVector{Ti}(A.colptr), oneVector{Ti}(A.rowval), oneVector{Tv}(A.nzval), size(t) + ) +end +function oneSparseMatrixCSR(t::Adjoint{Tv, <:SparseMatrixCSC{Tv, Ti}}) where {Tv, Ti} + A = parent(t) + return oneSparseMatrixCSR{Tv, Ti}( + oneVector{Ti}(A.colptr), oneVector{Ti}(A.rowval), oneVector{Tv}(conj.(A.nzval)), size(t) + ) +end +const SparseMatrixCSCAdjOrTrans = Union{Transpose{<:Any, <:SparseMatrixCSC}, Adjoint{<:Any, <:SparseMatrixCSC}} +oneSparseMatrixCSC(t::SparseMatrixCSCAdjOrTrans) = oneSparseMatrixCSC(SparseMatrixCSC(t)) +oneSparseMatrixCOO(t::SparseMatrixCSCAdjOrTrans) = oneSparseMatrixCOO(SparseMatrixCSC(t)) + +for X in (:oneSparseMatrixCSR, :oneSparseMatrixCSC, :oneSparseMatrixCOO) + @eval begin + # sparse vectors are single-column matrices + $X(v::SparseVector) = $X(SparseMatrixCSC(v)) + # element type conversion + $X{Tv}(A::SparseMatrixCSC{<:Any, Ti}) where {Tv, Ti} = $X(SparseMatrixCSC{Tv, Ti}(A)) + $X{Tv, Ti}(A::SparseMatrixCSC) where {Tv, Ti} = $X(SparseMatrixCSC{Tv, Ti}(A)) + $X{Tv}(A::Union{SparseVector, SparseMatrixCSCAdjOrTrans}) where {Tv} = $X{Tv}(SparseMatrixCSC(A)) end end -SparseArrays.nnz(A::oneAbstractSparseMatrix) = A.nnz -SparseArrays.nonzeros(A::oneAbstractSparseMatrix) = A.nzVal +# GPU to CPU (GPUArrays provides the CSC method) +function SparseArrays.SparseMatrixCSC(A::oneSparseMatrixCSR) + m, n = size(A) + At = SparseMatrixCSC(n, m, Array(A.rowPtr), Array(A.colVal), Array(A.nzVal)) + return SparseMatrixCSC(transpose(At)) +end +SparseArrays.SparseMatrixCSC(A::oneSparseMatrixCOO) = + sparse(Array(A.rowInd), Array(A.colInd), Array(A.nzVal), size(A)...) + + +## adapt + +Adapt.adapt_storage(::Type{oneArray}, xs::SparseMatrixCSC) = oneSparseMatrixCSC(xs) +Adapt.adapt_storage(::Type{<:oneArray{T}}, xs::SparseMatrixCSC) where {T} = oneSparseMatrixCSC{T}(xs) +Adapt.adapt_storage(::Type{oneArray}, xs::oneAbstractSparseMatrix) = xs +Adapt.adapt_storage(::Type{Array}, xs::oneAbstractSparseMatrix) = SparseMatrixCSC(xs) + +# device-side counterparts for use in kernels +Adapt.adapt_structure(to::KernelAdaptor, A::oneSparseMatrixCSR{Tv, Ti}) where {Tv, Ti} = + GPUArrays.GPUSparseDeviceMatrixCSR{ + Tv, Ti, oneDeviceVector{Ti, AS.CrossWorkgroup}, oneDeviceVector{Tv, AS.CrossWorkgroup}, AS.CrossWorkgroup, +}(adapt(to, A.rowPtr), adapt(to, A.colVal), adapt(to, A.nzVal), A.dims, A.nnz) +Adapt.adapt_structure(to::KernelAdaptor, A::oneSparseMatrixCSC{Tv, Ti}) where {Tv, Ti} = + GPUArrays.GPUSparseDeviceMatrixCSC{ + Tv, Ti, oneDeviceVector{Ti, AS.CrossWorkgroup}, oneDeviceVector{Tv, AS.CrossWorkgroup}, AS.CrossWorkgroup, +}(adapt(to, A.colPtr), adapt(to, A.rowVal), adapt(to, A.nzVal), A.dims, A.nnz) +Adapt.adapt_structure(to::KernelAdaptor, A::oneSparseMatrixCOO{Tv, Ti}) where {Tv, Ti} = + GPUArrays.GPUSparseDeviceMatrixCOO{ + Tv, Ti, oneDeviceVector{Ti, AS.CrossWorkgroup}, oneDeviceVector{Tv, AS.CrossWorkgroup}, AS.CrossWorkgroup, +}(adapt(to, A.rowInd), adapt(to, A.colInd), adapt(to, A.nzVal), A.dims, A.nnz) + + +## input/output for (gpu, cpu) in [:oneSparseMatrixCSR => :SparseMatrixCSC, :oneSparseMatrixCSC => :SparseMatrixCSC, diff --git a/lib/mkl/oneMKL.jl b/lib/mkl/oneMKL.jl index 64f54358..7e243f62 100644 --- a/lib/mkl/oneMKL.jl +++ b/lib/mkl/oneMKL.jl @@ -27,6 +27,7 @@ include("utils.jl") include("wrappers_blas.jl") include("wrappers_lapack.jl") include("wrappers_sparse.jl") +include("sparse_conversions.jl") include("linalg.jl") include("interfaces.jl") include("fft.jl") diff --git a/lib/mkl/sparse_conversions.jl b/lib/mkl/sparse_conversions.jl new file mode 100644 index 00000000..ffac7636 --- /dev/null +++ b/lib/mkl/sparse_conversions.jl @@ -0,0 +1,206 @@ +# conversions between sparse formats, transposition, and sparse addition +# +# Everything here is built from generic GPUArrays operations (broadcast, sortperm, gather), so +# it runs on the device for any element type and does not involve oneMKL. Sorting uses plain +# integer keys, which is the configuration AcceleratedKernels' sortperm is reliable for. + +# the first index (row for CSR, column for CSC) of every stored entry, from a pointer array +function _expand_ptr(ptr::oneVector{Ti}, nnzA::Integer) where {Ti} + nnzA == 0 && return similar(ptr, 0) + return Ti.(searchsortedlast.(Ref(ptr), oneVector{Ti}(1:nnzA))) +end + +# compress (rows, cols, vals) triplets of an nrows×ncols matrix into a pointer-array layout: +# returns (ptr, idx, vals) with the entries sorted by row, then by column +function _compress( + rows::oneVector{Ti}, cols::oneVector{Ti}, vals::oneVector, nrows::Integer, ncols::Integer + ) where {Ti} + nnzA = length(vals) + if nnzA == 0 + return _empty_ptr(Ti, nrows + 1), similar(cols, 0), similar(vals, 0) + end + key = (Int.(rows) .- 1) .* Int(ncols) .+ Int.(cols) + perm = sortperm(key) + sorted_rows = rows[perm] + ptr = Ti.(searchsortedfirst.(Ref(sorted_rows), oneVector{Ti}(1:(nrows + 1)))) + return ptr, cols[perm], vals[perm] +end + +# transpose an m×n matrix in pointer-array layout (CSR: ptr=rowPtr, idx=colVal), yielding the +# n×m transpose in the same layout. Since CSC(A) has the same layout as CSR(Aᵀ), this also +# converts between the CSR and CSC formats. +function _transpose_csr(ptr::oneVector{Ti}, idx::oneVector{Ti}, vals::oneVector, m::Integer, n::Integer) where {Ti} + return _compress(idx, _expand_ptr(ptr, length(vals)), vals, n, m) +end + + +## conversions + +oneSparseMatrixCSR(A::oneSparseMatrixCSR) = A +oneSparseMatrixCSC(A::oneSparseMatrixCSC) = A +oneSparseMatrixCOO(A::oneSparseMatrixCOO) = A + +# conversion of the element and index types, on the device +_convert_vector(::Type{T}, x::oneVector) where {T} = eltype(x) === T ? x : T.(x) +function _with_types(A::oneAbstractSparseMatrix, ::Type{Tv}, ::Type{Ti}) where {Tv, Ti} + p, i, v = _storage(A) + (eltype(v) === Tv && eltype(p) === Ti) && return A + return GPUArrays.sparse_array_type(A){Tv, Ti}( + _convert_vector(Ti, p), _convert_vector(Ti, i), _convert_vector(Tv, v), size(A) + ) +end + +# typed constructors convert the format and the types (GPUArrays uses these, e.g. in `triu`) +for X in (:oneSparseMatrixCSR, :oneSparseMatrixCSC, :oneSparseMatrixCOO) + @eval begin + $X{Tv, Ti}(A::oneAbstractSparseMatrix) where {Tv, Ti} = _with_types($X(A), Tv, Ti) + $X{Tv}(A::oneAbstractSparseMatrix{<:Any, Ti}) where {Tv, Ti} = _with_types($X(A), Tv, Ti) + end +end + +function oneSparseMatrixCSC(A::oneSparseMatrixCSR{Tv, Ti}) where {Tv, Ti} + m, n = size(A) + return oneSparseMatrixCSC{Tv, Ti}(_transpose_csr(A.rowPtr, A.colVal, A.nzVal, m, n)..., (m, n)) +end +function oneSparseMatrixCSR(A::oneSparseMatrixCSC{Tv, Ti}) where {Tv, Ti} + m, n = size(A) + return oneSparseMatrixCSR{Tv, Ti}(_transpose_csr(A.colPtr, A.rowVal, A.nzVal, n, m)..., (m, n)) +end + +function oneSparseMatrixCOO(A::oneSparseMatrixCSR{Tv, Ti}) where {Tv, Ti} + rowInd = _expand_ptr(A.rowPtr, nnz(A)) + return oneSparseMatrixCOO{Tv, Ti}(rowInd, copy(A.colVal), copy(A.nzVal), size(A)) +end +# via CSR so that the entries end up in row-major order +oneSparseMatrixCOO(A::oneSparseMatrixCSC) = oneSparseMatrixCOO(oneSparseMatrixCSR(A)) + +function oneSparseMatrixCSR(A::oneSparseMatrixCOO{Tv, Ti}) where {Tv, Ti} + m, n = size(A) + return oneSparseMatrixCSR{Tv, Ti}(_compress(A.rowInd, A.colInd, A.nzVal, m, n)..., (m, n)) +end +function oneSparseMatrixCSC(A::oneSparseMatrixCOO{Tv, Ti}) where {Tv, Ti} + m, n = size(A) + return oneSparseMatrixCSC{Tv, Ti}(_compress(A.colInd, A.rowInd, A.nzVal, n, m)..., (m, n)) +end + + +## transposition + +function GPUArrays._sptranspose(A::oneSparseMatrixCSR{Tv, Ti}) where {Tv, Ti} + m, n = size(A) + return oneSparseMatrixCSR{Tv, Ti}(_transpose_csr(A.rowPtr, A.colVal, A.nzVal, m, n)..., (n, m)) +end +function GPUArrays._spadjoint(A::oneSparseMatrixCSR{Tv, Ti}) where {Tv, Ti} + m, n = size(A) + return oneSparseMatrixCSR{Tv, Ti}(_transpose_csr(A.rowPtr, A.colVal, conj.(A.nzVal), m, n)..., (n, m)) +end + +function GPUArrays._sptranspose(A::oneSparseMatrixCSC{Tv, Ti}) where {Tv, Ti} + m, n = size(A) + return oneSparseMatrixCSC{Tv, Ti}(_transpose_csr(A.colPtr, A.rowVal, A.nzVal, n, m)..., (n, m)) +end +function GPUArrays._spadjoint(A::oneSparseMatrixCSC{Tv, Ti}) where {Tv, Ti} + m, n = size(A) + return oneSparseMatrixCSC{Tv, Ti}(_transpose_csr(A.colPtr, A.rowVal, conj.(A.nzVal), n, m)..., (n, m)) +end + +function GPUArrays._sptranspose(A::oneSparseMatrixCOO{Tv, Ti}) where {Tv, Ti} + m, n = size(A) + # re-sort so that the result is in row-major order + return oneSparseMatrixCOO(oneSparseMatrixCSR{Tv, Ti}(_compress(A.colInd, A.rowInd, A.nzVal, n, m)..., (n, m))) +end +function GPUArrays._spadjoint(A::oneSparseMatrixCOO{Tv, Ti}) where {Tv, Ti} + m, n = size(A) + return oneSparseMatrixCOO(oneSparseMatrixCSR{Tv, Ti}(_compress(A.colInd, A.rowInd, conj.(A.nzVal), n, m)..., (n, m))) +end + +# materializing lazy wrappers of device matrices +for X in (:oneSparseMatrixCSR, :oneSparseMatrixCSC, :oneSparseMatrixCOO) + @eval begin + $X(t::Transpose{<:Any, <:oneAbstractSparseMatrix}) = $X(GPUArrays._sptranspose(parent(t))) + $X(t::Adjoint{<:Any, <:oneAbstractSparseMatrix}) = $X(GPUArrays._spadjoint(parent(t))) + end +end + + +## addition and subtraction + +# GPUArrays implements these through broadcasting over matrices of the same format (e.g. for +# `issymmetric`, which computes `A - transpose(A)`), so wrapped operands are materialized first. +_materialize(A::oneAbstractSparseMatrix) = A +_materialize(t::Transpose{<:Any, <:oneAbstractSparseMatrix}) = GPUArrays._sptranspose(parent(t)) +_materialize(t::Adjoint{<:Any, <:oneAbstractSparseMatrix}) = GPUArrays._spadjoint(parent(t)) + +for op in (:+, :-), X in (:oneSparseMatrixCSR, :oneSparseMatrixCSC) + @eval begin + Base.$op(A::$X, B::$X) = broadcast($op, A, B) + Base.$op(A::$X, B::Union{Transpose{<:Any, <:$X}, Adjoint{<:Any, <:$X}}) = + broadcast($op, A, _materialize(B)) + Base.$op(A::Union{Transpose{<:Any, <:$X}, Adjoint{<:Any, <:$X}}, B::$X) = + broadcast($op, _materialize(A), B) + end +end + + +## structural operations on COO matrices +# +# GPUArrays implements triu/tril/kron/reshape/droptol! for CSR and CSC matrices by converting +# to the COO format and back, so the actual work happens here. + +# keep the entries of `A` selected by `mask` +function _select(A::oneSparseMatrixCOO{Tv, Ti}, mask::AbstractVector{Bool}) where {Tv, Ti} + return oneSparseMatrixCOO{Tv, Ti}(A.rowInd[mask], A.colInd[mask], A.nzVal[mask], size(A)) +end + +LinearAlgebra.triu(A::oneSparseMatrixCOO, k::Integer = 0) = _select(A, A.rowInd .+ k .<= A.colInd) +LinearAlgebra.tril(A::oneSparseMatrixCOO, k::Integer = 0) = _select(A, A.rowInd .+ k .>= A.colInd) + +function SparseArrays.droptol!(A::oneSparseMatrixCOO, tol::Real) + mask = abs.(A.nzVal) .> tol + _invalidate_handle!(A) + A.rowInd = A.rowInd[mask] + A.colInd = A.colInd[mask] + A.nzVal = A.nzVal[mask] + A.nnz = length(A.nzVal) + return A +end + +function Base.reshape(A::oneSparseMatrixCOO{Tv, Ti}, dims::Dims{2}) where {Tv, Ti} + prod(dims) == length(A) || throw(DimensionMismatch("new dimensions $dims must be consistent with array size $(length(A))")) + m = size(A, 1) + m2 = dims[1] + # column-major linear index of every entry, re-interpreted in the new shape + linear = (Int.(A.colInd) .- 1) .* m .+ Int.(A.rowInd) + rowInd = Ti.(mod1.(linear, m2)) + colInd = Ti.(fld1.(linear, m2)) + return oneSparseMatrixCOO{Tv, Ti}(rowInd, colInd, copy(A.nzVal), dims) +end + +# Diagonal matrices as COO, for kron +function oneSparseMatrixCOO(D::Diagonal{Tv}) where {Tv} + n = size(D, 1) + ind = oneVector{Int}(1:n) + return oneSparseMatrixCOO{Tv, Int}(ind, copy(ind), oneVector{Tv}(D.diag), (n, n)) +end + +function LinearAlgebra.kron(A::oneSparseMatrixCOO{Tv, Ti}, B::oneSparseMatrixCOO{Tv, Ti}) where {Tv, Ti} + mA, nA = size(A) + mB, nB = size(B) + nnzA, nnzB = Int(nnz(A)), Int(nnz(B)) + # every entry of A is paired with every entry of B; entries of A vary slowest + Ar = repeat(A.rowInd; inner = nnzB) + Ac = repeat(A.colInd; inner = nnzB) + Av = repeat(A.nzVal; inner = nnzB) + Br = repeat(B.rowInd; outer = nnzA) + Bc = repeat(B.colInd; outer = nnzA) + Bv = repeat(B.nzVal; outer = nnzA) + rowInd = (Ar .- one(Ti)) .* Ti(mB) .+ Br + colInd = (Ac .- one(Ti)) .* Ti(nB) .+ Bc + return oneSparseMatrixCOO{Tv, Ti}(rowInd, colInd, Av .* Bv, (mA * mB, nA * nB)) +end +function LinearAlgebra.kron(A::oneSparseMatrixCOO{Tv, Ti}, B::oneSparseMatrixCOO) where {Tv, Ti} + T = promote_type(Tv, eltype(B)) + return kron(_with_types(A, T, Ti), _with_types(B, T, Ti)) +end +LinearAlgebra.kron(A::oneSparseMatrixCOO, D::Diagonal) = kron(A, oneSparseMatrixCOO(D)) +LinearAlgebra.kron(D::Diagonal, A::oneSparseMatrixCOO) = kron(oneSparseMatrixCOO(D), A) diff --git a/lib/mkl/wrappers_sparse.jl b/lib/mkl/wrappers_sparse.jl index 3da9aba9..5af480c4 100644 --- a/lib/mkl/wrappers_sparse.jl +++ b/lib/mkl/wrappers_sparse.jl @@ -6,14 +6,18 @@ const _deferred_sparse_handles = Vector{matrix_handle_t}() const _deferred_sparse_handles_lock = ReentrantLock() -function sparse_release_matrix_handle(A::oneAbstractSparseMatrix) - return if A.handle !== nothing - lock(_deferred_sparse_handles_lock) do - push!(_deferred_sparse_handles, A.handle) - end +function _defer_release(handle::matrix_handle_t) + return lock(_deferred_sparse_handles_lock) do + push!(_deferred_sparse_handles, handle) end end +function sparse_release_matrix_handle(A::oneAbstractSparseMatrix) + handle = A.handle + handle === nothing || _defer_release(handle) + return +end + function flush_deferred_sparse_releases() handles = lock(_deferred_sparse_handles_lock) do if isempty(_deferred_sparse_handles) @@ -50,6 +54,71 @@ function _check_csc_support() ) end + +## lazy oneMKL matrix handles + +# oneMKL operations are only defined for these element and index types; matrices with other +# types can still be created and used with the generic GPUArrays functionality. +const onemklSparseFloat = onemklFloat +const onemklSparseInt = Union{Int32, Int64} + +# matrices without stored entries never get a handle: the operations short-circuit instead +_mkl_empty(A::oneAbstractSparseMatrix) = nnz(A) == 0 || any(iszero, size(A)) + +# result of op(A) * x for an empty A +function _scale_output!(beta::Number, y::AbstractArray) + if iszero(beta) + fill!(y, zero(eltype(y))) + else + y .*= beta + end + return y +end + +# drop the cached oneMKL handle, e.g. because the storage vectors are about to be replaced +function _invalidate_handle!(A::oneAbstractSparseMatrix) + handle = A.handle + handle === nothing && return A + A.handle = nothing + _defer_release(handle) + return A +end + +""" + sparse_matrix_handle(A::oneAbstractSparseMatrix) + +Return the oneMKL matrix handle describing `A`, creating it on first use. The handle refers to +the storage vectors of `A`, so those must not be replaced or resized afterwards (use `copyto!`, +which takes care of this, or construct a new matrix). +""" +function sparse_matrix_handle(A::oneAbstractSparseMatrix) + handle = A.handle + handle === nothing || return handle + _mkl_empty(A) && throw(ArgumentError("cannot create a oneMKL handle for an empty sparse matrix")) + + flush_deferred_sparse_releases() + Support._check_sparse_abi() + handle_ptr = Ref{matrix_handle_t}() + onemklXsparse_init_matrix_handle(handle_ptr) + try + _set_matrix_data!(A, handle_ptr[]) + catch + _defer_release(handle_ptr[]) + rethrow() + end + A.handle = handle_ptr[] + return handle_ptr[] +end + +function _set_matrix_data!(A::oneAbstractSparseMatrix, handle::matrix_handle_t) + throw( + ArgumentError( + "oneMKL sparse operations only support Float32, Float64, ComplexF32 and ComplexF64 " * + "matrices with Int32 or Int64 indices, got $(typeof(A))" + ) + ) +end + for (fname, elty, intty) in ((:onemklSsparse_set_csr_data , :Float32 , :Int32), (:onemklSsparse_set_csr_data_64, :Float32 , :Int64), (:onemklDsparse_set_csr_data , :Float64 , :Int32), @@ -59,81 +128,20 @@ for (fname, elty, intty) in ((:onemklSsparse_set_csr_data , :Float32 , :Int3 (:onemklZsparse_set_csr_data , :ComplexF64, :Int32), (:onemklZsparse_set_csr_data_64, :ComplexF64, :Int64)) @eval begin - - function oneSparseMatrixCSR( - rowPtr::oneVector{$intty}, colVal::oneVector{$intty}, - nzVal::oneVector{$elty}, dims::NTuple{2, Int} - ) - flush_deferred_sparse_releases() - handle_ptr = Ref{matrix_handle_t}() - onemklXsparse_init_matrix_handle(handle_ptr) - m, n = dims - nnzA = length(nzVal) - queue = global_queue(context(nzVal), device(nzVal)) - # Don't update handle if matrix is empty - if m != 0 && n != 0 - Support._check_sparse_abi() - $fname(sycl_queue(queue), handle_ptr[], m, n, nnzA, 'O', rowPtr, colVal, nzVal) - dA = oneSparseMatrixCSR{$elty, $intty}(handle_ptr[], rowPtr, colVal, nzVal, (m, n), nnzA) - finalizer(sparse_release_matrix_handle, dA) - else - dA = oneSparseMatrixCSR{$elty, $intty}(nothing, rowPtr, colVal, nzVal, (m, n), nnzA) - end - return dA - end - - function oneSparseMatrixCSC( - colPtr::oneVector{$intty}, rowVal::oneVector{$intty}, - nzVal::oneVector{$elty}, dims::NTuple{2, Int} - ) - flush_deferred_sparse_releases() - queue = global_queue(context(nzVal), device(nzVal)) - handle_ptr = Ref{matrix_handle_t}() - onemklXsparse_init_matrix_handle(handle_ptr) - m, n = dims - nnzA = length(nzVal) - # Don't update handle if matrix is empty - if m != 0 && n != 0 - Support._check_sparse_abi() - _check_csc_support() - $fname(sycl_queue(queue), handle_ptr[], n, m, nnzA, 'O', colPtr, rowVal, nzVal) # CSC of A is CSR of Aᵀ - dA = oneSparseMatrixCSC{$elty, $intty}(handle_ptr[], colPtr, rowVal, nzVal, (m, n), nnzA) - finalizer(sparse_release_matrix_handle, dA) - else - dA = oneSparseMatrixCSC{$elty, $intty}(nothing, colPtr, rowVal, nzVal, (m, n), nnzA) - end - return dA - end - - - function oneSparseMatrixCSR(A::SparseMatrixCSC{$elty, $intty}) + function _set_matrix_data!(A::oneSparseMatrixCSR{$elty, $intty}, handle::matrix_handle_t) m, n = size(A) - At = SparseMatrixCSC(A |> transpose) - rowPtr = oneVector{$intty}(At.colptr) - colVal = oneVector{$intty}(At.rowval) - nzVal = oneVector{$elty}(At.nzval) - return oneSparseMatrixCSR(rowPtr, colVal, nzVal, (m, n)) - end - - function SparseArrays.SparseMatrixCSC(A::oneSparseMatrixCSR{$elty, $intty}) - handle_ptr = Ref{matrix_handle_t}() - At = SparseMatrixCSC(reverse(A.dims)..., Vector(A.rowPtr), Vector(A.colVal), Vector(A.nzVal)) - A_csc = SparseMatrixCSC(At |> transpose) - return A_csc + queue = global_queue(context(A.nzVal), device(A.nzVal)) + $fname(sycl_queue(queue), handle, m, n, nnz(A), 'O', A.rowPtr, A.colVal, A.nzVal) + return end - function oneSparseMatrixCSC(A::SparseMatrixCSC{$elty, $intty}) + function _set_matrix_data!(A::oneSparseMatrixCSC{$elty, $intty}, handle::matrix_handle_t) + _check_csc_support() m, n = size(A) - colPtr = oneVector{$intty}(A.colptr) - rowVal = oneVector{$intty}(A.rowval) - nzVal = oneVector{$elty}(A.nzval) - return oneSparseMatrixCSC(colPtr, rowVal, nzVal, (m, n)) - end - - function SparseArrays.SparseMatrixCSC(A::oneSparseMatrixCSC{$elty, $intty}) - handle_ptr = Ref{matrix_handle_t}() - A_csc = SparseMatrixCSC(A.dims..., Vector(A.colPtr), Vector(A.rowVal), Vector(A.nzVal)) - return A_csc + queue = global_queue(context(A.nzVal), device(A.nzVal)) + # CSC of A is CSR of Aᵀ + $fname(sycl_queue(queue), handle, n, m, nnz(A), 'O', A.colPtr, A.rowVal, A.nzVal) + return end end end @@ -147,36 +155,24 @@ for (fname, elty, intty) in ((:onemklSsparse_set_coo_data , :Float32 , :Int3 (:onemklZsparse_set_coo_data , :ComplexF64, :Int32), (:onemklZsparse_set_coo_data_64, :ComplexF64, :Int64)) @eval begin - function oneSparseMatrixCOO(A::SparseMatrixCSC{$elty, $intty}) - flush_deferred_sparse_releases() - handle_ptr = Ref{matrix_handle_t}() - onemklXsparse_init_matrix_handle(handle_ptr) + function _set_matrix_data!(A::oneSparseMatrixCOO{$elty, $intty}, handle::matrix_handle_t) m, n = size(A) - row, col, val = findnz(A) - rowInd = oneVector{$intty}(row) - colInd = oneVector{$intty}(col) - nzVal = oneVector{$elty}(val) - nnzA = length(val) - queue = global_queue(context(nzVal), device(nzVal)) - if m != 0 && n != 0 - Support._check_sparse_abi() - $fname(sycl_queue(queue), handle_ptr[], m, n, nnzA, 'O', rowInd, colInd, nzVal) - dA = oneSparseMatrixCOO{$elty, $intty}(handle_ptr[], rowInd, colInd, nzVal, (m, n), nnzA) - finalizer(sparse_release_matrix_handle, dA) - else - dA = oneSparseMatrixCOO{$elty, $intty}(nothing, rowInd, colInd, nzVal, (m, n), nnzA) - end - return dA - end - - function SparseArrays.SparseMatrixCSC(A::oneSparseMatrixCOO{$elty, $intty}) - handle_ptr = Ref{matrix_handle_t}() - A = sparse(Vector(A.rowInd), Vector(A.colInd), Vector(A.nzVal), A.dims...) - return A + queue = global_queue(context(A.nzVal), device(A.nzVal)) + $fname(sycl_queue(queue), handle, m, n, nnz(A), 'O', A.rowInd, A.colInd, A.nzVal) + return end end end +function oneAPI.unsafe_free!(A::oneAbstractSparseMatrix) + _invalidate_handle!(A) + foreach(unsafe_free!, _storage(A)) + return +end + + +## operations + for SparseMatrix in (:oneSparseMatrixCSR, :oneSparseMatrixCOO) for (fname, elty) in ((:onemklSsparse_gemv, :Float32), (:onemklDsparse_gemv, :Float64), @@ -190,8 +186,9 @@ for SparseMatrix in (:oneSparseMatrixCSR, :oneSparseMatrixCOO) beta::Number, y::oneStridedVector{$elty}) + _mkl_empty(A) && return _scale_output!(beta, y) queue = global_queue(context(x), device(x)) - $fname(sycl_queue(queue), trans, alpha, A.handle, x, beta, y) + $fname(sycl_queue(queue), trans, alpha, sparse_matrix_handle(A), x, beta, y) y end end @@ -199,8 +196,9 @@ for SparseMatrix in (:oneSparseMatrixCSR, :oneSparseMatrixCOO) @eval begin function sparse_optimize_gemv!(trans::Char, A::$SparseMatrix) + _mkl_empty(A) && return A queue = global_queue(context(A.nzVal), device(A.nzVal)) - onemklXsparse_optimize_gemv(sycl_queue(queue), trans, A.handle) + onemklXsparse_optimize_gemv(sycl_queue(queue), trans, sparse_matrix_handle(A)) return A end end @@ -225,10 +223,11 @@ for SparseMatrix in (:oneSparseMatrixCSC,) beta::Number, y::oneStridedVector{$elty}) + _mkl_empty(A) && return _scale_output!(beta, y) queue = global_queue(context(x), device(x)) m, n = size(A) if m != 0 && n != 0 - $fname(sycl_queue(queue), flip_trans(trans), alpha, A.handle, x, beta, y) + $fname(sycl_queue(queue), flip_trans(trans), alpha, sparse_matrix_handle(A), x, beta, y) end y end @@ -250,6 +249,7 @@ for SparseMatrix in (:oneSparseMatrixCSC,) y::oneStridedVector{$elty} ) + _mkl_empty(A) && return _scale_output!(beta, y) # Compute A^H*x via identity: # conj(y_new) = conj(alpha) * (A^T) * conj(x) + conj(beta) * conj(y) # Since S=A^T and op='N' computes S*x = A^T*x, we can realize this with one call. @@ -262,7 +262,7 @@ for SparseMatrix in (:oneSparseMatrixCSC,) end queue = global_queue(context(x), device(x)) - $fname(sycl_queue(queue), flip_trans(trans), alpha, A.handle, x, beta, y) + $fname(sycl_queue(queue), flip_trans(trans), alpha, sparse_matrix_handle(A), x, beta, y) if trans == 'C' y .= conj.(y) @@ -275,9 +275,10 @@ for SparseMatrix in (:oneSparseMatrixCSC,) end @eval begin function sparse_optimize_gemv!(trans::Char, A::$SparseMatrix) + _mkl_empty(A) && return A # complex 'C' case is implemented using op='N' on S=A^T with conjugation trick queue = global_queue(context(A.nzVal), device(A.nzVal)) - onemklXsparse_optimize_gemv(sycl_queue(queue), flip_trans(trans), A.handle) + onemklXsparse_optimize_gemv(sycl_queue(queue), flip_trans(trans), sparse_matrix_handle(A)) return A end end @@ -298,6 +299,7 @@ for (fname, elty) in ((:onemklSsparse_gemm, :Float32), C::oneStridedMatrix{$elty} ) + _mkl_empty(A) && return _scale_output!(beta, C) mB, nB = size(B) mC, nC = size(C) (nB != nC) && (transb == 'N') && throw(ArgumentError("B and C must have the same number of columns.")) @@ -306,21 +308,23 @@ for (fname, elty) in ((:onemklSsparse_gemm, :Float32), ldb = max(1,stride(B,2)) ldc = max(1,stride(C,2)) queue = global_queue(context(C), device(C)) - $fname(sycl_queue(queue), 'C', transa, transb, alpha, A.handle, B, nrhs, ldb, beta, C, ldc) + $fname(sycl_queue(queue), 'C', transa, transb, alpha, sparse_matrix_handle(A), B, nrhs, ldb, beta, C, ldc) C end end end function sparse_optimize_gemm!(trans::Char, A::oneSparseMatrixCSR) + _mkl_empty(A) && return A queue = global_queue(context(A.nzVal), device(A.nzVal)) - onemklXsparse_optimize_gemm(sycl_queue(queue), trans, A.handle) + onemklXsparse_optimize_gemm(sycl_queue(queue), trans, sparse_matrix_handle(A)) return A end function sparse_optimize_gemm!(trans::Char, transB::Char, nrhs::Int, A::oneSparseMatrixCSR) + _mkl_empty(A) && return A queue = global_queue(context(A.nzVal), device(A.nzVal)) - onemklXsparse_optimize_gemm_advanced(sycl_queue(queue), 'C', trans, transB, A.handle, nrhs) + onemklXsparse_optimize_gemm_advanced(sycl_queue(queue), 'C', trans, transB, sparse_matrix_handle(A), nrhs) return A end @@ -335,6 +339,7 @@ for (fname, elty) in ((:onemklSsparse_gemm, :Float32), beta::Number, C::oneStridedMatrix{$elty}) + _mkl_empty(A) && return _scale_output!(beta, C) mB, nB = size(B) mC, nC = size(C) (nB != nC) && (transb == 'N') && throw(ArgumentError("B and C must have the same number of columns.")) @@ -343,7 +348,7 @@ for (fname, elty) in ((:onemklSsparse_gemm, :Float32), ldb = max(1,stride(B,2)) ldc = max(1,stride(C,2)) queue = global_queue(context(C), device(C)) - $fname(sycl_queue(queue), 'C', flip_trans(transa), transb, alpha, A.handle, B, nrhs, ldb, beta, C, ldc) + $fname(sycl_queue(queue), 'C', flip_trans(transa), transb, alpha, sparse_matrix_handle(A), B, nrhs, ldb, beta, C, ldc) C end end @@ -365,6 +370,7 @@ for (fname, elty) in ( C::oneStridedMatrix{$elty} ) + _mkl_empty(A) && return _scale_output!(beta, C) # Map op(A) to op(S) where S = A^T stored as CSR in the handle # transa: 'N' -> op(S)='T'; 'T' -> op(S)='N'; 'C' -> # real: op(S)='N' (since A^H == A^T) @@ -408,7 +414,7 @@ for (fname, elty) in ( transb_eff = transb end - $fname(sycl_queue(queue), 'C', flip_trans(transa), transb_eff, alpha, A.handle, B, nrhs, ldb, beta, C, ldc) + $fname(sycl_queue(queue), 'C', flip_trans(transa), transb_eff, alpha, sparse_matrix_handle(A), B, nrhs, ldb, beta, C, ldc) # Undo conjugation to obtain C_new if transa == 'C' @@ -424,14 +430,16 @@ for (fname, elty) in ( end function sparse_optimize_gemm!(trans::Char, A::oneSparseMatrixCSC) + _mkl_empty(A) && return A queue = global_queue(context(A.nzVal), device(A.nzVal)) - onemklXsparse_optimize_gemm(sycl_queue(queue), flip_trans(trans), A.handle) + onemklXsparse_optimize_gemm(sycl_queue(queue), flip_trans(trans), sparse_matrix_handle(A)) return A end function sparse_optimize_gemm!(trans::Char, transB::Char, nrhs::Int, A::oneSparseMatrixCSC) + _mkl_empty(A) && return A queue = global_queue(context(A.nzVal), device(A.nzVal)) - onemklXsparse_optimize_gemm_advanced(sycl_queue(queue), 'C', flip_trans(trans), transB, A.handle, nrhs) + onemklXsparse_optimize_gemm_advanced(sycl_queue(queue), 'C', flip_trans(trans), transB, sparse_matrix_handle(A), nrhs) return A end @@ -447,8 +455,9 @@ for (fname, elty) in ((:onemklSsparse_symv, :Float32), beta::Number, y::oneStridedVector{$elty}) + _mkl_empty(A) && return _scale_output!(beta, y) queue = global_queue(context(y), device(y)) - $fname(sycl_queue(queue), uplo, alpha, A.handle, x, beta, y) + $fname(sycl_queue(queue), uplo, alpha, sparse_matrix_handle(A), x, beta, y) y end end @@ -467,8 +476,9 @@ for (fname, elty) in ((:onemklSsparse_symv, :Float32), beta::Number, y::oneStridedVector{$elty}) + _mkl_empty(A) && return _scale_output!(beta, y) queue = global_queue(context(y), device(y)) - $fname(sycl_queue(queue), flip_uplo(uplo), alpha, A.handle, x, beta, y) + $fname(sycl_queue(queue), flip_uplo(uplo), alpha, sparse_matrix_handle(A), x, beta, y) y end end @@ -488,16 +498,18 @@ for (fname, elty) in ((:onemklSsparse_trmv, :Float32), beta::Number, y::oneStridedVector{$elty}) + _mkl_empty(A) && return _scale_output!(beta, y) queue = global_queue(context(y), device(y)) - $fname(sycl_queue(queue), uplo, trans, diag, alpha, A.handle, x, beta, y) + $fname(sycl_queue(queue), uplo, trans, diag, alpha, sparse_matrix_handle(A), x, beta, y) y end end end function sparse_optimize_trmv!(uplo::Char, trans::Char, diag::Char, A::oneSparseMatrixCSR) + _mkl_empty(A) && return A queue = global_queue(context(A.nzVal), device(A.nzVal)) - onemklXsparse_optimize_trmv(sycl_queue(queue), uplo, trans, diag, A.handle) + onemklXsparse_optimize_trmv(sycl_queue(queue), uplo, trans, diag, sparse_matrix_handle(A)) return A end @@ -520,6 +532,7 @@ for (fname, elty) in ( y::oneStridedVector{$elty} ) + _mkl_empty(A) && return _scale_output!(beta, y) # Intel oneAPI sparse trmv only supports nontrans operations. # Since CSC(A) is stored as CSR(A^T), we cannot map CSC operations # to CSR operations for triangular operations without transpose support. @@ -531,13 +544,14 @@ for (fname, elty) in ( ) ) queue = global_queue(context(y), device(y)) - $fname(sycl_queue(queue), uplo, flip_trans(trans), diag, alpha, A.handle, x, beta, y) + $fname(sycl_queue(queue), uplo, flip_trans(trans), diag, alpha, sparse_matrix_handle(A), x, beta, y) return y end end end function sparse_optimize_trmv!(uplo::Char, trans::Char, diag::Char, A::oneSparseMatrixCSC) + _mkl_empty(A) && return A throw( ArgumentError( "sparse_optimize_trmv! is not supported for oneSparseMatrixCSC due to Intel oneAPI limitations. " * @@ -546,7 +560,7 @@ function sparse_optimize_trmv!(uplo::Char, trans::Char, diag::Char, A::oneSparse ) ) queue = global_queue(context(A.nzVal), device(A.nzVal)) - onemklXsparse_optimize_trmv(sycl_queue(queue), uplo, flip_trans(trans), diag, A.handle) + onemklXsparse_optimize_trmv(sycl_queue(queue), uplo, flip_trans(trans), diag, sparse_matrix_handle(A)) return A end @@ -563,16 +577,18 @@ for (fname, elty) in ((:onemklSsparse_trsv, :Float32), x::oneStridedVector{$elty}, y::oneStridedVector{$elty}) + _mkl_empty(A) && throw(ArgumentError("cannot perform a triangular solve with an empty sparse matrix")) queue = global_queue(context(y), device(y)) - $fname(sycl_queue(queue), uplo, trans, diag, alpha, A.handle, x, y) + $fname(sycl_queue(queue), uplo, trans, diag, alpha, sparse_matrix_handle(A), x, y) y end end end function sparse_optimize_trsv!(uplo::Char, trans::Char, diag::Char, A::oneSparseMatrixCSR) + _mkl_empty(A) && return A queue = global_queue(context(A.nzVal), device(A.nzVal)) - onemklXsparse_optimize_trsv(sycl_queue(queue), uplo, trans, diag, A.handle) + onemklXsparse_optimize_trsv(sycl_queue(queue), uplo, trans, diag, sparse_matrix_handle(A)) return A end @@ -593,6 +609,7 @@ for (fname, elty) in ( y::oneStridedVector{$elty} ) + _mkl_empty(A) && throw(ArgumentError("cannot perform a triangular solve with an empty sparse matrix")) throw( ArgumentError( "sparse_trsv! is not supported for oneSparseMatrixCSC due to Intel oneAPI limitations. " * @@ -601,13 +618,14 @@ for (fname, elty) in ( ) ) queue = global_queue(context(y), device(y)) - onemklXsparse_optimize_trsv(sycl_queue(queue), uplo, flip_trans(trans), diag, A.handle) + onemklXsparse_optimize_trsv(sycl_queue(queue), uplo, flip_trans(trans), diag, sparse_matrix_handle(A)) return A end end end function sparse_optimize_trsv!(uplo::Char, trans::Char, diag::Char, A::oneSparseMatrixCSC) + _mkl_empty(A) && return A throw( ArgumentError( "sparse_optimize_trsv! is not supported for oneSparseMatrixCSC due to Intel oneAPI limitations. " * @@ -616,7 +634,7 @@ function sparse_optimize_trsv!(uplo::Char, trans::Char, diag::Char, A::oneSparse ) ) queue = global_queue(context(A.nzVal), device(A.nzVal)) - onemklXsparse_optimize_trsv(sycl_queue(queue), uplo, flip_trans(trans), diag, A.handle) + onemklXsparse_optimize_trsv(sycl_queue(queue), uplo, flip_trans(trans), diag, sparse_matrix_handle(A)) return A end @@ -634,6 +652,7 @@ for (fname, elty) in ((:onemklSsparse_trsm, :Float32), X::oneStridedMatrix{$elty}, Y::oneStridedMatrix{$elty}) + _mkl_empty(A) && throw(ArgumentError("cannot perform a triangular solve with an empty sparse matrix")) mX, nX = size(X) mY, nY = size(Y) (mX != mY) && (transX == 'N') && throw(ArgumentError("X and Y must have the same number of rows.")) @@ -644,21 +663,23 @@ for (fname, elty) in ((:onemklSsparse_trsm, :Float32), ldx = max(1,stride(X,2)) ldy = max(1,stride(Y,2)) queue = global_queue(context(Y), device(Y)) - $fname(sycl_queue(queue), 'C', transA, transX, uplo, diag, alpha, A.handle, X, nrhs, ldx, Y, ldy) + $fname(sycl_queue(queue), 'C', transA, transX, uplo, diag, alpha, sparse_matrix_handle(A), X, nrhs, ldx, Y, ldy) Y end end end function sparse_optimize_trsm!(uplo::Char, trans::Char, diag::Char, A::oneSparseMatrixCSR) + _mkl_empty(A) && return A queue = global_queue(context(A.nzVal), device(A.nzVal)) - onemklXsparse_optimize_trsm(sycl_queue(queue), uplo, trans, diag, A.handle) + onemklXsparse_optimize_trsm(sycl_queue(queue), uplo, trans, diag, sparse_matrix_handle(A)) return A end function sparse_optimize_trsm!(uplo::Char, trans::Char, diag::Char, nrhs::Int, A::oneSparseMatrixCSR) + _mkl_empty(A) && return A queue = global_queue(context(A.nzVal), device(A.nzVal)) - onemklXsparse_optimize_trsm_advanced(sycl_queue(queue), 'C', uplo, trans, diag, A.handle, nrhs) + onemklXsparse_optimize_trsm_advanced(sycl_queue(queue), 'C', uplo, trans, diag, sparse_matrix_handle(A), nrhs) return A end @@ -682,6 +703,7 @@ for (fname, elty) in ( Y::oneStridedMatrix{$elty} ) + _mkl_empty(A) && throw(ArgumentError("cannot perform a triangular solve with an empty sparse matrix")) # Intel oneAPI sparse trsm only supports nontrans operations for the matrix A. # Since CSC(A) is stored as CSR(A^T), we cannot map CSC operations # to CSR operations for triangular solve operations without transpose support. @@ -703,13 +725,14 @@ for (fname, elty) in ( ldx = max(1, stride(X, 2)) ldy = max(1, stride(Y, 2)) queue = global_queue(context(Y), device(Y)) - $fname(sycl_queue(queue), 'C', flip_trans(transA), transX, uplo, diag, alpha, A.handle, X, nrhs, ldx, Y, ldy) + $fname(sycl_queue(queue), 'C', flip_trans(transA), transX, uplo, diag, alpha, sparse_matrix_handle(A), X, nrhs, ldx, Y, ldy) return Y end end end function sparse_optimize_trsm!(uplo::Char, trans::Char, diag::Char, A::oneSparseMatrixCSC) + _mkl_empty(A) && return A throw( ArgumentError( "sparse_optimize_trsm! is not supported for oneSparseMatrixCSC due to Intel oneAPI limitations. " * @@ -718,11 +741,12 @@ function sparse_optimize_trsm!(uplo::Char, trans::Char, diag::Char, A::oneSparse ) ) queue = global_queue(context(A.nzVal), device(A.nzVal)) - onemklXsparse_optimize_trsm(sycl_queue(queue), uplo, trans, diag, A.handle) + onemklXsparse_optimize_trsm(sycl_queue(queue), uplo, trans, diag, sparse_matrix_handle(A)) return A end function sparse_optimize_trsm!(uplo::Char, trans::Char, diag::Char, nrhs::Int, A::oneSparseMatrixCSC) + _mkl_empty(A) && return A throw( ArgumentError( "sparse_optimize_trsm! is not supported for oneSparseMatrixCSC due to Intel oneAPI limitations. " * @@ -731,6 +755,6 @@ function sparse_optimize_trsm!(uplo::Char, trans::Char, diag::Char, nrhs::Int, A ) ) queue = global_queue(context(A.nzVal), device(A.nzVal)) - onemklXsparse_optimize_trsm_advanced(sycl_queue(queue), 'C', uplo, trans, diag, A.handle, nrhs) + onemklXsparse_optimize_trsm_advanced(sycl_queue(queue), 'C', uplo, trans, diag, sparse_matrix_handle(A), nrhs) return A end diff --git a/src/oneAPIKernels.jl b/src/oneAPIKernels.jl index c35cc629..7b90d2ca 100644 --- a/src/oneAPIKernels.jl +++ b/src/oneAPIKernels.jl @@ -37,6 +37,15 @@ Adapt.adapt_storage(::oneAPIBackend, a::AbstractArray) = Adapt.adapt(oneArray, a Adapt.adapt_storage(::oneAPIBackend, a::oneArray) = a Adapt.adapt_storage(::KA.CPU, a::oneArray) = convert(Array, a) +# sparse arrays (oneMKL is only available on Linux) +@static if Sys.islinux() + import GPUArrays, SparseArrays + KA.get_backend(::oneAPI.oneMKL.oneAbstractSparseMatrix) = oneAPIBackend() + # without this, `adapt_storage(::oneAPIBackend, ::AbstractArray)` would densify sparse arrays + Adapt.adapt_storage(::oneAPIBackend, a::GPUArrays.AbstractGPUSparseArray) = a + Adapt.adapt_storage(::KA.CPU, a::oneAPI.oneMKL.oneAbstractSparseMatrix) = SparseArrays.SparseMatrixCSC(a) +end + ## Memory Operations diff --git a/test/onemkl.jl b/test/onemkl.jl index 0f743f69..f5a3faf5 100644 --- a/test/onemkl.jl +++ b/test/onemkl.jl @@ -1105,8 +1105,6 @@ end end @testset "oneSparseMatrixCSC" begin - csc_supported || continue - (T isa Complex) && continue for S in (Int32, Int64) A = sprand(T, 20, 10, 0.5) A = SparseMatrixCSC{T, S}(A) @@ -1127,9 +1125,148 @@ end B = oneSparseMatrixCOO(A) A2 = SparseMatrixCSC(B) @test A == A2 + C = oneSparseMatrixCOO(B.rowInd, B.colInd, B.nzVal, size(B)) + @test SparseMatrixCSC(C) == A + D = oneSparseMatrixCOO(oneVector(S[]), oneVector(S[]), oneVector(T[]), (0, 0)) # empty matrix end end + + sparse_matrices = (oneSparseMatrixCSR, oneSparseMatrixCSC, oneSparseMatrixCOO) + + @testset "GPUArrays interface" begin + @testset "$SparseMatrix" for SparseMatrix in sparse_matrices + A = SparseMatrixCSC{T, Int32}(sprand(T, 20, 10, 0.5)) + B = SparseMatrix(A) + @test B isa GPUArrays.AbstractGPUSparseMatrix{T, Int32} + @test size(B) == (20, 10) + @test size(B, 1) == 20 + @test size(B, 2) == 10 + @test size(B, 3) == 1 + @test nnz(B) == nnz(A) + @test collect(B) == collect(A) + @test Array(B) == Array(A) + @test Array(oneArray(B)) == Array(A) + @test SparseMatrixCSC(copy(B)) == A + @test similar(B) isa SparseMatrix{T, Int32} + @test nnz(similar(B)) == nnz(A) + @test similar(B, Float32) isa SparseMatrix{Float32, Int32} + @test similar(B, Float32, 3, 4) isa SparseMatrix{Float32, Int32} + @test size(similar(B, Float32, (3, 4))) == (3, 4) + @test SparseMatrixCSC(similar(B, (3, 4))) == spzeros(T, 3, 4) + @test similar(B, T, (5,)) isa oneVector{T} + @test SparseMatrixCSC(SparseMatrix(sparsevec(A[:, 1]))) == A[:, 1:1] + @test SparseMatrix{T, Int64}(A) isa SparseMatrix{T, Int64} + @test SparseMatrix{T}(SparseMatrixCSC{T, Int64}(A)) isa SparseMatrix{T, Int64} + @allowscalar begin + @test B[2, 3] == A[2, 3] + @test all(B[i, j] == A[i, j] for i in 1:20, j in 1:10) + @test Array(B[:, 2]) == A[:, 2] + end + @test_throws BoundsError B[0, 1] + @test_throws BoundsError B[21, 1] + end + A = SparseMatrixCSC{T, Int32}(sprand(T, 20, 10, 0.5)) + @test adapt(oneArray, A) isa oneSparseMatrixCSC{T, Int32} + @test SparseMatrixCSC(adapt(oneArray, A)) == A + @test adapt(oneArray{T}, SparseMatrixCSC{T, Int64}(A)) isa oneSparseMatrixCSC{T, Int64} + @test adapt(Array, oneSparseMatrixCSR(A)) == A + end + + @testset "conversions" begin + for S in (Int32, Int64) + A = SparseMatrixCSC{T, S}(sprand(T, 20, 10, 0.3)) + # include empty rows and columns + A[3, :] .= 0 + A[:, 5] .= 0 + dropzeros!(A) + @testset "$src -> $dst" for src in sparse_matrices, dst in sparse_matrices + B = dst(src(A)) + @test B isa dst{T, S} + @test SparseMatrixCSC(B) == A + end + @testset "transpose $SparseMatrix" for SparseMatrix in sparse_matrices + B = SparseMatrix(A) + @test SparseMatrixCSC(GPUArrays._sptranspose(B)) == transpose(A) + @test SparseMatrixCSC(GPUArrays._spadjoint(B)) == adjoint(A) + @test SparseMatrixCSC(SparseMatrix(transpose(B))) == transpose(A) + @test SparseMatrixCSC(SparseMatrix(adjoint(B))) == adjoint(A) + @test SparseMatrixCSC(SparseMatrix(transpose(A))) == transpose(A) + @test SparseMatrixCSC(SparseMatrix(adjoint(A))) == adjoint(A) + end + Z = spzeros(T, S, 5, 7) + @testset "empty $src -> $dst" for src in sparse_matrices, dst in sparse_matrices + @test SparseMatrixCSC(dst(src(Z))) == Z + end + end + end + + @testset "broadcast and reductions" begin + @testset "$SparseMatrix" for SparseMatrix in (oneSparseMatrixCSR, oneSparseMatrixCSC) + A = SparseMatrixCSC{T, Int32}(sprand(T, 20, 10, 0.3)) + B = SparseMatrix(A) + C = B .* T(2) + @test C isa SparseMatrix{T, Int32} + @test SparseMatrixCSC(C) == A .* T(2) + D = B .+ T(1) + @test D isa oneMatrix{T} + @test Array(D) ≈ Array(A .+ T(1)) + @test SparseMatrixCSC(B .* B) ≈ A .* A + @test sum(B) ≈ sum(A) + @test Array(sum(B; dims = 1)) ≈ sum(A; dims = 1) + @test Array(sum(B; dims = 2)) ≈ sum(A; dims = 2) + @test opnorm(B, 1) ≈ opnorm(A, 1) + @test opnorm(B, Inf) ≈ opnorm(A, Inf) + @test !iszero(B) + @test iszero(B .* T(0)) + I, J, V = findnz(B) + @test (Array(I), Array(J), Array(V)) == findnz(A) + + Asq = SparseMatrixCSC{T, Int32}(sprand(T, 10, 10, 0.3)) + Bsq = SparseMatrix(Asq) + @test SparseMatrixCSC(Bsq + Bsq) == Asq + Asq + @test SparseMatrixCSC(Bsq - transpose(Bsq)) == Asq - transpose(Asq) + @test SparseMatrixCSC(adjoint(Bsq) + Bsq) == adjoint(Asq) + Asq + @test !issymmetric(Bsq) + @test issymmetric(SparseMatrix(Asq + transpose(Asq))) + @test SparseMatrixCSC(triu(Bsq)) == triu(Asq) + @test SparseMatrixCSC(tril(Bsq, -1)) == tril(Asq, -1) + end + end + + @testset "lazy oneMKL handles" begin + A = SparseMatrixCSC{T, Int32}(sprand(T, 20, 10, 0.3)) + x = oneArray(rand(T, 10)) + @testset "$SparseMatrix" for SparseMatrix in csr_csc_matrices + B = SparseMatrix(A) + @test B.handle === nothing + @test Array(B * x) ≈ A * Array(x) + @test B.handle !== nothing + # matrices produced by generic code + C = B .* T(2) + @test Array(C * x) ≈ (A .* T(2)) * Array(x) + @test Array(SparseMatrix(oneSparseMatrixCOO(A)) * x) ≈ A * Array(x) + # copyto! replaces the storage, so the handle must be recreated + copyto!(C, B) + @test C.handle === nothing + @test Array(C * x) ≈ A * Array(x) + # empty matrices do not need a handle + E = similar(B, T, 20, 10) + @test Array(E * x) == zeros(T, 20) + y = oneArray(ones(T, 20)) + @test Array(mul!(y, E, x, true, T(2))) == fill(T(2), 20) + oneAPI.unsafe_free!(B) + @test B.handle === nothing + end + if !csc_supported + @test_throws ErrorException oneSparseMatrixCSC(A) * x + end + # any element type can be stored, but oneMKL only operates on BLAS types + F = oneSparseMatrixCSR(sprand(Float16, 10, 10, 0.3)) + @test F isa oneSparseMatrixCSR{Float16, Int} + @test_throws ArgumentError oneMKL.sparse_matrix_handle(F) + end + @testset "sparse gemv" begin @testset "$SparseMatrix" for SparseMatrix in coo_csr_csc_matrices @testset "transa = $transa" for (transa, opa) in [('N', identity), ('T', transpose), ('C', adjoint)] diff --git a/test/runtests.jl b/test/runtests.jl index 045fd1d3..09a7f6c7 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -121,6 +121,19 @@ init_worker_code = quote end TestSuite.supported_eltypes(::Type{<:oneArray}) = eltypes + # run the GPUArrays sparse testsuite on the oneMKL sparse matrix types (COO is excluded + # because GPUArrays' generic sparse broadcast only supports vectors, CSR and CSC). + # oneMKL itself is only needed for linear algebra, so integer element types are the only + # ones to exclude: the testsuite's sprand/mapreduce checks do not make sense for them. + TestSuite.sparse_types(::Type{<:oneArray}) = (oneMKL.oneSparseMatrixCSR, oneMKL.oneSparseMatrixCSC) + function TestSuite.supported_eltypes(::Type{<:oneArray}, test) + typs = copy(eltypes) + if startswith(string(test), "test_sparse") + filter!(ET -> !(ET <: Integer || ET <: Complex{<:Integer}), typs) + end + return typs + end + const validation_layer = parse(Bool, get(ENV, "ZE_ENABLE_VALIDATION_LAYER", "false")) const parameter_validation = parse(Bool, get(ENV, "ZE_ENABLE_PARAMETER_VALIDATION", "false")) @@ -173,6 +186,7 @@ end init_code = quote using oneAPI, Adapt + using GPUArrays: GPUArrays, @allowscalar import ..TestSuite, ..testf import ..eltypes, ..float16_supported, ..float64_supported, From d24babb22bb7cbc0c2ac22ecb9e549392ab403df Mon Sep 17 00:00:00 2001 From: Michel Schanen Date: Thu, 24 Sep 2026 10:59:49 -0500 Subject: [PATCH 2/2] Fix sparse storage sharing, mixed-type copyto!, and empty unit-triangular ops - Every sparse matrix now holds its own reference to its storage vectors, so unsafe_free! on one matrix no longer frees vectors shared with another (single-input broadcast outputs, type conversions). - copyto! between sparse matrices of the same format accepts different element and index types, replacing the storage and dropping the handle instead of falling back to in-place or scalar GPUArrays methods. - Empty triangular matrices with a unit diagonal act as the identity in trmv/trsv/trsm; with a non-unit diagonal the solves throw SingularException. Co-Authored-By: Claude Opus 5.5 (1M context) --- lib/mkl/array.jl | 38 ++++++++++++--------- lib/mkl/wrappers_sparse.jl | 35 ++++++++++++++++---- test/onemkl.jl | 68 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 120 insertions(+), 21 deletions(-) diff --git a/lib/mkl/array.jl b/lib/mkl/array.jl index 7d85d718..c09a2151 100644 --- a/lib/mkl/array.jl +++ b/lib/mkl/array.jl @@ -14,6 +14,12 @@ import Adapt: adapt # convert sparse matrices freely without paying for oneMKL handle set-up, and so that the struct # can hold element types oneMKL does not support. +# The storage vectors are often shared between matrices (e.g. the output of a single-input +# broadcast reuses the pointer array of its input, and type conversions reuse the vectors whose +# type does not change). Every matrix therefore holds its own reference to the underlying +# memory, so that `unsafe_free!` on one matrix does not free the storage of another. +_own_ref(x::oneVector) = GPUArrays.derive(eltype(x), x, size(x), 0) + mutable struct oneSparseMatrixCSR{Tv, Ti} <: GPUArrays.AbstractGPUSparseMatrixCSR{Tv, Ti} handle::Union{Nothing, matrix_handle_t} rowPtr::oneVector{Ti} @@ -26,7 +32,7 @@ mutable struct oneSparseMatrixCSR{Tv, Ti} <: GPUArrays.AbstractGPUSparseMatrixCS rowPtr::oneVector{Ti}, colVal::oneVector{Ti}, nzVal::oneVector{Tv}, dims::NTuple{2, <:Integer} ) where {Tv, Ti <: Integer} - A = new{Tv, Ti}(nothing, rowPtr, colVal, nzVal, Int.(dims), Ti(length(nzVal))) + A = new{Tv, Ti}(nothing, _own_ref(rowPtr), _own_ref(colVal), _own_ref(nzVal), Int.(dims), Ti(length(nzVal))) return finalizer(sparse_release_matrix_handle, A) end end @@ -43,7 +49,7 @@ mutable struct oneSparseMatrixCSC{Tv, Ti} <: GPUArrays.AbstractGPUSparseMatrixCS colPtr::oneVector{Ti}, rowVal::oneVector{Ti}, nzVal::oneVector{Tv}, dims::NTuple{2, <:Integer} ) where {Tv, Ti <: Integer} - A = new{Tv, Ti}(nothing, colPtr, rowVal, nzVal, Int.(dims), Ti(length(nzVal))) + A = new{Tv, Ti}(nothing, _own_ref(colPtr), _own_ref(rowVal), _own_ref(nzVal), Int.(dims), Ti(length(nzVal))) return finalizer(sparse_release_matrix_handle, A) end end @@ -60,7 +66,7 @@ mutable struct oneSparseMatrixCOO{Tv, Ti} <: GPUArrays.AbstractGPUSparseMatrixCO rowInd::oneVector{Ti}, colInd::oneVector{Ti}, nzVal::oneVector{Tv}, dims::NTuple{2, <:Integer} ) where {Tv, Ti <: Integer} - A = new{Tv, Ti}(nothing, rowInd, colInd, nzVal, Int.(dims), Ti(length(nzVal))) + A = new{Tv, Ti}(nothing, _own_ref(rowInd), _own_ref(colInd), _own_ref(nzVal), Int.(dims), Ti(length(nzVal))) return finalizer(sparse_release_matrix_handle, A) end end @@ -175,33 +181,35 @@ Base.copy(A::oneSparseMatrixCSC{Tv, Ti}) where {Tv, Ti} = Base.copy(A::oneSparseMatrixCOO{Tv, Ti}) where {Tv, Ti} = oneSparseMatrixCOO{Tv, Ti}(copy(A.rowInd), copy(A.colInd), copy(A.nzVal), size(A)) +_copy_as(::Type{T}, x::oneVector) where {T} = eltype(x) === T ? copy(x) : T.(x) + # `copyto!` may change the sparsity structure, so it replaces the storage vectors (they may be # shared with other matrices, e.g. the output of a single-input broadcast reuses the pointer # array of its input) and drops the oneMKL handle, which refers to the old storage. -function Base.copyto!(dst::oneSparseMatrixCSR{Tv, Ti}, src::oneSparseMatrixCSR{Tv, Ti}) where {Tv, Ti} +function Base.copyto!(dst::oneSparseMatrixCSR{Tv, Ti}, src::oneSparseMatrixCSR) where {Tv, Ti} size(dst) == size(src) || throw(ArgumentError("Inconsistent Sparse Matrix size")) _invalidate_handle!(dst) - dst.rowPtr = copy(src.rowPtr) - dst.colVal = copy(src.colVal) - dst.nzVal = copy(src.nzVal) + dst.rowPtr = _copy_as(Ti, src.rowPtr) + dst.colVal = _copy_as(Ti, src.colVal) + dst.nzVal = _copy_as(Tv, src.nzVal) dst.nnz = src.nnz return dst end -function Base.copyto!(dst::oneSparseMatrixCSC{Tv, Ti}, src::oneSparseMatrixCSC{Tv, Ti}) where {Tv, Ti} +function Base.copyto!(dst::oneSparseMatrixCSC{Tv, Ti}, src::oneSparseMatrixCSC) where {Tv, Ti} size(dst) == size(src) || throw(ArgumentError("Inconsistent Sparse Matrix size")) _invalidate_handle!(dst) - dst.colPtr = copy(src.colPtr) - dst.rowVal = copy(src.rowVal) - dst.nzVal = copy(src.nzVal) + dst.colPtr = _copy_as(Ti, src.colPtr) + dst.rowVal = _copy_as(Ti, src.rowVal) + dst.nzVal = _copy_as(Tv, src.nzVal) dst.nnz = src.nnz return dst end -function Base.copyto!(dst::oneSparseMatrixCOO{Tv, Ti}, src::oneSparseMatrixCOO{Tv, Ti}) where {Tv, Ti} +function Base.copyto!(dst::oneSparseMatrixCOO{Tv, Ti}, src::oneSparseMatrixCOO) where {Tv, Ti} size(dst) == size(src) || throw(ArgumentError("Inconsistent Sparse Matrix size")) _invalidate_handle!(dst) - dst.rowInd = copy(src.rowInd) - dst.colInd = copy(src.colInd) - dst.nzVal = copy(src.nzVal) + dst.rowInd = _copy_as(Ti, src.rowInd) + dst.colInd = _copy_as(Ti, src.colInd) + dst.nzVal = _copy_as(Tv, src.nzVal) dst.nnz = src.nnz return dst end diff --git a/lib/mkl/wrappers_sparse.jl b/lib/mkl/wrappers_sparse.jl index 5af480c4..40ec2eaa 100644 --- a/lib/mkl/wrappers_sparse.jl +++ b/lib/mkl/wrappers_sparse.jl @@ -75,6 +75,29 @@ function _scale_output!(beta::Number, y::AbstractArray) return y end +# an empty triangular matrix is the identity if it has a unit diagonal, and singular otherwise +# (unless it has no rows or columns, in which case there is nothing to compute) +_empty_triangular_is_identity(diag::Char, A::oneAbstractSparseMatrix) = + diag == 'U' || any(iszero, size(A)) + +# result of op(A) * x for an empty triangular A +function _empty_trmv!(diag::Char, alpha::Number, x::AbstractVector, beta::Number, y::AbstractVector) + _scale_output!(beta, y) + diag == 'U' && (y .+= alpha .* x) + return y +end + +# result of op(A) \ alpha * op(X) for an empty triangular A +function _empty_trsm!( + diag::Char, A::oneAbstractSparseMatrix, alpha::Number, transX::Char, + X::AbstractArray, Y::AbstractArray + ) + _empty_triangular_is_identity(diag, A) || throw(LinearAlgebra.SingularException(1)) + opX = transX == 'N' ? X : transX == 'T' ? permutedims(X) : conj.(permutedims(X)) + Y .= alpha .* opX + return Y +end + # drop the cached oneMKL handle, e.g. because the storage vectors are about to be replaced function _invalidate_handle!(A::oneAbstractSparseMatrix) handle = A.handle @@ -498,7 +521,7 @@ for (fname, elty) in ((:onemklSsparse_trmv, :Float32), beta::Number, y::oneStridedVector{$elty}) - _mkl_empty(A) && return _scale_output!(beta, y) + _mkl_empty(A) && return _empty_trmv!(diag, alpha, x, beta, y) queue = global_queue(context(y), device(y)) $fname(sycl_queue(queue), uplo, trans, diag, alpha, sparse_matrix_handle(A), x, beta, y) y @@ -532,7 +555,7 @@ for (fname, elty) in ( y::oneStridedVector{$elty} ) - _mkl_empty(A) && return _scale_output!(beta, y) + _mkl_empty(A) && return _empty_trmv!(diag, alpha, x, beta, y) # Intel oneAPI sparse trmv only supports nontrans operations. # Since CSC(A) is stored as CSR(A^T), we cannot map CSC operations # to CSR operations for triangular operations without transpose support. @@ -577,7 +600,7 @@ for (fname, elty) in ((:onemklSsparse_trsv, :Float32), x::oneStridedVector{$elty}, y::oneStridedVector{$elty}) - _mkl_empty(A) && throw(ArgumentError("cannot perform a triangular solve with an empty sparse matrix")) + _mkl_empty(A) && return _empty_trsm!(diag, A, alpha, 'N', x, y) queue = global_queue(context(y), device(y)) $fname(sycl_queue(queue), uplo, trans, diag, alpha, sparse_matrix_handle(A), x, y) y @@ -609,7 +632,7 @@ for (fname, elty) in ( y::oneStridedVector{$elty} ) - _mkl_empty(A) && throw(ArgumentError("cannot perform a triangular solve with an empty sparse matrix")) + _mkl_empty(A) && return _empty_trsm!(diag, A, alpha, 'N', x, y) throw( ArgumentError( "sparse_trsv! is not supported for oneSparseMatrixCSC due to Intel oneAPI limitations. " * @@ -652,7 +675,7 @@ for (fname, elty) in ((:onemklSsparse_trsm, :Float32), X::oneStridedMatrix{$elty}, Y::oneStridedMatrix{$elty}) - _mkl_empty(A) && throw(ArgumentError("cannot perform a triangular solve with an empty sparse matrix")) + _mkl_empty(A) && return _empty_trsm!(diag, A, alpha, transX, X, Y) mX, nX = size(X) mY, nY = size(Y) (mX != mY) && (transX == 'N') && throw(ArgumentError("X and Y must have the same number of rows.")) @@ -703,7 +726,7 @@ for (fname, elty) in ( Y::oneStridedMatrix{$elty} ) - _mkl_empty(A) && throw(ArgumentError("cannot perform a triangular solve with an empty sparse matrix")) + _mkl_empty(A) && return _empty_trsm!(diag, A, alpha, transX, X, Y) # Intel oneAPI sparse trsm only supports nontrans operations for the matrix A. # Since CSC(A) is stored as CSR(A^T), we cannot map CSC operations # to CSR operations for triangular solve operations without transpose support. diff --git a/test/onemkl.jl b/test/onemkl.jl index f5a3faf5..253eba4d 100644 --- a/test/onemkl.jl +++ b/test/onemkl.jl @@ -1267,6 +1267,74 @@ end @test_throws ArgumentError oneMKL.sparse_matrix_handle(F) end + @testset "shared storage" begin + A = SparseMatrixCSC{T, Int32}(sprand(T, 20, 10, 0.3)) + x = oneArray(rand(T, 10)) + @testset "$SparseMatrix" for SparseMatrix in csr_csc_matrices + B = SparseMatrix(A) + @test Array(B * x) ≈ A * Array(x) + # a single-input broadcast reuses the pointer array of its input + C = B .* T(2) + oneAPI.unsafe_free!(C) + @test SparseMatrixCSC(B) == A + @test Array(B * x) ≈ A * Array(x) + # converting the index type reuses the values + D = SparseMatrix{T, Int64}(B) + oneAPI.unsafe_free!(D) + @test SparseMatrixCSC(B) == A + @test Array(B * x) ≈ A * Array(x) + end + end + + @testset "copyto! with different types" begin + A = SparseMatrixCSC{T, Int64}(sprand(T, 20, 10, 0.3)) + x = oneArray(rand(T, 10)) + # COO matrices have no `*` method, only the oneMKL wrapper + matvec(M, x) = oneMKL.sparse_gemv!('N', one(T), M, x, zero(T), similar(x, size(M, 1))) + @testset "$SparseMatrix" for SparseMatrix in coo_csr_csc_matrices + src = SparseMatrix(A) + dst = SparseMatrix(SparseMatrixCSC{T, Int32}(sprand(T, 20, 10, 0.1))) + matvec(dst, x) # create a handle + shared = typeof(dst)(oneMKL._storage(dst)..., size(dst)) + expected = SparseMatrixCSC(shared) + copyto!(dst, src) + @test dst isa SparseMatrix{T, Int32} + @test dst.handle === nothing + @test SparseMatrixCSC(dst) == A + @test Array(matvec(dst, x)) ≈ A * Array(x) + # matrices sharing the old storage are left untouched + @test SparseMatrixCSC(shared) == expected + end + end + + @testset "empty triangular matrices" begin + E = oneSparseMatrixCSR(spzeros(T, 10, 10)) + x = rand(T, 10) + y = rand(T, 10) + alpha = rand(T) + beta = rand(T) + # with a unit diagonal, an empty triangular matrix is the identity + dy = oneVector(y) + oneMKL.sparse_trmv!('L', 'N', 'U', alpha, E, oneVector(x), beta, dy) + @test collect(dy) ≈ alpha * x + beta * y + dy = oneVector(y) + oneMKL.sparse_trmv!('L', 'N', 'N', alpha, E, oneVector(x), beta, dy) + @test collect(dy) ≈ beta * y + dy = oneVector(y) + oneMKL.sparse_trsv!('L', 'N', 'U', alpha, E, oneVector(x), dy) + @test collect(dy) ≈ alpha * x + @test_throws SingularException oneMKL.sparse_trsv!('L', 'N', 'N', alpha, E, oneVector(x), dy) + X = rand(T, 10, 4) + dY = oneMatrix(zeros(T, 10, 4)) + oneMKL.sparse_trsm!('U', 'N', 'N', 'U', alpha, E, oneMatrix(X), dY) + @test collect(dY) ≈ alpha * X + dY = oneMatrix(zeros(T, 10, 4)) + oneMKL.sparse_trsm!('U', 'N', 'C', 'U', alpha, E, oneMatrix(collect(X')), dY) + @test collect(dY) ≈ alpha * X + @test_throws SingularException oneMKL.sparse_trsm!('U', 'N', 'N', 'N', alpha, E, oneMatrix(X), dY) + @test collect(UnitLowerTriangular(E) \ oneVector(x)) ≈ x + end + @testset "sparse gemv" begin @testset "$SparseMatrix" for SparseMatrix in coo_csr_csc_matrices @testset "transa = $transa" for (transa, opa) in [('N', identity), ('T', transpose), ('C', adjoint)]