Support in-place rmul/lmul for Diagonals (#2856)

Author: kshyatt

lib/cublas/linalg.jl

@@ -1,6 +1,6 @@
 # interfacing with LinearAlgebra standard library
 
-using LinearAlgebra: MulAddMul, AdjOrTrans, wrap, UpperOrLowerTriangular
+using LinearAlgebra: MulAddMul, AdjOrTrans, wrap, UpperOrLowerTriangular, rmul!, lmul!
 
 #
 # BLAS 1
@@ -413,6 +413,14 @@ function LinearAlgebra.mul!(C::CuMatrix{T}, A::Diagonal{T,<:CuVector}, B::Adjoin
     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
+
+function LinearAlgebra.rmul!(A::CuMatrix{T}, B::Diagonal{T,<:CuVector{T}}) where {T<:CublasFloat}
+    return dgmm!('R', A, B.diag, A)
+end
+
 # diagm
 
 LinearAlgebra.diagm(kv::Pair{<:Integer,<:CuVector}...) = _cuda_diagm(nothing, kv...)

test/libraries/cublas/extensions.jl

@@ -531,18 +531,30 @@ k = 13
             h_C = Array(d_C)
             @test C ≈ h_C
         end
-        @testset "diagonal -- mul!" begin
+        @testset "diagonal -- mul!, rmul!, lmul!" begin
             XA = rand(elty,m,n)
             d_XA = CuArray(XA)
             d_X = Diagonal(d_x)
             mul!(d_XA, d_X, d_A)
             Array(d_XA) ≈ Diagonal(x) * A
+            
+            XA = rand(elty,m,n)
+            d_XA = CuArray(XA)
+            d_X = Diagonal(d_x)
+            lmul!(d_X, d_XA)
+            Array(d_XA) ≈ Diagonal(x) * XA
 
             AY = rand(elty,m,n)
             d_AY = CuArray(AY)
             d_Y = Diagonal(d_y)
             mul!(d_AY, d_A, d_Y)
             Array(d_AY) ≈ A * Diagonal(y)
+            
+            AY = rand(elty,m,n)
+            d_AY = CuArray(AY)
+            d_Y = Diagonal(d_y)
+            rmul!(d_AY, d_Y)
+            Array(d_AY) ≈ AY * Diagonal(y)
 
             YA = rand(elty,n,m)
             d_YA = CuArray(YA)