Unify DataLayouts and simplify API

[dennisYatunin]

Purpose

This PR completes the first task outlined in #2468, unifying most of the structs in the DataLayouts module and simplifying their API to eliminate code duplication. This is a prerequisite for #2514.

All loops and reductions over layouts now run through two scope-dispatched primitives, foreach_slice and reduce_points, which work identically on CPUs (including nested multithreading) and GPUs (as single kernels). The rewrite is verified to compile for GPUs without runtime dispatch, and hot paths match the performance of main on CPUs.

Tested against the ClimaAtmos CI pipeline in this build (only a single MSE test failure due to roundoff error).

Content

• Rename AbstractData to DataLayout, and make it a subtype of AbstractArray
• Simplify all layouts to four structs: DataF, VIJHWithF, VIH1, and IH1JH2
  • Enable consistent treatment of static arrays (MArrays or cuda shmem arrays) by similar(::DataLayout, ...)
  • Allow any parent array dimension to be treated as an F axis
  • Make getindex and setindex! match standard AbstractArray behavior
  • Propagate linear indexing through the IndexStyle, instead of separate structs like NonExtrudedBroadcasted
• Replace dispatch based on parent arrays with a distinct layer of abstraction
  • Introduce the DataScope singleton, a more versatile generalization of a ClimaComms.AbstractDevice
  • Propagate this singleton through all operations on DataLayouts, including broadcasts and reductions
  • Define rules for combining DataScopes, as well as rules for partitioning into smaller DataScopes
  • Automatically divide column/slab/level/point views among DataScope partitions
• Define a minimal set of communication primitives through DataScope dispatch: foreach_slice and reduce_points
  • Implement the loops in fill!, copyto!, and fused_copyto! through the foreach_point primitive (a simple wrapper for foreach_slice)
  • Enable column masking in reduce and fused_copyto!
  • Make fused_copyto! treat all layouts consistently (and actually fuse the copyto! loops every time it is called)
• Simplify all BroadcastStyles to one concrete struct, still called DataStyle
  • Use DataStyle to propagate the layout type, ignoring DataF layouts when combining distinct types
  • Limit layout-specific broadcasting behavior to four methods of shape_params
• Make the broadcasting implementation consistent with the default methods for AbstractArrays in Base
  • Replicate the optimizations for broadcasting over Refs, and for simple identity function broadcasts
  • Formalize how single-valued Tuples in broadcasts are treated as Refs
  • Have multi-valued Tuples fall back to the default behavior for AbstractArrays instead of erroring
• Guarantee that the new primitives compile for GPUs without runtime dispatch
  • Add Utilities.stable_view, which constructs inference-stable SubArrays for all slice and property views (GPUArrays replaces contiguous CuArray views with uninferrable derived arrays), eliminating reshaped views on GPUs altogether
  • Keep 128-bit integers and runtime-built error strings out of kernel argument types and kernel code, which LLVM cannot compile to PTX
• Keep old code loadable during the transition
  • Add DataLayouts/deprecated.jl with exported aliases for AbstractData, IJFH, and IJHF (downstream packages and registered satellite packages still reference them)
  • Read legacy layout shapes from existing HDF5 files by reshaping them explicitly

Roadmap

A suggested order for reviewing the changes:

1. Core API definitions — the layout types, struct storage, and scopes:
   src/DataLayouts/DataLayouts.jl, src/DataLayouts/struct_storage.jl, src/DataLayouts/scopes.jl, src/DataLayouts/deprecated.jl
2. Indexing, broadcasting, and loop primitives — how values are accessed and iterated:
   src/DataLayouts/indexing.jl, src/DataLayouts/broadcast.jl, src/DataLayouts/masks.jl, src/DataLayouts/loops.jl, src/Utilities/safe_mapreduce.jl, src/Utilities/Utilities.jl
3. GPU implementations of the primitives:
   ext/cuda/scopes.jl, ext/cuda/loops.jl, ext/cuda/cuda_utils.jl, ext/cuda/data_layouts.jl, ext/cuda/adapt.jl
4. Name and indexing updates in consumers — mostly mechanical (4-D (v, i, j, h) indices, 5-D parent arrays, shape_params):
   src/Fields/, src/Spaces/, src/Grids/, src/Operators/, src/MatrixFields/, src/Limiters/, src/Topologies/dss*.jl, src/Remapping/, src/InputOutput/, remaining ext/cuda/ files, lib/ClimaCoreMakie, lib/ClimaCorePlots, examples/
5. New and updated tests:
   test/DataLayouts/unit_loops.jl, test/Utilities/unit_stable_view.jl, updated test/DataLayouts/, test/Fields/, test/Spaces/, test/Operators/, and GPU expectation updates (test/gpu/latency_benchmarks.jl, test/Spaces/opt_spaces.jl, test/Operators/finitedifference/opt_examples.jl)

---

• Code follows the style guidelines OR N/A.
• Unit tests are included OR N/A.
• Code is exercised in an integration test OR N/A.
• Documentation has been added/updated OR N/A.

[dennisYatunin]

When looking at this PR, one question also jumped to my mind, which is what kind of the scopes we expect to exist when we are done.

At the moment there seems to be a CPU hierarchy: ThisThreadPool -> ThisThread and GPU hierarchy ExternalDevice -> ThisKernelGrid -> ThisKernelBlock -> ThisThread

Is it more-or-less fixed or is it still work in progress?

