Merge branch 'master' into sds/rng_wavefront64

Author: simeonschaub

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/Project.toml

@@ -7,6 +7,7 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

test/device/launch.jl

@@ -131,7 +131,7 @@ end
     # end
 end
 
-if !iszero(AMDGPU.HIP.properties(AMDGPU.device()).cooperativeLaunch)
+if VERSION >= v"1.12-" && !iszero(AMDGPU.HIP.properties(AMDGPU.device()).cooperativeLaunch)
     @testset "Cooperative Groups" begin
         function test_kernel!(x)
             block_row, block_col = workgroupIdx().x, workgroupIdx().y
@@ -151,6 +151,8 @@ if !iszero(AMDGPU.HIP.properties(AMDGPU.device()).cooperativeLaunch)
                 end
                 AMDGPU.Device.sync_grid()
             end
+
+            return nothing
         end
 
         n = 4

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