Skip to content

Remove vertical chunking from aux vars and tendency terms #473

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sbrus89 merged 10 commits into from
Aug 27, 2026
Merged

Remove vertical chunking from aux vars and tendency terms #473

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
2e3dc04
Use nonunit diffusivities in tendency terms test
mwarusz Jul 23, 2026
53336f4
Remove vert chunking from most tendency terms
mwarusz May 23, 2026
40ad947
Remove vert chunking from aux vars
mwarusz May 23, 2026
a6e6ced
Add and use subviewUnmanaged
mwarusz Jul 22, 2026
e56ac36
Fix double mult by EddyDiff2
mwarusz Jul 23, 2026
64a93e2
Fix vertical bounds in tracer flux divergence
mwarusz Jul 23, 2026
0409bdc
Allocate scratch for pseudo-thickness cell aux
mwarusz Jul 24, 2026
939123d
Improve LinearIdxWrapper to better hide implementation details
mwarusz Jul 24, 2026
5a36da2
Update kinetic energy computation in vert mixing
mwarusz Aug 18, 2026
00f7a37
Pass by value in subviewUnmanaged
mwarusz Aug 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 21 additions & 6 deletions components/omega/src/infra/OmegaKokkos.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -202,12 +202,20 @@ bool arraysEqual(const ArrayTypeA &A, const ArrayTypeB &B) {
return Equal;
}

// Works the same as Kokkos::subview, but returns an unmanaged view
// This is safe when we know that the subview won't outlive the input view
// and can be slightly faster
template <class View, class... Args>
KOKKOS_FUNCTION auto subviewUnmanaged(const View &A, Args... args) {
return Kokkos::subview(View(A.data(), A.layout()), args...);
}

