Start using hierarchical parallelism to limit vertical loops extents

components/omega/doc/devGuide/AuxiliaryState.md

@@ -24,10 +24,9 @@ OMEGA::AuxiliaryState* DefAuxState = OMEGA::AuxiliaryState::getDefault();
 
 ## Creation of non-default auxiliary states
 
-A non-default auxiliary state can be created from a string `Name`, horizontal mesh `Mesh`, a number of
-vertical layers `NVertLayers`, and a number of Tracers `NTracers`:
+A non-default auxiliary state can be created from a string `Name`, horizontal mesh `Mesh`, halo layer `MeshHalo`, vertical coordinate `VCoord`, and a number of Tracers `NTracers`:
 ```c++
-OMEGA::AuxiliaryState*  NewAuxState = OMEGA::AuxiliaryState::create(Name, Mesh, NVertLayers, NTracers);
+OMEGA::AuxiliaryState*  NewAuxState = OMEGA::AuxiliaryState::create(Name, Mesh, MeshHalo, VCoord, NTracers);
 ```
 For conveniece, this returns a pointer to the newly created state. Given its name, a pointer to a named auxiliary state
 can be obtained at any time by calling the static `get` method:
@@ -37,9 +36,9 @@ OMEGA::AuxiliaryState* NewAuxState = OMEGA::AuxiliaryState::get(Name);
 
 ## Computation of auxiliary variables
 To compute all auxiliary variables stored in an auxiliary state `AuxState`,
-given ocean state `State` and time level `TimeLevel` do:
+given ocean state `State`, an array of tracers `TracerArray`, and time level `TimeLevel` do:
 ```c++
-AuxState.computeAll(State, TimeLevel);
+AuxState.computeAll(State, TracerArray, TimeLevel);
 ```
 
 ## Removal of auxiliary states

components/omega/doc/devGuide/Tendencies.md

@@ -26,16 +26,15 @@ tendency terms, which each store constant mesh information as private member var
 ## Creation of non-default tendencies
 
 A non-default tendency group can be created with or without custom tendencies.
-Without custom tendencies, it is created from a string `Name`, horizontal mesh `Mesh`, number of
-vertical layers `NVertLayers`, number of tracers `NTracers`, and a configuration `Options`:
+Without custom tendencies, it is created from a string `Name`, horizontal mesh `Mesh`, vertical coordinate `VCoord`, number of tracers `NTracers`, and a configuration `Options`:
 ```c++
-OMEGA::Tendencies*  NewTendencies = OMEGA::Tendencies::create(Name, Mesh, NVertLayers, NTracers, Options);
+OMEGA::Tendencies*  NewTendencies = OMEGA::Tendencies::create(Name, Mesh, VCoord, NTracers, Options);
 ```
 For convenience, this returns a pointer to the newly created instance.
 To allow the user to provide custom tendencies, the `create` function can take two additional arguments
 `CustomThicknessTend` and `CustomVelocityTend`
 ```c++
-OMEGA::Tendencies*  NewTendencies = OMEGA::Tendencies::create(Name, Mesh, NVertLayers, NTracers, Options, CustomThicknessTend, CustomVelocityTend);
+OMEGA::Tendencies*  NewTendencies = OMEGA::Tendencies::create(Name, Mesh, VCoord, NTracers, Options, CustomThicknessTend, CustomVelocityTend);
 ```
 The two custom tendency arguments need to be callable objects that take a Kokkos array `Tend`, ocean state `State`,
 auxiliary state `AuxState`, two integers: `ThickTimeLevel` and `VelTimeLevel`, and time instant `Time`.
@@ -49,7 +48,7 @@ OMEGA::Tendencies* NewTendencies = OMEGA::Tendencies::get(Name);
 ## Computation of tendencies
 To compute all tendencies for layer thickness, normal velocity, and tracer equations,
 given ocean state, `State`, a group of auxiliary variables, `AuxState`, an array of tracers,
