Merge pull request #2522 from CliMA/dy/data_layouts_refactor

Unify DataLayouts and simplify API

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

@@ -26,6 +26,7 @@ steps:
       # Instantiate and dev install ClimaCore
       - julia --project=$COUPLER_PATH/experiments/AMIP/ -e 'using Pkg; Pkg.instantiate(;verbose=true)'
       - julia --project=$COUPLER_PATH/experiments/AMIP/ -e 'using Pkg; Pkg.develop(path=".")'
+      - julia --project=$COUPLER_PATH/experiments/AMIP/ -e 'using Pkg; Pkg.develop(path="lib/ClimaCoreMakie")'
       - julia --project=$COUPLER_PATH/experiments/AMIP/ -e 'using Pkg; Pkg.add("MPI")'
       - julia --project=$COUPLER_PATH/experiments/AMIP/ -e 'using Pkg; Pkg.precompile()'
       - julia --project=$COUPLER_PATH/experiments/AMIP/ -e 'using Pkg; Pkg.status()'

.buildkite/perf/pipeline.yml

@@ -104,76 +104,41 @@ 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
+      - 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_copyto.jl"
+        command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_loops.jl"
 
-      - label: "Unit: cartesian_field_index"
-        key: unit_data_cartesian_field_index
+      - label: "Unit: data loops (4 threads)"
+        key: threaded_unit_data_loops
         retry: *retry_policy
-        command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_cartesian_field_index.jl"
+        command: "julia --threads=4 --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_loops.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 +161,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 +198,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:
 
@@ -1629,7 +1583,7 @@ steps:
         retry: *retry_policy
         command: "julia --color=yes --project=.buildkite test/DataLayouts/benchmark_fill.jl"
 
-      - label: "Perf: DataLayouts copyto!"
+      - label: "Perf: DataLayouts copyto"
         key: "cpu_datalayouts_copyto"
         retry: *retry_policy
         command: "julia --color=yes --project=.buildkite test/DataLayouts/benchmark_copyto.jl"
@@ -2093,7 +2047,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 +2368,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"

.buildkite/pipeline.yml

@@ -18,7 +18,11 @@ jobs:
       - name: Install dependencies
         run: >
           julia --project=docs -e 'using Pkg; Pkg.develop(path=".");
-          Pkg.develop(path="lib/ClimaCoreVTK"); 
+          Pkg.develop(path="lib/ClimaCoreVTK");
+          Pkg.develop(path="lib/ClimaCoreTempestRemap");
+          Pkg.develop(path="lib/ClimaCoreMakie");
+          Pkg.develop(path="lib/ClimaCorePlots");
+          Pkg.develop(path="lib/ClimaCoreSpectra");
           Pkg.instantiate(;verbose=true)'
       - name: Build and deploy
         env:

.github/workflows/docs.yml

@@ -4,6 +4,17 @@ ClimaCore.jl Release Notes
 main
 -------
 
+- Refactor DataLayouts module [2522](https://github.com/CliMA/ClimaCore.jl/pull/2522)
+  - All data layout types are unified into a single `DataLayout` type, and all
+    loops over data go through two communication primitives (`foreach_slice` and
+    `reduce_points`) that work the same way on CPUs, multi-threaded CPUs, and
+    GPUs. Runtime performance on CPUs and GPUs matches the previous release.
+  - Data in a `Field` is now indexed in the order `[v, i, j, h]`, and every
+    parent array has a 5-dimensional shape, with one axis used to store struct
+    fields. Code that indexes into `parent(field)` arrays must be updated.
+  - Several old names (`AbstractData`, `IJFH`, `IJHF`) are available as aliases,
+    and `[i, j, f, v, h]` indexing is still allowed for backward compatibility.
+
 v0.14.55
 -------
 

NEWS.md

@@ -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/index_swapping.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.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(ξ)

benchmarks/scripts/thermo_bench_bw.jl

@@ -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