Share one foreach_axis/sum_axes pair between the CPU and GPU operators

The axis loop was written two ways: the CPU operators called
`unrolled_foreach(axis_vals(op))` inline at five sites, each repeating the
`@inline` annotation and a comment explaining it, while the CUDA extension had
its own `sum_over_axes` plus a bare `unrolled_foreach(axis_vals(op))` in the
shared-memory fills. Both now go through `foreach_axis` and `sum_axes`, defined
once next to `axis_vals`, so the axis-tuple plumbing and the reason the per-axis
accumulators are kept separate are stated in one place.

The `@inline` stays in each closure body rather than moving into the helpers:
only a method, not an anonymous function, can carry it, and forcing the inline
from inside the helper -- by wrapping `f` in another closure to annotate its
call site -- reintroduces the per-slab heap allocation it exists to prevent (I
measured every operator kernel allocating again). The helpers' docstring now
carries that explanation, along with the matching note about `@inbounds`.

Verified: CPU results bitwise unchanged (74/74 expressions) with allocations
back at zero for all 14 benchmark kernels; GPU results bitwise identical for all
70 supported expressions. Tests pass on both devices.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0148rCSSVswyipJTZjqpEUMd

Author: imreddyTeja

ext/cuda/operators_sem_shmem.jl

@@ -7,6 +7,7 @@ import ClimaCore.Operators: form_jacobian, form_weighted_arg
 import ClimaCore.Operators:
     axis_vals,
     axis_index,
+    foreach_axis,
     contravariant,
     curl_covariant_components,
     curl_uses_component,