The CPU hierarchy is fixed, and I doubt it will ever need to get more complicated. I was initially hoping to keep the GPU hierarchy simple as well, but I had to add a bit more complexity to match ClimaCore's current behavior.

Aside from some naming simplifications (ExternalDevice is now ThisHost, ThisKernelGrid is now ThisKernel, and ThisKernelBlock is now ThisBlock), I've introduced a new layer between blocks and threads called ThisSubBlock{N}. The type parameter N denotes the number of threads, which can be any value between 2 and 32 (with 32 corresponding to one warp).

When processing slice views in parallel across ThisKernel, the scope assigned to each slice will be determined according to its size:

• For single-point views, foreach_point will assign ThisThread to each point.
• For columns with 64 levels, foreach_column will assign ThisBlock to each column.
• For slabs with 16 nodal points, foreach_slab will assign ThisSubBlock{16} to each slab.

This will match the current implementation of copyto! in ClimaCore's CUDA extension:

• For Broadcasted expressions, copyto! assigns one thread to each point.
• For StencilBroadcasted expressions with 64 levels, copyto! assigns one block to each column.
• For SpectralBroadcasted expressions with 16 nodal points, copyto! assigns one half-warp to each slab.
  • In particular, spectral_partition assigns one block to Ni * Nj * Nvthreads points, where Nvthreads is roughly n_max_threads ÷ (Ni * Nj). When Ni * Nj is 16, this is equivalent to assigning one half-warp to each slab, and grouping multiple half-warps into blocks of size n_max_threads.

To minimize the amount of code we'll need in the CUDA extension, I've handled all of these cases in a single foreach_slice primitive, which foreach_point/foreach_level/foreach_slab/foreach_column just call with a particular slice operator (view/level/slab/column, respectively). The most important part of foreach_slice is the slice_subscope(scope, op, args...) function, which is defined as follows:

• "By default, this is the smallest subset of scope that does not require any thread to process more than one point from the largest slice returned by op. When no such subset is available, the largest subset is used in order to minimize the number of points per thread."

If the default slice_subscope is not optimal in some kernels, it should be straightforward to override for special cases by adding new methods to the CUDA extension, which can be specialized on the current scope, the slice operator, or the types of DataLayout/LazyDataLayout arguments. This could potentially be helpful for optimizing fused broadcast expressions with significant register pressure or shared memory requirements. Aside from any such special cases, though, we'll only need to add CUDA extension code for launching kernels from ThisHost, and we'll be able to eliminate a lot of our current code duplication.

Let me know if any of this can be changed to further simplify things for you, @Mikolaj-A-Kowalski.

[Copilot]

Copilot wasn't able to review this pull request because it exceeds the maximum number of lines (20,000). Try reducing the number of changed lines and requesting a review from Copilot again.

[tapios]

This looks good as far as I can tell. It's nice to see the simplifications realized here.

I will send a list of detailed code comments from an audit separately, which identified two apparent reduction bugs.

[tapios]

One more thing: Please add a smoke test for the DG Bickley jet (just run it a few steps). I checked it works with this branch, but numerical_flux.jl is not covered by tests; we want to make sure the DG part remains maintained (and soon build out that option).

[imreddyTeja]

Have you looked at gpu benchmarks compared to main? It looks like the end to end test is a bit slower

[ph-kev]

Why are the y and z dimensions in the launch configuration not supported?

[dennisYatunin]

This is a limitation of the current design, since nonlinear grid/block dimensions make it much harder to implement automatic coarsening (looping over multiple indices per thread) and automatic launch config adjustment (using the largest possible blocks while making sure they can all run at the same time). Support for y and z dimensions can be added in a future PR, but only if the measurable performance benefits outweigh the increase in complexity.

[ph-kev]

I think this might be a better name?

Suggested change

	const MaybeFusedDataLayoutBroadcast = Union{LazyDataLayout, FusedMultiBroadcast}
	const LazyOrFusedDataLayoutBroadcast = Union{LazyDataLayout, FusedMultiBroadcast}

[dennisYatunin]

I actually used that name in my first draft of this PR. I ended up replacing it with the current name because LazyDataLayout and FusedMultiBroadcast are both "lazy" objects, so calling the union "lazy or fused" is a bit confusing.

[imreddyTeja]

I checked the performance of this PR at commit bd53967 vs main. Here are my findings:

From profiling prognostic edmf with 1M microphysics in the coupler:

1. Overall, each step takes about 1.25x longer
2. There is a negligible speed-up to spectral operations
3. There is no difference in finite difference operator performance. This is expected.
4. the total time spent in Pointwise operations is ~1.5x greater than before. This is where the slowdown comes from.

When looking at simple point wise expressions in isolation, I found:

1. This branch is always faster or the same as main when not using a VIJFH with V==64.
2. when using a VIJFH with V==64, this branch is slower than main, and uses more registers. This could be from the "tail effect", or because the specialized kernel for VIJFH_64 uses the block/grid indices for indexing.

The slowdown is not ideal, but I do not think it is enough of a degradation to prevent merging this. The "tail effect" issue should be an easy fix, and any potential differences in indexing computation should be less impactful for non-trivial point-wise kernels. At the moment, over half the point wise kernels in ClimaAtmos take less than 25 microseconds. That is short enough to where a few integer divisions will be noticeable. Hopefully, fusing more point wise operations will make this a non-problem.