// Takes a functor that uses multidimensional indexing
// and converts it into one that also accepts linear index
template <class F, int Rank> struct LinearIdxWrapper : F {
template <class F, int Rank> class LinearIdxWrapper : private F {
static_assert(Rank >= 1 && Rank <= 5, "LinearIdxWrapper supports ranks 1-5");
using F::operator();

public:
template <class Array>
LinearIdxWrapper(F &&Functor, Array &&Bounds) : F(std::move(Functor)) {
computeStrides(std::forward<Array>(Bounds));
Expand All @@ -227,13 +235,19 @@ template <class F, int Rank> struct LinearIdxWrapper : F {
}
}

template <int N = Rank, class... Args>
KOKKOS_FORCEINLINE_FUNCTION std::enable_if_t<N == 1>
operator()(int Idx, Args &&...OtherArgs) const {
F::operator()(Idx, std::forward<Args>(OtherArgs)...);
}

template <int N = Rank, class... Args>
KOKKOS_FORCEINLINE_FUNCTION std::enable_if_t<N == 2>
operator()(int Idx, Args &&...OtherArgs) const {
const int I1 = Idx / Strides[0];
const int I2 = Idx - I1 * Strides[0];

(*this)(I1, I2, std::forward<Args>(OtherArgs)...);
F::operator()(I1, I2, std::forward<Args>(OtherArgs)...);
}

template <int N = Rank, class... Args>
Expand All @@ -244,7 +258,7 @@ template <class F, int Rank> struct LinearIdxWrapper : F {
const int I2 = Idx / Strides[1];
const int I3 = Idx - I2 * Strides[1];

(*this)(I1, I2, I3, std::forward<Args>(OtherArgs)...);
F::operator()(I1, I2, I3, std::forward<Args>(OtherArgs)...);
}

template <int N = Rank, class... Args>
Expand All @@ -257,7 +271,7 @@ template <class F, int Rank> struct LinearIdxWrapper : F {
const int I3 = Idx / Strides[2];
const int I4 = Idx - I3 * Strides[2];

(*this)(I1, I2, I3, I4, std::forward<Args>(OtherArgs)...);
F::operator()(I1, I2, I3, I4, std::forward<Args>(OtherArgs)...);
}

template <int N = Rank, class... Args>
Expand All @@ -272,9 +286,10 @@ template <class F, int Rank> struct LinearIdxWrapper : F {
const int I4 = Idx / Strides[3];
const int I5 = Idx - I4 * Strides[3];

(*this)(I1, I2, I3, I4, I5, std::forward<Args>(OtherArgs)...);
F::operator()(I1, I2, I3, I4, I5, std::forward<Args>(OtherArgs)...);
}

private:
// SYCL doesn't allow 0-length arrays so add one extra element even though
// it is not needed
#ifdef KOKKOS_ENABLE_SYCL
Expand Down
112 changes: 31 additions & 81 deletions components/omega/src/ocn/AuxiliaryState.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -141,32 +141,22 @@ void AuxiliaryState::computeMomAux(const OceanState *State,

Pacer::start("AuxState:vertexAuxState1", 2);
parallelForOuter(
"vertexAuxState1", {Mesh->NVerticesAll},
"vertexAuxState1",
LaunchConfig({Mesh->NVerticesAll},
TeamScratch<Real>(2 * VCoord->NVertLayers)),
KOKKOS_LAMBDA(int IVertex, const TeamMember &Team) {
const int KMin = MinLayerVertexTop(IVertex);
const int KMax = MaxLayerVertexBot(IVertex);
const int KRange = vertRangeChunked(KMin, KMax);

parallelForInner(
Team, KRange, INNER_LAMBDA(int KChunk) {
LocVorticityAux.computeVarsOnVertex(
IVertex, KChunk, PseudoThickCell, NormalVelEdge);
});
LocVorticityAux.computeVarsOnVertex(Team, IVertex, PseudoThickCell,
NormalVelEdge);
});
Pacer::stop("AuxState:vertexAuxState1", 2);

Pacer::start("AuxState:cellAuxState1", 2);
parallelForOuter(
"cellAuxState1", {Mesh->NCellsAll},
"cellAuxState1",
LaunchConfig({Mesh->NCellsAll},
TeamScratch<Real>(2 * VCoord->NVertLayers)),
KOKKOS_LAMBDA(int ICell, const TeamMember &Team) {
const int KMin = MinLayerCell(ICell);
const int KMax = MaxLayerCell(ICell);
const int KRange = vertRangeChunked(KMin, KMax);

parallelForInner(
Team, KRange, INNER_LAMBDA(int KChunk) {
LocKineticAux.computeVarsOnCell(ICell, KChunk, NormalVelEdge);
});
LocKineticAux.computeVarsOnCell(Team, ICell, NormalVelEdge);
});
Pacer::stop("AuxState:cellAuxState1", 2);

Expand All @@ -183,64 +173,36 @@ void AuxiliaryState::computeMomAux(const OceanState *State,
parallelForOuter(
"edgeAuxState2", {Mesh->NEdgesAll},
KOKKOS_LAMBDA(int IEdge, const TeamMember &Team) {
const int KMin = MinLayerEdgeBot(IEdge);
const int KMax = MaxLayerEdgeTop(IEdge);
const int KRange = vertRangeChunked(KMin, KMax);

parallelForInner(
Team, KRange, INNER_LAMBDA(int KChunk) {
LocPseudoThicknessAux.computeVarsOnEdge(
IEdge, KChunk, PseudoThickCell, NormalVelEdge);
LocVelocityDel2Aux.computeVarsOnEdge(
IEdge, KChunk, VelocityDivCell, RelVortVertex);
});
LocPseudoThicknessAux.computeVarsOnEdge(Team, IEdge, PseudoThickCell,
NormalVelEdge);

LocVelocityDel2Aux.computeVarsOnEdge(Team, IEdge, VelocityDivCell,
RelVortVertex);
});

parallelForOuter(
"edgeAuxState2", {Mesh->NEdgesAll},
KOKKOS_LAMBDA(int IEdge, const TeamMember &Team) {
const int KMin = MinLayerEdgeTop(IEdge);
const int KMax = MaxLayerEdgeBot(IEdge);
const int KRange = vertRangeChunked(KMin, KMax);

parallelForInner(
Team, KRange, INNER_LAMBDA(int KChunk) {
LocVorticityAux.computeVarsOnEdge(IEdge, KChunk);
});
LocVorticityAux.computeVarsOnEdge(Team, IEdge);
});
Pacer::stop("AuxState:edgeAuxState2", 2);

Pacer::start("AuxState:vertexAuxState2", 2);
parallelForOuter(
"vertexAuxState2", {Mesh->NVerticesAll},
"vertexAuxState2",
LaunchConfig({Mesh->NVerticesAll},
TeamScratch<Real>(VCoord->NVertLayers)),
KOKKOS_LAMBDA(int IVertex, const TeamMember &Team) {
// Del2RelVortVertex is computed over the full vertex valid range
// [MinLayerVertexTop, MaxLayerVertexBot] so that boundary-vertex
// layers read by the biharmonic velocity tendency are valid rather
// than fill values (see VelocityDel2AuxVars::computeVarsOnVertex).
const int KMin = MinLayerVertexTop(IVertex);
const int KMax = MaxLayerVertexBot(IVertex);
const int KRange = vertRangeChunked(KMin, KMax);

parallelForInner(
Team, KRange, INNER_LAMBDA(int KChunk) {
LocVelocityDel2Aux.computeVarsOnVertex(IVertex, KChunk);
});
LocVelocityDel2Aux.computeVarsOnVertex(Team, IVertex);
});
Pacer::stop("AuxState:vertexAuxState2", 2);

Pacer::start("AuxState:cellAuxState2", 2);
parallelForOuter(
"cellAuxState2", {Mesh->NCellsAll},
"cellAuxState2",
LaunchConfig({Mesh->NCellsAll}, TeamScratch<Real>(VCoord->NVertLayers)),
KOKKOS_LAMBDA(int ICell, const TeamMember &Team) {
const int KMin = MinLayerCell(ICell);
const int KMax = MaxLayerCell(ICell);
const int KRange = vertRangeChunked(KMin, KMax);

parallelForInner(
Team, KRange, INNER_LAMBDA(int KChunk) {
LocVelocityDel2Aux.computeVarsOnCell(ICell, KChunk);
});
LocVelocityDel2Aux.computeVarsOnCell(Team, ICell);
});
Pacer::stop("AuxState:cellAuxState2", 2);

Expand Down Expand Up @@ -281,36 +243,24 @@ void AuxiliaryState::computeAll(const OceanState *State,

Pacer::start("AuxState:cellAuxState3", 2);
parallelForOuter(
"cellAuxState3", {Mesh->NCellsAll},
"cellAuxState3",
LaunchConfig({Mesh->NCellsAll}, TeamScratch<Real>(VCoord->NVertLayers)),
KOKKOS_LAMBDA(int ICell, const TeamMember &Team) {
const int KMin = MinLayerCell(ICell);
const int KMax = MaxLayerCell(ICell);
const int KRange = vertRangeChunked(KMin, KMax);

parallelForInner(
Team, KRange, INNER_LAMBDA(int KChunk) {
LocPseudoThicknessAux.computeVarsOnCells(
ICell, KChunk, PseudoThickCell, NormalVelEdge,
TimeStepSeconds);
});
LocPseudoThicknessAux.computeVarsOnCells(
Team, ICell, PseudoThickCell, NormalVelEdge, TimeStepSeconds);
Comment thread
mwarusz marked this conversation as resolved.
});
Pacer::stop("AuxState:cellAuxState3", 2);

const auto &MeanPseudoThickEdge = PseudoThicknessAux.MeanPseudoThickEdge;

Pacer::start("AuxState:cellAuxState4", 2);
parallelForOuter(
"cellAuxState4", {NTracers, Mesh->NCellsAll},
"cellAuxState4",
LaunchConfig({NTracers, Mesh->NCellsAll},
TeamScratch<Real>(VCoord->NVertLayers)),
KOKKOS_LAMBDA(int LTracer, int ICell, const TeamMember &Team) {
const int KMin = MinLayerCell(ICell);
const int KMax = MaxLayerCell(ICell);
const int KRange = vertRangeChunked(KMin, KMax);

parallelForInner(
Team, KRange, INNER_LAMBDA(int KChunk) {
LocTracerAux.computeVarsOnCells(
LTracer, ICell, KChunk, MeanPseudoThickEdge, TracerArray);
});
LocTracerAux.computeVarsOnCells(Team, LTracer, ICell,
MeanPseudoThickEdge, TracerArray);
});
Pacer::stop("AuxState:cellAuxState4", 2);

Expand Down
Loading
Loading