Unify DataLayouts and simplify API

Author: dennisYatunin

.buildkite/pipeline.yml

@@ -104,76 +104,31 @@ steps:
   - group: "Unit: DataLayouts"
     steps:
 
-      - label: "Unit: data0d"
-        key: unit_data0d
+      - label: "Unit: data fill and copyto (1 thread)"
+        key: unit_data_fill_and_copyto
         retry: *retry_policy
-        command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/data0d.jl"
+        command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_fill_and_copyto.jl"
 
-      - label: "Unit: data_fill"
-        key: unit_data_fill
+      - label: "Unit: data fill and copyto (4 threads)"
+        key: threaded_unit_data_fill_and_copyto
         retry: *retry_policy
-        command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_fill.jl"
+        command: "julia --threads=4 --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_fill_and_copyto.jl"
 
-      - label: "Unit: data_copyto"
-        key: unit_data_copyto
-        retry: *retry_policy
-        command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_copyto.jl"
-
-      - label: "Unit: cartesian_field_index"
-        key: unit_data_cartesian_field_index
-        retry: *retry_policy
-        command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_cartesian_field_index.jl"
-
-      - label: "Unit: non_extruded_broadcast"
-        key: unit_non_extruded_broadcast
-        retry: *retry_policy
-        command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_non_extruded_broadcast.jl"
-
-      - label: "Unit: mapreduce"
+      - label: "Unit: mapreduce (1 thread)"
         key: unit_data_mapreduce
         retry: *retry_policy
         command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_mapreduce.jl"
 
+      - label: "Unit: mapreduce (4 threads)"
+        key: threaded_unit_data_mapreduce
+        retry: *retry_policy
+        command: "julia --threads=4 --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_mapreduce.jl"
+
       - label: "Unit: data_opt_similar"
         key: data_opt_similar
         retry: *retry_policy
         command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/opt_similar.jl"
 
-      - label: "Unit: opt_universal_size"
-        key: opt_universal_size
-        retry: *retry_policy
-        command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/opt_universal_size.jl"
-
-      - label: "Unit: data_ndims"
-        key: unit_data_ndims
-        retry: *retry_policy
-        command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_ndims.jl"
-
-      - label: "Unit: unit_data2array"
-        key: unit_data2array
-        retry: *retry_policy
-        command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_data2array.jl"
-
-      - label: "Unit: data1d"
-        key: unit_data1d
-        retry: *retry_policy
-        command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/data1d.jl"
-
-      - label: "Unit: data2d"
-        key: unit_data2d
-        retry: *retry_policy
-        command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/data2d.jl"
-
-      - label: "Unit: data1dx"
-        key: unit_data1dx
-        retry: *retry_policy
-        command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/data1dx.jl"
-
-      - label: "Unit: data2dx"
-        key: unit_data2dx
-        retry: *retry_policy
-        command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/data2dx.jl"
-
       - label: "Unit: data cuda"
         key: unit_data_cuda
         retry: *retry_policy
@@ -196,18 +151,18 @@ steps:
         agents:
           slurm_gpus: 1
 
-      - label: "Unit: data fill"
-        key: gpu_unit_data_fill
+      - label: "Unit: data fill and copyto (1 gpu)"
+        key: gpu_unit_data_fill_and_copyto
         retry: *retry_policy
         command:
           - "julia --project=.buildkite -e 'using CUDA; CUDA.versioninfo()'"
-          - "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_fill.jl"
+          - "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_fill_and_copyto.jl"
         env:
           CLIMACOMMS_DEVICE: "CUDA"
         agents:
           slurm_gpus: 1
 
-      - label: "Unit: data mapreduce"
+      - label: "Unit: data mapreduce (1 gpu)"
         key: gpu_unit_data_mapreduce
         retry: *retry_policy
         command:
@@ -233,17 +188,6 @@ steps:
            modules: mpiwrapper/2024_05_27 climacommon/2026_02_18
         soft_fail: true # remove this after library issues are fixed
 
