Propagate unified DataLayouts API and refine it after review

Updates all modules, the CUDA extension, and the test suite for the new
DataLayout types, 4-D (v, i, j, h) index convention, and loop primitives
introduced by the previous commit, and applies the revisions from review:

- Add Utilities.nested_view, which prevents unbounded nesting of SubArray
  and ReshapedArray wrappers from repeated views and reshapes; use it for
  slice views and struct field views to avoid type-depth widening and
  per-view allocations
- Use Base's pairwise mapreduce over linear positions for reductions, and
  re-evaluate impure scalar broadcasts at every point
- Simplify BroadcastStyle combination, getproperty, view via rebuild, and
  the Style{Tuple} conversion; restore kwarg-based layout constructors
- Restrict the recursion_relation overrides to Core.kwcall methods, which
  keeps nested loop primitives allocation-free and cuts the compilation
  time of several unit tests by more than an order of magnitude
- Use fieldoffset for struct field views (prefix-tuple sizeof includes
  trailing padding) and branch-free thread-pool sizes for JET cleanliness
- Add unit tests for nested loops, reduction accuracy, masked reductions,
  0-dimensional broadcasts, scalar broadcast allocations, and nested_view

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: dennisYatunin

.buildkite/Manifest.toml

@@ -12,6 +12,7 @@ ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884"
 ClimaCorePlots = "cf7c7e5a-b407-4c48-9047-11a94a308626"
 ClimaCoreTempestRemap = "d934ef94-cdd4-4710-83d6-720549644b70"
 ClimaCoreVTK = "c8b6d40d-e815-466f-95ae-c48aefa668fa"
+ClimaInterpolations = "dd0f122e-fa3b-47f3-bcf0-93bbc60d885e"
 ClimaParams = "5c42b081-d73a-476f-9059-fd94b934656c"
 ClimaTimeSteppers = "595c0a79-7f3d-439a-bc5a-b232dc3bde79"
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"

.buildkite/Project.toml

@@ -114,6 +114,16 @@ steps:
         retry: *retry_policy
         command: "julia --threads=4 --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_fill_and_copyto.jl"
 
+      - label: "Unit: data loops (1 thread)"
+        key: unit_data_loops
+        retry: *retry_policy
+        command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_loops.jl"
+
+      - label: "Unit: data loops (4 threads)"
+        key: threaded_unit_data_loops
+        retry: *retry_policy
+        command: "julia --threads=4 --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_loops.jl"
+
       - label: "Unit: mapreduce (1 thread)"
         key: unit_data_mapreduce
         retry: *retry_policy

.buildkite/pipeline.yml

@@ -4,13 +4,85 @@
 CurrentModule = ClimaCore
 ```
 
+## Data layouts
+
 ```@docs
 DataLayouts
+DataLayouts.DataLayout
 DataLayouts.DataF
+DataLayouts.VIJHWithF
 DataLayouts.VIJFH
 DataLayouts.VIJHF
 DataLayouts.VIH1
 DataLayouts.IH1JH2
+```
+
+## Layout properties
+
+```@docs
+DataLayouts.layout_type
+DataLayouts.parent_type
+DataLayouts.f_dim
+DataLayouts.shape_params
+DataLayouts.inferred_size
+DataLayouts.has_inferred_size
+DataLayouts.vijh_params
+DataLayouts.nlevels
+DataLayouts.nquadpoints
+DataLayouts.nelems
+DataLayouts.ncomponents
+DataLayouts.layout_constructor
+DataLayouts.rebuild
+DataLayouts.reassign
+```
+
+## Data scopes
+
+```@docs
+DataLayouts.DataScope
+DataLayouts.ThisThread
+DataLayouts.ThisThreadPool
+DataLayouts.partition
+DataLayouts.is_subscope
+DataLayouts.num_threads
+DataLayouts.num_partitions
+DataLayouts.thread_rank
+DataLayouts.partition_rank
+DataLayouts.parallelize_over
+DataLayouts.synchronize
+DataLayouts.scoped_array
+DataLayouts.scoped_static_array
+DataLayouts.strided_access
+DataLayouts.subscope_indices
+```
+
+## Loops and reductions
+
+```@docs
+DataLayouts.each_slice_index
+DataLayouts.slice_subscope
+DataLayouts.foreach_slice
+DataLayouts.foreach_point
+DataLayouts.foreach_level
+DataLayouts.foreach_slab
+DataLayouts.foreach_column
+DataLayouts.reduce_points
+DataLayouts.column_reduce!
+```
+
+## Masks
+
+```@docs
+DataLayouts.DataMask
+DataLayouts.NoMask
+DataLayouts.IJHMask
+DataLayouts.set_mask_maps!
+DataLayouts.should_compute
+```
+
+## Struct storage
+
+```@docs
 DataLayouts.bitcast_struct
 DataLayouts.default_basetype
 DataLayouts.check_basetype
