Even more

Author: kshyatt

src/host/linalg.jl

@@ -275,23 +275,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

@@ -234,6 +234,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