Make the spectral-element operators dimension generic

[imreddyTeja]

Authorship: this PR was written primarily by Claude Code (see the
Co-Authored-By trailer on the commits). The refactors, the GPU bug fix, the new
tests and this description are the agent's work; I set the task and its scope.
Every measurement quoted below was produced by the agent rather than
independently reproduced, so please review with that in mind.

Rebased onto #2556. That PR unified the strong/weak variants behind a
FormType parameter while this one was open; the two refactors are orthogonal
(form vs. dimension) and compose, so this branch is rebuilt on top of its
form_* helpers rather than merged textually. One method per operator now
covers both dimensions and both forms. The numbers below are measured
against the post-#2556 main, so they are smaller than in the original
description: #2556 already removed the strong/weak half of the duplication.

What

The spectral-element operators had separate 1D ({(1,)}) and 2D ({(1, 2)}) implementations of the same tensor-product stencil, with each 2D method repeating its 1D counterpart once per axis. This was true on both sides:

• CPU: apply_operator in src/Operators/spectralelement.jl
• GPU: operator_evaluate in ext/cuda/operators_spectral_element.jl, plus operator_shmem/operator_fill_shmem! in ext/cuda/operators_sem_shmem.jl

Each of Divergence, SplitDivergence, WeakDivergence, Gradient, WeakGradient, Curl and WeakCurl now has a single implementation that loops over the axes it works over, with the axis index carried in the type domain so the loop unrolls and the node indices stay statically sized.

The last two commits then remove the rest of the per-dimension duplication in the same file — Interpolate/Restrict, tensor_product!, and the node accessors — described in its own section below.

This removes 78 lines from src + ext (850 added, 928 removed). Counting only code -- excluding blank lines, # comments and docstrings -- it is 389 lines, a 14% reduction (2752 -> 2363); the refactors account for -412 code lines and the GPU bug fix plus the Geometry move add 23 back. The operator methods themselves go from 37 to 18:

	methods before	after	code lines before	after	net
ext/cuda/operators_sem_shmem.jl (operator_shmem, operator_fill_shmem!)	16	8	235	133	-102
ext/cuda/operators_spectral_element.jl (operator_evaluate)	8	4	342	259	-83
src/Operators/spectralelement.jl (apply_operator)	13	6	1038	811	-227
5 files touched by the 1D fix and the Geometry move	--	--	1137	1160	+23
total	37	18	2752	2363	-389

The six remaining apply_operator methods are the four unified differential operators, one for Interpolate and Restrict together, and the zero-axes fallback.

Shared helpers (documented in a "Dimension-generic building blocks" section):

