Make single-multiple field solvers mask-aware (#2270)

Make reductions mask-aware

Author: kmdeck

ext/cuda/matrix_fields_multiple_field_solve.jl

@@ -21,6 +21,7 @@ NVTX.@annotate function multiple_field_solve!(
     Nnames = length(names)
     Ni, Nj, _, _, Nh = size(Fields.field_values(x1))
     sscache = Operators.strip_space(cache)
+    mask = Spaces.get_mask(axes(x1))
     ssx = Operators.strip_space(x)
     ssA = Operators.strip_space(A)
     ssb = Operators.strip_space(b)
@@ -33,7 +34,7 @@ NVTX.@annotate function multiple_field_solve!(
     device = ClimaComms.device(x[first(names)])
 
     us = UniversalSize(Fields.field_values(x1))
-    args = (device, caches, xs, As, bs, x1, us, Val(Nnames))
+    args = (device, caches, xs, As, bs, x1, us, mask, Val(Nnames))
 
     nitems = Ni * Nj * Nh * Nnames
     threads = threads_via_occupancy(multiple_field_solve_kernel!, args)
@@ -86,10 +87,12 @@ function multiple_field_solve_kernel!(
     bs,
     x1,
     us::UniversalSize,
+    mask,
     ::Val{Nnames},
 ) where {Nnames}
     @inbounds begin
         (I, iname) = multiple_field_solve_universal_index(us)
+        DataLayouts.should_compute(mask, I) || return nothing
         if multiple_field_solve_is_valid_index(I, us)
             (i, j, _, _, h) = I.I
             generated_single_field_solve!(

ext/cuda/matrix_fields_single_field_solve.jl

@@ -16,7 +16,8 @@ import ClimaCore.RecursiveApply: ⊠, ⊞, ⊟, rmap, rzero, rdiv
 function single_field_solve!(device::ClimaComms.CUDADevice, cache, x, A, b)
     Ni, Nj, _, _, Nh = size(Fields.field_values(A))
     us = UniversalSize(Fields.field_values(A))
-    args = (device, cache, x, A, b, us)
+    mask = Spaces.get_mask(axes(x))
+    args = (device, cache, x, A, b, us, mask)
     threads = threads_via_occupancy(single_field_solve_kernel!, args)
     nitems = Ni * Nj * Nh
     n_max_threads = min(threads, nitems)
@@ -30,8 +31,9 @@ function single_field_solve!(device::ClimaComms.CUDADevice, cache, x, A, b)
     call_post_op_callback() && post_op_callback(x, device, cache, x, A, b)
 end
 
-function single_field_solve_kernel!(device, cache, x, A, b, us)
+function single_field_solve_kernel!(device, cache, x, A, b, us, mask)
     I = columnwise_universal_index(us)
+    DataLayouts.should_compute(mask, I) || return nothing
     if columnwise_is_valid_index(I, us)
         (i, j, _, _, h) = I.I
         _single_field_solve!(

ext/cuda/operators_integral.jl

@@ -20,6 +20,10 @@ function column_reduce_device!(
 ) where {F, T}
     Ni, Nj, _, _, Nh = size(Fields.field_values(output))
     us = UniversalSize(Fields.field_values(output))
+    mask = Spaces.get_mask(space)
+    if !(mask isa DataLayouts.NoMask) && space isa Spaces.FiniteDifferenceSpace
+        error("Masks not supported for FiniteDifferenceSpace")
+    end
     args = (
         single_column_reduce!,
         f,
@@ -29,6 +33,7 @@ function column_reduce_device!(
         init,
         space,
         us,
+        mask,
     )
     nitems = Ni * Nj * Nh
     threads = threads_via_occupancy(bycolumn_kernel!, args)
@@ -57,6 +62,10 @@ function column_accumulate_device!(
     space,
 ) where {F, T}
     out_fv = Fields.field_values(output)
+    mask = Spaces.get_mask(space)
+    if !(mask isa DataLayouts.NoMask) && space isa Spaces.FiniteDifferenceSpace
+        error("Masks not supported for FiniteDifferenceSpace")
+    end
     us = UniversalSize(out_fv)
     args = (
         single_column_accumulate!,
@@ -67,6 +76,7 @@ function column_accumulate_device!(
         init,
         space,
         us,
+        mask,
     )
     (Ni, Nj, _, _, Nh) = DataLayouts.universal_size(us)
     nitems = Ni * Nj * Nh
@@ -81,7 +91,7 @@ function column_accumulate_device!(
     )
 end
 
-bycolumn_kernel!(
+function bycolumn_kernel!(
     single_column_function!::S,
     f::F,
     transform::T,
@@ -90,11 +100,13 @@ bycolumn_kernel!(
     init,
     space,
     us::DataLayouts.UniversalSize,
-) where {S, F, T} =
+    mask,
+) where {S, F, T}
     if space isa Spaces.FiniteDifferenceSpace
         single_column_function!(f, transform, output, input, init, space)
     else
         I = columnwise_universal_index(us)
+        DataLayouts.should_compute(mask, I) || return nothing
         if columnwise_is_valid_index(I, us)
             (i, j, _, _, h) = I.I
             single_column_function!(
@@ -107,3 +119,4 @@ bycolumn_kernel!(
             )
         end
     end
+end

src/Fields/indices.jl

@@ -244,3 +244,9 @@ function byslab(
         end
     end
 end
+
+universal_index(colidx::Fields.ColumnIndex{2}) =
+    CartesianIndex(colidx.ij[1], colidx.ij[2], 1, 1, colidx.h)
+
+universal_index(colidx::Fields.ColumnIndex{1}) =
+    CartesianIndex(colidx.ij[1], 1, 1, 1, colidx.h)

src/MatrixFields/single_field_solver.jl

@@ -1,3 +1,4 @@
+import ..DataLayouts
 dual_type(::Type{A}) where {A} = typeof(Geometry.dual(A.instance))
 
 inv_return_type(::Type{X}) where {X} = error(
@@ -54,11 +55,10 @@ function single_field_solve_diag_matrix_row!(
     A::ColumnwiseBandMatrixField,
     b,
 )
-    Aⱼs = unzip_tuple_field_values(Fields.field_values(A.entries))
-    b_vals = Fields.field_values(b)
-    x_vals = Fields.field_values(x)
-    (A₀,) = Aⱼs
-    @. x_vals = inv(A₀) ⊠ b_vals
+    # Use fields here, and not field values, so that this operation is
+    # mask-aware.
+    A₀ = A.entries.:1
+    @. x = inv(A₀) ⊠ b
 end
 single_field_solve!(_, x, A::ScalingFieldMatrixEntry, b) =
     x .= (inv(scaling_value(A)),) .* b
@@ -82,17 +82,22 @@ function _single_field_solve!(
     b,
 )
     space = axes(x)
+    mask = Spaces.get_mask(space)
     if space isa Spaces.FiniteDifferenceSpace
+        @assert mask isa DataLayouts.NoMask
         _single_field_solve_col!(device, cache, x, A, b)
     else
         Fields.bycolumn(space) do colidx
-            _single_field_solve_col!(
-                device,
-                cache[colidx],
-                x[colidx],
-                A[colidx],
-                b[colidx],
-            )
+            I = Fields.universal_index(colidx)
+            if DataLayouts.should_compute(mask, I)
+                _single_field_solve_col!(
+                    device,
+                    cache[colidx],
+                    x[colidx],
+                    A[colidx],
+                    b[colidx],
+                )
+            end
         end
     end
 end

src/Operators/integrals.jl

@@ -124,7 +124,7 @@ bottom of each column and moving upward, and the result of the final iteration
 is passed to the `transform` function before being stored in `output`. If `init`
 is specified, it is used as the initial value of the iteration; otherwise, the
 value at the bottom of each column in `input` is used as the initial value.
-    
+
 With `first_level` and `last_level` denoting the indices of the boundary levels
 of `input`, the reduction in each column can be summarized as follows:
   - If `init` is unspecified,
@@ -156,27 +156,35 @@ function column_reduce!(
     column_reduce_device!(device, f, transform, output, input, init, space)
 end
 
-column_reduce_device!(
+function column_reduce_device!(
     ::ClimaComms.AbstractCPUDevice,
     f::F,
     transform::T,
     output,
     input,
     init,
     space,
-) where {F, T} =
-    space isa Spaces.FiniteDifferenceSpace ?
-    single_column_reduce!(f, transform, output, input, init, space) :
-    Fields.bycolumn(space) do colidx
-        single_column_reduce!(
-            f,
-            transform,
-            output[colidx],
-            input[colidx],
-            init,
-            space[colidx],
-        )
+) where {F, T}
+    mask = Spaces.get_mask(space)
+    if space isa Spaces.FiniteDifferenceSpace
+        @assert mask isa DataLayouts.NoMask
+        single_column_reduce!(f, transform, output, input, init, space)
+    else
+        Fields.bycolumn(space) do colidx
+            I = Fields.universal_index(colidx)
+            if DataLayouts.should_compute(mask, I)
+                single_column_reduce!(
+                    f,
+                    transform,
+                    output[colidx],
+                    input[colidx],
+                    init,
+                    space[colidx],
+                )
+            end
+        end
     end
+end
 
 # On GPUs, input and output go through strip_space to become _input and _output.
 function single_column_reduce!(
@@ -214,7 +222,7 @@ from the bottom of each column and moving upward, and the result of each
 iteration is passed to the `transform` function before being stored in `output`.
 The `init` value is is optional for center-to-center, face-to-face, and
 face-to-center accumulation, but it is required for center-to-face accumulation.
-    
+
 With `first_level` and `last_level` denoting the indices of the boundary levels
 of `input`, the accumulation in each column can be summarized as follows:
   - For center-to-center and face-to-face accumulation with `init` unspecified,
@@ -276,27 +284,35 @@ function column_accumulate!(
     column_accumulate_device!(device, f, transform, output, input, init, space)
 end
 
-column_accumulate_device!(
+function column_accumulate_device!(
     ::ClimaComms.AbstractCPUDevice,
     f::F,
     transform::T,
     output,
     input,
     init,
     space,
-) where {F, T} =
-    space isa Spaces.FiniteDifferenceSpace ?
-    single_column_accumulate!(f, transform, output, input, init, space) :
-    Fields.bycolumn(space) do colidx
-        single_column_accumulate!(
-            f,
-            transform,
-            output[colidx],
-            input[colidx],
-            init,
-            space[colidx],
-        )
+) where {F, T}
+    mask = Spaces.get_mask(space)
+    if space isa Spaces.FiniteDifferenceSpace
+        @assert mask isa DataLayouts.NoMask
+        single_column_accumulate!(f, transform, output, input, init, space)
+    else
+        Fields.bycolumn(space) do colidx
+            I = Fields.universal_index(colidx)
+            if DataLayouts.should_compute(mask, I)
+                single_column_accumulate!(
+                    f,
+                    transform,
+                    output[colidx],
+                    input[colidx],
+                    init,
+                    space[colidx],
+                )
+            end
+        end
     end
+end
 
 # On GPUs, input and output go through strip_space to become _input and _output.
 function single_column_accumulate!(