Test GPU QR in testsuite

Author: kshyatt

Project.toml

@@ -30,6 +30,7 @@ GenericLinearAlgebra = "0.3.19"
 GenericSchur = "0.5.6"
 JET = "0.9, 0.10"
 LinearAlgebra = "1"
+Random = "1"
 SafeTestsets = "0.1"
 StableRNGs = "1"
 Test = "1"
@@ -43,11 +44,12 @@ Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "JET", "SafeTestsets", "Test", "TestExtras", "ChainRulesCore", "ChainRulesTestUtils", "StableRNGs", "Zygote", "CUDA", "AMDGPU", "GenericLinearAlgebra", "GenericSchur"]
+test = ["Aqua", "JET", "SafeTestsets", "Test", "TestExtras", "ChainRulesCore", "ChainRulesTestUtils", "StableRNGs", "Zygote", "CUDA", "AMDGPU", "GenericLinearAlgebra", "GenericSchur", "Random"]

src/implementations/qr.jl

@@ -270,10 +270,12 @@ function _gpu_unmqr!(
 end
 
 function _gpu_qr!(
-        A::AbstractMatrix, Q::AbstractMatrix, R::AbstractMatrix; positive = false, blocksize = 1
+        A::AbstractMatrix, Q::AbstractMatrix, R::AbstractMatrix; pivoted = false, positive = false, blocksize = 1
     )
     blocksize > 1 &&
         throw(ArgumentError("CUSOLVER/ROCSOLVER does not provide a blocked implementation for a QR decomposition"))
+    pivoted &&
+        throw(ArgumentError("CUSOLVER/ROCSOLVER does not provide a pivoted implementation for a QR decomposition"))
     m, n = size(A)
     minmn = min(m, n)
     computeR = length(R) > 0
@@ -309,10 +311,12 @@ function _gpu_qr!(
 end
 
 function _gpu_qr_null!(
-        A::AbstractMatrix, N::AbstractMatrix; positive = false, blocksize = 1
+        A::AbstractMatrix, N::AbstractMatrix; positive = false, blocksize = 1, pivoted = false
     )
     blocksize > 1 &&
         throw(ArgumentError("CUSOLVER/ROCSOLVER does not provide a blocked implementation for a QR decomposition"))
+    pivoted &&
+        throw(ArgumentError("CUSOLVER/ROCSOLVER does not provide a pivoted implementation for a QR decomposition"))
     m, n = size(A)
     minmn = min(m, n)
     fill!(N, zero(eltype(N)))

test/amd/qr.jl

@@ -3,6 +3,7 @@ using Test
 using TestExtras
 using StableRNGs
 using LinearAlgebra: diag, I, Diagonal
+using CUDA, AMDGPU
 
 BLASFloats = (Float32, Float64, ComplexF32, ComplexF64)
 GenericFloats = (Float16, BigFloat, Complex{BigFloat})
@@ -14,6 +15,14 @@ m = 54
 for T in BLASFloats, n in (37, m, 63)
     TestSuite.seed_rng!(123)
     TestSuite.test_qr(T, (m, n))
+    if CUDA.functional()
+        TestSuite.test_qr(CuMatrix{T}, (m, n); test_pivoted = false, test_blocksize = false)
+        TestSuite.test_qr(Diagonal{T, CuVector{T}}, m; test_pivoted = false, test_blocksize = false)
+    end
+    if AMDGPU.functional()
+        TestSuite.test_qr(ROCMatrix{T}, (m, n); test_pivoted = false, test_blocksize = false)
+        TestSuite.test_qr(Diagonal{T, ROCVector{T}}, m; test_pivoted = false, test_blocksize = false)
+    end
 end
 for T in (BLASFloats..., GenericFloats...)
     AT = Diagonal{T, Vector{T}}

test/cuda/qr.jl

@@ -57,9 +57,6 @@ end
 
 using CUDA
 if CUDA.functional()
-    @safetestset "CUDA QR" begin
-        include("cuda/qr.jl")
-    end
     @safetestset "CUDA LQ" begin
         include("cuda/lq.jl")
     end
@@ -85,9 +82,6 @@ end
 
 using AMDGPU
 if AMDGPU.functional()
-    @safetestset "AMDGPU QR" begin
-        include("amd/qr.jl")
-    end
     @safetestset "AMDGPU LQ" begin
         include("amd/lq.jl")
     end

test/qr.jl

@@ -13,6 +13,7 @@ using MatrixAlgebraKit
 using MatrixAlgebraKit: diagview
 using LinearAlgebra: Diagonal, norm, istriu
 using Random, StableRNGs
+using AMDGPU, CUDA
 
 const tests = Dict()
 
@@ -33,7 +34,11 @@ seed_rng!(seed) = Random.seed!(rng, seed)
 
 instantiate_matrix(::Type{T}, size) where {T <: Number} = randn(rng, T, size)
 instantiate_matrix(::Type{AT}, size) where {AT <: Array} = randn(rng, eltype(AT), size)
+instantiate_matrix(::Type{AT}, size) where {AT <: CuArray} = CuArray(randn(rng, eltype(AT), size))
+instantiate_matrix(::Type{AT}, size) where {AT <: ROCArray} = ROCArray(randn(rng, eltype(AT), size))
 instantiate_matrix(::Type{AT}, size) where {AT <: Diagonal} = Diagonal(randn(rng, eltype(AT), size))
+instantiate_matrix(::Type{AT}, size) where {T, AT <: Diagonal{T, <:CuVector}} = Diagonal(CuArray(randn(rng, eltype(AT), size)))
+instantiate_matrix(::Type{AT}, size) where {T, AT <: Diagonal{T, <:ROCVector}} = Diagonal(ROCArray(randn(rng, eltype(AT), size)))
 
 precision(::Type{T}) where {T <: Number} = sqrt(eps(real(T)))
 precision(::Type{T}) where {T} = precision(eltype(T))

test/runtests.jl

@@ -1,9 +1,15 @@
 function test_qr(T::Type, sz; kwargs...)
     summary_str = testargs_summary(T, sz)
     return @testset "qr $summary_str" begin
-        test_qr_compact(T, sz; kwargs...)
-        test_qr_full(T, sz; kwargs...)
-        test_qr_null(T, sz; kwargs...)
+        if T <: Union{CuArray, ROCArray}
+            test_qr_compact(T, sz; pivoted = false, kwargs...)
+            test_qr_full(T, sz; pivoted = false, kwargs...)
+            test_qr_null(T, sz; pivoted = false, kwargs...)
+        else
+            test_qr_compact(T, sz; kwargs...)
+            test_qr_full(T, sz; kwargs...)
+            test_qr_null(T, sz; kwargs...)
+        end
     end
 end