@@ -19,4 +91,14 @@ DataLayouts.num_basetypes
 DataLayouts.struct_field_view
 DataLayouts.set_struct!
 DataLayouts.get_struct
+DataLayouts.view_struct
+```
+
+## Broadcasting
+
+```@docs
+DataLayouts.DataStyle
+DataLayouts.LazyDataLayout
+DataLayouts.layout_args
+DataLayouts.modify_args
 ```

docs/src/APIs/datalayouts_api.md

@@ -13,6 +13,7 @@ Utilities.is_inferred_type
 Utilities.return_type
 Utilities.unsafe_eltype
 Utilities.safe_eltype
+Utilities.nested_view
 ```
 
 ## Utilities.PlusHalf

docs/src/APIs/utilities_api.md

@@ -91,16 +91,16 @@ if haskey(ENV, "RESTART_FILE")
     ᶠlocal_geometry = Fields.local_geometry_field(Y.f)
 else
     t_start = FT(0)
-    horizontal_layout_types = Dict()
-    horizontal_layout_types["VIJFH"] = DataLayouts.VIJFH
-    horizontal_layout_types["VIJHF"] = DataLayouts.VIJHF
-    horizontal_layout_type =
-        horizontal_layout_types[get(ENV, "horizontal_layout_type", "VIJFH")]
+    VIJHs = Dict()
+    VIJHs["VIJFH"] = DataLayouts.VIJFH
+    VIJHs["VIJHF"] = DataLayouts.VIJHF
+    VIJH =
+        VIJHs[get(ENV, "horizontal_layout_type", "VIJFH")]
     h_space = make_horizontal_space(
         horizontal_mesh,
         npoly,
         comms_ctx,
-        horizontal_layout_type,
+        VIJH,
     )
     center_space, face_space =
         make_hybrid_spaces(h_space, z_max, z_elem; z_stretch)

examples/hybrid/driver.jl

@@ -18,19 +18,19 @@ import ClimaCore.Utilities: half, new, return_type
 include(joinpath("cuda", "adapt.jl"))
 include(joinpath("cuda", "cuda_utils.jl"))
 include(joinpath("cuda", "data_layouts.jl"))
-# include(joinpath("cuda", "fields.jl"))
+include(joinpath("cuda", "fields.jl"))
 include(joinpath("cuda", "topologies_dss.jl"))
