Add reverse option to column integral operators

Add a `reverse` keyword to `column_reduce!` and `column_accumulate!` so
vertical reductions and accumulations can proceed from the top of a column
downward.

Extend the operator tests to cover forward and reverse traversal across
center-to-center, face-to-face, center-to-face, and face-to-center
accumulations, and document the new behavior.

This description was generated by Codex.

Author: cflsepyt

ext/cuda/operators_integral.jl

@@ -17,6 +17,7 @@ function column_reduce_device!(
     input,
     init,
     space,
+    reverse,
 ) where {F, T}
     Ni, Nj, _, _, Nh = size(Fields.field_values(output))
     us = UniversalSize(Fields.field_values(output))
@@ -36,6 +37,7 @@ function column_reduce_device!(
         us,
         mask,
         cart_inds,
+        reverse,
     )
     nitems = Ni * Nj * Nh
     threads = threads_via_occupancy(bycolumn_kernel!, args)
@@ -49,7 +51,7 @@ function column_reduce_device!(
     )
     call_post_op_callback() && post_op_callback(
         output,
-        (dev, f, transform, output, input, init, space),
+        (dev, f, transform, output, input, init, space, reverse),
         (;),
     )
 end
@@ -62,6 +64,7 @@ function column_accumulate_device!(
     input,
     init,
     space,
+    reverse,
 ) where {F, T}
     out_fv = Fields.field_values(output)
     mask = Spaces.get_mask(space)
@@ -81,6 +84,7 @@ function column_accumulate_device!(
         us,
         mask,
         cart_inds,
+        reverse,
     )
     (Ni, Nj, _, _, Nh) = DataLayouts.universal_size(us)
     nitems = Ni * Nj * Nh
@@ -106,9 +110,10 @@ function bycolumn_kernel!(
     us::DataLayouts.UniversalSize,
     mask,
     cart_inds,
+    reverse,
 ) where {S, F, T}
     if space isa Spaces.FiniteDifferenceSpace
-        single_column_function!(f, transform, output, input, init, space)
+        single_column_function!(f, transform, output, input, init, space, reverse)
     else
         tidx = linear_thread_idx()
         if linear_is_valid_index(tidx, us) && tidx ≤ length(unval(cart_inds))
@@ -123,6 +128,7 @@ function bycolumn_kernel!(
                 column(input, i, j, h),
                 init,
                 column(space, i, j, h),
+                reverse,
             )
         end
     end

src/Operators/integrals.jl

@@ -107,7 +107,7 @@ Analogue of `Base._InitialValue` for `column_reduce!` and `column_accumulate!`.
 struct UnspecifiedInit end
 
 """
-    column_reduce!(f, output, input; [init], [transform])
+    column_reduce!(f, output, input; [init], [transform], [reverse])
 
 Applies `reduce` to `input` along the vertical direction, storing the result in
 `output`. The `input` can be either a `Field` or an `AbstractBroadcasted` that
@@ -116,10 +116,12 @@ computed by iteratively applying `f` to the values in `input`, starting from the
 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.
+value at the starting boundary of each column in `input` is used as the initial
+value. By default, reduction starts at the bottom boundary and proceeds upward.
+When `reverse = true`, it starts at the top boundary and proceeds downward.
 
 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:
+of `input`, the default reduction in each column can be summarized as follows:
   - If `init` is unspecified,
     ```
     reduced_value = input[first_level]
@@ -143,10 +145,11 @@ function column_reduce!(
     input::Union{Fields.Field, PointwiseOrColumnwiseBroadcasted};
     init = UnspecifiedInit(),
     transform::T = identity,
+    reverse::Bool = false,
 ) where {F, T}
     device = ClimaComms.device(output)
     space = axes(input)
-    column_reduce_device!(device, f, transform, output, input, init, space)
+    column_reduce_device!(device, f, transform, output, input, init, space, reverse)
 end
 
 function column_reduce_device!(
@@ -157,11 +160,12 @@ function column_reduce_device!(
     input,
     init,
     space,
+    reverse,
 ) 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)
+        single_column_reduce!(f, transform, output, input, init, space, reverse)
     else
         Fields.bycolumn(space) do colidx
             I = Fields.universal_index(colidx)
@@ -173,6 +177,7 @@ function column_reduce_device!(
                     input[colidx],
                     init,
                     space[colidx],
+                    reverse,
                 )
             end
         end
@@ -187,37 +192,45 @@ function single_column_reduce!(
     _input,
     init,
     space,
+    reverse,
 ) where {F, T}
     first_level = left_idx(space)
     last_level = right_idx(space)
+    start_level, stop_level, direction =
+        reverse ? (last_level, first_level, -1) : (first_level, last_level, 1)
     @inbounds if init == UnspecifiedInit()
-        reduced_value = get_level_value(space, _input, first_level)
-        next_level = first_level + 1
+        reduced_value = get_level_value(space, _input, start_level)
+        next_level = start_level + direction
     else
         reduced_value = init
-        next_level = first_level
+        next_level = start_level
     end
-    @inbounds for level in next_level:last_level
+    n_steps = direction * (stop_level - next_level) + 1
+    @inbounds for i in 1:n_steps
+        level = next_level + direction * (i - 1)
         reduced_value = f(reduced_value, get_level_value(space, _input, level))
     end
     Fields.field_values(_output)[] = transform(reduced_value)
     return nothing
 end
 
 """
