Fix spdiagm with specified pairs (#2784)

* Fix spdiagm with specified pairs

* Should subtype keys

* Remove DS_store

Author: ErikQQY

lib/cusparse/array.jl

@@ -336,14 +336,55 @@ function SparseArrays.sparsevec(I::CuArray{Ti}, V::CuArray{Tv}, n::Integer) wher
     CuSparseVector(I, V, n)
 end
 
-function SparseArrays.spdiagm(v::CuVector{Tv}) where {Tv}
-    nzVal = v
-    N = Int32(length(nzVal))
-
-    colPtr = CuArray(one(Int32):(N + one(Int32)))
-    rowVal = CuArray(one(Int32):N)
-    dims = (N, N)
-    CuSparseMatrixCSC(colPtr, rowVal, nzVal, dims)
+SparseArrays.spdiagm(kv::Pair{<:Integer,<:CuVector}...) = _cuda_spdiagm(nothing, kv...)
+SparseArrays.spdiagm(m::Integer, n::Integer, kv::Pair{<:Integer,<:CuVector}...) = _cuda_spdiagm((Int(m),Int(n)), kv...)
+SparseArrays.spdiagm(v::CuVector) = _cuda_spdiagm(nothing, 0 => v)
+SparseArrays.spdiagm(m::Integer, n::Integer, v::CuVector) = _cuda_spdiagm((Int(m), Int(n)), 0 => v)
+
+function _cuda_spdiagm(size, kv::Pair{<:Integer, <:CuVector}...)
+    I, J, V, mmax, nmax = _cuda_spdiagm_internal(kv...)
+    mnmax = max(mmax, nmax)
+    m, n = something(size, (mnmax,mnmax))
+    (m ≥ mmax && n ≥ nmax) || throw(DimensionMismatch("invalid size=$size"))
+    return sparse(CuVector(I), CuVector(J), V, m, n)
+end
+
+function _cuda_spdiagm_internal(kv::Pair{T,<:CuVector}...) where {T<:Integer}
+    ncoeffs = 0
+    for p in kv
+        ncoeffs += SparseArrays._nnz(p.second)
+    end
+    I = Vector{T}(undef, ncoeffs)
+    J = Vector{T}(undef, ncoeffs)
+    V = CuArray{promote_type(map(x -> eltype(x.second), kv)...)}(undef, ncoeffs)
+    i = 0
+    m = 0
+    n = 0
+    for p in kv
+        k = p.first
+        v = p.second
+        if k < 0
+            row = -k
+            col = 0
+        elseif k > 0
+            row = 0
+            col = k
+        else
+            row = 0
+            col = 0
+        end
+        numel = SparseArrays._nnz(v)
+        r = 1+i:numel+i
+        I_r, J_r = SparseArrays._indices(v, row, col)
+        copyto!(view(I, r), I_r)
+        copyto!(view(J, r), J_r)
+        copyto!(view(V, r), v)
+        veclen = length(v)
+        m = max(m, row + veclen)
+        n = max(n, col + veclen)
+        i += numel
+    end
+    return I, J, V, m, n
 end
 
 LinearAlgebra.issymmetric(M::Union{CuSparseMatrixCSC,CuSparseMatrixCSR}) = size(M, 1) == size(M, 2) ? norm(M - transpose(M), Inf) == 0 : false

test/libraries/cusparse/interfaces.jl

@@ -581,7 +581,23 @@ end
         cuda_vec = CuVector(ref_vec)
 
         ref_spdiagm = spdiagm(ref_vec) # SparseArrays
-        cuda_spdiagm = spdiagm(cuda_vec) # CuSparseMatrixCSC
+        cuda_spdiagm = spdiagm(cuda_vec)
+
+        ref_cuda_sparse = CuSparseMatrixCSC(ref_spdiagm)
+
+        @test ref_cuda_sparse.rowVal == cuda_spdiagm.rowVal
+
+        @test ref_cuda_sparse.nzVal == cuda_spdiagm.nzVal
+
+        @test ref_cuda_sparse.colPtr == cuda_spdiagm.colPtr
+    end
+
+    @testset "spdiagm(2 => CuVector{$elty})" for elty in [Float32, Float64, ComplexF32, ComplexF64]
+        ref_vec = collect(elty, 100:121)
+        cuda_vec = CuVector(ref_vec)
+
+        ref_spdiagm = spdiagm(2 => ref_vec) # SparseArrays
+        cuda_spdiagm = spdiagm(2 => cuda_vec)
 
         ref_cuda_sparse = CuSparseMatrixCSC(ref_spdiagm)