interpolate_array: add support for GPUs

CliMA · Nov 1, 2023 · 91b6267 · 91b6267
1 parent 28bbd23
commit 91b6267
Show file tree

Hide file tree

Showing 5 changed files with 736 additions and 324 deletions.
diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
@@ -231,12 +231,25 @@ steps:
         command: "srun julia --color=yes --check-bounds=yes --project=test test/Remapping/distributed_remapping.jl"
         env:
           CLIMACOMMS_CONTEXT: "MPI"
+          CLIMACOMMS_DEVICE: "CPU"
         agents:
           slurm_ntasks: 2
 
       - label: "Unit: distributed remapping (1 process)"
         key: distributed_remapping_1proc
         command: "julia --color=yes --check-bounds=yes --project=test test/Remapping/distributed_remapping.jl"
+        env:
+          CLIMACOMMS_DEVICE: "CPU"
+
+      - label: "Unit: distributed remapping with CUDA"
+        key: distributed_remapping_gpu
+        command: "srun julia --color=yes --check-bounds=yes --project=test test/Remapping/distributed_remapping.jl"
+        env:
+          CLIMACOMMS_CONTEXT: "MPI"
+          CLIMACOMMS_DEVICE: "CUDA"
+        agents:
+          slurm_ntasks: 2
+          slurm_gpus: 1
 
       - label: "Unit: distributed gather"
         key: unit_distributed_gather4

diff --git a/src/Remapping/Remapping.jl b/src/Remapping/Remapping.jl
@@ -13,6 +13,7 @@ import ..DataLayouts,
     ..Hypsography
 
 using ..RecursiveApply
+using CUDA
 
 include("interpolate_array.jl")
 include("distributed_remapping.jl")

diff --git a/src/Remapping/distributed_remapping.jl b/src/Remapping/distributed_remapping.jl
@@ -239,10 +239,17 @@ function interpolate(
 
     if length(remapper.target_zcoords) == 0
         out_local_array = zeros(FT, size(remapper.local_target_hcoords_bitmask))
-        interpolated_values = [
-            interpolate_slab(field, Fields.SlabIndex(nothing, gidx), weights) for (gidx, weights) in
-            zip(remapper.local_indices, remapper.interpolation_coeffs)
-        ]
+
+        interpolated_values = zeros(FT, length(remapper.local_indices))
+        slab_indices =
+            [Fields.SlabIndex(nothing, gidx) for gidx in remapper.local_indices]
+
+        interpolate_slab!(
+            interpolated_values,
+            field,
+            slab_indices,
+            remapper.interpolation_coeffs,
+        )
 
         # out_local_array[remapper.local_target_hcoords_bitmask] returns a view on space we
         # want to write on