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
-        retry: *retry_policy
-        command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/data0d.jl"
-
-      - label: "Unit: data_fill"
-        key: unit_data_fill
-        retry: *retry_policy
-        command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_fill.jl"
-
-      - label: "Unit: data_copyto"
+      - 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
+      - label: "Unit: data copyto threaded"
+        key: threaded_unit_data_copyto
         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"
+        command: "julia --threads=4 --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_copyto.jl"
 
       - label: "Unit: mapreduce"
         key: unit_data_mapreduce
         retry: *retry_policy
         command: "julia --color=yes --check-bounds=yes --project=.buildkite test/DataLayouts/unit_mapreduce.jl"
 
+      - label: "Unit: mapreduce threaded"
+        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,17 +151,6 @@ steps:
         agents:
           slurm_gpus: 1
 
-      - label: "Unit: data fill"
-        key: gpu_unit_data_fill
-        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"
-        env:
-          CLIMACOMMS_DEVICE: "CUDA"
-        agents:
-          slurm_gpus: 1
-
       - label: "Unit: data mapreduce"
         key: gpu_unit_data_mapreduce
         retry: *retry_policy
@@ -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

NEWS.md

@@ -523,7 +523,7 @@ v0.14.6
 
 - ![][badge-✨feature/enhancement] Added `array2field(::Field)` and `field2array(::Field)` convenience functions, to help facilitate use with RRTMGP. PR [#1768](https://github.com/CliMA/ClimaCore.jl/pull/1768)
 
-- ![][badge-🚀performance] `Nv` is now a type parameter in DataLayouts that have vertical levels. As a result, users can use `DataLayouts.nlevels(::AbstractData)` to obtain a compile-time constant for the number of vertical levels.
+- ![][badge-🚀performance] `Nv` is now a type parameter in DataLayouts that have vertical levels. As a result, users can use `DataLayouts.nlevels(::DataLayout)` to obtain a compile-time constant for the number of vertical levels.
 
 - ![][badge-✨feature/enhancement] Added `interpolate(field, target_hcoords,
   target_zcoord)` convenience function so that the `Remapper` does not have to

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]

docs/src/debugging.md

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

examples/hybrid/driver.jl

@@ -197,9 +197,9 @@ any(isnan, sol.u[end]) && error("NaNs found in result.")
 
 if is_distributed # replace sol.u on the root processor with the global sol.u
     global_Y_c_1 =
-        DataLayouts.gather(comms_ctx, Fields.field_values(sol.u[1].c))
+        ClimaComms.gather(comms_ctx, Fields.field_values(sol.u[1].c))
     global_Y_f_1 =
-        DataLayouts.gather(comms_ctx, Fields.field_values(sol.u[1].f))
+        ClimaComms.gather(comms_ctx, Fields.field_values(sol.u[1].f))
     if ClimaComms.iamroot(comms_ctx)
         global_h_space = make_horizontal_space(
             horizontal_mesh,
@@ -220,9 +220,9 @@ if is_distributed # replace sol.u on the root processor with the global sol.u
     end
     for i in 1:length(sol.u)
         global_Y_c =
-            DataLayouts.gather(comms_ctx, Fields.field_values(sol.u[i].c))
+            ClimaComms.gather(comms_ctx, Fields.field_values(sol.u[i].c))
         global_Y_f =
-            DataLayouts.gather(comms_ctx, Fields.field_values(sol.u[i].f))
+            ClimaComms.gather(comms_ctx, Fields.field_values(sol.u[i].f))
         if ClimaComms.iamroot(comms_ctx)
             global_sol_u[i] = Fields.FieldVector(
                 c = Fields.Field(global_Y_c, global_center_space),

examples/sphere/shallow_water.jl

@@ -14,8 +14,7 @@ import ClimaCore:
     Operators,
     Spaces,
     Quadratures,
-    Topologies,
-    DataLayouts
+    Topologies
 
 import QuadGK
 using OrdinaryDiffEqSSPRK: ODEProblem, init, solve!, SSPRK33
@@ -564,7 +563,7 @@ function shallow_water_driver(ARGS, ::Type{FT}) where {FT}
     if !usempi
         Y0_global = deepcopy(Y)
     else
-        Y0_global_values = DataLayouts.gather(context, Fields.field_values(Y))
+        Y0_global_values = ClimaComms.gather(context, Fields.field_values(Y))
         if ClimaComms.iamroot(context)
             Y0_global = Fields.Field(Y0_global_values, global_space)
         end
@@ -602,7 +601,7 @@ function shallow_water_driver(ARGS, ::Type{FT}) where {FT}
     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)

ext/ClimaCoreCUDAExt.jl

@@ -12,7 +12,6 @@ using CUDA: threadIdx, blockIdx, blockDim
 import StaticArrays: SVector, SMatrix, SArray
 import ClimaCore.DebugOnly: call_post_op_callback, post_op_callback
 import ClimaCore.DataLayouts: mapreduce_cuda
-import ClimaCore.DataLayouts: ToCUDA
 import ClimaCore.DataLayouts: NoMask, IJHMask
 import ClimaCore.DataLayouts: slab, column
 import ClimaCore.Utilities: half, new, cart_ind, linear_ind

ext/cuda/adapt.jl

@@ -54,12 +54,3 @@ Adapt.adapt_structure(
     lim.rtol,
     Limiters.NoConvergenceStats(),
 )
-
-Adapt.adapt_structure(to::CUDA.KernelAdaptor, mask::DataLayouts.IJHMask) =
-    DataLayouts.IJHMask(
-        Adapt.adapt(to, mask.is_active),
-        nothing,
-        Adapt.adapt(to, mask.i_map),
-        Adapt.adapt(to, mask.j_map),
-        Adapt.adapt(to, mask.h_map),
-    )