Remove CUDA from the test extras

wip Use get_extension instead of eval-parse Add composed function stencil benchmark
CliMA · Jun 28, 2024 · 5ed9017 · 5ed9017
1 parent 13c6176
commit 5ed9017
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 7 deletions.
diff --git a/Project.toml b/Project.toml
@@ -93,7 +93,6 @@ AssociatedLegendrePolynomials = "2119f1ac-fb78-50f5-8cc0-dda848ebdb19"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
 CountFlops = "1db9610d-79e1-487a-8d40-77f3295c7593"
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898"
 Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7"
@@ -109,5 +108,5 @@ TerminalLoggers = "5d786b92-1e48-4d6f-9151-6b4477ca9bed"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "ArgParse", "AssociatedLegendrePolynomials", "BenchmarkTools", "Combinatorics", "CountFlops", "CUDA", "Dates", "FastBroadcast", "Krylov", "JET", "Logging", "MPI", "OrderedCollections", "PrettyTables", "Random", "SafeTestsets", "StatsBase", "TerminalLoggers", "Test"]
+test = ["Aqua", "ArgParse", "AssociatedLegendrePolynomials", "BenchmarkTools", "Combinatorics", "CountFlops", "Dates", "FastBroadcast", "Krylov", "JET", "Logging", "MPI", "OrderedCollections", "PrettyTables", "Random", "SafeTestsets", "StatsBase", "TerminalLoggers", "Test"]
 
diff --git a/test/MatrixFields/field_matrix_solvers.jl b/test/MatrixFields/field_matrix_solvers.jl
@@ -53,9 +53,10 @@ function test_field_matrix_solver(; test_name, alg, A, b, use_rel_error = false)
 
         # In addition to ignoring the type instabilities from CUDA, ignore those
         # from CUBLAS (norm), KrylovKit (eigsolve), and CoreLogging (@debug).
+        cuda_mod = Base.get_extension(ClimaComms, :ClimaCommsCUDAExt).CUDA
         ignored = (
             ignore_cuda...,
-            using_cuda ? AnyFrameModule(CUDA.CUBLAS) :
+            using_cuda ? AnyFrameModule(cuda_mod.CUBLAS) :
             AnyFrameModule(MatrixFields.KrylovKit),
             AnyFrameModule(Base.CoreLogging),
         )

diff --git a/test/MatrixFields/matrix_field_test_utils.jl b/test/MatrixFields/matrix_field_test_utils.jl
@@ -1,9 +1,9 @@
 using Test
 using JET
-import CUDA
 import Random: seed!
 
 import ClimaComms
+ClimaComms.@import_required_backends
 import ClimaCore:
     Geometry,
     Domains,
@@ -45,7 +45,10 @@ const comms_device = ClimaComms.device()
 # comms_device = ClimaComms.CPUSingleThreaded()
 @show comms_device
 const using_cuda = comms_device isa ClimaComms.CUDADevice
-const ignore_cuda = using_cuda ? (AnyFrameModule(CUDA),) : ()
+cuda_module(ext) = using_cuda ? ext.CUDA : ext
+const cuda_mod = cuda_module(Base.get_extension(ClimaComms, :ClimaCommsCUDAExt))
+const ignore_cuda = using_cuda ? (AnyFrameModule(cuda_mod),) : ()
+const invalid_ir_error = using_cuda ? cuda_mod.InvalidIRError : ErrorException
 
 # Test the allocating and non-allocating versions of a field broadcast against
 # a reference non-allocating implementation. Ensure that they are performant,
@@ -63,7 +66,7 @@ function test_field_broadcast(;
 ) where {F1, F2, F3}
     @testset "$test_name" begin
         if test_broken_with_cuda && using_cuda
-            @test_throws CUDA.InvalidIRError get_result()
+            @test_throws invalid_ir_error get_result()
             @warn "$test_name:\n\tCUDA.InvalidIRError"
             return
         end
@@ -133,7 +136,7 @@ function test_field_broadcast_against_array_reference(;
 ) where {F1, F2, F3}
     @testset "$test_name" begin
         if test_broken_with_cuda && using_cuda
-            @test_throws CUDA.InvalidIRError get_result()
+            @test_throws invalid_ir_error get_result()
             @warn "$test_name:\n\tCUDA.InvalidIRError"
             return
         end