More diag mul methods (#631)

Author: kshyatt

src/host/linalg.jl

@@ -266,8 +266,8 @@ function LinearAlgebra.mul!(C::Diagonal{<:Any, <:AbstractGPUArray},
 end
 
 function LinearAlgebra.mul!(C::Diagonal{<:Any, <:AbstractGPUArray},
-                            A::AbstractGPUArray,
-                            B::AbstractGPUArray)
+                            A::Union{AbstractGPUArray, Adjoint{T,<:AbstractGPUArray{T}}, Transpose{T,<:AbstractGPUArray{T}}},
+                            B::Union{AbstractGPUArray, Adjoint{T,<:AbstractGPUArray{T}}, Transpose{T,<:AbstractGPUArray{T}}}) where {T}
     dc = C.diag
     d  = length(dc)
     m, n   = size(A, 1), size(A, 2)
@@ -282,23 +282,22 @@ end
 
 function LinearAlgebra.mul!(B::AbstractGPUVecOrMat,
                             D::Diagonal{<:Any, <:AbstractGPUArray},
-                            A::AbstractGPUVecOrMat)
+                            A::Union{AbstractGPUArray, Adjoint{T,<:AbstractGPUArray{T}}, Transpose{T,<:AbstractGPUArray{T}}}) where {T}
     dd = D.diag
     d = length(dd)
     m, n = size(A, 1), size(A, 2)
     m′, n′ = size(B, 1), size(B, 2)
     m == d || throw(DimensionMismatch("right hand side has $m rows but D is $d by $d"))
     (m, n) == (m′, n′) || throw(DimensionMismatch("expect output to be $m by $n, but got $m′ by $n′"))
     @. B = dd * A
-
     B
 end
 
 function LinearAlgebra.mul!(B::AbstractGPUVecOrMat,
                             D::Diagonal{<:Any, <:AbstractGPUArray},
-                            A::AbstractGPUVecOrMat,
+                            A::Union{AbstractGPUArray, Adjoint{T,<:AbstractGPUArray{T}}, Transpose{T,<:AbstractGPUArray{T}}},
                             α::Number,
-                            β::Number)
+                            β::Number) where {T}
     dd = D.diag
     d = length(dd)
     m, n = size(A, 1), size(A, 2)

test/testsuite/linalg.jl

@@ -238,6 +238,9 @@
                 mul!(X, D, B)
                 mul!(Y, Diagonal(collect(d)), collect(B))
                 @test collect(X) ≈ Y
+                mul!(X, D, adjoint(B))
+                mul!(Y, Diagonal(collect(d)), collect(adjoint(B)))
+                @test collect(X) ≈ Y
                 mul!(X, D, B, α, β)
                 mul!(Y, Diagonal(collect(d)), collect(B), α, β)
                 @test collect(X) ≈ Y
@@ -259,6 +262,11 @@
                 C = Diagonal(d)
                 mul!(C, a, b)
                 @test collect(C) ≈ Diagonal(collect(a) * collect(b)) 
+                a = transpose(AT(diagm(rand(elty, n))))
+                b = adjoint(AT(diagm(rand(elty, n))))
+                C = Diagonal(d)
+                mul!(C, a, b)
+                @test collect(C) ≈ Diagonal(collect(a) * collect(b)) 
             end
         end