CUBLAS: Add support for diagm (#2786)

Author: ErikQQY

lib/cublas/linalg.jl

@@ -385,6 +385,40 @@ function LinearAlgebra.mul!(C::CuMatrix{T}, A::Diagonal{T,<:CuVector}, B::Adjoin
     return C
 end
 
+# diagm
+
+LinearAlgebra.diagm(kv::Pair{<:Integer,<:CuVector}...) = _cuda_diagm(nothing, kv...)
+LinearAlgebra.diagm(m::Integer, n::Integer, kv::Pair{<:Integer,<:CuVector}...) = _cuda_diagm((Int(m),Int(n)), kv...)
+LinearAlgebra.diagm(v::CuVector) = LinearAlgebra.diagm(0 => v)
+LinearAlgebra.diagm(m::Integer, n::Integer, v::CuVector) = LinearAlgebra.diagm(m, n, 0 => v)
+
+function _cuda_diagm(size, kv::Pair{<:Integer,<:CuVector}...)
+    A = LinearAlgebra.diagm_container(size, kv...)
+    for p in kv
+        inds = LinearAlgebra.diagind(A, p.first)
+        copyto!(view(A, inds), p.second)
+    end
+    return A
+end
+
+function LinearAlgebra.diagm_container(size, kv::Pair{<:Integer,<:CuVector}...)
+    T = promote_type(map(x -> eltype(x.second), kv)...)
+    U = promote_type(T, typeof(zero(T)))
+    return cu(zeros(U, LinearAlgebra.diagm_size(size, kv...)...))
+end
+
+function LinearAlgebra.diagm_size(size::Nothing, kv::Pair{<:Integer,<:CuVector}...)
+    mnmax = mapreduce(x -> length(x.second) + abs(Int(x.first)), max, kv; init=0)
+    return mnmax, mnmax
+end
+function LinearAlgebra.diagm_size(size::Tuple{Int,Int}, kv::Pair{<:Integer,<:CuVector}...)
+    mmax = mapreduce(x -> length(x.second) - min(0,Int(x.first)), max, kv; init=0)
+    nmax = mapreduce(x -> length(x.second) + max(0,Int(x.first)), max, kv; init=0)
+    m, n = size
+    (m ≥ mmax && n ≥ nmax) || throw(DimensionMismatch(lazy"invalid size=$size"))
+    return m, n
+end
+
 # symmetric mul!
 
 op_wrappers = ((identity, T -> 'N', identity),

test/libraries/cublas/level3.jl

@@ -467,6 +467,23 @@ k = 13
             h_C = triu(h_C)
             @test C ≈ h_C
         end
+        @testset "diagm" begin
+            A = rand(elty, m)
+            B = rand(elty, n)
+            # move to device
+            d_A = CuArray(A)
+            d_B = CuArray(B)
+            diagA = diagm(d_A)
+            diagB = diagm(2 => d_B)
+            # move back to host and compare
+            diagind_A = diagind(diagA, 0)
+            diagind_B = diagind(diagB, 2)
+            h_A = Array(diagA[diagind_A])
+            h_B = Array(diagB[diagind_B])
+
+            @test A ≈ h_A
+            @test B ≈ h_B
+        end
         if elty <: Complex
             @testset "herk!" begin
                 alpha = rand(real(elty))