Reduce latency in stencil operators

Author: charleskawczynski

.buildkite/Manifest.toml

@@ -1,6 +1,6 @@
 # This file is machine-generated - editing it directly is not advised
 
-julia_version = "1.10.8"
+julia_version = "1.10.9"
 manifest_format = "2.0"
 project_hash = "6ab89829ea190189b0319a6f8e22b3515e5283c2"
 
@@ -327,7 +327,7 @@ weakdeps = ["CUDA", "MPI"]
     ClimaCommsMPIExt = "MPI"
 
 [[deps.ClimaCore]]
-deps = ["Adapt", "BandedMatrices", "BlockArrays", "ClimaComms", "CubedSphere", "DataStructures", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "KrylovKit", "LinearAlgebra", "MultiBroadcastFusion", "NVTX", "PkgVersion", "RecursiveArrayTools", "RootSolvers", "SparseArrays", "StaticArrays", "Statistics", "UnrolledUtilities"]
+deps = ["Adapt", "BandedMatrices", "BlockArrays", "ClimaComms", "CubedSphere", "DataStructures", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "KrylovKit", "LinearAlgebra", "MultiBroadcastFusion", "NVTX", "PkgVersion", "PrecompileTools", "RecursiveArrayTools", "RootSolvers", "SparseArrays", "StaticArrays", "Statistics", "UnrolledUtilities"]
 path = ".."
 uuid = "d414da3d-4745-48bb-8d80-42e94e092884"
 version = "0.14.26"
@@ -1538,7 +1538,7 @@ version = "0.3.23+4"
 [[deps.OpenLibm_jll]]
 deps = ["Artifacts", "Libdl"]
 uuid = "05823500-19ac-5b8b-9628-191a04bc5112"
-version = "0.8.1+2"
+version = "0.8.1+4"
 
 [[deps.OpenMPI_jll]]
 deps = ["Artifacts", "CompilerSupportLibraries_jll", "Hwloc_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "MPIPreferences", "TOML", "Zlib_jll"]

Project.toml

@@ -21,6 +21,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MultiBroadcastFusion = "c3c07f87-98de-43f2-a76f-835b330b2cbb"
 NVTX = "5da4648a-3479-48b8-97b9-01cb529c0a1f"
 PkgVersion = "eebad327-c553-4316-9ea0-9fa01ccd7688"
+PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 RootSolvers = "7181ea78-2dcb-4de3-ab41-2b8ab5a31e74"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"

src/ClimaCore.jl

@@ -29,5 +29,6 @@ include("CommonSpaces/CommonSpaces.jl")
 
 include("deprecated.jl")
 include("to_device.jl")
+include("precompile.jl")
 
 end # module

src/DataLayouts/struct.jl

@@ -179,6 +179,13 @@ Base.@propagate_inbounds @generated function get_struct(
     ::Val{D},
     start_index::CartesianIndex,
 ) where {T, S, D}
+    # recursion base case: hit array type is the same as the struct leaf type
+    if T === S # Use Union-splitting for better latency
+        return quote
+            Base.@_propagate_inbounds_meta
+            @inbounds return array[start_index]
+        end
+    end
     tup = :(())
     for i in 1:fieldcount(S)
         push!(
@@ -201,16 +208,6 @@ Base.@propagate_inbounds @generated function get_struct(
     end
 end
 
-# recursion base case: hit array type is the same as the struct leaf type
-Base.@propagate_inbounds function get_struct(
-    array::AbstractArray{S},
-    ::Type{S},
-    ::Val{D},
-    start_index::CartesianIndex,
-) where {S, D}
-    @inbounds return array[start_index]
-end
-
 """
     set_struct!(array, val::S, Val(D), start_index)