Wu/reverse accumulate

ext/cuda/operators_integral.jl

@@ -62,6 +62,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 +82,7 @@ function column_accumulate_device!(
         us,
         mask,
         cart_inds,
+        reverse,
     )
     (Ni, Nj, _, _, Nh) = DataLayouts.universal_size(us)
     nitems = Ni * Nj * Nh
@@ -106,9 +108,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 +126,7 @@ function bycolumn_kernel!(
                 column(input, i, j, h),
                 init,
                 column(space, i, j, h),
+                reverse,
             )
         end
     end

src/Operators/integrals.jl

@@ -215,6 +215,8 @@ 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:
@@ -267,14 +269,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 +288,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 +305,7 @@ function column_accumulate_device!(
                     input[colidx],
                     init,
                     space[colidx],
+                    reverse,
                 )
             end
         end
@@ -315,6 +320,7 @@ function single_column_accumulate!(
     _input,
     init,
     space,
+    reverse,
 ) where {F, T}
     device = ClimaComms.device(space)
     first_level = left_idx(space)
@@ -323,25 +329,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
 
@@ -139,19 +143,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)