Remove vertical chunking from aux vars and tendency terms by mwarusz · Pull Request #473 · E3SM-Project/Omega

mwarusz · 2026-07-20T16:05:45Z

This PR removes vertical chunking from all auxiliary variables and tendency terms computations. Vertical chunking was originally introduced to help with vectorization on CPUs. However, in its current form, it doesn't seem to improve CPU performance, and often makes things slower. Removing the chunking makes the code cleaner and makes it much simpler to limit vertical bounds to active layers.

I still need to finish running a couple of benchmarks to prove that this is beneficial, but I am opening this PR to show examples of how to write operators without vertical chunking.

This is just the first from a planned series of PRs to remove vertical chunking from all of Omega.

Checklist

mwarusz · 2026-07-22T19:54:47Z

Benchmarks

Single node runs using the baroclinic channel polaris test case on a 2 km mesh. The test case setup was modified to use 80 vertical layers and turn on most of the tendency terms. The following tables shows the Omega:RunLoop timer after integrating for 500 time steps. For CPU benchmarks, Omega was compiled with OMEGA_VECTOR_LENGTH 8 or 16 (on aurora).

CPU runs

Machine (compiler)	`develop`	this PR
pm-cpu (intel)	76.126 s	72.012 s
pm-cpu (gnu*)	81.295 s	74.967 s
frontier (craygnu)	105.795 s	97.231 s
aurora (oneapi-ifx*)	89.385 s	71.553 s

* I had to modify cime compiler flags to enable vectorization with these compilers

GPU runs

Machine (compiler)	`develop`	this PR
pm-gpu (gnugpu)	5.903 s	5.807 s
frontier (craygnu-mphipcc)	7.514 s	6.798 s
aurora (oneapi-ifxgpu)	5.182 s	6.328 s

In addition to these benchmarks, on aurora CPU I measured the fraction of vectorized instructions using the aps tool from Intel with the following result:

develop: 38.80%
this PR: 51.60%

Summary

Except for aurora GPU, everything is faster without vertical chunking.

mwarusz · 2026-07-23T19:18:47Z

Testing

CTest unit tests

Machine: aurora / frontier / frontier
Compiler: oneapi-ifx / craygnu / craygnu-mphipcc
Build type: Release
Result: All tests passed (on all machines)

Polaris `omega_pr` suite

Baseline workdir: /lus/flare/projects/E3SM_Dec/mwaruszewski/omega-pr-testing/remove-chunking-tend-aux/baseline-ifx/
Baseline build: /lus/flare/projects/E3SM_Dec/mwaruszewski/omega-pr-testing/remove-chunking-tend-aux/baseline-ifx/build
PR build: /lus/flare/projects/E3SM_Dec/mwaruszewski/omega-pr-testing/remove-chunking-tend-aux/build-pr-ifx
PR workdir: /lus/flare/projects/E3SM_Dec/mwaruszewski/omega-pr-testing/remove-chunking-tend-aux/pr-ifx
Machine: aurora
Compiler: oneapi-ifx
Build type: Release
Log: not found
Result: All tests passed

Polaris `omega_pr` suite

Baseline workdir: /lustre/orion/cli115/scratch/mwaruszewski/omega-pr-testing/remove-chunking-tend-aux/baseline-craygnu-mphipcc/
Baseline build: /lustre/orion/cli115/scratch/mwaruszewski/omega-pr-testing/remove-chunking-tend-aux/baseline-craygnu-mphipcc/build
PR build: /lustre/orion/cli115/scratch/mwaruszewski/omega-pr-testing/remove-chunking-tend-aux/build-pr-craygnu-mphipcc
PR workdir: /lustre/orion/cli115/scratch/mwaruszewski/omega-pr-testing/remove-chunking-tend-aux/pr-craygnu-mphipcc
Machine: frontier
Partition: batch
Compiler: craygnu-mphipcc
Build type: Release
Log: /lustre/orion/cli115/scratch/mwaruszewski/omega-pr-testing/remove-chunking-tend-aux/pr-craygnu-mphipcc/polaris_omega_pr.o5060812
Result: All tests passed