-    column_accumulate!(f, output, input; [init], [transform])
+    column_accumulate!(f, output, input; [init], [transform], [reverse])
 
 Applies `accumulate` to `input` along the vertical direction, storing the result
 in `output`. The `input` can be either a `Field` or an `AbstractBroadcasted`
-that performs pointwise or columnwise operations on `Field`s. Each accumulated
-value is computed by iteratively applying `f` to the values in `input`, starting
-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`.
+that performs pointwise or columnwise operations on `Field`s. By default, each
+accumulated value is computed by iteratively applying `f` to the values in
+`input`, starting 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.
+When `reverse = true`, accumulation starts at the top boundary and proceeds
+downward, with the corresponding staggered boundary offsets reversed.
 
 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:
+of `input`, the default accumulation in each column can be summarized as follows:
   - For center-to-center and face-to-face accumulation with `init` unspecified,
     ```
     accumulated_value = input[first_level]
@@ -267,14 +280,15 @@ function column_accumulate!(
     input::Union{Fields.Field, PointwiseOrColumnwiseBroadcasted};
     init = UnspecifiedInit(),
     transform::T = identity,
+    reverse::Bool = false,
 ) where {F, T}
     device = ClimaComms.device(output)
     space = axes(input)
     init == UnspecifiedInit() &&
         Spaces.staggering(space) == Spaces.CellCenter() &&
         Spaces.staggering(axes(output)) == Spaces.CellFace() &&
         error("init must be specified for center-to-face accumulation")
-    column_accumulate_device!(device, f, transform, output, input, init, space)
+    column_accumulate_device!(device, f, transform, output, input, init, space, reverse)
 end
 
 function column_accumulate_device!(
@@ -285,11 +299,12 @@ function column_accumulate_device!(
     input,
     init,
     space,
+    reverse,
 ) 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)
+        single_column_accumulate!(f, transform, output, input, init, space, reverse)
     else
         Fields.bycolumn(space) do colidx
             I = Fields.universal_index(colidx)
@@ -301,6 +316,7 @@ function column_accumulate_device!(
                     input[colidx],
                     init,
                     space[colidx],
+                    reverse,
                 )
             end
         end
@@ -315,6 +331,7 @@ function single_column_accumulate!(
     _input,
     init,
     space,
+    reverse,
 ) where {F, T}
     device = ClimaComms.device(space)
     first_level = left_idx(space)