@@ -91,7 +92,8 @@ Base.@propagate_inbounds function operator_fill_shmem!(
     local_geometry = get_local_geometry(space, ij, slabidx)
     node = shmem_index(op, ij)
     jacobian = form_jacobian(F(), local_geometry)
-    unrolled_foreach(axis_vals(op)) do vd
+    foreach_axis(op) do vd
+        @inline
         @inbounds Jv[axis_index(vd)][node] =
             jacobian * contravariant(vd, arg, local_geometry)
     end
@@ -124,7 +126,8 @@ Base.@propagate_inbounds function operator_fill_shmem!(
     local_geometry = get_local_geometry(space, ij, slabidx)
     node = shmem_index(op, ij)
     (; J) = local_geometry
-    unrolled_foreach(axis_vals(op)) do vd
+    foreach_axis(op) do vd
+        @inline
         @inbounds work[axis_index(vd)][node] =
             J * contravariant(vd, arg1, local_geometry)
     end

ext/cuda/operators_spectral_element.jl

@@ -14,8 +14,9 @@ import ClimaCore.Operators:
     covariant_components,
     covariant_vector,
     curl_term,
-    replace_index
-import UnrolledUtilities: unrolled_map, unrolled_reduce
+    replace_index,
+    sum_axes
+import UnrolledUtilities: unrolled_map
 import Base.Broadcast: Broadcasted
 
 """
@@ -194,9 +195,8 @@ end
 # form_deriv_entry supplies the derivative matrix entries (transposed and
 # sign-flipped for the weak form), form_jacobian_rescale or form_weight_rescale
 # divides the form's Jacobian or quadrature-weight factor back out of the result,
-# and the loop over axes is unrolled with one accumulator per dimension, combined
-# once at the end, which keeps the accumulation loops as independent dependency
-# chains.
+# and `sum_axes` unrolls the loop over axes, keeping one accumulator per
+# dimension.
 
 """
     apply_stencil(form, D, w, node, ::Val{d}, i, Nq)
@@ -242,9 +242,6 @@ end
 Base.@propagate_inbounds shmem_components(work, node) =
     unrolled_map(w_k -> isnothing(w_k) ? nothing : (@inbounds w_k[node]), work)
 
-# Sum of the per-axis contributions of an operator, unrolled over its axes.
-@inline sum_over_axes(f, op) = unrolled_reduce(+, unrolled_map(f, axis_vals(op)))
-
 Base.@propagate_inbounds function operator_evaluate(
     op::Divergence{I, F},
     Jv,
@@ -260,7 +257,8 @@ Base.@propagate_inbounds function operator_evaluate(
     local_geometry = get_local_geometry(space, ij, slabidx)
     node = shmem_index(op, ij)
 
-    DJv = sum_over_axes(op) do vd
+    DJv = sum_axes(op) do vd
+        @inline
         d = axis_index(vd)
         @inbounds apply_stencil(F(), D, Jv[d], node, vd, ij[d], Nq)
     end
@@ -283,7 +281,8 @@ Base.@propagate_inbounds function operator_evaluate(
     node = shmem_index(op, ij)
     psi = last(work)
 
-    result = sum_over_axes(op) do vd
+    result = sum_axes(op) do vd
+        @inline
         d = axis_index(vd)
         i = ij[d]
         Juᵈ = work[d]
@@ -347,7 +346,8 @@ Base.@propagate_inbounds function operator_evaluate(
     local_geometry = get_local_geometry(space, ij, slabidx)
     node = shmem_index(op, ij)
 
-    result = sum_over_axes(op) do vd
+    result = sum_axes(op) do vd
+        @inline
         d = axis_index(vd)
         @inbounds curl_stencil(F(), D, work, node, vd, ij[d], Nq)
     end

src/Operators/spectralelement.jl

@@ -1,4 +1,5 @@
-import UnrolledUtilities: unrolled_map, unrolled_foreach, unrolled_any
+import UnrolledUtilities:
+    unrolled_map, unrolled_foreach, unrolled_reduce, unrolled_any
 import ..Utilities.Unrolled: unrolled_map_with_inbounds
 
 abstract type AbstractSpectralStyle <: Fields.AbstractFieldStyle end
@@ -670,10 +671,11 @@ Base.Broadcast.BroadcastStyle(
 # axis index in the type domain so that the loop over axes unrolls and the node
 # indices stay statically sized.
 #
-# The per-axis closures below carry two annotations that matter for performance:
-# `@inline`, because the mutable slab temporary they write to would otherwise
-# escape and be heap-allocated once per slab, and `@inbounds`, which a closure
-# body does not inherit from the enclosing `@inbounds` block.
+# Loops over axes go through `foreach_axis` and `sum_axes`, which own the
+# inlining discipline the per-axis closures depend on. Their bodies still need
+# their own `@inbounds`: a closure body does not inherit it from the enclosing
+# `@inbounds` block, and only a method (not an anonymous function) can opt into
+# propagating it.
 
 """
     axis_vals(op)
@@ -691,6 +693,32 @@ The axis index `d` itself: `ij[axis_index(vd)]` is the node index along axis `d`
 """
 @inline axis_index(::Val{d}) where {d} = d
 
+"""
+    foreach_axis(f, op)
+    sum_axes(f, op)
+
+Call `f(Val(d))` for each axis index `d` that `op` works over, unrolled;
+`sum_axes` sums the results, keeping one accumulator per axis and combining them
+once at the end so that the per-axis accumulation loops stay independent
+dependency chains.
+
+Two annotations belong in the body of `f`, and neither can be supplied from here,
+because only a method -- not an anonymous function -- can carry them:
+
+  - `@inline`, as the first statement. Without it the closure is not inlined, so
+    the mutable slab temporary it writes to escapes and is heap-allocated once per
+    slab; that measured 1.1-3.8x slower across the operator kernels. (Forcing the
+    inline from inside these helpers does not work: wrapping `f` in another
+    closure to annotate its call site reintroduces the allocation.)
+  - `@inbounds`, since a closure body does not inherit it from an enclosing
+    `@inbounds` block.
+"""
+@inline foreach_axis(f::F, op::SpectralElementOperator) where {F} =
+    unrolled_foreach(f, axis_vals(op))
+
+@inline sum_axes(f::F, op::SpectralElementOperator) where {F} =
+    unrolled_reduce(+, unrolled_map(f, axis_vals(op)))
+
 """
     slab_dims(op, Nq)
 
@@ -819,8 +847,8 @@ Base.@propagate_inbounds function apply_operator(
     @inbounds for ij in node_indices(space)
         local_geometry = get_local_geometry(space, ij, slabidx)
         v = get_node(space, arg, ij, slabidx)
-        unrolled_foreach(axis_vals(op)) do vd
-            @inline # `out` heap-allocates per slab if this closure is not inlined
+        foreach_axis(op) do vd
+            @inline
             @inbounds begin
                 i = ij[axis_index(vd)]
                 Jvᵈ =
@@ -958,8 +986,8 @@ Base.@propagate_inbounds function apply_operator(
         node = slab_node_index(ij)
         local_geometry = get_local_geometry(space, ij, slabidx)
         u = get_node(space, arg1, ij, slabidx)
-        unrolled_foreach(axis_vals(op)) do vd
-            @inline # `Ju` heap-allocates per slab if this closure is not inlined
+        foreach_axis(op) do vd
+            @inline
             @inbounds Ju[axis_index(vd)][node] =
                 local_geometry.J * contravariant(vd, u, local_geometry)
         end
@@ -970,8 +998,8 @@ Base.@propagate_inbounds function apply_operator(
     fill!(parent(out), zero(FT))
     @inbounds for ij in node_indices(space)
         node = slab_node_index(ij)
-        unrolled_foreach(axis_vals(op)) do vd
-            @inline # `out` heap-allocates per slab if this closure is not inlined
+        foreach_axis(op) do vd
+            @inline
             @inbounds begin
                 i = ij[axis_index(vd)]
                 Juᵈ = Ju[axis_index(vd)]
@@ -1046,8 +1074,8 @@ Base.@propagate_inbounds function apply_operator(
     @inbounds for ij in node_indices(space)
         local_geometry = get_local_geometry(space, ij, slabidx)
         x = form_weighted_arg(form, local_geometry, get_node(space, arg, ij, slabidx))
-        unrolled_foreach(axis_vals(op)) do vd
-            @inline # `out` heap-allocates per slab if this closure is not inlined
+        foreach_axis(op) do vd
+            @inline
             @inbounds begin
                 i = ij[axis_index(vd)]
                 for k in 1:Nq
@@ -1195,8 +1223,8 @@ Base.@propagate_inbounds function apply_operator(
         local_geometry = get_local_geometry(space, ij, slabidx)
         v = get_node(space, arg, ij, slabidx)
         u = curl_covariant_components(op, form, v, local_geometry)
-        unrolled_foreach(axis_vals(op)) do vd
-            @inline # `out` heap-allocates per slab if this closure is not inlined
+        foreach_axis(op) do vd
+            @inline
             @inbounds begin
                 i = ij[axis_index(vd)]
                 for k in 1:Nq