-# include(joinpath("cuda", "operators_finite_difference.jl"))
-# include(joinpath("cuda", "remapping_distributed.jl"))
-# include(joinpath("cuda", "operators_integral.jl"))
-# include(joinpath("cuda", "remapping_interpolate_array.jl"))
-# include(joinpath("cuda", "limiters.jl"))
-# include(joinpath("cuda", "operators_sem_shmem.jl"))
-# include(joinpath("cuda", "operators_fd_shmem_common.jl"))
-# include(joinpath("cuda", "operators_fd_shmem.jl"))
-# include(joinpath("cuda", "operators_columnwise.jl"))
-# include(joinpath("cuda", "matrix_fields_single_field_solve.jl"))
-# include(joinpath("cuda", "matrix_fields_multiple_field_solve.jl"))
-# include(joinpath("cuda", "operators_spectral_element.jl"))
+include(joinpath("cuda", "operators_finite_difference.jl"))
+include(joinpath("cuda", "remapping_distributed.jl"))
+include(joinpath("cuda", "operators_integral.jl"))
+include(joinpath("cuda", "remapping_interpolate_array.jl"))
+include(joinpath("cuda", "limiters.jl"))
+include(joinpath("cuda", "operators_sem_shmem.jl"))
+include(joinpath("cuda", "operators_fd_shmem_common.jl"))
+include(joinpath("cuda", "operators_fd_shmem.jl"))
+include(joinpath("cuda", "operators_columnwise.jl"))
+include(joinpath("cuda", "matrix_fields_single_field_solve.jl"))
+include(joinpath("cuda", "matrix_fields_multiple_field_solve.jl"))
+include(joinpath("cuda", "operators_spectral_element.jl"))
 
 end

ext/ClimaCoreCUDAExt.jl

@@ -5,4 +5,4 @@ import ClimaCore: DataLayouts
 
 include("scopes.jl")
 include("loops.jl")
-# include("data_layouts_threadblock.jl")
+include("data_layouts_threadblock.jl")

ext/cuda/data_layouts.jl

@@ -4,37 +4,41 @@ maximum_allowable_threads() = (
     CUDA.attribute(CUDA.device(), CUDA.DEVICE_ATTRIBUTE_MAX_BLOCK_DIM_Z),
 )
 
+# Wrappers for sizes and indices that are passed to kernels as type parameters.
+@inline unval(x) = x
+@inline unval(::Val{x}) where {x} = x
+
 """
-    universal_index(::DataLayout)
+    linear_thread_idx()
 
-Returns a universal cartesian index,
-computed from CUDA's `threadIdx`,
-`blockIdx` and `blockDim`.
+Returns the linear index of the current thread across all blocks, computed from
+CUDA's `threadIdx`, `blockIdx` and `blockDim`.
 """
-function universal_index end
+@inline linear_thread_idx() =
+    (CUDA.blockIdx().x - Int32(1)) * CUDA.blockDim().x + CUDA.threadIdx().x
 
 ##### Masked
-@inline cartesian_indices_mask(us::DataLayouts.UniversalSize, mask::IJHMask) =
-    cartesian_indices_mask(us, typeof(mask), mask.N[1])
+@inline cartesian_indices_mask(data, mask::IJHMask) =
+    cartesian_indices_mask(data, typeof(mask), mask.N[1])
 
 @inline function cartesian_indices_mask(
-    us::DataLayouts.UniversalSize,
+    data,
     ::Type{<:IJHMask},
     n_active_columns::Integer,
 )
-    (Ni, _, _, Nv, Nh) = DataLayouts.universal_size(us)
+    Nv = size(data, 1)
     return CartesianIndices((1:Nv, 1:n_active_columns))
 end
-@inline masked_partition(mask::IJHMask, n_max_threads, us) =
-    masked_partition(typeof(mask), mask.N[1], n_max_threads, us)
+@inline masked_partition(mask::IJHMask, n_max_threads, data) =
+    masked_partition(typeof(mask), mask.N[1], n_max_threads, data)
 
 @inline function masked_partition(
     ::Type{<:IJHMask},
     n_active_columns,
     n_max_threads,
-    us,
+    data,
 )
-    (_, _, _, Nv, _) = DataLayouts.universal_size(us)
+    Nv = size(data, 1)
     nitems = n_active_columns * Nv
     return linear_partition(nitems, n_max_threads)
 end
@@ -60,39 +64,28 @@ end
     blocks = cld(nitems, threads)
     return (; threads, blocks)
 end
-@inline function cartesian_indices(us::UniversalSize)
-    inds = DataLayouts.universal_size(us)
-    return CartesianIndices(map(Base.OneTo, inds))
-end
-@inline linear_is_valid_index(i::Integer, us::UniversalSize) =
-    1 ≤ i ≤ DataLayouts.get_N(us)
+@inline cartesian_indices(data) =
+    CartesianIndices(map(Base.OneTo, size(data)))
+@inline linear_is_valid_index(i::Integer, data) = 1 ≤ i ≤ length(data)
 
 ##### Column-wise