@@ -323,25 +340,32 @@ function single_column_accumulate!(
     is_c2c_or_f2f = Spaces.staggering(space) == Spaces.staggering(axes(output))
     is_c2f = !is_c2c_or_f2f && Spaces.staggering(space) == Spaces.CellCenter()
     is_f2c = !is_c2c_or_f2f && !is_c2f
+    # With `reverse = true`, start at the top boundary, step downward, and
+    # reverse the center/face half-level offset used by staggered accumulation.
+    start_level, stop_level, direction =
+        reverse ? (last_level, first_level, -1) : (first_level, last_level, 1)
+    stagger = reverse ? -half : half
     @inbounds if init == UnspecifiedInit()
         @assert !is_c2f
-        accumulated_value = get_level_value(space, _input, first_level)
-        next_level = first_level + 1
-        init_output_level = is_c2c_or_f2f ? first_level : nothing
+        accumulated_value = get_level_value(space, _input, start_level)
+        next_level = start_level + direction
+        init_output_level = is_c2c_or_f2f ? start_level : nothing
     else
         accumulated_value =
-            is_f2c ? f(init, get_level_value(space, _input, first_level)) : init
-        next_level = is_f2c ? first_level + 1 : first_level
-        init_output_level = is_c2f ? first_level - half : nothing
+            is_f2c ? f(init, get_level_value(space, _input, start_level)) : init
+        next_level = is_f2c ? start_level + direction : start_level
+        init_output_level = is_c2f ? start_level - stagger : nothing
     end
     @inbounds if !isnothing(init_output_level)
         Fields.level(output, init_output_level)[] = transform(accumulated_value)
     end
-    @inbounds for level in next_level:last_level
+    n_steps = direction * (stop_level - next_level) + 1
+    @inbounds for i in 1:n_steps
+        level = next_level + direction * (i - 1)
         accumulated_value =
             f(accumulated_value, get_level_value(space, _input, level))
         output_level =
-            is_c2c_or_f2f ? level : (is_c2f ? level + half : level - half)
+            is_c2c_or_f2f ? level : (is_c2f ? level + stagger : level - stagger)
         Fields.level(output, output_level)[] = transform(accumulated_value)
     end
 end

test/Operators/integrals.jl

@@ -106,9 +106,13 @@ end
 function test_column_reduce_and_accumulate!(center_space)
     face_space = center_to_face_space(center_space)
     ᶜwhole_number = ones(center_space)
+    ᶜwhole_number_reverse = ones(center_space)
     column_accumulate!(+, ᶜwhole_number, ᶜwhole_number) # 1:Nv per column
+    column_accumulate!(+, ᶜwhole_number_reverse, ᶜwhole_number_reverse; reverse = true)
     ᶠwhole_number = ones(face_space)
     column_accumulate!(+, ᶠwhole_number, ᶠwhole_number) # 1:(Nv + 1) per column
+    ᶠwhole_number_reverse = ones(face_space)
+    column_accumulate!(+, ᶠwhole_number_reverse, ᶠwhole_number_reverse; reverse = true) # 1:(Nv + 1) per column
 
     safe_binomial(n, k) = binomial(Int32(n), Int32(k)) # GPU-compatible binomial
 
@@ -125,12 +129,19 @@ function test_column_reduce_and_accumulate!(center_space)
     init = (1, 0) # m₀ = 1, m₋₁ = 0 (m₋₁ can be set to any finite value)
     transform = first # Get mₙ from each (mₙ, mₙ₋₁) pair before saving to output.
 
-    for input in (ᶜwhole_number, ᶠwhole_number)
-        last_input_level = Fields.level(input, Operators.right_idx(axes(input)))
-        output = similar(last_input_level)
-        reference_output = motzkin_number.(last_input_level)
+    for (input, reverse) in (
+        (ᶜwhole_number, false),
+        (ᶠwhole_number, false),
+        (ᶜwhole_number_reverse, true),
+        (ᶠwhole_number_reverse, true),
+    )
+        final_input_idx =
+            reverse ? Operators.left_idx(axes(input)) : Operators.right_idx(axes(input))
+        final_input_level = Fields.level(input, final_input_idx)
+        output = similar(final_input_level)
+        reference_output = motzkin_number.(final_input_level)
 
-        set_output! = () -> column_reduce!(f, output, input; init, transform)
+        set_output! = () -> column_reduce!(f, output, input; init, transform, reverse)
         set_output!()
         @test output == reference_output
         @test_opt ignored_modules = CUDA_FRAMES set_output!()
@@ -139,19 +150,26 @@ function test_column_reduce_and_accumulate!(center_space)
 
     ᶜoutput = similar(ᶜwhole_number)
     ᶠoutput = similar(ᶠwhole_number)
-    for (input, output, reference_output) in (
-        (ᶜwhole_number, ᶜoutput, motzkin_number.(ᶜwhole_number)),
-        (ᶠwhole_number, ᶠoutput, motzkin_number.(ᶠwhole_number)),
-        (ᶠwhole_number, ᶜoutput, motzkin_number.(ᶜwhole_number .+ 1)),
-        (ᶜwhole_number, ᶠoutput, motzkin_number.(ᶠwhole_number .- 1)),
+    for (input, output, reference_output, reverse) in (
+        (ᶜwhole_number, ᶜoutput, motzkin_number.(ᶜwhole_number), false),
+        (ᶠwhole_number, ᶠoutput, motzkin_number.(ᶠwhole_number), false),
+        (ᶠwhole_number, ᶜoutput, motzkin_number.(ᶜwhole_number .+ 1), false),
+        (ᶜwhole_number, ᶠoutput, motzkin_number.(ᶠwhole_number .- 1), false),
+        (ᶜwhole_number_reverse, ᶜoutput, motzkin_number.(ᶜwhole_number_reverse), true),
+        (ᶠwhole_number_reverse, ᶠoutput, motzkin_number.(ᶠwhole_number_reverse), true),
+        (ᶠwhole_number_reverse, ᶜoutput, motzkin_number.(ᶜwhole_number_reverse .+ 1), true),
+        (ᶜwhole_number_reverse, ᶠoutput, motzkin_number.(ᶠwhole_number_reverse .- 1), true),
     )
+
         set_output! =
-            () -> column_accumulate!(f, output, input; init, transform)
+            () -> column_accumulate!(f, output, input; init, transform, reverse)
         set_output!()
         @test output == reference_output
         @test_opt ignored_modules = CUDA_FRAMES set_output!()
         test_allocs(@allocated set_output!())
     end
+
+
 end
 
 function test_fubinis_theorem(space)