More tests and bugfixing for CUSPARSE (#2725)

Author: kshyatt

lib/cusparse/conversions.jl

@@ -587,6 +587,7 @@ end
 
 function CuSparseMatrixCSR{Tv}(coo::CuSparseMatrixCOO{Tv}; index::SparseChar='O') where {Tv}
     m,n = size(coo)
+    nnz(coo) == 0 && return CuSparseMatrixCSR{Tv}(CUDA.ones(Cint, m+1), coo.colInd, nonzeros(coo), size(coo))
     coo = sort_coo(coo, 'R')
     csrRowPtr = CuVector{Cint}(undef, m+1)
     cusparseXcoo2csr(handle(), coo.rowInd, nnz(coo), m, csrRowPtr, index)
@@ -595,6 +596,7 @@ end
 
 function CuSparseMatrixCOO{Tv}(csr::CuSparseMatrixCSR{Tv}; index::SparseChar='O') where {Tv}
     m,n = size(csr)
+    nnz(csr) == 0 && return CuSparseMatrixCOO{Tv}(CUDA.zeros(Cint, 0), CUDA.zeros(Cint, 0), nonzeros(csr), size(csr))
     cooRowInd = CuVector{Cint}(undef, nnz(csr))
     cusparseXcsr2coo(handle(), csr.rowPtr, nnz(csr), m, cooRowInd, index)
     CuSparseMatrixCOO{Tv}(cooRowInd, csr.colVal, nonzeros(csr), size(csr))
@@ -604,6 +606,7 @@ end
 
 function CuSparseMatrixCSC{Tv}(coo::CuSparseMatrixCOO{Tv}; index::SparseChar='O') where {Tv}
     m,n = size(coo)
+    nnz(coo) == 0 && return CuSparseMatrixCSC{Tv}(CUDA.ones(Cint, n+1), coo.rowInd, nonzeros(coo), size(coo))
     coo = sort_coo(coo, 'C')
     cscColPtr = CuVector{Cint}(undef, n+1)
     cusparseXcoo2csr(handle(), coo.colInd, nnz(coo), n, cscColPtr, index)
@@ -612,6 +615,7 @@ end
 
 function CuSparseMatrixCOO{Tv}(csc::CuSparseMatrixCSC{Tv}; index::SparseChar='O') where {Tv}
     m,n = size(csc)
+    nnz(csc) == 0 && return CuSparseMatrixCOO{Tv}(CUDA.zeros(Cint, 0), CUDA.zeros(Cint, 0), nonzeros(csc), size(csc))
     cooColInd = CuVector{Cint}(undef, nnz(csc))
     cusparseXcsr2coo(handle(), csc.colPtr, nnz(csc), n, cooColInd, index)
     coo = CuSparseMatrixCOO{Tv}(csc.rowVal, cooColInd, nonzeros(csc), size(csc))

lib/cusparse/linalg.jl

@@ -57,15 +57,15 @@ function Base.reshape(A::CuSparseMatrixCOO, dims::Dims)
     sparse(new_row, new_col, A.nzVal, dims[1], length(dims) == 1 ? 1 : dims[2], fmt = :coo)
 end
 
-function LinearAlgebra.kron(A::CuSparseMatrixCOO{T}, B::CuSparseMatrixCOO{T}) where {T}
+function LinearAlgebra.kron(A::CuSparseMatrixCOO{T, Ti}, B::CuSparseMatrixCOO{T, Ti}) where {Ti, T}
     mA,nA = size(A)
     mB,nB = size(B)
     out_shape = (mA * mB, nA * nB)
     Annz = Int64(A.nnz)
     Bnnz = Int64(B.nnz)
 
     if Annz == 0 || Bnnz == 0
-        return CuSparseMatrixCOO(CuVector{T}(undef, 0), CuVector{T}(undef, 0), CuVector{T}(undef, 0), out_shape)
+        return CuSparseMatrixCOO(CuVector{Ti}(undef, 0), CuVector{Ti}(undef, 0), CuVector{T}(undef, 0), out_shape)
     end
 
     row = (A.rowInd .- 1) .* mB
@@ -82,15 +82,15 @@ function LinearAlgebra.kron(A::CuSparseMatrixCOO{T}, B::CuSparseMatrixCOO{T}) wh
     sparse(row, col, data, out_shape..., fmt = :coo)
 end
 
-function LinearAlgebra.kron(A::CuSparseMatrixCOO{T}, B::Diagonal) where {T}
+function LinearAlgebra.kron(A::CuSparseMatrixCOO{T, Ti}, B::Diagonal) where {Ti, T}
     mA,nA = size(A)
     mB,nB = size(B)
     out_shape = (mA * mB, nA * nB)
     Annz = Int64(A.nnz)
     Bnnz = nB
 
     if Annz == 0 || Bnnz == 0
-        return CuSparseMatrixCOO(CuVector{T}(undef, 0), CuVector{T}(undef, 0), CuVector{T}(undef, 0), out_shape)
+        return CuSparseMatrixCOO(CuVector{Ti}(undef, 0), CuVector{Ti}(undef, 0), CuVector{T}(undef, 0), out_shape)
     end
 
     row = (A.rowInd .- 1) .* mB
@@ -107,15 +107,15 @@ function LinearAlgebra.kron(A::CuSparseMatrixCOO{T}, B::Diagonal) where {T}
     sparse(row, col, data, out_shape..., fmt = :coo)
 end
 
-function LinearAlgebra.kron(A::Diagonal, B::CuSparseMatrixCOO{T}) where {T}
+function LinearAlgebra.kron(A::Diagonal, B::CuSparseMatrixCOO{T, Ti}) where {Ti, T}
     mA,nA = size(A)
     mB,nB = size(B)
     out_shape = (mA * mB, nA * nB)
     Annz = nA
     Bnnz = Int64(B.nnz)
 
     if Annz == 0 || Bnnz == 0
-        return CuSparseMatrixCOO(CuVector{T}(undef, 0), CuVector{T}(undef, 0), CuVector{T}(undef, 0), out_shape)
+        return CuSparseMatrixCOO(CuVector{Ti}(undef, 0), CuVector{Ti}(undef, 0), CuVector{T}(undef, 0), out_shape)
     end
 
     row = (0:nA-1) .* mB

test/libraries/cusparse.jl

@@ -199,6 +199,14 @@ blockdim = 5
     d_x = LowerTriangular(CuSparseMatrixCSC(x))
     @test !istriu(d_x)
     @test istril(d_x)
+    if VERSION >= v"1.11.2"
+        d_x = UpperTriangular(transpose(CuSparseMatrixCSC(x)))
+        @test istriu(d_x)
+        @test !istril(d_x)
+        d_x = LowerTriangular(transpose(CuSparseMatrixCSC(x)))
+        @test !istriu(d_x)
+        @test istril(d_x)
+    end
     d_x = UpperTriangular(triu(transpose(CuSparseMatrixCSC(x)), 1))
     @test istriu(d_x)
     @test !istril(d_x)

test/libraries/cusparse/linalg.jl

@@ -43,13 +43,23 @@ end
     b = sprand(ComplexF32, 100, 100, 0.1)
     A = typ(a)
     B = typ(b)
+    za = spzeros(ComplexF32, 100, 100)
+    ZA = typ(za)
     @test collect(kron(A, B)) ≈ kron(a, b)
     @test collect(kron(transpose(A), B)) ≈ kron(transpose(a), b)
     @test collect(kron(A, transpose(B))) ≈ kron(a, transpose(b))
     @test collect(kron(transpose(A), transpose(B))) ≈ kron(transpose(a), transpose(b))
     @test collect(kron(A', B)) ≈ kron(a', b)
     @test collect(kron(A, B')) ≈ kron(a, b')
     @test collect(kron(A', B')) ≈ kron(a', b')
+    
+    @test collect(kron(ZA, B)) ≈ kron(za, b)
+    @test collect(kron(transpose(ZA), B)) ≈ kron(transpose(za), b)
+    @test collect(kron(ZA, transpose(B))) ≈ kron(za, transpose(b))
+    @test collect(kron(transpose(ZA), transpose(B))) ≈ kron(transpose(za), transpose(b))
+    @test collect(kron(ZA', B)) ≈ kron(za', b)
+    @test collect(kron(ZA, B')) ≈ kron(za, b')
+    @test collect(kron(ZA', B')) ≈ kron(za', b')
 
     C = I(50)
     @test collect(kron(A, C)) ≈ kron(a, C)
@@ -88,4 +98,5 @@ end
     A2 = typ(A)
 
     @test dot(x, A, y) ≈ dot(x2, A2, y2)
+    @test_throws DimensionMismatch("dimensions must match") dot(CUDA.rand(elty, N1+1), A2, y2)
 end