From 695823e22cfa7971b2c397879c8020016f8f2d8d Mon Sep 17 00:00:00 2001 From: Hung-Chih Wu Date: Mon, 29 Jun 2026 10:28:42 -0700 Subject: [PATCH 1/6] add option to reverse column_accumulate --- ext/cuda/operators_integral.jl | 6 +++++- src/Operators/integrals.jl | 15 ++++++++++++--- 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/ext/cuda/operators_integral.jl b/ext/cuda/operators_integral.jl index e6a7ec8c8a..e13a4d85f9 100644 --- a/ext/cuda/operators_integral.jl +++ b/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 diff --git a/src/Operators/integrals.jl b/src/Operators/integrals.jl index 20faaa8ed9..7f650d2176 100644 --- a/src/Operators/integrals.jl +++ b/src/Operators/integrals.jl @@ -267,6 +267,7 @@ 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) @@ -274,7 +275,7 @@ function column_accumulate!( 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 = From a9b4c2615390cd73e4e32ed4db8f8c0697cce139 Mon Sep 17 00:00:00 2001 From: Hung-Chih Wu Date: Tue, 30 Jun 2026 11:01:21 -0700 Subject: [PATCH 2/6] add test for reverse column_accumulate --- test/Operators/integrals.jl | 24 +++++++++++++++++------- 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/test/Operators/integrals.jl b/test/Operators/integrals.jl index 2d5fda91be..fb62eb55b6 100644 --- a/test/Operators/integrals.jl +++ b/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) From 15a620945aab73b4957b5676f7e9c64e853ac4a9 Mon Sep 17 00:00:00 2001 From: Hung-Chih Wu Date: Wed, 15 Jul 2026 09:50:26 -0700 Subject: [PATCH 3/6] Apply JuliaFormatter --- ext/cuda/operators_integral.jl | 4 ++-- src/Operators/integrals.jl | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/ext/cuda/operators_integral.jl b/ext/cuda/operators_integral.jl index e13a4d85f9..ad7c87f21d 100644 --- a/ext/cuda/operators_integral.jl +++ b/ext/cuda/operators_integral.jl @@ -62,7 +62,7 @@ function column_accumulate_device!( input, init, space, - reverse + reverse, ) where {F, T} out_fv = Fields.field_values(output) mask = Spaces.get_mask(space) @@ -126,7 +126,7 @@ function bycolumn_kernel!( column(input, i, j, h), init, column(space, i, j, h), - reverse + reverse, ) end end diff --git a/src/Operators/integrals.jl b/src/Operators/integrals.jl index 7f650d2176..79b6ba0807 100644 --- a/src/Operators/integrals.jl +++ b/src/Operators/integrals.jl @@ -286,7 +286,7 @@ function column_accumulate_device!( input, init, space, - reverse + reverse, ) where {F, T} mask = Spaces.get_mask(space) if space isa Spaces.FiniteDifferenceSpace @@ -303,7 +303,7 @@ function column_accumulate_device!( input[colidx], init, space[colidx], - reverse + reverse, ) end end From 06e07611635141cb917337f8132fc2d79a120124 Mon Sep 17 00:00:00 2001 From: Hung-Chih Wu Date: Wed, 15 Jul 2026 16:31:24 -0700 Subject: [PATCH 4/6] Address reviewer comments. --- src/Operators/integrals.jl | 30 +++++++++++++++++------------- test/Operators/integrals.jl | 1 + 2 files changed, 18 insertions(+), 13 deletions(-) diff --git a/src/Operators/integrals.jl b/src/Operators/integrals.jl index 79b6ba0807..5dc3561fa5 100644 --- a/src/Operators/integrals.jl +++ b/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: @@ -327,30 +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 - indices = if reverse - last_level:-1:next_level - else - next_level:last_level - end - @inbounds for level in indices + 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 diff --git a/test/Operators/integrals.jl b/test/Operators/integrals.jl index fb62eb55b6..2c46898a30 100644 --- a/test/Operators/integrals.jl +++ b/test/Operators/integrals.jl @@ -110,6 +110,7 @@ function test_column_reduce_and_accumulate!(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 From 979f527f98a6368356d8fc157d02e2419849a93d Mon Sep 17 00:00:00 2001 From: Hung-Chih Wu Date: Fri, 17 Jul 2026 14:19:43 -0700 Subject: [PATCH 5/6] Add reverse option to column_reduce. --- ext/cuda/operators_integral.jl | 4 +++- src/Operators/integrals.jl | 28 +++++++++++++++++++--------- test/Operators/integrals.jl | 17 ++++++++++++----- 3 files changed, 34 insertions(+), 15 deletions(-) diff --git a/ext/cuda/operators_integral.jl b/ext/cuda/operators_integral.jl index ad7c87f21d..5308b69193 100644 --- a/ext/cuda/operators_integral.jl +++ b/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 diff --git a/src/Operators/integrals.jl b/src/Operators/integrals.jl index 5dc3561fa5..deb5433e37 100644 --- a/src/Operators/integrals.jl +++ b/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,17 +192,22 @@ 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) diff --git a/test/Operators/integrals.jl b/test/Operators/integrals.jl index 2c46898a30..40ff409498 100644 --- a/test/Operators/integrals.jl +++ b/test/Operators/integrals.jl @@ -129,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!() From 7ee6dc2b984eb97cb1ee707267ceefcc264738cd Mon Sep 17 00:00:00 2001 From: Hung-Chih Wu Date: Fri, 17 Jul 2026 14:25:56 -0700 Subject: [PATCH 6/6] Update column_accumulate docstring. --- src/Operators/integrals.jl | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/src/Operators/integrals.jl b/src/Operators/integrals.jl index deb5433e37..655e0a2c5f 100644 --- a/src/Operators/integrals.jl +++ b/src/Operators/integrals.jl @@ -215,21 +215,22 @@ function single_column_reduce!( 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]