-@inline function cartesian_indices_columnwise(us::DataLayouts.UniversalSize)
-    (Ni, Nj, _, _, Nh) = DataLayouts.universal_size(us)
-    inds = (Ni, Nj, Nh)
-    return CartesianIndices(map(Base.OneTo, inds))
+@inline function cartesian_indices_columnwise(data)
+    (_, Ni, Nj, Nh) = size(data)
+    return CartesianIndices(map(Base.OneTo, (Ni, Nj, Nh)))
 end
 
 ##### Element-wise (e.g., limiters)
 # TODO
 
 ##### Multiple-field solve partition
-@inline function cartesian_indices_multiple_field_solve(
-    us::DataLayouts.UniversalSize;
-    Nnames,
-)
-    (Ni, Nj, _, _, Nh) = DataLayouts.universal_size(us)
-    inds = (Ni, Nj, Nh, Nnames)
-    return CartesianIndices(map(Base.OneTo, inds))
+@inline function cartesian_indices_multiple_field_solve(data; Nnames)
+    (_, Ni, Nj, Nh) = size(data)
+    return CartesianIndices(map(Base.OneTo, (Ni, Nj, Nh, Nnames)))
 end
 
 ##### spectral kernel partition
-@inline function spectral_partition(
-    us::DataLayouts.UniversalSize,
-    n_max_threads::Integer = 256;
-)
-    (Ni, Nj, _, Nv, Nh) = DataLayouts.universal_size(us)
+@inline function spectral_partition(data, n_max_threads::Integer = 256)
+    (Nv, Ni, Nj, Nh) = size(data)
     Nvthreads = min(fld(n_max_threads, Ni * Nj), maximum_allowable_threads()[3])
     Nvblocks = cld(Nv, Nvthreads)
     @assert prod((Ni, Nj, Nvthreads)) ≤ n_max_threads "threads,n_max_threads=($(prod((Ni, Nj, Nvthreads))),$n_max_threads)"
@@ -121,11 +114,11 @@ end
 
 ##### shmem fd kernel partition
 @inline function fd_shmem_stencil_partition(
-    us::DataLayouts.UniversalSize,
+    data,
     n_face_levels::Integer,
     n_max_threads::Integer = 256;
 )
-    (Ni, Nj, _, Nv, Nh) = DataLayouts.universal_size(us)
+    (Nv, Ni, Nj, Nh) = size(data)
     Nvthreads = n_face_levels
     @assert Nvthreads <= maximum_allowable_threads()[1] "Number of vertical face levels cannot exceed $(maximum_allowable_threads()[1])"
     Nvblocks = cld(Nv, Nvthreads) # +1 may be needed to guarantee that shared memory is populated at the last cell face
@@ -135,19 +128,15 @@ end
         Nvthreads,
     )
 end
-@inline function fd_shmem_stencil_universal_index(
-    space::Spaces.AbstractSpace,
-    us,
-)
+@inline function fd_shmem_stencil_universal_index(space::Spaces.AbstractSpace, data)
     (tv,) = CUDA.threadIdx()
     (h, bv, ij) = CUDA.blockIdx()
     v = tv + (bv - 1) * CUDA.blockDim().x
-    (Ni, Nj, _, _, _) = DataLayouts.universal_size(us)
+    (_, Ni, Nj, _) = size(data)
     if Ni * Nj < ij
         return CartesianIndex((-1, -1, -1, -1))
     end
     @inbounds (i, j) = CartesianIndices((Ni, Nj))[ij].I
     return CartesianIndex((v, i, j, h))
 end
-@inline fd_shmem_stencil_is_valid_index(I, us::UniversalSize) =
-    1 ≤ I[5] ≤ DataLayouts.get_Nh(us)
+@inline fd_shmem_stencil_is_valid_index(I, data) = 1 ≤ I[4] ≤ size(data, 4)