-      - label: "Unit: data copyto"
-        key: gpu_unit_data_copyto
-        retry: *retry_policy
-        command:
-          - "julia --project=.buildkite -e 'using CUDA; CUDA.versioninfo()'"
-          - "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_copyto.jl"
-        env:
-          CLIMACOMMS_DEVICE: "CUDA"
-        agents:
-          slurm_gpus: 1
-
   - group: "Unit: Geometry"
     steps:
 
@@ -1624,27 +1568,11 @@ steps:
   - group: "Perf: DataLayouts"
     steps:
 
-      - label: "Perf: DataLayouts fill"
-        key: "cpu_datalayouts_fill"
-        retry: *retry_policy
-        command: "julia --color=yes --project=.buildkite test/DataLayouts/benchmark_fill.jl"
-
       - label: "Perf: DataLayouts copyto!"
         key: "cpu_datalayouts_copyto"
         retry: *retry_policy
         command: "julia --color=yes --project=.buildkite test/DataLayouts/benchmark_copyto.jl"
 
-      - label: "Perf: DataLayouts fill"
-        key: "gpu_datalayouts_fill"
-        retry: *retry_policy
-        command:
-          - "julia --project=.buildkite -e 'using CUDA; CUDA.versioninfo()'"
-          - "julia --color=yes --project=.buildkite test/DataLayouts/benchmark_fill.jl"
-        env:
-          CLIMACOMMS_DEVICE: "CUDA"
-        agents:
-          slurm_gpus: 1
-
       - label: "Perf: DataLayouts copyto"
         key: "gpu_datalayouts_copyto"
         retry: *retry_policy
@@ -2093,7 +2021,7 @@ steps:
         env:
           TEST_NAME: "sphere/baroclinic_wave_rhoe_hf"
           FLOAT_TYPE: "Float32"
-          horizontal_layout_type: "IJHF"
+          horizontal_layout_type: "VIJHF"
           CLIMACOMMS_DEVICE: "CUDA"
         agents:
           slurm_gpus: 1
@@ -2414,7 +2342,7 @@ steps:
         env:
           TEST_NAME: "sphere/baroclinic_wave_rhoe_hf"
           FLOAT_TYPE: "Float64"
-          horizontal_layout_type: "IJHF"
+          horizontal_layout_type: "VIJHF"
 
       - label: ":computer: 3D sphere baroclinic wave (ρe)"
         key: "cpu_baroclinic_wave_rho_e"

benchmarks/scripts/index_swapping.jl

@@ -106,7 +106,7 @@ function custom_kernel_knl_bc_1swap!(y1, bc, us)
         if tidx ≤ get_N(us)
             n = (get_Nij(us), get_Nij(us), 1, get_Nv(us), get_Nh(us))
             GCI = CartesianIndices(map(x -> Base.OneTo(x), n))[tidx]
-            # Perform index swap (as in `getindex(::AbstractData, ::CartesianIndex)`)
+            # Perform index swap (as in `getindex(::DataLayout, ::CartesianIndex)`)
             i, j, _, v, h = GCI.I
             CI = CartesianIndex(v, i, j, 1, h)
             y1[CI] = bc[CI]
@@ -140,7 +140,7 @@ function custom_kernel_knl_bc_2swap!(y1, bc, us)
             (v, i, j, _, h) = CIK.I
             GCI = CartesianIndex(i, j, 1, v, h)
 
-            # Swap again (in `getindex(::AbstractData, ::CartesianIndex)`)
+            # Swap again (in `getindex(::DataLayout, ::CartesianIndex)`)
             (i, j, _, v, h) = GCI.I
             CI = CartesianIndex(v, i, j, 1, h)
             y1[CI] = bc[CI]

benchmarks/scripts/thermo_bench.jl

@@ -154,7 +154,7 @@ using Test
     )
     x = fill((; ts = nt_ts, nt_core...), cspace)
     xv = fill((; ts = nt_ts, nt_core...), cspace)
-    (_, Nij, _, Nv, Nh) = size(Fields.field_values(x.ts))
+    (Nv, Nij, _, Nh) = size(Fields.field_values(x.ts))
     us = TB.UniversalSizesStatic(Nv, Nij, Nh)
     function to_vec(ξ)
         pns = propertynames(ξ)

