Use ClimaCartesianIndices

charleskawczynski · charleskawczynski · commit 8c80b51981bd · 2025-04-21T11:18:39.000-04:00
diff --git a/Project.toml b/Project.toml
@@ -7,6 +7,7 @@ version = "0.14.31"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
+ClimaCartesianIndices = "afc79dc6-9922-4b50-ae92-da7b5116c2c8"
 ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
 CubedSphere = "7445602f-e544-4518-8976-18f8e8ae6cdb"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
diff --git a/ext/cuda/data_layouts_threadblock.jl b/ext/cuda/data_layouts_threadblock.jl
@@ -1,6 +1,8 @@
 const CI5 = CartesianIndex{5}
-# using ClimaCartesianIndices: FastCartesianIndices
-FastCartesianIndices(x) = CartesianIndices(x)
+import ClimaCartesianIndices
+FastCartesianIndices(x) =
+    length(prod(x)) > typemax(Int32) ? CartesianIndices(x) :
+    ClimaCartesianIndices.FastCartesianIndices(x)
 
 maximum_allowable_threads() = (
     CUDA.attribute(CUDA.device(), CUDA.DEVICE_ATTRIBUTE_MAX_BLOCK_DIM_X),
