Refactor gpu tests into gpu/ folder

JuliaSmoothOptimizers · Oct 16, 2024 · b8ea9bc · b8ea9bc
1 parent 6dc3b23
commit b8ea9bc
Show file tree

Hide file tree

Showing 7 changed files with 31 additions and 11 deletions.
diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
@@ -11,5 +11,6 @@ steps:
       using Pkg
       Pkg.add("CUDA")
       Pkg.instantiate()
+      include("test/gpu/test_S_kwarg.jl")
       include("test/gpu/nvidia.jl")'
     timeout_in_minutes: 30
diff --git a/test/gpu/amdgpu.jl b/test/gpu/amdgpu.jl
@@ -0,0 +1,15 @@
+using Test, LinearAlgebra, SparseArrays
+using LinearOperators, AMDGPU
+
+@testset "AMDGPU -- AMDGPU.jl" begin
+  A = ROCArray(rand(Float32, 5, 5))
+  B = ROCArray(rand(Float32, 10, 10))
+  C = ROCArray(rand(Float32, 20, 20))
+  M = BlockDiagonalOperator(A, B, C)
+
+  v = ROCArray(rand(Float32, 35))
+  y = M * v
+  @test y isa ROCArray{Float32}
+
+  @testset "AMDGPU S kwarg" test_S_kwarg(arrayType = ROCArray)
+end
diff --git a/test/gpu/jlarrays.jl b/test/gpu/jlarrays.jl
@@ -0,0 +1 @@
+test_S_kwarg(arrayType = JLArray)
diff --git a/test/gpu/metal.jl b/test/gpu/metal.jl
@@ -0,0 +1,3 @@
+if Sys.isapple() && occursin("arm64", Sys.MACHINE)
+  test_S_kwarg(arrayType = MtlArray, notMetal = false)
+end
diff --git a/test/gpu/nvidia.jl b/test/gpu/nvidia.jl
@@ -13,4 +13,5 @@ using LinearOperators, CUDA, CUDA.CUSPARSE, CUDA.CUSOLVER
   v = CUDA.rand(35)
   y = M * v
   @test y isa CuVector{Float32}
+  @testset "Nvidia S kwarg" test_S_kwarg(arrayType = CuArray)
 end
diff --git a/test/test_S_kwarg.jl → test/gpu/test_S_kwarg.jl b/test/test_S_kwarg.jl → test/gpu/test_S_kwarg.jl
@@ -1,4 +1,6 @@
-function test_S_kwarg(; arrayType = JLArray, notMetal = true)
+using Test, LinearOperators, LinearAlgebra
+
+function test_S_kwarg(; arrayType, notMetal = true)
   mat = arrayType(rand(Float32, 32, 32))
   vec = arrayType(rand(Float32, 32))
   vecT = typeof(vec)
@@ -8,7 +10,7 @@ function test_S_kwarg(; arrayType = JLArray, notMetal = true)
     vecTother = typeof(arrayType(rand(Float32, 32)))
   end
 
-  @testset ExtendedTestSet "S Kwarg with arrayType $(arrayType)" begin
+  @testset "S Kwarg with arrayType $(arrayType)" begin
     @test vecT == LinearOperators.storage_type(mat)
 
     # constructors.jl
@@ -37,9 +39,4 @@ function test_S_kwarg(; arrayType = JLArray, notMetal = true)
     notMetal && @test LinearOperators.storage_type(BlockDiagonalOperator(mat, mat; S = vecTother)) == vecTother
   end
 
-end
-
-test_S_kwarg()
-if Sys.isapple() && occursin("arm64", Sys.MACHINE)
-    test_S_kwarg(arrayType = MtlArray, notMetal = false)
 end
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -1,9 +1,6 @@
 using Arpack, Test, TestSetExtensions, LinearOperators
 using LinearAlgebra, LDLFactorizations, SparseArrays, JLArrays
 using Zygote
-if Sys.isapple() && occursin("arm64", Sys.MACHINE)
-    using Metal
-end
 include("test_aux.jl")
 
 include("test_linop.jl")
@@ -19,4 +16,9 @@ include("test_normest.jl")
 include("test_diag.jl")
 include("test_chainrules.jl")
 include("test_solve_shifted_system.jl")
-include("test_S_kwarg.jl")
+include("gpu/test_S_kwarg.jl")
+include("gpu/jlarrays.jl")
+if Sys.isapple() && occursin("arm64", Sys.MACHINE)
+    using Metal
+    include("gpu/metal.jl")
+end