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..c09a2151 100644 --- a/lib/mkl/array.jl +++ b/lib/mkl/array.jl @@ -1,49 +1,304 @@ 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. + +# 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} 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, _own_ref(rowPtr), _own_ref(colVal), _own_ref(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, _own_ref(colPtr), _own_ref(rowVal), _own_ref(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, _own_ref(rowInd), _own_ref(colInd), _own_ref(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)) + +_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) where {Tv, Ti} + size(dst) == size(src) || throw(ArgumentError("Inconsistent Sparse Matrix size")) + _invalidate_handle!(dst) + 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) where {Tv, Ti} + size(dst) == size(src) || throw(ArgumentError("Inconsistent Sparse Matrix size")) + _invalidate_handle!(dst) + 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) where {Tv, Ti} + size(dst) == size(src) || throw(ArgumentError("Inconsistent Sparse Matrix size")) + _invalidate_handle!(dst) + 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 + +# 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..40ec2eaa 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,94 @@ 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 + +# 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 + 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 +151,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 +178,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 +209,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 +219,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 +246,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 +272,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 +285,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 +298,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 +322,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 +331,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 +362,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 +371,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 +393,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 +437,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 +453,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 +478,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 +499,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 +521,18 @@ for (fname, elty) in ((:onemklSsparse_trmv, :Float32), beta::Number, y::oneStridedVector{$elty}) + _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, 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 +555,7 @@ for (fname, elty) in ( y::oneStridedVector{$elty} ) + _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. @@ -531,13 +567,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 +583,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 +600,18 @@ for (fname, elty) in ((:onemklSsparse_trsv, :Float32), x::oneStridedVector{$elty}, y::oneStridedVector{$elty}) + _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, 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 +632,7 @@ for (fname, elty) in ( y::oneStridedVector{$elty} ) + _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. " * @@ -601,13 +641,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 +657,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 +675,7 @@ for (fname, elty) in ((:onemklSsparse_trsm, :Float32), X::oneStridedMatrix{$elty}, Y::oneStridedMatrix{$elty}) + _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.")) @@ -644,21 +686,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 +726,7 @@ for (fname, elty) in ( Y::oneStridedMatrix{$elty} ) + _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. @@ -703,13 +748,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 +764,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 +778,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..253eba4d 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,216 @@ 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 "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)] 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,