Support mul!(Diagonal, A, B) (#2977)

Author: kshyatt

lib/cublas/linalg.jl

@@ -537,6 +537,12 @@ function LinearAlgebra.mul!(C::CuMatrix{T}, A::Diagonal{T,<:CuVector}, B::Adjoin
     return C
 end
 
+function LinearAlgebra.mul!(C::Diagonal{T, <:CuVector}, A::Union{<:CuMatrix{T}, Adjoint{T, <:CuMatrix}, Transpose{T, <:CuMatrix}}, B::Union{<:CuMatrix{T}, Adjoint{T, <:CuMatrix}, Transpose{T, <:CuMatrix}}) where {T<:CublasFloat}
+    Cfull   = A*B
+    C.diag .= diag(Cfull)
+    return C
+end
+
 function LinearAlgebra.lmul!(A::Diagonal{T,<:CuVector{T}}, B::CuMatrix{T}) where {T<:CublasFloat}
     return dgmm!('L', B, A.diag, B)
 end

test/libraries/cublas/extensions.jl

@@ -585,6 +585,15 @@ k = 13
             @test Array(d_AX) ≈ A' * Diagonal(x)
 
             @test Array(d_X) == Diagonal(Array(d_x))
+            
+            d_X = Diagonal(copy(d_x))
+            diagA = diagm(rand(elty, m))
+            d_diagA = CuArray(diagA)
+            diagB = diagm(rand(elty, m))
+            d_diagB = CuArray(diagB)
+            diagAdiagB  = diagA * diagB'
+            mul!(d_X, d_diagA, d_diagB')
+            @test Diagonal(collect(d_X.diag)) ≈ Diagonal(diagAdiagB)
         end
     end # extensions