Copilot

Pull request overview

This PR removes vertical chunking from Omega ocean auxiliary-variable and
tendency-term computations, switching operators to team-parallel (outer) +
layer-range (inner) patterns that better align with active-layer bounds.

Changes:

Refactor ocean tendency-term operators to accept TeamMember and iterate over
active-layer Range{KMin, KMax} instead of chunk indices.
Refactor auxiliary-variable computations (vorticity, del2, tracer, pseudo
thickness, kinetic aux) to the same non-chunked, active-layer approach.
Update unit tests and add subviewUnmanaged to reduce overhead in hot paths.

Reviewed changes

Copilot reviewed 17 out of 17 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
components/omega/test/ocn/TendencyTermsTest.cpp	Updates test kernels to call the new team-based tendency operators.
components/omega/test/ocn/AuxiliaryVarsTest.cpp	Updates test kernels to call the new team-based auxiliary-variable operators.
components/omega/src/ocn/TendencyTerms.h	Removes chunked operators and rewrites tendency terms using active-layer ranges and team scratch.
components/omega/src/ocn/TendencyTerms.cpp	Wires new members (e.g., `NVertLayers`, min/max layer arrays) into tendency-term constructors.
components/omega/src/ocn/Tendencies.cpp	Updates tendency assembly to launch non-chunked team kernels (often with per-team scratch).
components/omega/src/ocn/auxiliaryVars/VorticityAuxVars.h	Refactors vorticity aux computations to team + active-layer ranges.
components/omega/src/ocn/auxiliaryVars/VorticityAuxVars.cpp	Adds `NVertLayers` initialization needed for scratch-sized temporaries.
components/omega/src/ocn/auxiliaryVars/VelocityDel2AuxVars.h	Refactors velocity del2 aux computations to team + active-layer ranges with scratch.
components/omega/src/ocn/auxiliaryVars/VelocityDel2AuxVars.cpp	Adds `NVertLayers` initialization for scratch allocation sizing.
components/omega/src/ocn/auxiliaryVars/TracerAuxVars.h	Refactors tracer aux computations to team + active-layer ranges with scratch and unmanaged subviews.
components/omega/src/ocn/auxiliaryVars/TracerAuxVars.cpp	Adds `NVertLayers` initialization needed by refactored tracer aux kernels.
components/omega/src/ocn/auxiliaryVars/PseudoThicknessAuxVars.h	Refactors pseudo-thickness aux computations to team + active-layer ranges with scratch.
components/omega/src/ocn/auxiliaryVars/PseudoThicknessAuxVars.cpp	Adds `NVertLayers` initialization for scratch-sized temporaries.
components/omega/src/ocn/auxiliaryVars/KineticAuxVars.h	Refactors kinetic aux computations to team + active-layer ranges with scratch.
components/omega/src/ocn/auxiliaryVars/KineticAuxVars.cpp	Adds `NVertLayers` initialization for scratch-sized temporaries.
components/omega/src/ocn/AuxiliaryState.cpp	Updates auxiliary-state assembly to call new team-based aux kernels.
components/omega/src/infra/OmegaKokkos.h	Adds `subviewUnmanaged` helper used by refactored kernels.

Copilot

Pull request overview

Copilot reviewed 17 out of 17 changed files in this pull request and generated 2 comments.

Comments suppressed due to low confidence (2)

components/omega/src/ocn/AuxiliaryState.cpp:152

