Specialize LinearAlgebra.mul! on Diagonal for GPUArrays (#416)

Author: vpuri3

src/host/linalg.jl

@@ -173,6 +173,36 @@ if VERSION < v"1.8-"
         return B
     end
 else
+    function LinearAlgebra.mul!(B::AbstractGPUVecOrMat,
+                                D::Diagonal{<:Any, <:AbstractGPUArray},
+                                A::AbstractGPUVecOrMat)
+        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,
+                                α::Number,
+                                β::Number)
+        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
+
+        B
+    end
+
     function LinearAlgebra.ldiv!(B::AbstractGPUVecOrMat,
                                  D::Diagonal{<:Any, <:AbstractGPUArray},
                                  A::AbstractGPUVecOrMat)

test/testsuite/linalg.jl

@@ -153,6 +153,23 @@
             @test_throws SingularException D \ B
         end
 
+        @testset "mul! + Diagonal" begin
+            n = 128
+            d = AT(rand(Float32, n))
+            D = Diagonal(d)
+            B = AT(rand(Float32, n, n))
+            X = AT(zeros(Float32, n, n))
+            Y = zeros(Float32, n, n)
+            α = rand(Float32)
+            β = rand(Float32)
+            mul!(X, D, B)
+            mul!(Y, Diagonal(collect(d)), collect(B))
+            @test collect(X) ≈ Y
+            mul!(X, D, B, α, β)
+            mul!(Y, Diagonal(collect(d)), collect(B), α, β)
+            @test collect(X) ≈ Y
+        end
+
         @testset "ldiv! + Diagonal" begin
             n = 128
             d = AT(rand(Float32, n))