Merge pull request #821 from kshyatt/ksh/lrmul

Author: luraess

src/blas/highlevel.jl

@@ -330,3 +330,41 @@ if VERSION ≥ v"1.12-"
     LinearAlgebra.copytrito!(B::Matrix{T}, A::ROCMatrix{T}, uplo::AbstractChar) where {T <: ROCBLASFloat} = 
         invoke(LinearAlgebra.copytrito!, Tuple{AbstractMatrix, AbstractMatrix, AbstractChar}, B, A, uplo)
 end
+
+function LinearAlgebra.lmul!(A::Diagonal{T,<:ROCVector{T}}, B::ROCMatrix{T}) where {T<:ROCBLASFloat}
+    return dgmm!('L', B, A.diag, B)
+end
+
+function LinearAlgebra.rmul!(A::ROCMatrix{T}, B::Diagonal{T,<:ROCVector{T}}) where {T<:ROCBLASFloat}
+    return dgmm!('R', A, B.diag, A)
+end
+
+# eltypes do not match
+function LinearAlgebra.lmul!(A::Diagonal{T,<:ROCVector{T}}, B::ROCMatrix) where {T<:ROCBLASFloat}
+    @. B = A.diag * B
+    return B
+end
+function LinearAlgebra.lmul!(A::Diagonal{Td,<:ROCVector{Td}}, B::Transpose{Tt, <:ROCMatrix{Tt}}) where {Td<:ROCBLASFloat, Tt<:ROCBLASFloat}
+    @. B = A.diag * B
+    return B
+end
+function LinearAlgebra.lmul!(A::Diagonal{Td,<:ROCVector{Td}}, B::Adjoint{Tt, <:ROCMatrix{Tt}}) where {Td<:ROCBLASFloat, Tt<:ROCBLASFloat}
+    @. B = A.diag * B
+    return B
+end
+# eltypes do not match
+function LinearAlgebra.rmul!(A::ROCMatrix, B::Diagonal{T,<:ROCVector{T}}) where {T<:ROCBLASFloat}
+    At = transpose(A)
+    @. At = B.diag * At
+    return A
+end
+function LinearAlgebra.rmul!(A::Transpose{Tt, <:ROCMatrix{Tt}}, B::Diagonal{Td,<:ROCVector{Td}}) where {Td<:ROCBLASFloat, Tt<:ROCBLASFloat}
+    At = parent(A)
+    @. At = B.diag * At
+    return transpose(At)
+end
+function LinearAlgebra.rmul!(A::Adjoint{Tt, <:ROCMatrix{Tt}}, B::Diagonal{Td,<:ROCVector{Td}}) where {Td<:ROCBLASFloat, Tt<:ROCBLASFloat}
+    At = parent(A)
+    @. At = adjoint(B.diag) * At
+    return adjoint(At)
+end

test/rocarray/blas.jl

@@ -288,4 +288,56 @@ end
     end
 end
 
+@testset "Extension" begin
+    for TA in (Float32, Float64, ComplexF32, ComplexF64), TB in (Float32, Float64)
+        x    = rand(TB, m)
+        d_x  = ROCArray(x)
+        XA   = rand(TA, m, n)
+        d_XA = ROCArray(XA)
+        d_X  = Diagonal(d_x)
+        lmul!(d_X, d_XA)
+        @test Array(d_XA) ≈ Diagonal(x) * XA
+
+        x    = rand(TB, m)
+        d_x  = ROCArray(x)
+        XA   = rand(TA, n, m)
+        d_AX = transpose(ROCArray(XA))
+        d_X  = Diagonal(d_x)
+        lmul!(d_X, d_AX)
+        @test Array(d_AX) ≈ Diagonal(x) * transpose(XA)
+
+        x    = rand(TB, m)
+        d_x  = ROCArray(x)
+        XA   = rand(TA, n, m)
+        d_AX = adjoint(ROCArray(XA))
+        d_X  = Diagonal(d_x)
+        lmul!(d_X, d_AX)
+        @test Array(d_AX) ≈ Diagonal(x) * adjoint(XA)
+
+        y    = rand(TB, n)
+        d_y  = ROCArray(y)
+        AY   = rand(TA, m, n)
+        d_AY = ROCArray(AY)
+        d_Y  = Diagonal(d_y)
+        rmul!(d_AY, d_Y)
+        @test Array(d_AY) ≈ AY * Diagonal(y)
+
+        y    = rand(TB, n)
+        d_y  = ROCArray(y)
+        AY   = rand(TA, n, m)
+        d_YA = transpose(ROCArray(AY))
+        d_Y  = Diagonal(d_y)
+        d_YA = rmul!(d_YA, d_Y)
+        @test Array(d_YA) ≈ transpose(AY) * Diagonal(y)
+
+        y    = rand(TB, n)
+        d_y  = ROCArray(y)
+        AY   = rand(TA, n, m)
+        d_YA = adjoint(ROCArray(AY))
+        d_Y  = Diagonal(d_y)
+        d_YA = rmul!(d_YA, d_Y)
+        @test Array(d_YA) ≈ adjoint(AY) * Diagonal(y)
+    end
+end
+
 end