The LaunchConfig scratch sizing uses TeamScratch(2 * VCoord->NVertLayers), but LocKineticAux::computeVarsOnCell allocates two ScratchArray1DReal views (KineticEnergyCellTmp and VelocityDivCellTmp). Use a 2-view TeamScratch calculation to ensure enough scratch bytes (per-view padding/alignment can make this larger than a single 2*N allocation).

       LaunchConfig({Mesh->NCellsAll},
                    TeamScratch<Real>(2 * VCoord->NVertLayers)),
       KOKKOS_LAMBDA(int ICell, const TeamMember &Team) {

components/omega/test/ocn/AuxiliaryVarsTest.cpp:468

This kernel launches with TeamScratch(2 * VCoord->NVertLayers), but VorticityAuxVars::computeVarsOnVertex allocates two separate scratch views (PseudoThickVertex and RelVortVertexTmp). Use TeamScratch<Real, Real>(VCoord->NVertLayers, VCoord->NVertLayers) so the scratch allocation matches the sum of both views.

   parallelForOuter(
       LaunchConfig({Decomp->NVerticesHaloH(0)},
                    TeamScratch<Real>(2 * VCoord->NVertLayers)),
       KOKKOS_LAMBDA(int IVertex, const TeamMember &Team) {
          VorticityAux.computeVarsOnVertex(Team, IVertex, PseudoThickCell,
                                           NormalVelEdge);
       });

sbrus89 · 2026-08-12T17:59:47Z

@brian-oneill, can you take a look at this when you get a chance? Thanks!

Vertical chunking is on its way out of Omega (PR #473 removes it from the auxiliary variables and tendency terms), so this branch should not add more of it. The EOS is not touched by that PR, so remove chunking here across the board rather than leaving the new specific volume derivative code to be converted later. The functors now follow the same pattern as the de-chunked auxiliary variables: each takes a TeamMember and a cell index and loops over the active layers with parallelForInner(Team, Range{KMin, KMax}, INNER_LAMBDA(int K) { ... }); so the callers in Eos.cpp reduce to a single call inside parallelForOuter, with no vertRangeChunked and no inner loop over chunks. Dropping the chunk loop also removes VecLength from the TEOS-10 polynomial helpers. calcPCoeffs, calcPCoeffsDTt, calcPCoeffsDSs, calcDelta, calcDeltaDeriv and calcDeltaDP took arrays sized 6 * VecLength or 5 * VecLength together with a KVec index into them; they now take plain [6] and [5] arrays and no index. That storage was never shared between layers -- the coefficients are recomputed for every layer -- so nothing is lost, and the point-wise calcSpecVolAndDerivsAtPoint and the calcAlpha and calcBeta helpers no longer allocate VecLength times more stack than they use. With the chunk loop gone, the array-level derivative routine is exactly calcSpecVolAndDerivsAtPoint evaluated at each cell and layer, so it now calls it instead of repeating the polynomial evaluation. That leaves a single implementation of the TEOS-10 derivatives. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Vertical chunking is on its way out of Omega (PR E3SM-Project#473 removes it from the auxiliary variables and tendency terms), so this branch should not add more of it. The EOS is not touched by that PR, so remove chunking here across the board rather than leaving the new specific volume derivative code to be converted later. The functors now follow the same pattern as the de-chunked auxiliary variables: each takes a TeamMember and a cell index and loops over the active layers with parallelForInner(Team, Range{KMin, KMax}, INNER_LAMBDA(int K) { ... }); so the callers in Eos.cpp reduce to a single call inside parallelForOuter, with no vertRangeChunked and no inner loop over chunks. Dropping the chunk loop also removes VecLength from the TEOS-10 polynomial helpers. calcPCoeffs, calcPCoeffsDTt, calcPCoeffsDSs, calcDelta, calcDeltaDeriv and calcDeltaDP took arrays sized 6 * VecLength or 5 * VecLength together with a KVec index into them; they now take plain [6] and [5] arrays and no index. That storage was never shared between layers -- the coefficients are recomputed for every layer -- so nothing is lost, and the point-wise calcSpecVolAndDerivsAtPoint and the calcAlpha and calcBeta helpers no longer allocate VecLength times more stack than they use. With the chunk loop gone, the array-level derivative routine is exactly calcSpecVolAndDerivsAtPoint evaluated at each cell and layer, so it now calls it instead of repeating the polynomial evaluation. That leaves a single implementation of the TEOS-10 derivatives. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

sbrus89 · 2026-08-25T20:37:04Z

@mwarusz, when building on pm-gpu (gnugpu) I see the following errors:

/global/cfs/cdirs/e3sm/sbrus/polaris_worktrees/testing/e3sm_submodules/Omega/components/omega/src/ocn/auxiliaryVars/TracerAuxVars.h(27): error: identifier "Kokkos::ALL" is undefined in device code

/global/cfs/cdirs/e3sm/sbrus/polaris_worktrees/testing/e3sm_submodules/Omega/components/omega/src/ocn/auxiliaryVars/TracerAuxVars.h(29): error: identifier "Kokkos::ALL" is undefined in device code

/global/cfs/cdirs/e3sm/sbrus/polaris_worktrees/testing/e3sm_submodules/Omega/components/omega/src/ocn/auxiliaryVars/TracerAuxVars.h(27): error: identifier "Kokkos::ALL" is undefined in device code

/global/cfs/cdirs/e3sm/sbrus/polaris_worktrees/testing/e3sm_submodules/Omega/components/omega/src/ocn/auxiliaryVars/TracerAuxVars.h(29): error: identifier "Kokkos::ALL" is undefined in device code

mwarusz · 2026-08-25T22:13:09Z

@sbrus89

I pushed a fix and rebased on develop. With these changes I was able to build on pm-gpu.

sbrus89 · 2026-08-26T20:02:37Z

Passes CTests and omega_pr on pm-cpu (gnu) and pm-gpu (gnugpu):

Polaris `omega_pr` suite

Baseline workdir: /global/homes/s/sbrus/scratch/polaris_baseline_omega_pr_omega_gnugpu_mpich/
Baseline build: /global/homes/s/sbrus/scratch/polaris_baseline_omega_pr_omega_gnugpu_mpich/build
PR build: /global/homes/s/sbrus/scratch/polaris_remove_chunking_omega_pr_omega_gnugpu_mpich/build
PR workdir: /global/homes/s/sbrus/scratch/polaris_remove_chunking_omega_pr_omega_gnugpu_mpich
Machine: pm-gpu
Compiler: gnugpu
Build type: Release
Log: /global/homes/s/sbrus/scratch/polaris_remove_chunking_omega_pr_omega_gnugpu_mpich/polaris_omega_pr.o57627707
Result: All tests passed

Polaris `omega_pr` suite

Baseline workdir: /global/homes/s/sbrus/scratch/polaris_baseline_omega_pr_omega_gnu_mpich/
Baseline build: /global/homes/s/sbrus/scratch/polaris_baseline_omega_pr_omega_gnu_mpich/build
PR build: /global/homes/s/sbrus/scratch/polaris_remove_chunking_omega_pr_omega_gnu_mpich/build
PR workdir: /global/homes/s/sbrus/scratch/polaris_remove_chunking_omega_pr_omega_gnu_mpich
Machine: pm-cpu
Compiler: gnu
Build type: Release
Log: not found
Result: All tests passed

sbrus89 · 2026-08-26T22:12:50Z

Passes omega_pr on Frontier with craygnu-mphipcc as noted above, however there are diffs with craygnu. This is similar to what happened for #485, so we'll need a polaris submodule update to bless the diffs.

Polaris `omega_pr` suite

Baseline workdir: /ccs/home/brus/run/polaris_baseline_omega_pr_omega_craygnu-mphipcc_mpich/
Baseline build: /ccs/home/brus/run/polaris_baseline_omega_pr_omega_craygnu-mphipcc_mpich/build
PR build: /ccs/home/brus/run/polaris_remove_chunking_omega_pr_omega_craygnu-mphipcc_mpich/build
PR workdir: /ccs/home/brus/run/polaris_remove_chunking_omega_pr_omega_craygnu-mphipcc_mpich
Machine: frontier
Partition: batch
Compiler: craygnu-mphipcc
Build type: Release
Log: not found
Result: All tests passed

Polaris `omega_pr` suite

Baseline workdir: /ccs/home/brus/run/polaris_baseline_omega_pr_omega_craygnu_mpich/
Baseline build: /ccs/home/brus/run/polaris_baseline_omega_pr_omega_craygnu_mpich/build
PR build: /ccs/home/brus/run/polaris_remove_chuncking_omega_pr_omega_craygnu_mpich/build
PR workdir: /ccs/home/brus/run/polaris_remove_chuncking_omega_pr_omega_craygnu_mpich
Machine: frontier
Partition: batch
Compiler: craygnu
Build type: Release
Log: not found
Result:
- Diffs (19 of 22):
  - ocean/column/vmix_stable
  - ocean/planar/baroclinic_channel/10km/decomp
  - ocean/planar/baroclinic_channel/10km/restart
  - ocean/planar/baroclinic_channel/10km/threads
  - ocean/planar/barotropic_gyre/munk/free-slip
  - ocean/planar/manufactured_solution/convergence_both/default
  - ocean/planar/manufactured_solution/convergence_both/del2
  - ocean/planar/manufactured_solution/convergence_both/del4
  - ocean/planar/merry_go_round/default
  - ocean/planar/overflow/linear/zstar/smoke_test_horiz_adv_order_2_del4
  - ocean/planar/overflow/linear/zstar/smoke_test_horiz_adv_order_3
  - ocean/planar/overflow/linear/zstar/smoke_test_horiz_adv_order_4
  - ocean/planar/overflow/nonlinear/pstar/smoke_test_horiz_adv_order_4_del4
  - ocean/planar/seamount/linear/zstar/short
  - ocean/planar/seamount/nonlinear/sigma/short
  - ocean/spherical/icos/cosine_bell/decomp
  - ocean/spherical/icos/cosine_bell/restart
  - ocean/column/ekman
  - ocean/column/inertial

Sample diffs from ocean/planar/baroclinic_channel/10km/decomp:

temperature          Time index: 0
0:  l1: 2.21689333557151e-12  l2: 8.45217389650821e-14  linf: 1.42108547152020e-14
  ^[[91mFAIL^[[0m /ccs/home/brus/run/polaris_remove_chuncking_omega_pr_omega_craygnu_mpich/ocean/planar/baroclinic_channel/10km/decomp/4proc/output.nc
       /ccs/home/brus/run/polaris_baseline_omega_pr_omega_craygnu_mpich/ocean/planar/baroclinic_channel/10km/decomp/4proc/output.nc
salinity             Time index: 0
0:  l1: 7.21911419532262e-12  l2: 2.87396750538799e-13  linf: 2.84217094304040e-14
  ^[[91mFAIL^[[0m /ccs/home/brus/run/polaris_remove_chuncking_omega_pr_omega_craygnu_mpich/ocean/planar/baroclinic_channel/10km/decomp/4proc/output.nc
       /ccs/home/brus/run/polaris_baseline_omega_pr_omega_craygnu_mpich/ocean/planar/baroclinic_channel/10km/decomp/4proc/output.nc
normalVelocity       Time index: 0
0:  l1: 9.65391540298509e-12  l2: 1.55866843995814e-13  linf: 1.13492115011438e-14
  ^[[91mFAIL^[[0m /ccs/home/brus/run/polaris_remove_chuncking_omega_pr_omega_craygnu_mpich/ocean/planar/baroclinic_channel/10km/decomp/4proc/output.nc
       /ccs/home/brus/run/polaris_baseline_omega_pr_omega_craygnu_mpich/ocean/planar/baroclinic_channel/10km/decomp/4proc/output.nc
          baseline comp.:   ^[[91mFAIL^[[0m
          runtime:          ^[[94m0:00:01^[[0m

mwarusz force-pushed the omega/remove-chunking branch from ae5ec57 to 5375e80 Compare July 21, 2026 15:38

mwarusz force-pushed the omega/remove-chunking branch 2 times, most recently from 8f6b9d3 to 031a705 Compare July 23, 2026 19:13

mwarusz requested a review from Copilot July 23, 2026 19:20

Copilot started reviewing on behalf of mwarusz July 23, 2026 19:20 View session

Copilot AI reviewed Jul 23, 2026

View reviewed changes

Comment thread components/omega/src/ocn/AuxiliaryState.cpp

Comment thread components/omega/src/ocn/TendencyTerms.h Outdated

Comment thread components/omega/src/ocn/TendencyTerms.h Outdated

mwarusz force-pushed the omega/remove-chunking branch from 031a705 to af042eb Compare July 24, 2026 22:02

mwarusz requested a review from Copilot July 24, 2026 22:04

Copilot started reviewing on behalf of mwarusz July 24, 2026 22:05 View session

Copilot AI reviewed Jul 24, 2026

View reviewed changes

Comment thread components/omega/src/ocn/AuxiliaryState.cpp

Comment thread components/omega/test/ocn/AuxiliaryVarsTest.cpp

mwarusz requested a review from brian-oneill July 27, 2026 15:42

mwarusz force-pushed the omega/remove-chunking branch 2 times, most recently from bcbf20e to 90f0f6f Compare August 18, 2026 16:41

xylar mentioned this pull request Aug 25, 2026

Add velocity-component reconstruction for I/O #525

Draft

26 tasks

sbrus89 reviewed Aug 25, 2026

View reviewed changes

Comment thread components/omega/src/ocn/auxiliaryVars/TracerAuxVars.h

mwarusz added 9 commits August 25, 2026 16:07

Use nonunit diffusivities in tendency terms test

2e3dc04

Remove vert chunking from most tendency terms

53336f4

Remove vert chunking from aux vars

40ad947

Add and use subviewUnmanaged

a6e6ced

Fix double mult by EddyDiff2

e56ac36

Fix vertical bounds in tracer flux divergence

64a93e2

Allocate scratch for pseudo-thickness cell aux

0409bdc

Improve LinearIdxWrapper to better hide implementation details

939123d

Update kinetic energy computation in vert mixing

5a36da2

Pass by value in subviewUnmanaged

00f7a37

mwarusz force-pushed the omega/remove-chunking branch from 90f0f6f to 00f7a37 Compare August 25, 2026 22:10

sbrus89 merged commit c72e3cf into E3SM-Project:develop Aug 27, 2026
1 check passed

sbrus89 mentioned this pull request Aug 27, 2026

Update Omega submodule E3SM-Project/polaris#725

Merged

15 tasks

Navigation Menu

Uh oh!

Remove vertical chunking from aux vars and tendency terms - #473

Remove vertical chunking from aux vars and tendency terms#473
sbrus89 merged 10 commits into
E3SM-Project:developfrom
mwarusz:omega/remove-chunking

mwarusz commented Jul 20, 2026 •

edited by sbrus89

Loading

Uh oh!

mwarusz commented Jul 22, 2026

Uh oh!

mwarusz commented Jul 23, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

sbrus89 commented Aug 12, 2026

Uh oh!

sbrus89 commented Aug 25, 2026

Uh oh!

Uh oh!

mwarusz commented Aug 25, 2026

Uh oh!

sbrus89 commented Aug 26, 2026

Uh oh!

sbrus89 commented Aug 26, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

mwarusz commented Jul 20, 2026 • edited by sbrus89 Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mwarusz commented Jul 22, 2026

Benchmarks

CPU runs

GPU runs

Summary

Uh oh!

mwarusz commented Jul 23, 2026

Testing

CTest unit tests

Polaris omega_pr suite

Polaris omega_pr suite

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

sbrus89 commented Aug 12, 2026

Uh oh!

sbrus89 commented Aug 25, 2026

Uh oh!

Uh oh!

mwarusz commented Aug 25, 2026

Uh oh!

sbrus89 commented Aug 26, 2026

Polaris omega_pr suite

Polaris omega_pr suite

Uh oh!

sbrus89 commented Aug 26, 2026

Polaris omega_pr suite

Polaris omega_pr suite

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

mwarusz commented Jul 20, 2026 •

edited by sbrus89

Loading

Polaris `omega_pr` suite

Polaris `omega_pr` suite

Polaris `omega_pr` suite

Polaris `omega_pr` suite

Polaris `omega_pr` suite

Polaris `omega_pr` suite