benchmarks/scripts/thermo_bench_bw.jl

@@ -180,7 +180,7 @@ using Test
     )
     x = fill((; ts = zero(TBB.PhaseEquil{FT}), nt_core...), cspace)
     xv = fill((; ts = nt_ts, nt_core...), cspace)
-    (_, Nij, _, Nv, Nh) = size(Fields.field_values(x.ts))
+    (Nv, Nij, _, Nh) = size(Fields.field_values(x.ts))
     us = TBB.UniversalSizesStatic(Nv, Nij, Nh)
     function to_vec(ξ)
         pns = propertynames(ξ)

docs/clima_core_specific.md

@@ -41,7 +41,7 @@ ClimaCore.jl provides the dynamical core infrastructure for [CliMA](https://clim
 1. **`Field`** (`src/Fields/`) — the primary data type. A field wraps data on a space and supports broadcast, reductions, and operator application.
 2. **`Space`** (`src/Spaces/`) — represents a discretized function space (spectral element, finite-difference, or extruded hybrid). Constructed from a grid and a quadrature rule.
 3. **Operators** (`src/Operators/`) — lazy differential operators (gradient, divergence, curl, interpolation, restriction) that compose via Julia's broadcast system.
-4. **`DataLayout`** (`src/DataLayouts/`) — the storage backends (IJFH, VIJFH, VF, etc.) that determine memory layout for CPU vs GPU performance.
+4. **`DataLayout`** (`src/DataLayouts/`) — the storage backends (VIJFH, VIJHF, etc.) that determine memory layout for CPU vs GPU performance.
 5. **`MatrixFields`** (`src/MatrixFields/`) — banded-matrix field algebra used for implicit vertical solvers and Jacobian construction.
 
 ## Test groups

docs/src/APIs/datalayouts_api.md

@@ -7,17 +7,10 @@ CurrentModule = ClimaCore
 ```@docs
 DataLayouts
 DataLayouts.DataF
-DataLayouts.IF
-DataLayouts.IJF
-DataLayouts.VF
-DataLayouts.IFH
-DataLayouts.IJFH
-DataLayouts.VIFH
 DataLayouts.VIJFH
-DataLayouts.IHF
-DataLayouts.IJHF
-DataLayouts.VIHF
 DataLayouts.VIJHF
+DataLayouts.VIH1
+DataLayouts.IH1JH2
 DataLayouts.bitcast_struct
 DataLayouts.default_basetype
 DataLayouts.check_basetype
@@ -26,6 +19,4 @@ DataLayouts.num_basetypes
 DataLayouts.struct_field_view
 DataLayouts.set_struct!
 DataLayouts.get_struct
-DataLayouts.parent_array_type
-DataLayouts.promote_parent_array_type
 ```

docs/src/APIs/utilities_api.md

@@ -9,6 +9,8 @@ Utilities.unionall_type
 Utilities.replace_type_parameter
 Utilities.fieldtype_vals
 Utilities.new
+Utilities.is_inferred_type
+Utilities.return_type
 Utilities.unsafe_eltype
 Utilities.safe_eltype
 ```

docs/src/debugging.md

@@ -140,13 +140,13 @@ only when the `condition` is true (in this case `has_nans || has_inf`).
 Now, when we run our example, we will see
 ```julia
 julia> renormalized_energy(myrho, myP, myu)
-Infiltrating post_op_callback(::ClimaCore.DataLayouts.IJFH{Float64, 4, Array{Float64, 4}}, ::ClimaCore.DataLayouts.IJFH{Float64, 4, Array{Float64, 4}}, ::Vararg{Any}; kwargs::@Kwargs{})
+Infiltrating post_op_callback(::ClimaCore.DataLayouts.VIJFH{...}, ::ClimaCore.DataLayouts.VIJFH{...}, ::Vararg{Any}; kwargs::@Kwargs{})
   at REPL[40]:4
 infil>
 ```
 Here, we are dropped into a new REPL with full access to the variables in the scope where the `NaN` occurred. However, because of how `post_op_callback`, this is at a low level within `ClimaCore`, which is typically not useful. Hence, the next step is to type `@trace`, which prints out
 ```julia
-[1] post_op_callback(::ClimaCore.DataLayouts.IJFH{…}, ::ClimaCore.DataLayouts.IJFH{…}, ::Vararg{…}; kwargs::@Kwargs{})
+[1] post_op_callback(::ClimaCore.DataLayouts.VIJFH{…}, ::ClimaCore.DataLayouts.VIJFH{…}, ::Vararg{…}; kwargs::@Kwargs{})
     at REPL[40]:4
 [2] post_op_callback
     at REPL[40]:1
@@ -250,7 +250,7 @@ parts of the broadcasted object contains NaNs:
 ```julia
 using StructuredPrinting
 import ClimaCore: DataLayouts
-highlight_nans(x::DataLayouts.AbstractData) = any(y->isnan(y), parent(x));
+highlight_nans(x::DataLayouts.DataLayout) = any(y->isnan(y), parent(x));
 highlight_nans(_) = false;
 bc = Infiltrator.safehouse.args[2]; # we know that argument 2 is the broadcasted object
 (; result) = Infiltrator.safehouse; # get the result

docs/src/masks.md

@@ -169,60 +169,6 @@ operations of mask-aware and mask-unaware:
    This was a design implementation detail, users should not generally depend on the results where `mask == 0`, in case this is changed in the future. 
  - internal array operations (`fill!(parent(field), 0)`) mask-unaware.
 
-## Temporary work-arounds
-
-We can perform mask-aware reductions with the following work-around
-
-```julia
-using ClimaComms
-ClimaComms.@import_required_backends
-import ClimaCore: Spaces, Fields, DataLayouts, Geometry, Operators
-using ClimaCore.CommonSpaces
-using Test
-
-FT = Float64
-ᶜspace = ExtrudedCubedSphereSpace(FT;
-    z_elem = 10,
-    z_min = 0,
-    z_max = 1,
-    radius = 10,
-    h_elem = 10,
-    n_quad_points = 4,
-    staggering = CellCenter(),
-    enable_mask = true,
-)
-ᶠspace = Spaces.face_space(ᶜspace)
-ᶠcoords = Fields.coordinate_field(ᶠspace)
-
-# Set the mask
-Spaces.set_mask!(ᶜspace) do coords
-    coords.lat > 0.5
-end
-
-# get the mask
-mask = Spaces.get_mask(ᶜspace)
-
-# make a field of ones
-ᶜf = ones(ᶜspace) # ignores mask
-
-# bitmask spanning datalayout
-bm = DataLayouts.full_bitmask(mask, Fields.field_values(ᶜf));
-
-# mask-unaware integral (includes jacobian weighting)
-@show sum(ᶜf)
-
-# mask-unaware sum (excludes jacobian weighting)
-@show sum(Fields.field_values(ᶜf))
-
-# mask-aware sum (excludes jacobian)
-@show sum(parent(ᶜf)[bm])
-
-# level mask
-ᶜf_lev = Fields.level(ᶜf, 1);
-bm_lev = DataLayouts.full_bitmask(mask, Fields.field_values(ᶜf_lev));
-@show sum(parent(ᶜf_lev)[bm_lev])
-```
-
 ## Developer docs
 
 In order to support masks, we define their types in `DataLayouts`, since 

examples/bickleyjet/bickleyjet_cg_invariant_hypervisc.jl

@@ -9,8 +9,7 @@ import ClimaCore:
     Operators,
     Spaces,
     Topologies,
-    Quadratures,
-    DataLayouts
+    Quadratures
 using OrdinaryDiffEqSSPRK: ODEProblem, solve, SSPRK33
 
 using Logging
@@ -163,7 +162,7 @@ sol_global = []
 if usempi
     for sol_step in sol.u
         sol_step_values_global =
-            DataLayouts.gather(context, Fields.field_values(sol_step))
+            ClimaComms.gather(context, Fields.field_values(sol_step))
         if ClimaComms.iamroot(context)
             sol_step_global = Fields.Field(sol_step_values_global, global_space)
             push!(sol_global, sol_step_global)