Fixes for cuda kernels with empty fields

CliMA · Jun 10, 2024 · 3276b75 · 3276b75
1 parent 6e37848
commit 3276b75
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/ext/cuda/cuda_utils.jl b/ext/cuda/cuda_utils.jl
@@ -4,8 +4,9 @@ import ClimaCore.DataLayouts
 import ClimaCore.DataLayouts: empty_kernel_stats
 
 get_n_items(field::Fields.Field) =
-    prod(size(parent(Fields.field_values(field))))
-get_n_items(data::DataLayouts.AbstractData) = prod(size(parent(data)))
+    prod(DataLayouts.universal_size(Fields.field_values(field)))
+get_n_items(data::DataLayouts.AbstractData) =
+    prod(DataLayouts.universal_size(parent(data)))
 get_n_items(arr::AbstractArray) = prod(size(parent(arr)))
 get_n_items(tup::Tuple) = prod(tup)
 

diff --git a/ext/cuda/fill.jl b/ext/cuda/fill.jl
@@ -11,8 +11,8 @@ function knl_fill_flat!(dest::AbstractData, val)
 end
 
 function cuda_fill!(dest::AbstractData, val)
-    (_, _, _, Nv, Nh) = size(dest)
-    if Nv > 0 && Nh > 0
+    (_, _, Nf, Nv, Nh) = DataLayouts.universal_size(dest)
+    if Nv > 0 && Nh > 0 && Nf > 0
         auto_launch!(knl_fill_flat!, (dest, val), dest; auto = true)
     end
     return dest