-`TracerArray`, thickness time level 'ThickTimeLevel', velocity time level `VelTimeLevel`,
+`TracerArray`, thickness time level `ThickTimeLevel`, velocity time level `VelTimeLevel`,
 and time instant `Time` do:
 ```c++
 Tendencies.computeAllTendencies(State, AuxState, TracerArray, ThickTimeLevel, VelTimeLevel, Time);

components/omega/src/infra/OmegaKokkos.h

@@ -359,8 +359,7 @@ KOKKOS_FUNCTION void parallelForInner(const TeamMember &Team, int UpperBound,
 template <class T, class Enable = void> struct AccumTypeHelper;
 
 template <class T>
-struct AccumTypeHelper<
-    T, std::enable_if_t<std::is_arithmetic_v<std::remove_reference_t<T>>>> {
+struct AccumTypeHelper<T, std::enable_if_t<std::is_arithmetic_v<T>>> {
    using Type = T;
 };
 
@@ -386,7 +385,8 @@ inline void parallelReduceOuter(const std::string &Label,
    auto Policy = TeamPolicy(LinBound, OMEGA_TEAMSIZE);
    Kokkos::parallel_reduce(
        Label, Policy,
-       KOKKOS_LAMBDA(const TeamMember &Team, AccumType<R> &...Accums) {
+       KOKKOS_LAMBDA(const TeamMember &Team,
+                     AccumType<std::remove_reference_t<R>> &...Accums) {
           const int TeamId = Team.league_rank();
           LinFunctor(TeamId, Team, Accums...);
        },

components/omega/src/ocn/AuxiliaryState.cpp

@@ -18,15 +18,15 @@ static std::string stripDefault(const std::string &Name) {
 // Constructor. Constructs the member auxiliary variables and registers their
 // fields with IOStreams
 AuxiliaryState::AuxiliaryState(const std::string &Name, const HorzMesh *Mesh,
-                               const VertCoord *VCoord, Halo *MeshHalo,
-                               int NVertLayers, int NTracers)
-    : Mesh(Mesh), VCoord(VCoord), MeshHalo(MeshHalo), Name(stripDefault(Name)),
-      KineticAux(stripDefault(Name), Mesh, NVertLayers),
-      LayerThicknessAux(stripDefault(Name), Mesh, NVertLayers),
-      VorticityAux(stripDefault(Name), Mesh, NVertLayers),
-      VelocityDel2Aux(stripDefault(Name), Mesh, VCoord, NVertLayers),
+                               Halo *MeshHalo, const VertCoord *VCoord,
+                               int NTracers)
+    : Mesh(Mesh), MeshHalo(MeshHalo), VCoord(VCoord), Name(stripDefault(Name)),
+      KineticAux(stripDefault(Name), Mesh, VCoord),
+      LayerThicknessAux(stripDefault(Name), Mesh, VCoord),
+      VorticityAux(stripDefault(Name), Mesh, VCoord),
+      VelocityDel2Aux(stripDefault(Name), Mesh, VCoord),
       WindForcingAux(stripDefault(Name), Mesh),
-      TracerAux(stripDefault(Name), Mesh, VCoord, NVertLayers, NTracers) {
+      TracerAux(stripDefault(Name), Mesh, VCoord, NTracers) {
 
    GroupName = "AuxiliaryState";
    if (Name != "Default") {
@@ -65,31 +65,53 @@ void AuxiliaryState::computeMomAux(const OceanState *State, int ThickTimeLevel,
    State->getLayerThickness(LayerThickCell, ThickTimeLevel);
    State->getNormalVelocity(NormalVelEdge, VelTimeLevel);
 
-   const int NVertLayers = LayerThickCell.extent_int(1);
-   const int NChunks     = NVertLayers / VecLength;
-
    OMEGA_SCOPE(LocKineticAux, KineticAux);
    OMEGA_SCOPE(LocLayerThicknessAux, LayerThicknessAux);
    OMEGA_SCOPE(LocVorticityAux, VorticityAux);
    OMEGA_SCOPE(LocVelocityDel2Aux, VelocityDel2Aux);
    OMEGA_SCOPE(LocWindForcingAux, WindForcingAux);
 
+   OMEGA_SCOPE(MinLayerCell, VCoord->MinLayerCell);
+   OMEGA_SCOPE(MaxLayerCell, VCoord->MaxLayerCell);
+   OMEGA_SCOPE(MinLayerVertexBot, VCoord->MinLayerVertexBot);
+   OMEGA_SCOPE(MinLayerVertexTop, VCoord->MinLayerVertexTop);
+   OMEGA_SCOPE(MaxLayerVertexBot, VCoord->MaxLayerVertexBot);
+   OMEGA_SCOPE(MaxLayerVertexTop, VCoord->MaxLayerVertexTop);
+   OMEGA_SCOPE(MinLayerEdgeTop, VCoord->MinLayerEdgeTop);
+   OMEGA_SCOPE(MinLayerEdgeBot, VCoord->MinLayerEdgeBot);
+   OMEGA_SCOPE(MaxLayerEdgeBot, VCoord->MaxLayerEdgeBot);
+   OMEGA_SCOPE(MaxLayerEdgeTop, VCoord->MaxLayerEdgeTop);
+
    Pacer::start("AuxState:computeMomAux", 1);
 
    Pacer::start("AuxState:vertexAuxState1", 2);
-   parallelFor(
-       "vertexAuxState1", {Mesh->NVerticesAll, NChunks},
-       KOKKOS_LAMBDA(int IVertex, int KChunk) {
-          LocVorticityAux.computeVarsOnVertex(IVertex, KChunk, LayerThickCell,
-                                              NormalVelEdge);
+   parallelForOuter(
+       "vertexAuxState1", {Mesh->NVerticesAll},
+       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, LayerThickCell, NormalVelEdge);
+              });
        });
    Pacer::stop("AuxState:vertexAuxState1", 2);
 
    Pacer::start("AuxState:cellAuxState1", 2);
-   parallelFor(
-       "cellAuxState1", {Mesh->NCellsAll, NChunks},
-       KOKKOS_LAMBDA(int ICell, int KChunk) {
-          LocKineticAux.computeVarsOnCell(ICell, KChunk, NormalVelEdge);
+   parallelForOuter(
+       "cellAuxState1", {Mesh->NCellsAll},
+       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);
+              });
        });
    Pacer::stop("AuxState:cellAuxState1", 2);
 
@@ -104,39 +126,79 @@ void AuxiliaryState::computeMomAux(const OceanState *State, int ThickTimeLevel,
    Pacer::stop("AuxState:edgeAuxState1", 2);
 
    Pacer::start("AuxState:edgeAuxState2", 2);
-   parallelFor(
-       "edgeAuxState2", {Mesh->NEdgesAll, NChunks},
-       KOKKOS_LAMBDA(int IEdge, int KChunk) {
-          LocVorticityAux.computeVarsOnEdge(IEdge, KChunk);
-          LocLayerThicknessAux.computeVarsOnEdge(IEdge, KChunk, LayerThickCell,
-                                                 NormalVelEdge);
-          LocVelocityDel2Aux.computeVarsOnEdge(IEdge, KChunk, VelocityDivCell,
-                                               RelVortVertex);
+   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) {
+                 LocLayerThicknessAux.computeVarsOnEdge(
+                     IEdge, KChunk, LayerThickCell, NormalVelEdge);
+                 LocVelocityDel2Aux.computeVarsOnEdge(
+                     IEdge, KChunk, 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);
+              });
        });
    Pacer::stop("AuxState:edgeAuxState2", 2);
 
    Pacer::start("AuxState:vertexAuxState2", 2);
-   parallelFor(
-       "vertexAuxState2", {Mesh->NVerticesAll, NChunks},
-       KOKKOS_LAMBDA(int IVertex, int KChunk) {
-          LocVelocityDel2Aux.computeVarsOnVertex(IVertex, KChunk);
+   parallelForOuter(
+       "vertexAuxState2", {Mesh->NVerticesAll},
+       KOKKOS_LAMBDA(int IVertex, const TeamMember &Team) {
+          const int KMin   = MinLayerVertexBot(IVertex);
+          const int KMax   = MaxLayerVertexTop(IVertex);
+          const int KRange = vertRangeChunked(KMin, KMax);
+
+          parallelForInner(
+              Team, KRange, INNER_LAMBDA(int KChunk) {
+                 LocVelocityDel2Aux.computeVarsOnVertex(IVertex, KChunk);
+              });
        });
    Pacer::stop("AuxState:vertexAuxState2", 2);
 
    Pacer::start("AuxState:cellAuxState2", 2);
-   parallelFor(
-       "cellAuxState2", {Mesh->NCellsAll, NChunks},
-       KOKKOS_LAMBDA(int ICell, int KChunk) {
-          LocVelocityDel2Aux.computeVarsOnCell(ICell, KChunk);
+   parallelForOuter(
+       "cellAuxState2", {Mesh->NCellsAll},
+       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);
+              });
        });
    Pacer::stop("AuxState:cellAuxState2", 2);
 
    Pacer::start("AuxState:cellAuxState3", 2);
-   parallelFor(
-       "cellAuxState3", {Mesh->NCellsAll, NChunks},
-       KOKKOS_LAMBDA(int ICell, int KChunk) {
-          LocLayerThicknessAux.computeVarsOnCells(ICell, KChunk,
-                                                  LayerThickCell);
+   parallelForOuter(
+       "cellAuxState3", {Mesh->NCellsAll},
+       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) {
+                 LocLayerThicknessAux.computeVarsOnCells(ICell, KChunk,
+                                                         LayerThickCell);
+              });
        });
    Pacer::stop("AuxState:cellAuxState3", 2);
 
@@ -152,33 +214,49 @@ void AuxiliaryState::computeAll(const OceanState *State,
    State->getLayerThickness(LayerThickCell, ThickTimeLevel);
    State->getNormalVelocity(NormalVelEdge, VelTimeLevel);
 
-   const int NVertLayers = LayerThickCell.extent_int(1);
-   const int NChunks     = NVertLayers / VecLength;
-   const int NTracers    = TracerArray.extent_int(0);
+   const int NTracers = TracerArray.extent_int(0);
 
    OMEGA_SCOPE(LocTracerAux, TracerAux);
+   OMEGA_SCOPE(MinLayerCell, VCoord->MinLayerCell);
+   OMEGA_SCOPE(MaxLayerCell, VCoord->MaxLayerCell);
+   OMEGA_SCOPE(MinLayerEdgeBot, VCoord->MinLayerEdgeBot);
+   OMEGA_SCOPE(MaxLayerEdgeTop, VCoord->MaxLayerEdgeTop);
 
    Pacer::start("AuxState:computeAll", 1);
 
    computeMomAux(State, ThickTimeLevel, VelTimeLevel);
 
    Pacer::start("AuxState:edgeAuxState4", 2);
-   parallelFor(
-       "edgeAuxState4", {NTracers, Mesh->NEdgesAll, NChunks},
-       KOKKOS_LAMBDA(int LTracer, int IEdge, int KChunk) {
-          LocTracerAux.computeVarsOnEdge(LTracer, IEdge, KChunk, NormalVelEdge,
-                                         LayerThickCell, TracerArray);
+   parallelForOuter(
+       "edgeAuxState4", {NTracers, Mesh->NEdgesAll},
+       KOKKOS_LAMBDA(int LTracer, 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) {
+                 LocTracerAux.computeVarsOnEdge(LTracer, IEdge, KChunk,
+                                                NormalVelEdge, LayerThickCell,
+                                                TracerArray);
+              });
        });
    Pacer::stop("AuxState:edgeAuxState4", 2);
 
    const auto &MeanLayerThickEdge = LayerThicknessAux.MeanLayerThickEdge;
 
    Pacer::start("AuxState:cellAuxState4", 2);
-   parallelFor(
-       "cellAuxState4", {NTracers, Mesh->NCellsAll, NChunks},
-       KOKKOS_LAMBDA(int LTracer, int ICell, int KChunk) {
-          LocTracerAux.computeVarsOnCells(LTracer, ICell, KChunk,
-                                          MeanLayerThickEdge, TracerArray);
+   parallelForOuter(
+       "cellAuxState4", {NTracers, Mesh->NCellsAll},
+       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, MeanLayerThickEdge, TracerArray);
+              });
        });
    Pacer::stop("AuxState:cellAuxState4", 2);
 
@@ -193,9 +271,9 @@ void AuxiliaryState::computeAll(const OceanState *State,
 
 // Create a non-default auxiliary state
 AuxiliaryState *AuxiliaryState::create(const std::string &Name,
-                                       const HorzMesh *Mesh,
-                                       const VertCoord *VCoord, Halo *MeshHalo,
-                                       int NVertLayers, const int NTracers) {
+                                       const HorzMesh *Mesh, Halo *MeshHalo,
+                                       const VertCoord *VCoord,
+                                       const int NTracers) {
    if (AllAuxStates.find(Name) != AllAuxStates.end()) {
       LOG_ERROR("Attempted to create a new AuxiliaryState with name {} but it "
                 "already exists",
@@ -204,7 +282,7 @@ AuxiliaryState *AuxiliaryState::create(const std::string &Name,
    }
 
    auto *NewAuxState =
-       new AuxiliaryState(Name, Mesh, VCoord, MeshHalo, NVertLayers, NTracers);
+       new AuxiliaryState(Name, Mesh, MeshHalo, VCoord, NTracers);
    AllAuxStates.emplace(Name, NewAuxState);
 
    return NewAuxState;
@@ -214,14 +292,13 @@ AuxiliaryState *AuxiliaryState::create(const std::string &Name,
 // Halo have been initialized.
 void AuxiliaryState::init() {
    const HorzMesh *DefMesh    = HorzMesh::getDefault();
-   const VertCoord *DefVCoord = VertCoord::getDefault();
    Halo *DefHalo              = Halo::getDefault();
+   const VertCoord *DefVCoord = VertCoord::getDefault();
 
-   int NVertLayers = VertCoord::getDefault()->NVertLayers;
-   int NTracers    = Tracers::getNumTracers();
+   int NTracers = Tracers::getNumTracers();
 
-   AuxiliaryState::DefaultAuxState = AuxiliaryState::create(
-       "Default", DefMesh, DefVCoord, DefHalo, NVertLayers, NTracers);
+   AuxiliaryState::DefaultAuxState =
+       AuxiliaryState::create("Default", DefMesh, DefHalo, DefVCoord, NTracers);
 
    Config *OmegaConfig = Config::getOmegaConfig();
    DefaultAuxState->readConfigOptions(OmegaConfig);

components/omega/src/ocn/AuxiliaryState.h

@@ -50,8 +50,8 @@ class AuxiliaryState {
 
    // Create a non-default auxiliary state
    static AuxiliaryState *create(const std::string &Name, const HorzMesh *Mesh,
-                                 const VertCoord *VCoord, Halo *MeshHalo,
-                                 int NVertLayers, int NTracers);
+                                 Halo *MeshHalo, const VertCoord *VCoord,
+                                 int NTracers);
 
    /// Get the default auxiliary state
    static AuxiliaryState *getDefault();
@@ -83,16 +83,15 @@ class AuxiliaryState {
                    int TimeLevel) const;
 
  private:
-   AuxiliaryState(const std::string &Name, const HorzMesh *Mesh,
-                  const VertCoord *VCoord, Halo *MeshHalo, int NVertLayers,
-                  int NTracers);
+   AuxiliaryState(const std::string &Name, const HorzMesh *Mesh, Halo *MeshHalo,
+                  const VertCoord *VCoord, int NTracers);
 
    AuxiliaryState(const AuxiliaryState &) = delete;
    AuxiliaryState(AuxiliaryState &&)      = delete;
 
    const HorzMesh *Mesh;
-   const VertCoord *VCoord;
    Halo *MeshHalo;
+   const VertCoord *VCoord;
    static AuxiliaryState *DefaultAuxState;
    static std::map<std::string, std::unique_ptr<AuxiliaryState>> AllAuxStates;
 };