From 5ed90172dccdacd2be23f720284cf08103bb332a Mon Sep 17 00:00:00 2001
From: Charles Kawczynski <kawczynski.charles@gmail.com>
Date: Mon, 10 Jun 2024 15:34:46 -0400
Subject: [PATCH] Remove CUDA from the test extras

wip

Use get_extension instead of eval-parse

Add composed function stencil benchmark
---
 Project.toml                                 |  3 +--
 test/MatrixFields/field_matrix_solvers.jl    |  3 ++-
 test/MatrixFields/matrix_field_test_utils.jl | 11 +++++++----
 3 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/Project.toml b/Project.toml
index d47e3ac08e..979be8512f 100644
--- 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
index 28474db552..0caa9329f5 100644
--- 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
index e60b157109..e98f7fa2b9 100644
--- 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