Fix rmul for transpose/adjoint (#2871)

Author: kshyatt

lib/cublas/linalg.jl

@@ -441,9 +441,21 @@ function LinearAlgebra.lmul!(A::Diagonal{Td,<:CuVector{Td}}, B::Adjoint{Tt, <:Cu
     return B
 end
 # eltypes do not match
-LinearAlgebra.rmul!(A::CuMatrix, B::Diagonal{T,<:CuVector{T}}) where {T<:CublasFloat} = lmul!(B, transpose(A))
-LinearAlgebra.rmul!(A::Transpose{Tt, <:CuMatrix{Tt}}, B::Diagonal{Td,<:CuVector{Td}}) where {Td<:CublasFloat, Tt<:CublasFloat} = lmul!(B, A)
-LinearAlgebra.rmul!(A::Adjoint{Tt, <:CuMatrix{Tt}}, B::Diagonal{Td,<:CuVector{Td}}) where {Td<:CublasFloat, Tt<:CublasFloat} = conj!(lmul!(B, conj!(A)))
+function LinearAlgebra.rmul!(A::CuMatrix, B::Diagonal{T,<:CuVector{T}}) where {T<:CublasFloat}
+    At = transpose(A)
+    @. At = B.diag * At
+    return A
+end
+function LinearAlgebra.rmul!(A::Transpose{Tt, <:CuMatrix{Tt}}, B::Diagonal{Td,<:CuVector{Td}}) where {Td<:CublasFloat, Tt<:CublasFloat}
+    At = parent(A)
+    @. At = B.diag * At
+    return transpose(At)
+end
+function LinearAlgebra.rmul!(A::Adjoint{Tt, <:CuMatrix{Tt}}, B::Diagonal{Td,<:CuVector{Td}}) where {Td<:CublasFloat, Tt<:CublasFloat}
+    At = parent(A)
+    @. At = adjoint(B.diag) * At
+    return adjoint(At)
+end
 
 # diagm
 

test/libraries/cublas/extensions.jl

@@ -594,19 +594,51 @@ k = 13
 end # elty
 
 @testset "rmul/lmul with mixed eltypes ($Tr, $Tc)" for (Tr, Tc) in ((Float32, ComplexF32), (Float64, ComplexF64))
-    x = rand(Tr,m)
-    d_x = CuArray(x)
-    XA = rand(Tc,m,n)
+    x    = rand(Tr,m)
+    d_x  = CuArray(x)
+    XA   = rand(Tc,m,n)
     d_XA = CuArray(XA)
-    d_X = Diagonal(d_x)
+    d_X  = Diagonal(d_x)
     lmul!(d_X, d_XA)
     @test Array(d_XA) ≈ Diagonal(x) * XA
-            
-    y = rand(Tr,n)
-    d_y = CuArray(y)
-    AY = rand(Tc,m,n)
+
+    x    = rand(Tr,m)
+    d_x  = CuArray(x)
+    XA   = rand(Tc,n,m)
+    d_AX = transpose(CuArray(XA))
+    d_X  = Diagonal(d_x)
+    lmul!(d_X, d_AX)
+    @test Array(d_AX) ≈ Diagonal(x) * transpose(XA)
+
+    x    = rand(Tr,m)
+    d_x  = CuArray(x)
+    XA   = rand(Tc,n,m)
+    d_AX = adjoint(CuArray(XA))
+    d_X  = Diagonal(d_x)
+    lmul!(d_X, d_AX)
+    @test Array(d_AX) ≈ Diagonal(x) * adjoint(XA)
+
+    y    = rand(Tr,n)
+    d_y  = CuArray(y)
+    AY   = rand(Tc,m,n)
     d_AY = CuArray(AY)
-    d_Y = Diagonal(d_y)
+    d_Y  = Diagonal(d_y)
     rmul!(d_AY, d_Y)
     @test Array(d_AY) ≈ AY * Diagonal(y)
+
+    y    = rand(Tr,n)
+    d_y  = CuArray(y)
+    AY   = rand(Tc,n,m)
+    d_YA = transpose(CuArray(AY))
+    d_Y  = Diagonal(d_y)
+    d_YA = rmul!(d_YA, d_Y)
+    @test Array(d_YA) ≈ transpose(AY) * Diagonal(y)
+
+    y    = rand(Tr,n)
+    d_y  = CuArray(y)
+    AY   = rand(Tc,n,m)
+    d_YA = adjoint(CuArray(AY))
+    d_Y  = Diagonal(d_y)
+    d_YA = rmul!(d_YA, d_Y)
+    @test Array(d_YA) ≈ adjoint(AY) * Diagonal(y)
 end