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.

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
  • Add .claude/skills/fix-ci/, with a device-free @test_compilation checker that catches CPU and GPU compilation issues locally, plus the CI-fixing workflow built around it
• 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/stable_view.jl, src/Utilities/Utilities.jl (safe_mapreduce)
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)
6. Tooling — the CI-fixing skill and its device-free compilation checker:
   .claude/skills/fix-ci/ (entirely AI-generated)

---

• 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.