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

@@ -267,14 +267,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 +286,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 +303,7 @@ function column_accumulate_device!(
                     input[colidx],
                     init,
                     space[colidx],
+                    reverse,
                 )
             end
         end
@@ -315,6 +318,7 @@ function single_column_accumulate!(
     _input,
     init,
     space,
+    reverse,
 ) where {F, T}
     device = ClimaComms.device(space)
     first_level = left_idx(space)
@@ -337,7 +341,12 @@ function single_column_accumulate!(
     @inbounds if !isnothing(init_output_level)
         Fields.level(output, init_output_level)[] = transform(accumulated_value)
     end
-    @inbounds for level in next_level:last_level
+    indices = if reverse
+        last_level:-1:next_level
+    else
+        next_level:last_level
+    end
+    @inbounds for level in indices
         accumulated_value =
             f(accumulated_value, get_level_value(space, _input, level))
         output_level =

test/Operators/integrals.jl

@@ -106,9 +106,12 @@ 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 +142,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)