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.

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 combing DataScopes, as well as rules for partitioning into smaller DataScopes
  • Automatically assign column/slab/level views to partitions, and assign point views to individual threads
• Define a minimal set of communication primitives through DataScope dispatch: foreach_point and reduce
  • Implement the loops in fill!, copyto!, and fused_copyto! through the foreach_point primitive
  • 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 struct, still called DataStyle
  • Use DataStyle to propagate the layout type, ignoring DataF layouts when combining inconsistent types
  • Limit all 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

---

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