Merge pull request #1791 from CliMA/ck/test_rm_cuda

Remove CUDA from the test extras

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"]
 

test/MatrixFields/field_matrix_solvers.jl

@@ -54,8 +54,8 @@ 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).
         ignored = (
-            ignore_cuda...,
-            using_cuda ? AnyFrameModule(CUDA.CUBLAS) :
+            cuda_frames...,
+            cublas_frames...,
             AnyFrameModule(MatrixFields.KrylovKit),
             AnyFrameModule(Base.CoreLogging),
         )

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,11 @@ 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 cuda_frames = using_cuda ? (AnyFrameModule(cuda_mod),) : ()
+const cublas_frames = using_cuda ? (AnyFrameModule(cuda_mod.CUBLAS),) : ()
+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 +67,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
@@ -101,14 +105,14 @@ function test_field_broadcast(;
         # Test get_result and set_result! for type instabilities, and test
         # set_result! for allocations. Ignore the type instabilities in CUDA and
         # the allocations they incur.
-        @test_opt ignored_modules = ignore_cuda get_result()
-        @test_opt ignored_modules = ignore_cuda set_result!(result)
+        @test_opt ignored_modules = cuda_frames get_result()
+        @test_opt ignored_modules = cuda_frames set_result!(result)
         using_cuda || @test (@allocated set_result!(result)) == 0
 
         if !isnothing(ref_set_result!)
             # Test ref_set_result! for type instabilities and allocations to
             # ensure that the performance comparison is fair.
-            @test_opt ignored_modules = ignore_cuda ref_set_result!(ref_result)
+            @test_opt ignored_modules = cuda_frames ref_set_result!(ref_result)
             using_cuda || @test (@allocated ref_set_result!(ref_result)) == 0
         end
     end
@@ -133,7 +137,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
@@ -183,13 +187,13 @@ function test_field_broadcast_against_array_reference(;
         # Test get_result and set_result! for type instabilities, and test
         # set_result! for allocations. Ignore the type instabilities in CUDA and
         # the allocations they incur.
-        @test_opt ignored_modules = ignore_cuda get_result()
-        @test_opt ignored_modules = ignore_cuda set_result!(result)
+        @test_opt ignored_modules = cuda_frames get_result()
+        @test_opt ignored_modules = cuda_frames set_result!(result)
         using_cuda || @test (@allocated set_result!(result)) == 0
 
         # Test ref_set_result! for type instabilities and allocations to ensure
         # that the performance comparison is fair.
-        @test_opt ignored_modules = ignore_cuda call_ref_set_result!()
+        @test_opt ignored_modules = cuda_frames call_ref_set_result!()
         using_cuda || @test (@allocated call_ref_set_result!()) == 0
     end
 end

test/MatrixFields/operator_matrices.jl

@@ -321,7 +321,7 @@ end
         )),
     )
         get_result()
-        @test_opt ignored_modules = ignore_cuda get_result()
+        @test_opt ignored_modules = cuda_frames get_result()
     end
 
     test_field_broadcast(;