From 9fa2d3949863d54d3cceb6fbc3038060c07ce153 Mon Sep 17 00:00:00 2001
From: Charles Kawczynski <kawczynski.charles@gmail.com>
Date: Tue, 14 Nov 2023 14:08:26 -0800
Subject: [PATCH] Fixes

---
 test/Spaces/ddss1.jl             | 12 ++++++---
 test/Spaces/distributed/ddss2.jl | 12 ++++++---
 test/Spaces/distributed/ddss3.jl | 44 +++++++++++++++++---------------
 test/Spaces/distributed/ddss4.jl |  2 ++
 4 files changed, 42 insertions(+), 28 deletions(-)

diff --git a/test/Spaces/ddss1.jl b/test/Spaces/ddss1.jl
index 908fe35d3e..0fcb1346c3 100644
--- a/test/Spaces/ddss1.jl
+++ b/test/Spaces/ddss1.jl
@@ -1,5 +1,7 @@
 using Logging
 using Test
+import CUDA
+CUDA.allowscalar(false)
 
 import ClimaCore:
     Domains,
@@ -66,10 +68,12 @@ init_state_vector(local_geometry, p) = Geometry.Covariant12Vector(1.0, -1.0)
 
     @test Topologies.nlocalelems(Spaces.topology(space)) == 4
 
-    @test Topologies.local_neighboring_elements(space.topology, 1) == [2, 4]
-    @test Topologies.local_neighboring_elements(space.topology, 2) == [1, 3]
-    @test Topologies.local_neighboring_elements(space.topology, 3) == [2, 4]
-    @test Topologies.local_neighboring_elements(space.topology, 4) == [1, 3]
+    CUDA.@allowscalar begin
+        @test Topologies.local_neighboring_elements(space.topology, 1) == [2, 4]
+        @test Topologies.local_neighboring_elements(space.topology, 2) == [1, 3]
+        @test Topologies.local_neighboring_elements(space.topology, 3) == [2, 4]
+        @test Topologies.local_neighboring_elements(space.topology, 4) == [1, 3]
+    end
 
     y0 = init_state_scalar.(Fields.local_geometry_field(space), Ref(nothing))
     nel = Topologies.nlocalelems(Spaces.topology(space))
diff --git a/test/Spaces/distributed/ddss2.jl b/test/Spaces/distributed/ddss2.jl
index f2ce3d7dbe..2158e40064 100644
--- a/test/Spaces/distributed/ddss2.jl
+++ b/test/Spaces/distributed/ddss2.jl
@@ -1,3 +1,5 @@
+import CUDA
+CUDA.allowscalar(false)
 include("ddss_setup.jl")
 
 #=
@@ -17,11 +19,13 @@ include("ddss_setup.jl")
 
     @test Topologies.nlocalelems(Spaces.topology(space)) == 2
 
-    @test Topologies.local_neighboring_elements(space.topology, 1) == [2]
-    @test Topologies.local_neighboring_elements(space.topology, 2) == [1]
+    CUDA.@allowscalar begin
+        @test Topologies.local_neighboring_elements(space.topology, 1) == [2]
+        @test Topologies.local_neighboring_elements(space.topology, 2) == [1]
 
-    @test Topologies.ghost_neighboring_elements(space.topology, 1) == [2]
-    @test Topologies.ghost_neighboring_elements(space.topology, 2) == [1]
+        @test Topologies.ghost_neighboring_elements(space.topology, 1) == [2]
+        @test Topologies.ghost_neighboring_elements(space.topology, 2) == [1]
+    end
 
     init_state(local_geometry, p) = (ρ = 1.0)
     y0 = init_state.(Fields.local_geometry_field(space), Ref(nothing))
diff --git a/test/Spaces/distributed/ddss3.jl b/test/Spaces/distributed/ddss3.jl
index acdbf994b9..64dc6c7724 100644
--- a/test/Spaces/distributed/ddss3.jl
+++ b/test/Spaces/distributed/ddss3.jl
@@ -1,3 +1,5 @@
+import CUDA
+CUDA.allowscalar(false)
 include("ddss_setup.jl")
 
 #=
@@ -37,26 +39,28 @@ partition numbers
     space, comms_ctx = distributed_space((4, 4), (false, false), (Nq, 1, 1))
 
     @test Topologies.nlocalelems(Spaces.topology(space)) == (pid == 1 ? 6 : 5)
-    if pid == 1
-        # gidx 1
-        @test Topologies.local_neighboring_elements(space.topology, 1) ==
-              [2, 5, 6]
-        @test Topologies.ghost_neighboring_elements(space.topology, 1) == []
-        # gidx 6
-        @test Topologies.local_neighboring_elements(space.topology, 6) ==
-              [1, 2, 3, 5]
-        @test space.topology.recv_elem_gidx[Topologies.ghost_neighboring_elements(
-            space.topology,
-            6,
-        )] == [7, 9, 10, 11]
-    elseif pid == 2
-        # gidx 7
-        @test Topologies.local_neighboring_elements(space.topology, 1) ==
-              [2, 4, 5]
-        @test space.topology.recv_elem_gidx[Topologies.ghost_neighboring_elements(
-            space.topology,
-            1,
-        )] == [2, 3, 4, 6, 12]
+    CUDA.@allowscalar begin
+        if pid == 1
+            # gidx 1
+            @test Topologies.local_neighboring_elements(space.topology, 1) ==
+                  [2, 5, 6]
+            @test Topologies.ghost_neighboring_elements(space.topology, 1) == []
+            # gidx 6
+            @test Topologies.local_neighboring_elements(space.topology, 6) ==
+                  [1, 2, 3, 5]
+            @test space.topology.recv_elem_gidx[Topologies.ghost_neighboring_elements(
+                space.topology,
+                6,
+            )] == [7, 9, 10, 11]
+        elseif pid == 2
+            # gidx 7
+            @test Topologies.local_neighboring_elements(space.topology, 1) ==
+                  [2, 4, 5]
+            @test space.topology.recv_elem_gidx[Topologies.ghost_neighboring_elements(
+                space.topology,
+                1,
+            )] == [2, 3, 4, 6, 12]
+        end
     end
 
     init_state(local_geometry, p) = (ρ = 1.0)
diff --git a/test/Spaces/distributed/ddss4.jl b/test/Spaces/distributed/ddss4.jl
index bbfacf4a3d..0e1f08f5c6 100644
--- a/test/Spaces/distributed/ddss4.jl
+++ b/test/Spaces/distributed/ddss4.jl
@@ -1,3 +1,5 @@
+import CUDA
+CUDA.allowscalar(false)
 include("ddss_setup.jl")
 
 #=