Start using hierarchical parallelism to limit vertical loops extents

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,14 +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,
-                               Halo *MeshHalo, int NVertLayers, int NTracers)
-    : Mesh(Mesh), MeshHalo(MeshHalo), Name(stripDefault(Name)),
-      KineticAux(stripDefault(Name), Mesh, NVertLayers),
-      LayerThicknessAux(stripDefault(Name), Mesh, NVertLayers),
-      VorticityAux(stripDefault(Name), Mesh, NVertLayers),
-      VelocityDel2Aux(stripDefault(Name), Mesh, 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, NVertLayers, NTracers) {
+      TracerAux(stripDefault(Name), Mesh, VCoord, NTracers) {
 
    GroupName = "AuxiliaryState";
    if (Name != "Default") {
@@ -64,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);
 
@@ -103,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);
 
@@ -151,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,7 +272,8 @@ void AuxiliaryState::computeAll(const OceanState *State,
 // Create a non-default auxiliary state
 AuxiliaryState *AuxiliaryState::create(const std::string &Name,
                                        const HorzMesh *Mesh, Halo *MeshHalo,
-                                       int NVertLayers, const int NTracers) {
+                                       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",
@@ -202,7 +282,7 @@ AuxiliaryState *AuxiliaryState::create(const std::string &Name,
    }
 
    auto *NewAuxState =
-       new AuxiliaryState(Name, Mesh, MeshHalo, NVertLayers, NTracers);
+       new AuxiliaryState(Name, Mesh, MeshHalo, VCoord, NTracers);
    AllAuxStates.emplace(Name, NewAuxState);
 
    return NewAuxState;
@@ -211,14 +291,14 @@ AuxiliaryState *AuxiliaryState::create(const std::string &Name,
 // Create the default auxiliary state. Assumes that HorzMesh, VertCoord and
 // Halo have been initialized.
 void AuxiliaryState::init() {
-   const HorzMesh *DefMesh = HorzMesh::getDefault();
-   Halo *DefHalo           = Halo::getDefault();
+   const HorzMesh *DefMesh    = HorzMesh::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, 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,7 +50,8 @@ class AuxiliaryState {
 
    // Create a non-default auxiliary state
    static AuxiliaryState *create(const std::string &Name, const HorzMesh *Mesh,
-                                 Halo *MeshHalo, int NVertLayers, int NTracers);
+                                 Halo *MeshHalo, const VertCoord *VCoord,
+                                 int NTracers);
 
    /// Get the default auxiliary state
    static AuxiliaryState *getDefault();
@@ -83,13 +84,14 @@ class AuxiliaryState {
 
  private:
    AuxiliaryState(const std::string &Name, const HorzMesh *Mesh, Halo *MeshHalo,
-                  int NVertLayers, int NTracers);
+                  const VertCoord *VCoord, int NTracers);
 
    AuxiliaryState(const AuxiliaryState &) = delete;
    AuxiliaryState(AuxiliaryState &&)      = delete;
 
    const HorzMesh *Mesh;
    Halo *MeshHalo;
+   const VertCoord *VCoord;
    static AuxiliaryState *DefaultAuxState;
    static std::map<std::string, std::unique_ptr<AuxiliaryState>> AllAuxStates;
 };