Merge pull request #2004 from CliMA/ck/type_stability

Fix type instability in slab_index and vindex
CliMA · Sep 25, 2024 · 65d0e30 · 65d0e30
2 parents 9cb8366 + 2df6cf4
commit 65d0e30
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 8 deletions.
diff --git a/lib/ClimaCorePlots/src/ClimaCorePlots.jl b/lib/ClimaCorePlots/src/ClimaCorePlots.jl
@@ -5,8 +5,10 @@ import TriplotBase
 
 import ClimaComms
 # Keep in sync with definition(s) in ClimaCore.DataLayouts.
-@inline slab_index(i, j) = CartesianIndex(i, j, 1, 1, 1)
-@inline slab_index(i) = CartesianIndex(i, 1, 1, 1, 1)
+@inline slab_index(i::T, j::T) where {T} =
+    CartesianIndex(i, j, T(1), T(1), T(1))
+@inline slab_index(i::T) where {T} = CartesianIndex(i, T(1), T(1), T(1), T(1))
+
 import ClimaCore:
     ClimaCore,
     DataLayouts,

diff --git a/lib/ClimaCoreTempestRemap/src/ClimaCoreTempestRemap.jl b/lib/ClimaCoreTempestRemap/src/ClimaCoreTempestRemap.jl
@@ -4,9 +4,10 @@ export write_exodus, rll_mesh, overlap_mesh, remap_weights, apply_remap
 export def_time_coord, def_space_coord
 
 # Keep in sync with definition in DataLayouts.
-@inline slab_index(i, j) = CartesianIndex(i, j, 1, 1, 1)
-@inline slab_index(i) = CartesianIndex(i, 1, 1, 1, 1)
-@inline vindex(v) = CartesianIndex(1, 1, 1, v, 1)
+@inline slab_index(i::T, j::T) where {T} =
+    CartesianIndex(i, j, T(1), T(1), T(1))
+@inline slab_index(i::T) where {T} = CartesianIndex(i, T(1), T(1), T(1), T(1))
+@inline vindex(v::T) where {T} = CartesianIndex(T(1), T(1), T(1), v, T(1))
 
 using ClimaComms
 import ClimaCore

diff --git a/src/DataLayouts/DataLayouts.jl b/src/DataLayouts/DataLayouts.jl
@@ -1107,9 +1107,9 @@ type parameters.
 @inline farray_size(data::VIFH{S, Nv, Ni, Nh}) where {S, Nv, Ni, Nh} = (Nv, Ni, ncomponents(data), Nh)
 
 # Keep in sync with definition(s) in libs.
-@inline slab_index(i, j) = CartesianIndex(i, j, 1, 1, 1)
-@inline slab_index(i) = CartesianIndex(i, 1, 1, 1, 1)
-@inline vindex(v) = CartesianIndex(1, 1, 1, v, 1)
+@inline slab_index(i::T, j::T) where {T} = CartesianIndex(i, j, T(1), T(1), T(1))
+@inline slab_index(i::T) where {T} = CartesianIndex(i, T(1), T(1), T(1), T(1))
+@inline vindex(v::T) where {T} = CartesianIndex(T(1), T(1), T(1), v, T(1))
 
 """
     parent_array_type(data::AbstractData)