Merge pull request #2100 from CliMA/sk/swap_grid_dims

Swap CUDA grid dimensions for some partitions

Author: charleskawczynski

.buildkite/pipeline.yml

@@ -266,6 +266,7 @@ steps:
           - "julia --project=.buildkite -e 'using CUDA; CUDA.versioninfo()'"
           - "srun julia --color=yes --check-bounds=yes --project=.buildkite test/Spaces/unit_spaces.jl"
           - "srun julia --color=yes --check-bounds=yes --project=.buildkite test/Spaces/opt_spaces.jl"
+          - "julia --color=yes --check-bounds=yes --project=.buildkite test/Spaces/high_resolution_space.jl"
         env:
           CLIMACOMMS_DEVICE: "CUDA"
         agents:

ext/cuda/data_layouts_threadblock.jl

@@ -55,11 +55,11 @@ function is_valid_index end
     Nv_thread = min(Int(fld(n_max_threads, Nij * Nij)), Nv)
     Nv_blocks = cld(Nv, Nv_thread)
     @assert prod((Nv_thread, Nij, Nij)) ≤ n_max_threads "threads,n_max_threads=($(prod((Nv_thread, Nij, Nij))),$n_max_threads)"
-    return (; threads = (Nv_thread, Nij, Nij), blocks = (Nv_blocks, Nh))
+    return (; threads = (Nv_thread, Nij, Nij), blocks = (Nh, Nv_blocks))
 end
 @inline function universal_index(::Union{DataLayouts.VIJFH, DataLayouts.VIJHF})
     (tv, i, j) = CUDA.threadIdx()
-    (bv, h) = CUDA.blockIdx()
+    (h, bv) = CUDA.blockIdx()
     v = tv + (bv - 1) * CUDA.blockDim().x
     return CartesianIndex((i, j, 1, v, h))
 end
@@ -152,11 +152,11 @@ end
     Nv_thread = min(Int(fld(n_max_threads, Ni)), Nv)
     Nv_blocks = cld(Nv, Nv_thread)
     @assert prod((Nv_thread, Ni)) ≤ n_max_threads "threads,n_max_threads=($(prod((Nv_thread, Ni))),$n_max_threads)"
-    return (; threads = (Nv_thread, Ni), blocks = (Nv_blocks, Nh))
+    return (; threads = (Nv_thread, Ni), blocks = (Nh, Nv_blocks))
 end
 @inline function universal_index(::Union{DataLayouts.VIFH, DataLayouts.VIHF})
     (tv, i) = CUDA.threadIdx()
-    (bv, h) = CUDA.blockIdx()
+    (h, bv) = CUDA.blockIdx()
     v = tv + (bv - 1) * CUDA.blockDim().x
     return CartesianIndex((i, 1, 1, v, h))
 end

test/Spaces/high_resolution_space.jl

@@ -0,0 +1,22 @@
+using Test
+import ClimaComms
+ClimaComms.@import_required_backends
+import ClimaCore
+@testset "Construct high resolution space on GPU" begin
+    result = try
+        ClimaCore.CommonSpaces.ExtrudedCubedSphereSpace(
+            Float32;
+            radius = 1.0,
+            h_elem = 105,
+            z_elem = 10,
+            z_min = 1.0,
+            z_max = 2.0,
+            n_quad_points = 4,
+            staggering = ClimaCore.Grids.CellCenter(),
+        )
+    catch
+        println("unable to create center space")
+        false
+    end
+    @test result isa ClimaCore.Spaces.CenterExtrudedFiniteDifferenceSpace
+end