axis_vals(op)	(Val(1), Val(2)) for the axes op works over
slab_dims(op, Nq)	slab shape, one Nq per axis
replace_index(index, Val(d), k)	walk along axis d, or shrink a contracted axis' extent
contravariant/covariant(Val(d), u, lg)	component accessors by axis
covariant_vector/covariant_basis_vector	build a covariant vector over the operator's axes
curl_uses_component, curl_covariant_components, curl_term	Levi-Civita pieces; one helper serves Curl and WeakCurl, which just negates the coefficient
foreach_axis/sum_axes	the per-axis loop, shared by both devices
contract_axis!/contract_axes!/contract_temps	the sequential contraction fold (see below)
data_index, slab_level_index	node index padding and face/center level staggering
shmem_dims, sem_shmem, sem_shmem_per_axis, shmem_index	GPU shared memory
apply_stencil, curl_stencil	the two GPU stencil kernels (weak operators pass D')

Two annotations in the per-axis closures matter for performance and are commented in the code: @inline, without which the mutable slab temporary escapes and is heap-allocated once per slab (measured 1.1–3.8× slower), and @inbounds, which a closure body does not inherit from the enclosing block.

The rest of the duplication: contractions and node accessors

The differential operators above accumulate an independent contribution per axis, so one pass over a slab covers every axis. Interpolate, Restrict and tensor_product! instead contract their axes sequentially — contract axis 1 into a temporary, then contract axis 2 out of it — which is why they were left per-dimension in the first three commits. They fold over the axes rather than looping over them:

• contract_axis!(dst, M, src, post, Val(d), dims_out) contracts one axis against a matrix,
• contract_axes! folds it over a tuple of axes, each contraction reading what the previous one wrote and the last writing into the output,
• contract_temps allocates the length(I) - 1 intermediates, so a single-axis contraction gets no temporary at all — which is what the old 1D methods did by hand.

On top of that, Interpolate and Restrict differ only in three factors, in exactly the shape #2556 used for strong vs. weak: tensor_matrix (the matrix, transposed for Restrict), tensor_weighted_arg (WJ-weighted for Restrict) and tensor_rescale (divided by the output WJ for Restrict). So four methods become one. tensor_product! folds with the same primitives, which lets rmatmul1/rmatmul2 go (they had no other callers), and covers all three of its output layouts in one method instead of four.

The node accessors had the same 1D/2D split, differing only in padding the unused node axis with 1 and in repeating the face/center level staggering. Those are now data_index and slab_level_index:

	methods before	after
apply_operator for Interpolate/Restrict	4	1
tensor_product!	5	2
get_node	9	7
set_node!, get_local_geometry	2 each	1 each
node_indices	4	2

Two deliberate behaviour changes for downstream code, both in NEWS.md:

• The two-dimensional get_node and set_node! now accept a bare Center/FaceFiniteDifferenceSpace, which only their one-dimensional counterparts did. The asymmetry looks incidental (a 2D node index never arises over such a space) but the 1D case is load-bearing: a horizontal operator over an extracted column space reaches these accessors that way.
• Conversely, get_local_geometry and set_node! now raise "invalid space" for a space whose staggering is not recognised, instead of silently reading level slabidx.v. get_node always did. The full CPU and GPU suites pass with the strict version, which is the evidence that nothing relied on the permissive one.

tensor_product!'s invariants also move from dispatch into assertions (square node axes in and out; an IH1JH2 output takes a single level), since one method cannot express "square only if the input has a second node axis". A mismatched call now fails an assertion rather than raising MethodError.

Bug fix: SpectralElementSpace1D on the GPU

Spectral-element operators could not run on the GPU over a bare SpectralElementSpace1D at all — the space could not be passed to a kernel (KernelError: passing non-bitstype argument, since .grid holds an IntervalTopology with a mesh and domain). Unlike SpectralElementGrid2D, SpectralElementGrid1D had no device-side counterpart to convert to. This adds Grids.DeviceSpectralElementGrid1D and the corresponding Adapt rules.

Side effect: to_device now works on a SpectralElementSpace1D. It previously hit Adapt's identity fallback and returned the space unchanged.

Testing

New test/Operators/spectralelement/plane_cuda.jl (with a Unit: plane cuda pipeline entry) covers the 1D GPU path, which had no test before: all seven operators plus div∘grad and curl∘curl (which force two operators to share shared memory) over a bare SpectralElementSpace1D and an extruded 1D space, compared against the CPU, plus two absolute value checks. With only the fix reverted it fails 11/11 on the bare 1D space and passes 11/11 on the extruded one, isolating the bug.

New test/Operators/spectralelement/unit_tensor_product.jl covers tensor_product! against an explicit M * x * M' for all three of its output layouts and for the in-place two-argument form. The VIJFH output (used by ClimaCorePlots) and the in-place form had no coverage at all before, and both are rewritten here; only the IH1JH2 and VIH1 outputs were reached indirectly, through matrix_interpolate.

Verification beyond the test suites, comparing every operator over 7 spaces (1D/2D, bare/extruded, rectilinear/cubed sphere; 74 expressions) against results captured before the refactors:

• CPU: 73/74 bit-for-bit identical to post-Unify stong/weak forms with type parameter dispatch #2556 main, zero allocations across all 14 kernel/space combinations, kernel timings 0.985–1.018×. The one exception is Restrict on a SpectralElementSpace1D, where 3 of its 15 nodes move by at most 1 ulp (2.2e-16 on values of order 1). The order of summation is unchanged; LLVM simply contracts a different subset of the contraction's muladds into fmas once the loop over output nodes is written generically. Verified by reproducing every value with an explicit chain whose kth muladd is contracted iff mask[k]: one consistent mask reproduces all 15 values, which is what distinguishes "LLVM chose differently" from a reassociated sum. Two candidate fixes (applying the rescale in a second pass over the output; fusing it back into the last contraction) changed nothing, so the fused form was kept because it matches the structure of the code it replaces.
• GPU: 0 mismatches — of the 50 expressions that run on both main and this branch, 40 are bitwise identical and 10 (2D Divergence, SplitDivergence, Gradient, Curl) differ by at most 5.3e-15 on values of order 1–20, from reassociation and FMA selection after restructuring the accumulation loops; the GPU is deterministic run to run. The 20 bare-1D expressions that previously could not compile all agree with the CPU. All 70 supported expressions are bitwise unchanged across the last two commits, which touch the accessors the GPU kernels share.
• GPU kernel timings are within the ±7% session-to-session spread that the baseline shows against itself. A single before/after pair suggested wgrad_div was 9–11% slower and looked reproducible across two after-runs; sampling the baseline a second time showed the same kernel 1.07× slower than itself and div_grad 0.93×. minimum(trial).time suppresses within-trial noise but not between-session variation.

Test files run green on both devices: opt.jl (JET), rectilinear.jl, plane.jl, split_divergence.jl, unit_form_types.jl (added by #2556), unit_tensor_product.jl, covar_deriv_ops.jl, unit_diffusion2d.jl, sphere_{gradient,divergence,curl,diffusion,diffusion_vec,geometry,sphericalharmonics}.jl, unit_sphere_hyperdiffusion{,_vec}.jl, hybrid/unit_2d.jl, Fields/unit_field.jl, Geometry/tensors.jl; and on GPU plane_cuda.jl, rectilinear_cuda.jl, covar_deriv_ops.jl, hybrid/{unit_cuda,extruded_3dbox_cuda,extruded_sphere_cuda}.jl. detect_ambiguities is unchanged (40, all pre-existing, none naming anything added here).

Note: test/Spaces/unit_spaces.jl "2D spaces with mask" errors on GPU (a DivergenceF2C kernel throws DomainError from sqrt(-1) where the test expects NaNs). Reproduced identically with this branch's changes stashed, so it is pre-existing and unrelated.

Commits

The five commits are separately verified, in the order they were written:

1. 8804cf06 make the differential operators dimension generic (CPU and GPU)
2. 4ccf31f8 share one foreach_axis/sum_axes pair between the CPU and GPU operators
3. caeb50da move the generic covariant constructors into Geometry
4. 7d8a835a fold the tensor-product operators (Interpolate/Restrict, tensor_product!) over their axes
5. ef1be22e collapse the per-dimension node accessors and node_indices

🤖 Generated with Claude Code

https://claude.ai/code/session_0148rCSSVswyipJTZjqpEUMd