diff --git a/components/omega/src/infra/OmegaKokkos.h b/components/omega/src/infra/OmegaKokkos.h index e31896dbbe0a..df15c4f5aa55 100644 --- a/components/omega/src/infra/OmegaKokkos.h +++ b/components/omega/src/infra/OmegaKokkos.h @@ -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 +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 struct LinearIdxWrapper : F { +template class LinearIdxWrapper : private F { static_assert(Rank >= 1 && Rank <= 5, "LinearIdxWrapper supports ranks 1-5"); - using F::operator(); + public: template LinearIdxWrapper(F &&Functor, Array &&Bounds) : F(std::move(Functor)) { computeStrides(std::forward(Bounds)); @@ -227,13 +235,19 @@ template struct LinearIdxWrapper : F { } } + template + KOKKOS_FORCEINLINE_FUNCTION std::enable_if_t + operator()(int Idx, Args &&...OtherArgs) const { + F::operator()(Idx, std::forward(OtherArgs)...); + } + template KOKKOS_FORCEINLINE_FUNCTION std::enable_if_t operator()(int Idx, Args &&...OtherArgs) const { const int I1 = Idx / Strides[0]; const int I2 = Idx - I1 * Strides[0]; - (*this)(I1, I2, std::forward(OtherArgs)...); + F::operator()(I1, I2, std::forward(OtherArgs)...); } template @@ -244,7 +258,7 @@ template struct LinearIdxWrapper : F { const int I2 = Idx / Strides[1]; const int I3 = Idx - I2 * Strides[1]; - (*this)(I1, I2, I3, std::forward(OtherArgs)...); + F::operator()(I1, I2, I3, std::forward(OtherArgs)...); } template @@ -257,7 +271,7 @@ template struct LinearIdxWrapper : F { const int I3 = Idx / Strides[2]; const int I4 = Idx - I3 * Strides[2]; - (*this)(I1, I2, I3, I4, std::forward(OtherArgs)...); + F::operator()(I1, I2, I3, I4, std::forward(OtherArgs)...); } template @@ -272,9 +286,10 @@ template struct LinearIdxWrapper : F { const int I4 = Idx / Strides[3]; const int I5 = Idx - I4 * Strides[3]; - (*this)(I1, I2, I3, I4, I5, std::forward(OtherArgs)...); + F::operator()(I1, I2, I3, I4, I5, std::forward(OtherArgs)...); } + private: // SYCL doesn't allow 0-length arrays so add one extra element even though // it is not needed #ifdef KOKKOS_ENABLE_SYCL diff --git a/components/omega/src/ocn/AuxiliaryState.cpp b/components/omega/src/ocn/AuxiliaryState.cpp index 9a37a0739d98..7823a10590c7 100644 --- a/components/omega/src/ocn/AuxiliaryState.cpp +++ b/components/omega/src/ocn/AuxiliaryState.cpp @@ -141,32 +141,22 @@ void AuxiliaryState::computeMomAux(const OceanState *State, Pacer::start("AuxState:vertexAuxState1", 2); parallelForOuter( - "vertexAuxState1", {Mesh->NVerticesAll}, + "vertexAuxState1", + LaunchConfig({Mesh->NVerticesAll}, + TeamScratch(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(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); @@ -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(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(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); @@ -281,18 +243,11 @@ void AuxiliaryState::computeAll(const OceanState *State, Pacer::start("AuxState:cellAuxState3", 2); parallelForOuter( - "cellAuxState3", {Mesh->NCellsAll}, + "cellAuxState3", + LaunchConfig({Mesh->NCellsAll}, TeamScratch(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); }); Pacer::stop("AuxState:cellAuxState3", 2); @@ -300,17 +255,12 @@ void AuxiliaryState::computeAll(const OceanState *State, Pacer::start("AuxState:cellAuxState4", 2); parallelForOuter( - "cellAuxState4", {NTracers, Mesh->NCellsAll}, + "cellAuxState4", + LaunchConfig({NTracers, Mesh->NCellsAll}, + TeamScratch(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); diff --git a/components/omega/src/ocn/Tendencies.cpp b/components/omega/src/ocn/Tendencies.cpp index 5bb131f32ab6..6182720c4b3b 100644 --- a/components/omega/src/ocn/Tendencies.cpp +++ b/components/omega/src/ocn/Tendencies.cpp @@ -587,16 +587,11 @@ void Tendencies::computePseudoThicknessTendenciesOnly( if (LocThicknessFluxDiv.Enabled) { Pacer::start("Tend:thicknessFluxDiv", 2); parallelForOuter( - {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) { - LocThicknessFluxDiv(LocPseudoThicknessTend, ICell, KChunk, - ThickFluxEdge, NormalVelEdge); - }); + LaunchConfig({Mesh->NCellsAll}, + TeamScratch(VCoord->NVertLayers)), + KOKKOS_LAMBDA(int ICell, const TeamMember &Team) { + LocThicknessFluxDiv(Team, LocPseudoThicknessTend, ICell, + ThickFluxEdge, NormalVelEdge); }); Pacer::stop("Tend:thicknessFluxDiv", 2); } @@ -659,17 +654,12 @@ void Tendencies::computeVelocityTendenciesOnly( if (LocPotentialVortHAdv.Enabled) { Pacer::start("Tend:PotentialVortHAdv", 2); parallelForOuter( - {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) { - LocPotentialVortHAdv(LocNormalVelocityTend, IEdge, KChunk, - NormRVortEdge, NormFEdge, - FluxPseudoThickEdge, NormVelEdge); - }); + LaunchConfig({Mesh->NEdgesAll}, + TeamScratch(VCoord->NVertLayers)), + KOKKOS_LAMBDA(int IEdge, const TeamMember &Team) { + LocPotentialVortHAdv(Team, LocNormalVelocityTend, IEdge, + NormRVortEdge, NormFEdge, FluxPseudoThickEdge, + NormVelEdge); }); Pacer::stop("Tend:PotentialVortHAdv", 2); } @@ -680,13 +670,7 @@ void Tendencies::computeVelocityTendenciesOnly( Pacer::start("Tend:KEGrad", 2); parallelForOuter( {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) { - LocKEGrad(LocNormalVelocityTend, IEdge, KChunk, KECell); - }); + LocKEGrad(Team, LocNormalVelocityTend, IEdge, KECell); }); Pacer::stop("Tend:KEGrad", 2); } @@ -697,13 +681,7 @@ void Tendencies::computeVelocityTendenciesOnly( Pacer::start("Tend:SSHGrad", 2); parallelForOuter( {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) { - LocSSHGrad(LocNormalVelocityTend, IEdge, KChunk, SSHCell); - }); + LocSSHGrad(Team, LocNormalVelocityTend, IEdge, SSHCell); }); Pacer::stop("Tend:SSHGrad", 2); } @@ -715,14 +693,8 @@ void Tendencies::computeVelocityTendenciesOnly( Pacer::start("Tend:velocityDiffusion", 2); parallelForOuter( {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) { - LocVelocityDiffusion(LocNormalVelocityTend, IEdge, KChunk, - DivCell, RVortVertex); - }); + LocVelocityDiffusion(Team, LocNormalVelocityTend, IEdge, DivCell, + RVortVertex); }); Pacer::stop("Tend:velocityDiffusion", 2); } @@ -735,14 +707,8 @@ void Tendencies::computeVelocityTendenciesOnly( Pacer::start("Tend:velocityHyperDiff", 2); parallelForOuter( {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) { - LocVelocityHyperDiff(LocNormalVelocityTend, IEdge, KChunk, - Del2DivCell, Del2RVortVertex); - }); + LocVelocityHyperDiff(Team, LocNormalVelocityTend, IEdge, + Del2DivCell, Del2RVortVertex); }); Pacer::stop("Tend:velocityHyperDiff", 2); } @@ -763,14 +729,8 @@ void Tendencies::computeVelocityTendenciesOnly( Pacer::start("Tend:sfcStressForcing", 2); parallelForOuter( {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) { - LocSfcStressForcing(LocNormalVelocityTend, IEdge, KChunk, - NormalStressEdge, MeanPseudoThickEdge); - }); + LocSfcStressForcing(Team, LocNormalVelocityTend, IEdge, + NormalStressEdge, MeanPseudoThickEdge); }); Pacer::stop("Tend:sfcStressForcing", 2); } @@ -849,27 +809,17 @@ void Tendencies::computeTracerTendenciesOnly( if (LocTracerHorzAdv.Enabled) { Pacer::start("Tend:tracerHorzAdv", 2); parallelForOuter( - {NTracers, Mesh->NEdgesAll}, + LaunchConfig({NTracers, Mesh->NEdgesAll}, + TeamScratch(VCoord->NVertLayers)), KOKKOS_LAMBDA(int L, 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) { - LocTracerHorzAdv(L, IEdge, KChunk, TracerArray, - FluxPseudoThickEdge, NormalVelEdge); - }); + LocTracerHorzAdv(Team, L, IEdge, TracerArray, FluxPseudoThickEdge, + NormalVelEdge); }); parallelForOuter( - {NTracers, Mesh->NCellsAll}, + LaunchConfig({NTracers, Mesh->NCellsAll}, + TeamScratch(VCoord->NVertLayers)), KOKKOS_LAMBDA(int L, 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) { - LocTracerHorzAdv(LocTracerTend, L, ICell, KChunk); - }); + LocTracerHorzAdv(Team, LocTracerTend, L, ICell); }); Pacer::stop("Tend:tracerHorzAdv", 2); } @@ -880,17 +830,11 @@ void Tendencies::computeTracerTendenciesOnly( if (LocTracerDiffusion.Enabled) { Pacer::start("Tend:tracerDiffusion", 2); parallelForOuter( - {NTracers, Mesh->NCellsAll}, + LaunchConfig({NTracers, Mesh->NCellsAll}, + TeamScratch(VCoord->NVertLayers)), KOKKOS_LAMBDA(int L, 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) { - LocTracerDiffusion(LocTracerTend, L, ICell, KChunk, - TracerArray, MeanPseudoThickEdge); - }); + LocTracerDiffusion(Team, LocTracerTend, L, ICell, TracerArray, + MeanPseudoThickEdge); }); Pacer::stop("Tend:tracerDiffusion", 2); } @@ -900,17 +844,10 @@ void Tendencies::computeTracerTendenciesOnly( if (LocTracerHyperDiff.Enabled) { Pacer::start("Tend:tracerHyperDiff", 2); parallelForOuter( - {NTracers, Mesh->NCellsAll}, + LaunchConfig({NTracers, Mesh->NCellsAll}, + TeamScratch(VCoord->NVertLayers)), KOKKOS_LAMBDA(int L, 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) { - LocTracerHyperDiff(LocTracerTend, L, ICell, KChunk, - Del2TracersCell); - }); + LocTracerHyperDiff(Team, LocTracerTend, L, ICell, Del2TracersCell); }); Pacer::stop("Tend:tracerHyperDiff", 2); } @@ -971,15 +908,8 @@ void Tendencies::computePseudoThicknessTendencies( parallelForOuter( "computePseudoThickAux", {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) { - PseudoThicknessAux.computeVarsOnEdge( - IEdge, KChunk, PseudoThickCell, NormalVelEdge); - }); + PseudoThicknessAux.computeVarsOnEdge(Team, IEdge, PseudoThickCell, + NormalVelEdge); }); Pacer::stop("Tend:computePseudoThickAux", 2); @@ -1032,17 +962,12 @@ void Tendencies::computeTracerTendencies( AuxState->PseudoThicknessAux.MeanPseudoThickEdge; Pacer::start("Tend:computeTracerAuxCell", 2); parallelForOuter( - "computeTracerAuxCell", {NTracers, Mesh->NCellsAll}, + "computeTracerAuxCell", + LaunchConfig({NTracers, Mesh->NCellsAll}, + TeamScratch(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) { - TracerAux.computeVarsOnCells(LTracer, ICell, KChunk, - MeanPseudoThickEdge, TracerArray); - }); + TracerAux.computeVarsOnCells(Team, LTracer, ICell, + MeanPseudoThickEdge, TracerArray); }); Pacer::stop("Tend:computeTracerAuxCell", 2); diff --git a/components/omega/src/ocn/TendencyTerms.cpp b/components/omega/src/ocn/TendencyTerms.cpp index 5b142888142c..7e93f0fde0c5 100644 --- a/components/omega/src/ocn/TendencyTerms.cpp +++ b/components/omega/src/ocn/TendencyTerms.cpp @@ -22,16 +22,16 @@ PseudoThicknessFluxDivOnCell::PseudoThicknessFluxDivOnCell( const HorzMesh *Mesh, const VertCoord *VCoord) : NEdgesOnCell(Mesh->NEdgesOnCell), EdgesOnCell(Mesh->EdgesOnCell), DvEdge(Mesh->DvEdge), AreaCell(Mesh->AreaCell), - EdgeSignOnCell(Mesh->EdgeSignOnCell), MinLayerCell(VCoord->MinLayerCell), - MaxLayerCell(VCoord->MaxLayerCell), + EdgeSignOnCell(Mesh->EdgeSignOnCell), NVertLayers(VCoord->NVertLayers), + MinLayerCell(VCoord->MinLayerCell), MaxLayerCell(VCoord->MaxLayerCell), MinLayerEdgeBot(VCoord->MinLayerEdgeBot), MaxLayerEdgeTop(VCoord->MaxLayerEdgeTop) {} PotentialVortHAdvOnEdge::PotentialVortHAdvOnEdge(const HorzMesh *Mesh, const VertCoord *VCoord) : NEdgesOnEdge(Mesh->NEdgesOnEdge), EdgesOnEdge(Mesh->EdgesOnEdge), - WeightsOnEdge(Mesh->WeightsOnEdge), EdgeMask(VCoord->EdgeMask), - MinLayerEdgeBot(VCoord->MinLayerEdgeBot), + WeightsOnEdge(Mesh->WeightsOnEdge), NVertLayers(VCoord->NVertLayers), + EdgeMask(VCoord->EdgeMask), MinLayerEdgeBot(VCoord->MinLayerEdgeBot), MaxLayerEdgeTop(VCoord->MaxLayerEdgeTop) {} KEGradOnEdge::KEGradOnEdge(const HorzMesh *Mesh, const VertCoord *VCoord) @@ -63,7 +63,8 @@ VelocityHyperDiffOnEdge::VelocityHyperDiffOnEdge(const HorzMesh *Mesh, SfcStressForcingOnEdge::SfcStressForcingOnEdge(const HorzMesh *Mesh, const VertCoord *VCoord) : Enabled(false), EdgeMask(VCoord->EdgeMask), - MinLayerEdgeBot(VCoord->MinLayerEdgeBot) {} + MinLayerEdgeBot(VCoord->MinLayerEdgeBot), + MaxLayerEdgeTop(VCoord->MaxLayerEdgeTop) {} BottomDragOnEdge::BottomDragOnEdge(const HorzMesh *Mesh, const VertCoord *VCoord) @@ -88,15 +89,20 @@ TracerHorzAdvOnCell::TracerHorzAdvOnCell(const HorzMesh *Mesh, Mesh->NEdgesAll, VCoord->NVertLayers), NEdgesOnCell(Mesh->NEdgesOnCell), EdgesOnCell(Mesh->EdgesOnCell), CellsOnEdge(Mesh->CellsOnEdge), EdgeSignOnCell(Mesh->EdgeSignOnCell), - DvEdge(Mesh->DvEdge), AreaCell(Mesh->AreaCell) {} + DvEdge(Mesh->DvEdge), AreaCell(Mesh->AreaCell), + NVertLayers(VCoord->NVertLayers), MinLayerCell(VCoord->MinLayerCell), + MaxLayerCell(VCoord->MaxLayerCell), + MinLayerEdgeBot(VCoord->MinLayerEdgeBot), + MaxLayerEdgeTop(VCoord->MaxLayerEdgeTop) {} TracerDiffOnCell::TracerDiffOnCell(const HorzMesh *Mesh, const VertCoord *VCoord) : NEdgesOnCell(Mesh->NEdgesOnCell), EdgesOnCell(Mesh->EdgesOnCell), CellsOnEdge(Mesh->CellsOnEdge), EdgeSignOnCell(Mesh->EdgeSignOnCell), DvEdge(Mesh->DvEdge), DcEdge(Mesh->DcEdge), AreaCell(Mesh->AreaCell), - MeshScalingDel2(Mesh->MeshScalingDel2), EdgeMask(VCoord->EdgeMask), - MinLayerCell(VCoord->MinLayerCell), MaxLayerCell(VCoord->MaxLayerCell), + MeshScalingDel2(Mesh->MeshScalingDel2), NVertLayers(VCoord->NVertLayers), + EdgeMask(VCoord->EdgeMask), MinLayerCell(VCoord->MinLayerCell), + MaxLayerCell(VCoord->MaxLayerCell), MinLayerEdgeBot(VCoord->MinLayerEdgeBot), MaxLayerEdgeTop(VCoord->MaxLayerEdgeTop) {} @@ -105,8 +111,9 @@ TracerHyperDiffOnCell::TracerHyperDiffOnCell(const HorzMesh *Mesh, : NEdgesOnCell(Mesh->NEdgesOnCell), EdgesOnCell(Mesh->EdgesOnCell), CellsOnEdge(Mesh->CellsOnEdge), EdgeSignOnCell(Mesh->EdgeSignOnCell), DvEdge(Mesh->DvEdge), DcEdge(Mesh->DcEdge), AreaCell(Mesh->AreaCell), - MeshScalingDel4(Mesh->MeshScalingDel4), EdgeMask(VCoord->EdgeMask), - MinLayerCell(VCoord->MinLayerCell), MaxLayerCell(VCoord->MaxLayerCell), + MeshScalingDel4(Mesh->MeshScalingDel4), NVertLayers(VCoord->NVertLayers), + EdgeMask(VCoord->EdgeMask), MinLayerCell(VCoord->MinLayerCell), + MaxLayerCell(VCoord->MaxLayerCell), MinLayerEdgeBot(VCoord->MinLayerEdgeBot), MaxLayerEdgeTop(VCoord->MaxLayerEdgeTop) {} diff --git a/components/omega/src/ocn/TendencyTerms.h b/components/omega/src/ocn/TendencyTerms.h index 548832b290bf..a5582f3476c8 100644 --- a/components/omega/src/ocn/TendencyTerms.h +++ b/components/omega/src/ocn/TendencyTerms.h @@ -32,35 +32,37 @@ class PseudoThicknessFluxDivOnCell { /// The functor takes cell index, vertical chunk index, and pseudo-thickness /// flux array as inputs, outputs the tendency array - KOKKOS_FUNCTION void operator()(const Array2DReal &Tend, I4 ICell, I4 KChunk, + KOKKOS_FUNCTION void operator()(const TeamMember &Team, + const Array2DReal &Tend, I4 ICell, const Array2DReal &PseudoThicknessFlux, const Array2DReal &NormalVelEdge) const { - const I4 KStartCell = chunkStart(KChunk, MinLayerCell(ICell)); - const I4 KLenCell = chunkLength(KChunk, KStartCell, MaxLayerCell(ICell)); - const I4 KEndCell = KStartCell + KLenCell - 1; const Real InvAreaCell = 1._Real / AreaCell(ICell); - Real DivTmp[VecLength] = {0}; + ScratchArray1DReal DivTmp(teamScratch(Team), NVertLayers); + parallelForInner( + Team, NVertLayers, INNER_LAMBDA(int K) { DivTmp(K) = 0; }); for (int J = 0; J < NEdgesOnCell(ICell); ++J) { const I4 JEdge = EdgesOnCell(ICell, J); - const I4 KStartEdge = Kokkos::max(KStartCell, MinLayerEdgeBot(JEdge)); - const I4 KEndEdge = Kokkos::min(KEndCell, MaxLayerEdgeTop(JEdge)); + const int MinLyrEdgeBot = MinLayerEdgeBot(JEdge); + const int MaxLyrEdgeTop = MaxLayerEdgeTop(JEdge); - for (int K = KStartEdge; K <= KEndEdge; ++K) { - const I4 KVec = K - KStartCell; - DivTmp[KVec] -= DvEdge(JEdge) * EdgeSignOnCell(ICell, J) * - PseudoThicknessFlux(JEdge, K) * - NormalVelEdge(JEdge, K) * InvAreaCell; - } + parallelForInner( + Team, Range{MinLyrEdgeBot, MaxLyrEdgeTop}, INNER_LAMBDA(int K) { + DivTmp(K) -= DvEdge(JEdge) * EdgeSignOnCell(ICell, J) * + PseudoThicknessFlux(JEdge, K) * + NormalVelEdge(JEdge, K) * InvAreaCell; + }); } - for (int KVec = 0; KVec < KLenCell; ++KVec) { - const I4 K = KStartCell + KVec; - Tend(ICell, K) -= DivTmp[KVec]; - } + const int MinLyrCell = MinLayerCell(ICell); + const int MaxLyrCell = MaxLayerCell(ICell); + + parallelForInner( + Team, Range{MinLyrCell, MaxLyrCell}, + INNER_LAMBDA(int K) { Tend(ICell, K) -= DivTmp(K); }); } private: @@ -69,6 +71,7 @@ class PseudoThicknessFluxDivOnCell { Array1DReal DvEdge; Array1DReal AreaCell; Array2DReal EdgeSignOnCell; + I4 NVertLayers; Array1DI4 MinLayerCell; Array1DI4 MaxLayerCell; Array1DI4 MinLayerEdgeBot; @@ -88,40 +91,47 @@ class PotentialVortHAdvOnEdge { /// normalized relative vorticity, normalized planetary vorticity, layer /// thickness on edges, and normal velocity on edges as inputs, /// outputs the tendency array - KOKKOS_FUNCTION void operator()(const Array2DReal &Tend, I4 IEdge, I4 KChunk, + KOKKOS_FUNCTION void operator()(const TeamMember &Team, + const Array2DReal &Tend, I4 IEdge, const Array2DReal &NormRVortEdge, const Array2DReal &NormFEdge, const Array2DReal &FluxPseudoThickEdge, const Array2DReal &NormVelEdge) const { - const I4 KStart = chunkStart(KChunk, MinLayerEdgeBot(IEdge)); - const I4 KLen = chunkLength(KChunk, KStart, MaxLayerEdgeTop(IEdge)); - Real VortTmp[VecLength] = {0}; + ScratchArray1DReal VortTmp(teamScratch(Team), NVertLayers); + + const int KMin = MinLayerEdgeBot(IEdge); + const int KMax = MaxLayerEdgeTop(IEdge); + + parallelForInner( + Team, NVertLayers, INNER_LAMBDA(int K) { VortTmp(K) = 0; }); for (int J = 0; J < NEdgesOnEdge(IEdge); ++J) { I4 JEdge = EdgesOnEdge(IEdge, J); - for (int KVec = 0; KVec < KLen; ++KVec) { - const I4 K = KStart + KVec; - Real NormVort = (NormRVortEdge(IEdge, K) + NormFEdge(IEdge, K) + - NormRVortEdge(JEdge, K) + NormFEdge(JEdge, K)) * - 0.5_Real; - - VortTmp[KVec] += WeightsOnEdge(IEdge, J) * - FluxPseudoThickEdge(JEdge, K) * - NormVelEdge(JEdge, K) * NormVort; - } + parallelForInner( + Team, Range{KMin, KMax}, INNER_LAMBDA(int K) { + Real NormVort = + (NormRVortEdge(IEdge, K) + NormFEdge(IEdge, K) + + NormRVortEdge(JEdge, K) + NormFEdge(JEdge, K)) * + 0.5_Real; + + VortTmp(K) += WeightsOnEdge(IEdge, J) * + FluxPseudoThickEdge(JEdge, K) * + NormVelEdge(JEdge, K) * NormVort; + }); } - for (int KVec = 0; KVec < KLen; ++KVec) { - const I4 K = KStart + KVec; - Tend(IEdge, K) += EdgeMask(IEdge, K) * VortTmp[KVec]; - } + parallelForInner( + Team, Range{KMin, KMax}, INNER_LAMBDA(int K) { + Tend(IEdge, K) += EdgeMask(IEdge, K) * VortTmp(K); + }); } private: Array1DI4 NEdgesOnEdge; Array2DI4 EdgesOnEdge; Array2DReal WeightsOnEdge; + I4 NVertLayers; Array2DReal EdgeMask; Array1DI4 MinLayerEdgeBot; Array1DI4 MaxLayerEdgeTop; @@ -137,20 +147,24 @@ class KEGradOnEdge { /// The functor takes edge index, vertical chunk index, and kinetic energy /// array as inputs, outputs the tendency array - KOKKOS_FUNCTION void operator()(const Array2DReal &Tend, I4 IEdge, I4 KChunk, + KOKKOS_FUNCTION void operator()(const TeamMember &Team, + const Array2DReal &Tend, I4 IEdge, const Array2DReal &KECell) const { - const I4 KStart = chunkStart(KChunk, MinLayerEdgeBot(IEdge)); - const I4 KLen = chunkLength(KChunk, KStart, MaxLayerEdgeTop(IEdge)); const I4 JCell0 = CellsOnEdge(IEdge, 0); const I4 JCell1 = CellsOnEdge(IEdge, 1); + const Real InvDcEdge = 1._Real / DcEdge(IEdge); - for (int KVec = 0; KVec < KLen; ++KVec) { - const I4 K = KStart + KVec; - Tend(IEdge, K) -= EdgeMask(IEdge, K) * - (KECell(JCell1, K) - KECell(JCell0, K)) * InvDcEdge; - } + const int KMin = MinLayerEdgeBot(IEdge); + const int KMax = MaxLayerEdgeTop(IEdge); + + parallelForInner( + Team, Range{KMin, KMax}, INNER_LAMBDA(int K) { + Tend(IEdge, K) -= EdgeMask(IEdge, K) * + (KECell(JCell1, K) - KECell(JCell0, K)) * + InvDcEdge; + }); } private: @@ -173,20 +187,23 @@ class SSHGradOnEdge { /// The functor takes edge index, vertical chunk index, and array of /// pseudo-thickness/SSH, outputs tendency array - KOKKOS_FUNCTION void operator()(const Array2DReal &Tend, I4 IEdge, I4 KChunk, + KOKKOS_FUNCTION void operator()(const TeamMember &Team, + const Array2DReal &Tend, I4 IEdge, const Array1DReal &SshCell) const { - const I4 KStart = chunkStart(KChunk, MinLayerEdgeBot(IEdge)); - const I4 KLen = chunkLength(KChunk, KStart, MaxLayerEdgeTop(IEdge)); const I4 ICell0 = CellsOnEdge(IEdge, 0); const I4 ICell1 = CellsOnEdge(IEdge, 1); + const Real InvDcEdge = 1._Real / DcEdge(IEdge); - for (int KVec = 0; KVec < KLen; ++KVec) { - const I4 K = KStart + KVec; - Tend(IEdge, K) -= EdgeMask(IEdge, K) * Gravity * - (SshCell(ICell1) - SshCell(ICell0)) * InvDcEdge; - } + const int KMin = MinLayerEdgeBot(IEdge); + const int KMax = MaxLayerEdgeTop(IEdge); + + parallelForInner( + Team, Range{KMin, KMax}, INNER_LAMBDA(int K) { + Tend(IEdge, K) -= EdgeMask(IEdge, K) * Gravity * + (SshCell(ICell1) - SshCell(ICell0)) * InvDcEdge; + }); } private: @@ -210,12 +227,11 @@ class VelocityDiffusionOnEdge { /// The functor takes edge index, vertical chunk index, and arrays for /// divergence of horizontal velocity (defined at cell centers) and relative /// vorticity (defined at vertices), outputs tendency array - KOKKOS_FUNCTION void operator()(const Array2DReal &Tend, I4 IEdge, I4 KChunk, + KOKKOS_FUNCTION void operator()(const TeamMember &Team, + const Array2DReal &Tend, I4 IEdge, const Array2DReal &DivCell, const Array2DReal &RVortVertex) const { - const I4 KStart = chunkStart(KChunk, MinLayerEdgeBot(IEdge)); - const I4 KLen = chunkLength(KChunk, KStart, MaxLayerEdgeTop(IEdge)); const I4 ICell0 = CellsOnEdge(IEdge, 0); const I4 ICell1 = CellsOnEdge(IEdge, 1); @@ -225,16 +241,19 @@ class VelocityDiffusionOnEdge { const Real DcEdgeInv = 1._Real / DcEdge(IEdge); const Real DvEdgeInv = 1._Real / DvEdge(IEdge); - for (int KVec = 0; KVec < KLen; ++KVec) { - const I4 K = KStart + KVec; - const Real Del2U = - ((DivCell(ICell1, K) - DivCell(ICell0, K)) * DcEdgeInv - - (RVortVertex(IVertex1, K) - RVortVertex(IVertex0, K)) * - DvEdgeInv); + const int KMin = MinLayerEdgeBot(IEdge); + const int KMax = MaxLayerEdgeTop(IEdge); - Tend(IEdge, K) += - EdgeMask(IEdge, K) * ViscDel2 * MeshScalingDel2(IEdge) * Del2U; - } + parallelForInner( + Team, Range{KMin, KMax}, INNER_LAMBDA(int K) { + const Real Del2U = + ((DivCell(ICell1, K) - DivCell(ICell0, K)) * DcEdgeInv - + (RVortVertex(IVertex1, K) - RVortVertex(IVertex0, K)) * + DvEdgeInv); + + Tend(IEdge, K) += + EdgeMask(IEdge, K) * ViscDel2 * MeshScalingDel2(IEdge) * Del2U; + }); } private: @@ -262,12 +281,11 @@ class VelocityHyperDiffOnEdge { /// The functor takes the edge index, vertical chunk index, and arrays for /// the laplacian of divergence of horizontal velocity and the laplacian of /// the relative vorticity, outputs tendency array - KOKKOS_FUNCTION void operator()(const Array2DReal &Tend, I4 IEdge, I4 KChunk, + KOKKOS_FUNCTION void operator()(const TeamMember &Team, + const Array2DReal &Tend, I4 IEdge, const Array2DReal &Del2DivCell, const Array2DReal &Del2RVortVertex) const { - const I4 KStart = chunkStart(KChunk, MinLayerEdgeBot(IEdge)); - const I4 KLen = chunkLength(KChunk, KStart, MaxLayerEdgeTop(IEdge)); const I4 ICell0 = CellsOnEdge(IEdge, 0); const I4 ICell1 = CellsOnEdge(IEdge, 1); @@ -277,17 +295,22 @@ class VelocityHyperDiffOnEdge { const Real DcEdgeInv = 1._Real / DcEdge(IEdge); const Real DvEdgeInv = 1._Real / DvEdge(IEdge); - for (int KVec = 0; KVec < KLen; ++KVec) { - const I4 K = KStart + KVec; - const Real Del2U = - (DivFactor * (Del2DivCell(ICell1, K) - Del2DivCell(ICell0, K)) * - DcEdgeInv - - (Del2RVortVertex(IVertex1, K) - Del2RVortVertex(IVertex0, K)) * - DvEdgeInv); - - Tend(IEdge, K) -= - EdgeMask(IEdge, K) * ViscDel4 * MeshScalingDel4(IEdge) * Del2U; - } + const int KMin = MinLayerEdgeBot(IEdge); + const int KMax = MaxLayerEdgeTop(IEdge); + + parallelForInner( + Team, Range{KMin, KMax}, INNER_LAMBDA(int K) { + const Real Del2U = + (DivFactor * + (Del2DivCell(ICell1, K) - Del2DivCell(ICell0, K)) * + DcEdgeInv - + (Del2RVortVertex(IVertex1, K) - + Del2RVortVertex(IVertex0, K)) * + DvEdgeInv); + + Tend(IEdge, K) -= + EdgeMask(IEdge, K) * ViscDel4 * MeshScalingDel4(IEdge) * Del2U; + }); } private: @@ -311,21 +334,27 @@ class SfcStressForcingOnEdge { /// The functor takes the edge index, vertical chunk index, and arrays for /// normal surface stress and edge pseudo-thickness, outputs tendency array - KOKKOS_FUNCTION void operator()(const Array2DReal &Tend, I4 IEdge, I4 KChunk, + KOKKOS_FUNCTION void operator()(const TeamMember &Team, + const Array2DReal &Tend, I4 IEdge, const Array1DReal &NormalStressEdge, const Array2DReal &PseudoThickEdge) const { - if (KChunk == 0) { - const I4 K = MinLayerEdgeBot(IEdge); - - const Real InvThickEdge = 1._Real / PseudoThickEdge(IEdge, K); - Tend(IEdge, K) += EdgeMask(IEdge, K) * InvThickEdge * - NormalStressEdge(IEdge) / RhoSw; - } + const I4 KMin = MinLayerEdgeBot(IEdge); + const I4 KMax = MaxLayerEdgeTop(IEdge); + + parallelForInner( + Team, Range{KMin, KMax}, INNER_LAMBDA(int K) { + if (K == KMin) { + const Real InvThickEdge = 1._Real / PseudoThickEdge(IEdge, K); + Tend(IEdge, K) += EdgeMask(IEdge, K) * InvThickEdge * + NormalStressEdge(IEdge) / RhoSw; + } + }); } private: Array2DReal EdgeMask; Array1DI4 MinLayerEdgeBot; + Array1DI4 MaxLayerEdgeTop; }; /// Bottom drag @@ -381,62 +410,97 @@ class TracerHorzAdvOnCell { Real Coef3rdOrder = 0.25; TracerHorzAdvOnCell(const HorzMesh *Mesh, const VertCoord *VCoord); void init(); - KOKKOS_FUNCTION void operator()(const I4 L, const I4 IEdge, const I4 KChunk, + KOKKOS_FUNCTION void operator()(const TeamMember &Team, const I4 L, + const I4 IEdge, const Array3DReal &TracerCell, const Array2DReal &FluxPseudoThickEdge, const Array2DReal &NormVelEdge) const { - const I4 KStart = KChunk * VecLength; - const I4 KEnd = KStart + VecLength; - for (int K = KStart; K < KEnd; ++K) - HighOrderFlxHorz(L, IEdge, K) = 0; + + ScratchArray1DReal FlxTmp(teamScratch(Team), NVertLayers); + + const auto LTracerCell = + subviewUnmanaged(TracerCell, L, Kokkos::ALL, Kokkos::ALL); + + const int KMin = MinLayerEdgeBot(IEdge); + const int KMax = MaxLayerEdgeTop(IEdge); + + parallelForInner( + Team, NVertLayers, INNER_LAMBDA(int K) { FlxTmp(K) = 0; }); // Stay at low order at boundaries - for (int K = KStart; K < KEnd; ++K) { - const I4 JCell0 = CellsOnEdge(IEdge, 0); - const I4 JCell1 = CellsOnEdge(IEdge, 1); - const Real NormalThicknessFlux = - FluxPseudoThickEdge(IEdge, K) * NormVelEdge(IEdge, K); - const Real TracerWgt = DvEdge(IEdge) * 0.5_Real * NormalThicknessFlux; - HighOrderFlxHorz(L, IEdge, K) += - TracerWgt * (1._Real - AdvMaskHighOrder(IEdge, K)) * - (TracerCell(L, JCell1, K) + TracerCell(L, JCell0, K)); - } + parallelForInner( + Team, Range{KMin, KMax}, INNER_LAMBDA(int K) { + const I4 JCell0 = CellsOnEdge(IEdge, 0); + const I4 JCell1 = CellsOnEdge(IEdge, 1); + const Real NormalThicknessFlux = + FluxPseudoThickEdge(IEdge, K) * NormVelEdge(IEdge, K); + const Real TracerWgt = + DvEdge(IEdge) * 0.5_Real * NormalThicknessFlux; + FlxTmp(K) += TracerWgt * (1._Real - AdvMaskHighOrder(IEdge, K)) * + (LTracerCell(JCell1, K) + LTracerCell(JCell0, K)); + }); // High order (3rd or 4th) fluxes elsewhere when requested // - If HorzTracerFluxOrder = 2, NAdvCellsForEdge = 0 and // this loop is skipped. for (int I = 0; I < NAdvCellsForEdge(IEdge); ++I) { const I4 ICell = AdvCellsForEdge(IEdge, I); - for (int K = KStart; K < KEnd; ++K) { - const Real NormalThicknessFlux = - FluxPseudoThickEdge(IEdge, K) * NormVelEdge(IEdge, K); - const Real TracerWgt = - (AdvCoefs(I, IEdge) + - Coef3rdOrder * std::copysign(1._Real, NormalThicknessFlux) * - AdvCoefs3rd(I, IEdge)) * - NormalThicknessFlux; - HighOrderFlxHorz(L, IEdge, K) += TracerWgt * - TracerCell(L, ICell, K) * - AdvMaskHighOrder(IEdge, K); - } + parallelForInner( + Team, Range{KMin, KMax}, INNER_LAMBDA(int K) { + const Real NormalThicknessFlux = + FluxPseudoThickEdge(IEdge, K) * NormVelEdge(IEdge, K); + const Real TracerWgt = + (AdvCoefs(I, IEdge) + + Coef3rdOrder * + std::copysign(1._Real, NormalThicknessFlux) * + AdvCoefs3rd(I, IEdge)) * + NormalThicknessFlux; + FlxTmp(K) += TracerWgt * LTracerCell(ICell, K) * + AdvMaskHighOrder(IEdge, K); + }); } + + const auto LHighOrderFlxHorz = + subviewUnmanaged(HighOrderFlxHorz, L, Kokkos::ALL, Kokkos::ALL); + + parallelForInner( + Team, Range{KMin, KMax}, + INNER_LAMBDA(int K) { LHighOrderFlxHorz(IEdge, K) = FlxTmp(K); }); } - KOKKOS_FUNCTION void operator()(const Array3DReal &Tend, const I4 L, - const I4 ICell, const I4 KChunk) const { - const I4 KStart = KChunk * VecLength; - const I4 KEnd = KStart + VecLength; + KOKKOS_FUNCTION void operator()(const TeamMember &Team, + const Array3DReal &Tend, const I4 L, + const I4 ICell) const { + + const auto LTend = subviewUnmanaged(Tend, L, Kokkos::ALL, Kokkos::ALL); + const auto LHighOrderFlxHorz = + subviewUnmanaged(HighOrderFlxHorz, L, Kokkos::ALL, Kokkos::ALL); + const Real InvAreaCell = 1._Real / AreaCell(ICell); - for (int K = KStart; K < KEnd; ++K) - Tend(L, ICell, K) = 0; + + ScratchArray1DReal TendTmp(teamScratch(Team), NVertLayers); + parallelForInner( + Team, NVertLayers, INNER_LAMBDA(int K) { TendTmp(K) = 0; }); for (int I = 0; I < NEdgesOnCell(ICell); ++I) { const I4 IEdge = EdgesOnCell(ICell, I); - for (int K = KStart; K < KEnd; ++K) { - Tend(L, ICell, K) += EdgeSignOnCell(ICell, I) * - HighOrderFlxHorz(L, IEdge, K) * InvAreaCell; - } + + const int MinLyrEdgeBot = MinLayerEdgeBot(IEdge); + const int MaxLyrEdgeTop = MaxLayerEdgeTop(IEdge); + + parallelForInner( + Team, Range{MinLyrEdgeBot, MaxLyrEdgeTop}, INNER_LAMBDA(int K) { + TendTmp(K) += EdgeSignOnCell(ICell, I) * + LHighOrderFlxHorz(IEdge, K) * InvAreaCell; + }); } + + const int MinLyrCell = MinLayerCell(ICell); + const int MaxLyrCell = MaxLayerCell(ICell); + + parallelForInner( + Team, Range{MinLyrCell, MaxLyrCell}, + INNER_LAMBDA(int K) { LTend(ICell, K) += TendTmp(K); }); } private: @@ -455,6 +519,12 @@ class TracerHorzAdvOnCell { Array2DReal EdgeSignOnCell; Array1DReal DvEdge; Array1DReal AreaCell; + + I4 NVertLayers; + Array1DI4 MinLayerCell; + Array1DI4 MaxLayerCell; + Array1DI4 MinLayerEdgeBot; + Array1DI4 MaxLayerEdgeTop; }; // Tracer horizontal diffusion term @@ -467,21 +537,22 @@ class TracerDiffOnCell { TracerDiffOnCell(const HorzMesh *Mesh, const VertCoord *VCoord); KOKKOS_FUNCTION void - operator()(const Array3DReal &Tend, I4 L, I4 ICell, I4 KChunk, + operator()(const TeamMember &Team, const Array3DReal &Tend, I4 L, I4 ICell, const Array3DReal &TracerCell, const Array2DReal &MeanPseudoThickEdge) const { - const I4 KStartCell = chunkStart(KChunk, MinLayerCell(ICell)); - const I4 KLenCell = chunkLength(KChunk, KStartCell, MaxLayerCell(ICell)); - const I4 KEndCell = KStartCell + KLenCell - 1; + const auto LTend = subviewUnmanaged(Tend, L, Kokkos::ALL, Kokkos::ALL); + const auto LTracerCell = + subviewUnmanaged(TracerCell, L, Kokkos::ALL, Kokkos::ALL); + const Real InvAreaCell = 1._Real / AreaCell(ICell); - Real DiffTmp[VecLength] = {0}; + ScratchArray1DReal DiffTmp(teamScratch(Team), NVertLayers); + parallelForInner( + Team, NVertLayers, INNER_LAMBDA(int K) { DiffTmp(K) = 0; }); for (int J = 0; J < NEdgesOnCell(ICell); ++J) { - const I4 JEdge = EdgesOnCell(ICell, J); - const I4 KStartEdge = Kokkos::max(KStartCell, MinLayerEdgeBot(JEdge)); - const I4 KEndEdge = Kokkos::min(KEndCell, MaxLayerEdgeTop(JEdge)); + const I4 JEdge = EdgesOnCell(ICell, J); const I4 JCell0 = CellsOnEdge(JEdge, 0); const I4 JCell1 = CellsOnEdge(JEdge, 1); @@ -489,19 +560,26 @@ class TracerDiffOnCell { const Real RTemp = MeshScalingDel2(JEdge) * DvEdge(JEdge) / DcEdge(JEdge); - for (int K = KStartEdge; K <= KEndEdge; ++K) { - const I4 KVec = K - KStartCell; - const Real TracerGrad = - (TracerCell(L, JCell1, K) - TracerCell(L, JCell0, K)); + const int MinLyrEdgeBot = MinLayerEdgeBot(JEdge); + const int MaxLyrEdgeTop = MaxLayerEdgeTop(JEdge); - DiffTmp[KVec] -= EdgeMask(JEdge, K) * EdgeSignOnCell(ICell, J) * - RTemp * MeanPseudoThickEdge(JEdge, K) * TracerGrad; - } - } - for (int KVec = 0; KVec < KLenCell; ++KVec) { - const I4 K = KStartCell + KVec; - Tend(L, ICell, K) += EddyDiff2 * DiffTmp[KVec] * InvAreaCell; + parallelForInner( + Team, Range{MinLyrEdgeBot, MaxLyrEdgeTop}, INNER_LAMBDA(int K) { + const Real TracerGrad = + (LTracerCell(JCell1, K) - LTracerCell(JCell0, K)); + + DiffTmp(K) -= EdgeMask(JEdge, K) * EdgeSignOnCell(ICell, J) * + RTemp * MeanPseudoThickEdge(JEdge, K) * + TracerGrad; + }); } + const int MinLyrCell = MinLayerCell(ICell); + const int MaxLyrCell = MaxLayerCell(ICell); + + parallelForInner( + Team, Range{MinLyrCell, MaxLyrCell}, INNER_LAMBDA(int K) { + LTend(ICell, K) += EddyDiff2 * DiffTmp(K) * InvAreaCell; + }); } private: @@ -513,6 +591,7 @@ class TracerDiffOnCell { Array1DReal DcEdge; Array1DReal AreaCell; Array1DReal MeshScalingDel2; + I4 NVertLayers; Array2DReal EdgeMask; Array1DI4 MinLayerCell; Array1DI4 MaxLayerCell; @@ -529,21 +608,22 @@ class TracerHyperDiffOnCell { TracerHyperDiffOnCell(const HorzMesh *Mesh, const VertCoord *VCoord); - KOKKOS_FUNCTION void operator()(const Array3DReal &Tend, I4 L, I4 ICell, - I4 KChunk, + KOKKOS_FUNCTION void operator()(const TeamMember &Team, + const Array3DReal &Tend, I4 L, I4 ICell, const Array3DReal &TrDel2Cell) const { - const I4 KStartCell = chunkStart(KChunk, MinLayerCell(ICell)); - const I4 KLenCell = chunkLength(KChunk, KStartCell, MaxLayerCell(ICell)); - const I4 KEndCell = KStartCell + KLenCell - 1; + const auto LTend = subviewUnmanaged(Tend, L, Kokkos::ALL, Kokkos::ALL); + const auto LTrDel2Cell = + subviewUnmanaged(TrDel2Cell, L, Kokkos::ALL, Kokkos::ALL); + const Real InvAreaCell = 1._Real / AreaCell(ICell); - Real HypTmp[VecLength] = {0}; + ScratchArray1DReal HypTmp(teamScratch(Team), NVertLayers); + parallelForInner( + Team, NVertLayers, INNER_LAMBDA(int K) { HypTmp(K) = 0; }); for (int J = 0; J < NEdgesOnCell(ICell); ++J) { - const I4 JEdge = EdgesOnCell(ICell, J); - const I4 KStartEdge = Kokkos::max(KStartCell, MinLayerEdgeBot(JEdge)); - const I4 KEndEdge = Kokkos::min(KEndCell, MaxLayerEdgeTop(JEdge)); + const I4 JEdge = EdgesOnCell(ICell, J); const I4 JCell0 = CellsOnEdge(JEdge, 0); const I4 JCell1 = CellsOnEdge(JEdge, 1); @@ -551,19 +631,25 @@ class TracerHyperDiffOnCell { const Real RTemp = MeshScalingDel4(JEdge) * DvEdge(JEdge) / DcEdge(JEdge); - for (int K = KStartEdge; K <= KEndEdge; ++K) { - const I4 KVec = K - KStartCell; - const Real Del2TrGrad = - (TrDel2Cell(L, JCell1, K) - TrDel2Cell(L, JCell0, K)); + const int MinLyrEdgeBot = MinLayerEdgeBot(JEdge); + const int MaxLyrEdgeTop = MaxLayerEdgeTop(JEdge); - HypTmp[KVec] -= EdgeMask(JEdge, K) * EdgeSignOnCell(ICell, J) * - RTemp * Del2TrGrad; - } - } - for (int KVec = 0; KVec < KLenCell; ++KVec) { - const I4 K = KStartCell + KVec; - Tend(L, ICell, K) -= EddyDiff4 * HypTmp[KVec] * InvAreaCell; + parallelForInner( + Team, Range{MinLyrEdgeBot, MaxLyrEdgeTop}, INNER_LAMBDA(int K) { + const Real Del2TrGrad = + (LTrDel2Cell(JCell1, K) - LTrDel2Cell(JCell0, K)); + + HypTmp(K) -= EdgeMask(JEdge, K) * EdgeSignOnCell(ICell, J) * + RTemp * Del2TrGrad; + }); } + const int MinLyrCell = MinLayerCell(ICell); + const int MaxLyrCell = MaxLayerCell(ICell); + + parallelForInner( + Team, Range{MinLyrCell, MaxLyrCell}, INNER_LAMBDA(int K) { + LTend(ICell, K) -= EddyDiff4 * HypTmp(K) * InvAreaCell; + }); } private: @@ -575,6 +661,7 @@ class TracerHyperDiffOnCell { Array1DReal DcEdge; Array1DReal AreaCell; Array1DReal MeshScalingDel4; + I4 NVertLayers; Array2DReal EdgeMask; Array1DI4 MinLayerCell; Array1DI4 MaxLayerCell; diff --git a/components/omega/src/ocn/VertMix.cpp b/components/omega/src/ocn/VertMix.cpp index d9ab375387f6..35398e3ecf8a 100644 --- a/components/omega/src/ocn/VertMix.cpp +++ b/components/omega/src/ocn/VertMix.cpp @@ -701,17 +701,11 @@ void VertMix::VertMixImplicit(OceanState *State, AuxiliaryState *AuxState, Pacer::start("VertMix:computeKineticAuxForBottomDrag", 2); parallelForOuter( - "computeKineticAuxForBottomDrag", {Mesh->NCellsAll}, + "computeKineticAuxForBottomDrag", + LaunchConfig({Mesh->NCellsAll}, + TeamScratch(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("VertMix:computeKineticAuxForBottomDrag", 2); } diff --git a/components/omega/src/ocn/auxiliaryVars/KineticAuxVars.cpp b/components/omega/src/ocn/auxiliaryVars/KineticAuxVars.cpp index 5347699acdcc..f3e38428b9c2 100644 --- a/components/omega/src/ocn/auxiliaryVars/KineticAuxVars.cpp +++ b/components/omega/src/ocn/auxiliaryVars/KineticAuxVars.cpp @@ -15,7 +15,8 @@ KineticAuxVars::KineticAuxVars(const std::string &AuxStateSuffix, NEdgesOnCell(Mesh->NEdgesOnCell), EdgesOnCell(Mesh->EdgesOnCell), EdgeSignOnCell(Mesh->EdgeSignOnCell), DcEdge(Mesh->DcEdge), DvEdge(Mesh->DvEdge), AreaCell(Mesh->AreaCell), - MinLayerCell(VCoord->MinLayerCell), MaxLayerCell(VCoord->MaxLayerCell) {} + NVertLayers(VCoord->NVertLayers), MinLayerCell(VCoord->MinLayerCell), + MaxLayerCell(VCoord->MaxLayerCell) {} void KineticAuxVars::registerFields( const std::string &AuxGroupName, // name of Auxiliary field group diff --git a/components/omega/src/ocn/auxiliaryVars/KineticAuxVars.h b/components/omega/src/ocn/auxiliaryVars/KineticAuxVars.h index 850807702f64..34f36ebaf7fb 100644 --- a/components/omega/src/ocn/auxiliaryVars/KineticAuxVars.h +++ b/components/omega/src/ocn/auxiliaryVars/KineticAuxVars.h @@ -19,35 +19,42 @@ class KineticAuxVars { const VertCoord *VCoord); KOKKOS_FUNCTION void - computeVarsOnCell(int ICell, int KChunk, + computeVarsOnCell(const TeamMember &Team, int ICell, const Array2DReal &NormalVelEdge) const { - const int KStart = chunkStart(KChunk, MinLayerCell(ICell)); - const int KLen = chunkLength(KChunk, KStart, MaxLayerCell(ICell)); - const Real InvAreaCell = 1._Real / AreaCell(ICell); - Real KineticEnergyCellTmp[VecLength] = {0}; - Real VelocityDivCellTmp[VecLength] = {0}; + ScratchArray1DReal KineticEnergyCellTmp(teamScratch(Team), NVertLayers); + ScratchArray1DReal VelocityDivCellTmp(teamScratch(Team), NVertLayers); + + parallelForInner( + Team, NVertLayers, INNER_LAMBDA(int K) { + KineticEnergyCellTmp(K) = 0; + VelocityDivCellTmp(K) = 0; + }); + + const int KMin = MinLayerCell(ICell); + const int KMax = MaxLayerCell(ICell); for (int J = 0; J < NEdgesOnCell(ICell); ++J) { const int JEdge = EdgesOnCell(ICell, J); const Real AreaEdge = 0.5_Real * DvEdge(JEdge) * DcEdge(JEdge); - for (int KVec = 0; KVec < KLen; ++KVec) { - const int K = KStart + KVec; - KineticEnergyCellTmp[KVec] += AreaEdge * 0.5_Real * InvAreaCell * - NormalVelEdge(JEdge, K) * - NormalVelEdge(JEdge, K); - VelocityDivCellTmp[KVec] -= DvEdge(JEdge) * InvAreaCell * - EdgeSignOnCell(ICell, J) * - NormalVelEdge(JEdge, K); - } - } - for (int KVec = 0; KVec < KLen; ++KVec) { - const int K = KStart + KVec; - KineticEnergyCell(ICell, K) = KineticEnergyCellTmp[KVec]; - VelocityDivCell(ICell, K) = VelocityDivCellTmp[KVec]; + parallelForInner( + Team, Range{KMin, KMax}, INNER_LAMBDA(int K) { + KineticEnergyCellTmp(K) += AreaEdge * 0.5_Real * InvAreaCell * + NormalVelEdge(JEdge, K) * + NormalVelEdge(JEdge, K); + VelocityDivCellTmp(K) -= DvEdge(JEdge) * InvAreaCell * + EdgeSignOnCell(ICell, J) * + NormalVelEdge(JEdge, K); + }); } + + parallelForInner( + Team, Range{KMin, KMax}, INNER_LAMBDA(int K) { + KineticEnergyCell(ICell, K) = KineticEnergyCellTmp(K); + VelocityDivCell(ICell, K) = VelocityDivCellTmp(K); + }); } void registerFields(const std::string &AuxGroupName, @@ -61,6 +68,7 @@ class KineticAuxVars { Array1DReal DcEdge; Array1DReal DvEdge; Array1DReal AreaCell; + I4 NVertLayers; Array1DI4 MinLayerCell; Array1DI4 MaxLayerCell; }; diff --git a/components/omega/src/ocn/auxiliaryVars/PseudoThicknessAuxVars.cpp b/components/omega/src/ocn/auxiliaryVars/PseudoThicknessAuxVars.cpp index 627f6cd3177c..ed1dcb802c67 100644 --- a/components/omega/src/ocn/auxiliaryVars/PseudoThicknessAuxVars.cpp +++ b/components/omega/src/ocn/auxiliaryVars/PseudoThicknessAuxVars.cpp @@ -17,6 +17,7 @@ PseudoThicknessAuxVars::PseudoThicknessAuxVars( AreaCell(Mesh->AreaCell), DvEdge(Mesh->DvEdge), NEdgesOnCell(Mesh->NEdgesOnCell), EdgesOnCell(Mesh->EdgesOnCell), EdgeSignOnCell(Mesh->EdgeSignOnCell), CellsOnEdge(Mesh->CellsOnEdge), + NVertLayers(VCoord->NVertLayers), MinLayerEdgeBot(VCoord->MinLayerEdgeBot), MaxLayerEdgeTop(VCoord->MaxLayerEdgeTop), MinLayerCell(VCoord->MinLayerCell), MaxLayerCell(VCoord->MaxLayerCell) {} diff --git a/components/omega/src/ocn/auxiliaryVars/PseudoThicknessAuxVars.h b/components/omega/src/ocn/auxiliaryVars/PseudoThicknessAuxVars.h index d24ebfd50e8b..682d27267945 100644 --- a/components/omega/src/ocn/auxiliaryVars/PseudoThicknessAuxVars.h +++ b/components/omega/src/ocn/auxiliaryVars/PseudoThicknessAuxVars.h @@ -24,79 +24,83 @@ class PseudoThicknessAuxVars { const HorzMesh *Mesh, const VertCoord *VCoord); KOKKOS_FUNCTION void - computeVarsOnEdge(int IEdge, int KChunk, const Array2DReal &PseudoThickCell, + computeVarsOnEdge(const TeamMember &Team, int IEdge, + const Array2DReal &PseudoThickCell, const Array2DReal &NormalVelEdge) const { - const int KStart = chunkStart(KChunk, MinLayerEdgeBot(IEdge)); - const int KLen = chunkLength(KChunk, KStart, MaxLayerEdgeTop(IEdge)); const int JCell0 = CellsOnEdge(IEdge, 0); const int JCell1 = CellsOnEdge(IEdge, 1); - for (int KVec = 0; KVec < KLen; ++KVec) { - const int K = KStart + KVec; - MeanPseudoThickEdge(IEdge, K) = - 0.5_Real * - (PseudoThickCell(JCell0, K) + PseudoThickCell(JCell1, K)); - } + const int KMin = MinLayerEdgeBot(IEdge); + const int KMax = MaxLayerEdgeTop(IEdge); + + parallelForInner( + Team, Range{KMin, KMax}, INNER_LAMBDA(int K) { + MeanPseudoThickEdge(IEdge, K) = + 0.5_Real * + (PseudoThickCell(JCell0, K) + PseudoThickCell(JCell1, K)); + }); switch (FluxThickEdgeChoice) { case FluxThickEdgeOption::Center: - for (int KVec = 0; KVec < KLen; ++KVec) { - const int K = KStart + KVec; - FluxPseudoThickEdge(IEdge, K) = - 0.5_Real * - (PseudoThickCell(JCell0, K) + PseudoThickCell(JCell1, K)); - } + parallelForInner( + Team, Range{KMin, KMax}, INNER_LAMBDA(int K) { + FluxPseudoThickEdge(IEdge, K) = + 0.5_Real * + (PseudoThickCell(JCell0, K) + PseudoThickCell(JCell1, K)); + }); break; case FluxThickEdgeOption::Upwind: - for (int KVec = 0; KVec < KLen; ++KVec) { - const int K = KStart + KVec; - if (NormalVelEdge(IEdge, K) > 0) { - FluxPseudoThickEdge(IEdge, K) = PseudoThickCell(JCell0, K); - } else if (NormalVelEdge(IEdge, K) < 0) { - FluxPseudoThickEdge(IEdge, K) = PseudoThickCell(JCell1, K); - } else { - FluxPseudoThickEdge(IEdge, K) = Kokkos::max( - PseudoThickCell(JCell0, K), PseudoThickCell(JCell1, K)); - } - } + parallelForInner( + Team, Range{KMin, KMax}, INNER_LAMBDA(int K) { + if (NormalVelEdge(IEdge, K) > 0) { + FluxPseudoThickEdge(IEdge, K) = PseudoThickCell(JCell0, K); + } else if (NormalVelEdge(IEdge, K) < 0) { + FluxPseudoThickEdge(IEdge, K) = PseudoThickCell(JCell1, K); + } else { + FluxPseudoThickEdge(IEdge, K) = Kokkos::max( + PseudoThickCell(JCell0, K), PseudoThickCell(JCell1, K)); + } + }); break; } } - KOKKOS_FUNCTION void computeVarsOnCells(int ICell, int KChunk, + KOKKOS_FUNCTION void computeVarsOnCells(const TeamMember &Team, int ICell, const Array2DReal &PseudoThickCell, const Array2DReal &NormalVelEdge, const Real Dt) const { - const I4 KStartCell = chunkStart(KChunk, MinLayerCell(ICell)); - const I4 KLenCell = chunkLength(KChunk, KStartCell, MaxLayerCell(ICell)); - const I4 KEndCell = KStartCell + KLenCell - 1; + // Temporary for stacked shallow water + const int MinLyrCell = MinLayerCell(ICell); + const int MaxLyrCell = MaxLayerCell(ICell); - Real TmpProv[VecLength] = {0.}; + ScratchArray1DReal TmpProv(teamScratch(Team), NVertLayers); + parallelForInner( + Team, NVertLayers, INNER_LAMBDA(int K) { TmpProv(K) = 0; }); Real DtInvAreaCell = Dt / AreaCell(ICell); for (int J = 0; J < NEdgesOnCell(ICell); ++J) { - const I4 JEdge = EdgesOnCell(ICell, J); + const int JEdge = EdgesOnCell(ICell, J); - const I4 KStartEdge = Kokkos::max(KStartCell, MinLayerEdgeBot(JEdge)); - const I4 KEndEdge = Kokkos::min(KEndCell, MaxLayerEdgeTop(JEdge)); + const int MinLyrEdgeBot = MinLayerEdgeBot(JEdge); + const int MaxLyrEdgeTop = MaxLayerEdgeTop(JEdge); const Real Factor = DtInvAreaCell * DvEdge(JEdge) * EdgeSignOnCell(ICell, J); - for (int K = KStartEdge; K <= KEndEdge; ++K) { - const I4 KVec = K - KStartCell; - TmpProv[KVec] += Factor * FluxPseudoThickEdge(JEdge, K) * - NormalVelEdge(JEdge, K); - } + parallelForInner( + Team, Range{MinLyrEdgeBot, MaxLyrEdgeTop}, INNER_LAMBDA(int K) { + TmpProv(K) += Factor * FluxPseudoThickEdge(JEdge, K) * + NormalVelEdge(JEdge, K); + }); } - for (int KVec = 0; KVec < KLenCell; ++KVec) { - const int K = KStartCell + KVec; - ProvPseudoThickness(ICell, K) = - PseudoThickCell(ICell, K) + TmpProv[KVec]; - } + parallelForInner( + Team, Range{MinLyrCell, MaxLyrCell}, INNER_LAMBDA(int K) { + ProvPseudoThickness(ICell, K) = + PseudoThickCell(ICell, K) + TmpProv(K); + }); } void registerFields(const std::string &AuxGroupName, @@ -110,6 +114,7 @@ class PseudoThicknessAuxVars { Array2DI4 EdgesOnCell; Array2DReal EdgeSignOnCell; Array2DI4 CellsOnEdge; + I4 NVertLayers; Array1DI4 MinLayerEdgeBot; Array1DI4 MaxLayerEdgeTop; Array1DI4 MinLayerCell; diff --git a/components/omega/src/ocn/auxiliaryVars/TracerAuxVars.cpp b/components/omega/src/ocn/auxiliaryVars/TracerAuxVars.cpp index 6cbaffdf4dce..9f4c61ca0548 100644 --- a/components/omega/src/ocn/auxiliaryVars/TracerAuxVars.cpp +++ b/components/omega/src/ocn/auxiliaryVars/TracerAuxVars.cpp @@ -13,7 +13,8 @@ TracerAuxVars::TracerAuxVars(const std::string &AuxStateSuffix, NEdgesOnCell(Mesh->NEdgesOnCell), EdgesOnCell(Mesh->EdgesOnCell), CellsOnEdge(Mesh->CellsOnEdge), EdgeSignOnCell(Mesh->EdgeSignOnCell), DcEdge(Mesh->DcEdge), DvEdge(Mesh->DvEdge), AreaCell(Mesh->AreaCell), - EdgeMask(VCoord->EdgeMask), MinLayerEdgeBot(VCoord->MinLayerEdgeBot), + NVertLayers(VCoord->NVertLayers), EdgeMask(VCoord->EdgeMask), + MinLayerEdgeBot(VCoord->MinLayerEdgeBot), MaxLayerEdgeTop(VCoord->MaxLayerEdgeTop), MinLayerCell(VCoord->MinLayerCell), MaxLayerCell(VCoord->MaxLayerCell) {} diff --git a/components/omega/src/ocn/auxiliaryVars/TracerAuxVars.h b/components/omega/src/ocn/auxiliaryVars/TracerAuxVars.h index bbd0b81f55c4..7a5c13afedb0 100644 --- a/components/omega/src/ocn/auxiliaryVars/TracerAuxVars.h +++ b/components/omega/src/ocn/auxiliaryVars/TracerAuxVars.h @@ -19,17 +19,21 @@ class TracerAuxVars { const VertCoord *VCoord, const I4 NTracers); KOKKOS_FUNCTION void - computeVarsOnCells(int L, int ICell, int KChunk, + computeVarsOnCells(const TeamMember &Team, int L, int ICell, const Array2DReal &MeanPseudoThickEdge, const Array3DReal &TrCell) const { - const int KStartCell = chunkStart(KChunk, MinLayerCell(ICell)); - const int KLenCell = chunkLength(KChunk, KStartCell, MaxLayerCell(ICell)); - const int KEndCell = KStartCell + KLenCell - 1; + const auto LTrCell = + subviewUnmanaged(TrCell, L, Kokkos::ALL, Kokkos::ALL); + const auto LDel2TracersCell = + subviewUnmanaged(Del2TracersCell, L, Kokkos::ALL, Kokkos::ALL); const Real InvAreaCell = 1._Real / AreaCell(ICell); - Real Del2TrCellTmp[VecLength] = {0}; + ScratchArray1DReal Del2TrCellTmp(teamScratch(Team), NVertLayers); + + parallelForInner( + Team, NVertLayers, INNER_LAMBDA(int K) { Del2TrCellTmp(K) = 0; }); for (int J = 0; J < NEdgesOnCell(ICell); ++J) { const int JEdge = EdgesOnCell(ICell, J); @@ -39,21 +43,25 @@ class TracerAuxVars { const Real DvDcEdge = DvEdge(JEdge) / DcEdge(JEdge); - const int KStartEdge = Kokkos::max(KStartCell, MinLayerEdgeBot(JEdge)); - const int KEndEdge = Kokkos::min(KEndCell, MaxLayerEdgeTop(JEdge)); + const int MinLyrEdgeBot = MinLayerEdgeBot(JEdge); + const int MaxLyrEdgeTop = MaxLayerEdgeTop(JEdge); - for (int K = KStartEdge; K <= KEndEdge; ++K) { - const int KVec = K - KStartCell; - const Real TracerGrad = TrCell(L, JCell1, K) - TrCell(L, JCell0, K); - Del2TrCellTmp[KVec] -= EdgeMask(JEdge, K) * - EdgeSignOnCell(ICell, J) * DvDcEdge * - MeanPseudoThickEdge(JEdge, K) * TracerGrad; - } - } - for (int KVec = 0; KVec < KLenCell; ++KVec) { - const int K = KStartCell + KVec; - Del2TracersCell(L, ICell, K) = Del2TrCellTmp[KVec] * InvAreaCell; + parallelForInner( + Team, Range{MinLyrEdgeBot, MaxLyrEdgeTop}, INNER_LAMBDA(int K) { + const Real TracerGrad = LTrCell(JCell1, K) - LTrCell(JCell0, K); + Del2TrCellTmp(K) -= EdgeMask(JEdge, K) * + EdgeSignOnCell(ICell, J) * DvDcEdge * + MeanPseudoThickEdge(JEdge, K) * TracerGrad; + }); } + + const int MinLyrCell = MinLayerCell(ICell); + const int MaxLyrCell = MaxLayerCell(ICell); + + parallelForInner( + Team, Range{MinLyrCell, MaxLyrCell}, INNER_LAMBDA(int K) { + LDel2TracersCell(ICell, K) = Del2TrCellTmp(K) * InvAreaCell; + }); } void registerFields(const std::string &AuxGroupName, @@ -69,6 +77,7 @@ class TracerAuxVars { Array1DReal DvEdge; Array1DReal AreaCell; Array2DReal EdgeMask; + I4 NVertLayers; Array1DI4 MinLayerEdgeBot; Array1DI4 MaxLayerEdgeTop; Array1DI4 MinLayerCell; diff --git a/components/omega/src/ocn/auxiliaryVars/VelocityDel2AuxVars.cpp b/components/omega/src/ocn/auxiliaryVars/VelocityDel2AuxVars.cpp index e67da008c01b..217e5cb1ace2 100644 --- a/components/omega/src/ocn/auxiliaryVars/VelocityDel2AuxVars.cpp +++ b/components/omega/src/ocn/auxiliaryVars/VelocityDel2AuxVars.cpp @@ -22,6 +22,7 @@ VelocityDel2AuxVars::VelocityDel2AuxVars(const std::string &AuxStateSuffix, VerticesOnEdge(Mesh->VerticesOnEdge), EdgeMask(VCoord->EdgeMask), EdgeSignOnVertex(Mesh->EdgeSignOnVertex), AreaTriangle(Mesh->AreaTriangle), VertexDegree(Mesh->VertexDegree), + NVertLayers(VCoord->NVertLayers), MinLayerEdgeBot(VCoord->MinLayerEdgeBot), MaxLayerEdgeTop(VCoord->MaxLayerEdgeTop), MinLayerEdgeTop(VCoord->MinLayerEdgeTop), diff --git a/components/omega/src/ocn/auxiliaryVars/VelocityDel2AuxVars.h b/components/omega/src/ocn/auxiliaryVars/VelocityDel2AuxVars.h index 3e9e57782536..16702602c240 100644 --- a/components/omega/src/ocn/auxiliaryVars/VelocityDel2AuxVars.h +++ b/components/omega/src/ocn/auxiliaryVars/VelocityDel2AuxVars.h @@ -20,10 +20,9 @@ class VelocityDel2AuxVars { const VertCoord *VCoord); KOKKOS_FUNCTION void - computeVarsOnEdge(int IEdge, int KChunk, const Array2DReal &VelocityDivCell, + computeVarsOnEdge(const TeamMember &Team, int IEdge, + const Array2DReal &VelocityDivCell, const Array2DReal &RelVortVertex) const { - const int KStart = chunkStart(KChunk, MinLayerEdgeBot(IEdge)); - const int KLen = chunkLength(KChunk, KStart, MaxLayerEdgeTop(IEdge)); const int JCell0 = CellsOnEdge(IEdge, 0); const int JCell1 = CellsOnEdge(IEdge, 1); @@ -34,47 +33,54 @@ class VelocityDel2AuxVars { const Real InvDvEdge = 1._Real / Kokkos::max(DvEdge(IEdge), 0.25_Real * DcEdge(IEdge)); - for (int KVec = 0; KVec < KLen; ++KVec) { - const int K = KStart + KVec; - const Real GradDiv = - (VelocityDivCell(JCell1, K) - VelocityDivCell(JCell0, K)) * - InvDcEdge; - const Real CurlVort = - -(RelVortVertex(JVertex1, K) - RelVortVertex(JVertex0, K)) * - InvDvEdge; - Del2Edge(IEdge, K) = EdgeMask(IEdge, K) * GradDiv + CurlVort; - } + const int KMin = MinLayerEdgeBot(IEdge); + const int KMax = MaxLayerEdgeTop(IEdge); + + parallelForInner( + Team, Range{KMin, KMax}, INNER_LAMBDA(int K) { + const Real GradDiv = + (VelocityDivCell(JCell1, K) - VelocityDivCell(JCell0, K)) * + InvDcEdge; + const Real CurlVort = + -(RelVortVertex(JVertex1, K) - RelVortVertex(JVertex0, K)) * + InvDvEdge; + Del2Edge(IEdge, K) = EdgeMask(IEdge, K) * GradDiv + CurlVort; + }); } - KOKKOS_FUNCTION void computeVarsOnCell(int ICell, int KChunk) const { + KOKKOS_FUNCTION void computeVarsOnCell(const TeamMember &Team, + int ICell) const { const Real InvAreaCell = 1._Real / AreaCell(ICell); - const int KStartCell = chunkStart(KChunk, MinLayerCell(ICell)); - const int KLenCell = chunkLength(KChunk, KStartCell, MaxLayerCell(ICell)); - const int KEndCell = KStartCell + KLenCell - 1; - Real Del2DivCellTmp[VecLength] = {0}; + ScratchArray1DReal Del2DivCellTmp(teamScratch(Team), NVertLayers); + parallelForInner( + Team, NVertLayers, INNER_LAMBDA(int K) { Del2DivCellTmp(K) = 0; }); for (int J = 0; J < NEdgesOnCell(ICell); ++J) { const int JEdge = EdgesOnCell(ICell, J); const Real AreaEdge = 0.5_Real * DvEdge(JEdge) * DcEdge(JEdge); - const int KStartEdge = Kokkos::max(KStartCell, MinLayerEdgeBot(JEdge)); - const int KEndEdge = Kokkos::min(KEndCell, MaxLayerEdgeTop(JEdge)); + const int MinLyrEdgeBot = MinLayerEdgeBot(JEdge); + const int MaxLyrEdgeTop = MaxLayerEdgeTop(JEdge); - for (int K = KStartEdge; K <= KEndEdge; ++K) { - const int KVec = K - KStartCell; - Del2DivCellTmp[KVec] -= DvEdge(JEdge) * InvAreaCell * - EdgeSignOnCell(ICell, J) * - Del2Edge(JEdge, K); - } - } - for (int KVec = 0; KVec < KLenCell; ++KVec) { - const int K = KStartCell + KVec; - Del2DivCell(ICell, K) = Del2DivCellTmp[KVec]; + parallelForInner( + Team, Range{MinLyrEdgeBot, MaxLyrEdgeTop}, INNER_LAMBDA(int K) { + Del2DivCellTmp(K) -= DvEdge(JEdge) * InvAreaCell * + EdgeSignOnCell(ICell, J) * + Del2Edge(JEdge, K); + }); } + + const int MinLyrCell = MinLayerCell(ICell); + const int MaxLyrCell = MaxLayerCell(ICell); + + parallelForInner( + Team, Range{MinLyrCell, MaxLyrCell}, + INNER_LAMBDA(int K) { Del2DivCell(ICell, K) = Del2DivCellTmp(K); }); } - KOKKOS_FUNCTION void computeVarsOnVertex(int IVertex, int KChunk) const { + KOKKOS_FUNCTION void computeVarsOnVertex(const TeamMember &Team, + int IVertex) const { // Compute over the full vertex valid range [MinLayerVertexTop, // MaxLayerVertexBot] so that boundary-vertex layers (where only some // surrounding cells are active) receive a valid value. Each edge's @@ -82,32 +88,36 @@ class VelocityDel2AuxVars { // MaxLayerEdgeBot], where Del2Edge has been computed or zeroed; this // matches VorticityAuxVars::computeVarsOnVertex and avoids reading // uninitialized (fill-value) layers of Del2Edge for deeper edges. - const int KStartVertex = chunkStart(KChunk, MinLayerVertexTop(IVertex)); - const int KLenVertex = - chunkLength(KChunk, KStartVertex, MaxLayerVertexBot(IVertex)); - const int KEndVertex = KStartVertex + KLenVertex - 1; const Real InvAreaTriangle = 1._Real / AreaTriangle(IVertex); - Real Del2RelVortVertexTmp[VecLength] = {0}; + ScratchArray1DReal Del2RelVortVertexTmp(teamScratch(Team), NVertLayers); + + parallelForInner( + Team, NVertLayers, + INNER_LAMBDA(int K) { Del2RelVortVertexTmp(K) = 0; }); + + const int MinLyrVertexTop = MinLayerVertexTop(IVertex); + const int MaxLyrVertexBot = MaxLayerVertexBot(IVertex); for (int J = 0; J < VertexDegree; ++J) { const int JEdge = EdgesOnVertex(IVertex, J); - const int KStartEdge = - Kokkos::max(KStartVertex, MinLayerEdgeTop(JEdge)); - const int KEndEdge = Kokkos::min(KEndVertex, MaxLayerEdgeBot(JEdge)); - for (int K = KStartEdge; K <= KEndEdge; ++K) { - const int KVec = K - KStartVertex; - Del2RelVortVertexTmp[KVec] += InvAreaTriangle * DcEdge(JEdge) * - EdgeSignOnVertex(IVertex, J) * - Del2Edge(JEdge, K); - } - } - for (int KVec = 0; KVec < KLenVertex; ++KVec) { - const int K = KStartVertex + KVec; - Del2RelVortVertex(IVertex, K) = Del2RelVortVertexTmp[KVec]; + const int MinLyrEdgeTop = MinLayerEdgeTop(JEdge); + const int MaxLyrEdgeBot = MaxLayerEdgeBot(JEdge); + + parallelForInner( + Team, Range{MinLyrEdgeTop, MaxLyrEdgeBot}, INNER_LAMBDA(int K) { + Del2RelVortVertexTmp(K) += InvAreaTriangle * DcEdge(JEdge) * + EdgeSignOnVertex(IVertex, J) * + Del2Edge(JEdge, K); + }); } + + parallelForInner( + Team, Range{MinLyrVertexTop, MaxLyrVertexBot}, INNER_LAMBDA(int K) { + Del2RelVortVertex(IVertex, K) = Del2RelVortVertexTmp(K); + }); } void registerFields(const std::string &AuxGroupName, @@ -128,6 +138,7 @@ class VelocityDel2AuxVars { Array1DReal AreaTriangle; Array2DReal EdgeMask; I4 VertexDegree; + I4 NVertLayers; Array1DI4 MinLayerEdgeBot; Array1DI4 MaxLayerEdgeTop; Array1DI4 MinLayerEdgeTop; diff --git a/components/omega/src/ocn/auxiliaryVars/VorticityAuxVars.cpp b/components/omega/src/ocn/auxiliaryVars/VorticityAuxVars.cpp index 062aeac80cf6..20294980c3d3 100644 --- a/components/omega/src/ocn/auxiliaryVars/VorticityAuxVars.cpp +++ b/components/omega/src/ocn/auxiliaryVars/VorticityAuxVars.cpp @@ -24,7 +24,7 @@ VorticityAuxVars::VorticityAuxVars(const std::string &AuxStateSuffix, EdgeSignOnVertex(Mesh->EdgeSignOnVertex), DcEdge(Mesh->DcEdge), KiteAreasOnVertex(Mesh->KiteAreasOnVertex), AreaTriangle(Mesh->AreaTriangle), FVertex(Mesh->FVertex), - VerticesOnEdge(Mesh->VerticesOnEdge), + VerticesOnEdge(Mesh->VerticesOnEdge), NVertLayers(VCoord->NVertLayers), MinLayerVertexTop(VCoord->MinLayerVertexTop), MaxLayerVertexBot(VCoord->MaxLayerVertexBot), MinLayerCell(VCoord->MinLayerCell), MaxLayerCell(VCoord->MaxLayerCell), diff --git a/components/omega/src/ocn/auxiliaryVars/VorticityAuxVars.h b/components/omega/src/ocn/auxiliaryVars/VorticityAuxVars.h index e945bceb5417..6223001af693 100644 --- a/components/omega/src/ocn/auxiliaryVars/VorticityAuxVars.h +++ b/components/omega/src/ocn/auxiliaryVars/VorticityAuxVars.h @@ -23,74 +23,78 @@ class VorticityAuxVars { const VertCoord *VCoord); KOKKOS_FUNCTION void - computeVarsOnVertex(int IVertex, int KChunk, + computeVarsOnVertex(const TeamMember &Team, int IVertex, const Array2DReal &PseudoThickCell, const Array2DReal &NormalVelEdge) const { - const int KStartVertex = chunkStart(KChunk, MinLayerVertexTop(IVertex)); - const int KLenVertex = - chunkLength(KChunk, KStartVertex, MaxLayerVertexBot(IVertex)); - const int KEndVertex = KStartVertex + KLenVertex - 1; - const Real InvAreaTriangle = 1._Real / AreaTriangle(IVertex); - Real PseudoThickVertex[VecLength] = {0}; - Real RelVortVertexTmp[VecLength] = {0}; + ScratchArray1DReal PseudoThickVertex(teamScratch(Team), NVertLayers); + ScratchArray1DReal RelVortVertexTmp(teamScratch(Team), NVertLayers); + + parallelForInner( + Team, NVertLayers, INNER_LAMBDA(int K) { + PseudoThickVertex(K) = 0; + RelVortVertexTmp(K) = 0; + }); for (int J = 0; J < VertexDegree; ++J) { - const int JCell = CellsOnVertex(IVertex, J); - - const int KStartCell = Kokkos::max(KStartVertex, MinLayerCell(JCell)); - const int KEndCell = Kokkos::min(KEndVertex, MaxLayerCell(JCell)); - - for (int K = KStartCell; K <= KEndCell; ++K) { - const int KVec = K - KStartVertex; - PseudoThickVertex[KVec] += InvAreaTriangle * - KiteAreasOnVertex(IVertex, J) * - PseudoThickCell(JCell, K); - } - - const int JEdge = EdgesOnVertex(IVertex, J); - const int KStartEdge = - Kokkos::max(KStartVertex, MinLayerEdgeTop(JEdge)); - const int KEndEdge = Kokkos::min(KEndVertex, MaxLayerEdgeBot(JEdge)); - - for (int K = KStartEdge; K <= KEndEdge; ++K) { - const int KVec = K - KStartVertex; - RelVortVertexTmp[KVec] += InvAreaTriangle * DcEdge(JEdge) * - EdgeSignOnVertex(IVertex, J) * - NormalVelEdge(JEdge, K); - } + const int JCell = CellsOnVertex(IVertex, J); + const int MinLyrCell = MinLayerCell(JCell); + const int MaxLyrCell = MaxLayerCell(JCell); + + parallelForInner( + Team, Range{MinLyrCell, MaxLyrCell}, INNER_LAMBDA(int K) { + PseudoThickVertex(K) += InvAreaTriangle * + KiteAreasOnVertex(IVertex, J) * + PseudoThickCell(JCell, K); + }); + + const int JEdge = EdgesOnVertex(IVertex, J); + const int MinLyrEdgeTop = MinLayerEdgeTop(JEdge); + const int MaxLyrEdgeBot = MaxLayerEdgeBot(JEdge); + + parallelForInner( + Team, Range{MinLyrEdgeTop, MaxLyrEdgeBot}, INNER_LAMBDA(int K) { + RelVortVertexTmp(K) += InvAreaTriangle * DcEdge(JEdge) * + EdgeSignOnVertex(IVertex, J) * + NormalVelEdge(JEdge, K); + }); } - for (int KVec = 0; KVec < KLenVertex; ++KVec) { - const int K = KStartVertex + KVec; - const Real InvPseudoThickVertex = 1._Real / PseudoThickVertex[KVec]; + const int MinLyrVertexTop = MinLayerVertexTop(IVertex); + const int MaxLyrVertexBot = MaxLayerVertexBot(IVertex); - RelVortVertex(IVertex, K) = RelVortVertexTmp[KVec]; - NormRelVortVertex(IVertex, K) = - RelVortVertexTmp[KVec] * InvPseudoThickVertex; - NormPlanetVortVertex(IVertex, K) = - FVertex(IVertex) * InvPseudoThickVertex; - } + parallelForInner( + Team, Range{MinLyrVertexTop, MaxLyrVertexBot}, INNER_LAMBDA(int K) { + const Real InvPseudoThickVertex = 1._Real / PseudoThickVertex(K); + + RelVortVertex(IVertex, K) = RelVortVertexTmp(K); + NormRelVortVertex(IVertex, K) = + RelVortVertexTmp(K) * InvPseudoThickVertex; + NormPlanetVortVertex(IVertex, K) = + FVertex(IVertex) * InvPseudoThickVertex; + }); } - KOKKOS_FUNCTION void computeVarsOnEdge(int IEdge, int KChunk) const { - const int KStart = chunkStart(KChunk, MinLayerEdgeTop(IEdge)); - const int KLen = chunkLength(KChunk, KStart, MaxLayerEdgeBot(IEdge)); + KOKKOS_FUNCTION void computeVarsOnEdge(const TeamMember &Team, + int IEdge) const { const int JVertex0 = VerticesOnEdge(IEdge, 0); const int JVertex1 = VerticesOnEdge(IEdge, 1); - for (int KVec = 0; KVec < KLen; ++KVec) { - const int K = KStart + KVec; - NormRelVortEdge(IEdge, K) = - 0.5_Real * - (NormRelVortVertex(JVertex0, K) + NormRelVortVertex(JVertex1, K)); + const int KMin = MinLayerEdgeTop(IEdge); + const int KMax = MaxLayerEdgeBot(IEdge); - NormPlanetVortEdge(IEdge, K) = - 0.5_Real * (NormPlanetVortVertex(JVertex0, K) + - NormPlanetVortVertex(JVertex1, K)); - } + parallelForInner( + Team, Range{KMin, KMax}, INNER_LAMBDA(int K) { + NormRelVortEdge(IEdge, K) = + 0.5_Real * (NormRelVortVertex(JVertex0, K) + + NormRelVortVertex(JVertex1, K)); + + NormPlanetVortEdge(IEdge, K) = + 0.5_Real * (NormPlanetVortVertex(JVertex0, K) + + NormPlanetVortVertex(JVertex1, K)); + }); } void registerFields(const std::string &AuxGroupName, @@ -108,6 +112,7 @@ class VorticityAuxVars { Array2DI4 VerticesOnEdge; Array1DReal FVertex; + I4 NVertLayers; Array1DI4 MinLayerVertexTop; Array1DI4 MaxLayerVertexBot; Array1DI4 MinLayerCell; diff --git a/components/omega/test/ocn/AuxiliaryVarsTest.cpp b/components/omega/test/ocn/AuxiliaryVarsTest.cpp index 6260519f38ed..4f916d998c7d 100644 --- a/components/omega/test/ocn/AuxiliaryVarsTest.cpp +++ b/components/omega/test/ocn/AuxiliaryVarsTest.cpp @@ -345,9 +345,10 @@ int testKineticAuxVars(const Array2DReal &PseudoThicknessCell, KineticAuxVars KineticAux("", Mesh, VCoord); - parallelFor( - {Mesh->NCellsOwned, NVertLayers}, KOKKOS_LAMBDA(int ICell, int KLayer) { - KineticAux.computeVarsOnCell(ICell, KLayer, NormalVelocityEdge); + parallelForOuter( + LaunchConfig({Mesh->NCellsOwned}, TeamScratch(2 * NVertLayers)), + KOKKOS_LAMBDA(int ICell, const TeamMember &Team) { + KineticAux.computeVarsOnCell(Team, ICell, NormalVelocityEdge); }); const auto &NumKineticEnergyCell = KineticAux.KineticEnergyCell; const auto &NumVelocityDivCell = KineticAux.VelocityDivCell; @@ -392,9 +393,9 @@ int testPseudoThicknessAuxVars(const Array2DReal &PseudoThickCell, PseudoThicknessAuxVars PseudoThicknessAux("", Mesh, VCoord); PseudoThicknessAux.FluxThickEdgeChoice = FluxThickEdgeOption::Upwind; - parallelFor( - {Mesh->NEdgesOwned, NVertLayers}, KOKKOS_LAMBDA(int IEdge, int KLayer) { - PseudoThicknessAux.computeVarsOnEdge(IEdge, KLayer, PseudoThickCell, + parallelForOuter( + {Mesh->NEdgesOwned}, KOKKOS_LAMBDA(int IEdge, const TeamMember &Team) { + PseudoThicknessAux.computeVarsOnEdge(Team, IEdge, PseudoThickCell, NormalVelEdge); }); @@ -458,10 +459,11 @@ int testVorticityAuxVars(const Array2DReal &PseudoThickCell, // Compute numerical results for vertex variables - parallelFor( - {Decomp->NVerticesHaloH(0), NVertLayers}, - KOKKOS_LAMBDA(int IVertex, int KLayer) { - VorticityAux.computeVarsOnVertex(IVertex, KLayer, PseudoThickCell, + parallelForOuter( + LaunchConfig({Decomp->NVerticesHaloH(0)}, + TeamScratch(2 * VCoord->NVertLayers)), + KOKKOS_LAMBDA(int IVertex, const TeamMember &Team) { + VorticityAux.computeVarsOnVertex(Team, IVertex, PseudoThickCell, NormalVelEdge); }); @@ -511,9 +513,9 @@ int testVorticityAuxVars(const Array2DReal &PseudoThickCell, // Compute numerical results for vertex variables - parallelFor( - {Mesh->NEdgesOwned, NVertLayers}, KOKKOS_LAMBDA(int IEdge, int KLayer) { - VorticityAux.computeVarsOnEdge(IEdge, KLayer); + parallelForOuter( + {Mesh->NEdgesOwned}, KOKKOS_LAMBDA(int IEdge, const TeamMember &Team) { + VorticityAux.computeVarsOnEdge(Team, IEdge); }); const auto &NumNormRelVortEdge = VorticityAux.NormRelVortEdge; const auto &NumNormPlanetVortEdge = VorticityAux.NormPlanetVortEdge; @@ -575,10 +577,10 @@ int testVelocityDel2AuxVars(Real RTol) { // Compute numerical Del2 - parallelFor( - {Decomp->NEdgesHaloH(1), NVertLayers}, - KOKKOS_LAMBDA(int IEdge, int KLayer) { - VelocityDel2Aux.computeVarsOnEdge(IEdge, KLayer, ExactVelocityDivCell, + parallelForOuter( + {Decomp->NEdgesHaloH(1)}, + KOKKOS_LAMBDA(int IEdge, const TeamMember &Team) { + VelocityDel2Aux.computeVarsOnEdge(Team, IEdge, ExactVelocityDivCell, ExactRelVortVertex); }); const auto &NumDel2Edge = VelocityDel2Aux.Del2Edge; @@ -601,9 +603,10 @@ int testVelocityDel2AuxVars(Real RTol) { // Compute numerical Del2Div - parallelFor( - {Mesh->NCellsOwned, NVertLayers}, KOKKOS_LAMBDA(int ICell, int KLayer) { - VelocityDel2Aux.computeVarsOnCell(ICell, KLayer); + parallelForOuter( + LaunchConfig({Mesh->NCellsOwned}, TeamScratch(NVertLayers)), + KOKKOS_LAMBDA(int ICell, const TeamMember &Team) { + VelocityDel2Aux.computeVarsOnCell(Team, ICell); }); const auto &NumDel2DivCell = VelocityDel2Aux.Del2DivCell; @@ -625,10 +628,10 @@ int testVelocityDel2AuxVars(Real RTol) { // Compute numerical Del2RelVort - parallelFor( - {Mesh->NVerticesOwned, NVertLayers}, - KOKKOS_LAMBDA(int IVertex, int KLayer) { - VelocityDel2Aux.computeVarsOnVertex(IVertex, KLayer); + parallelForOuter( + LaunchConfig({Mesh->NVerticesOwned}, TeamScratch(NVertLayers)), + KOKKOS_LAMBDA(int IVertex, const TeamMember &Team) { + VelocityDel2Aux.computeVarsOnVertex(Team, IVertex); }); const auto &NumDel2RelVortVertex = VelocityDel2Aux.Del2RelVortVertex; @@ -682,10 +685,11 @@ int testTracerAuxVars(const Array2DReal &PseudoThickCell, // Compute numerical Del2TracerCell - parallelFor( - {NTracers, Mesh->NCellsOwned, NVertLayers}, - KOKKOS_LAMBDA(int L, int ICell, int KLayer) { - TracerAux.computeVarsOnCells(L, ICell, KLayer, PseudoThickEdge, + parallelForOuter( + LaunchConfig({NTracers, Mesh->NCellsOwned}, + TeamScratch(NVertLayers)), + KOKKOS_LAMBDA(int L, int ICell, const TeamMember &Team) { + TracerAux.computeVarsOnCells(Team, L, ICell, PseudoThickEdge, TracersOnCell); }); diff --git a/components/omega/test/ocn/TendencyTermsTest.cpp b/components/omega/test/ocn/TendencyTermsTest.cpp index e7f5038be0ce..8fe01ee4db59 100644 --- a/components/omega/test/ocn/TendencyTermsTest.cpp +++ b/components/omega/test/ocn/TendencyTermsTest.cpp @@ -139,8 +139,8 @@ struct TestSetupPlane { return 2. + std::cos(TwoPi * X / Lx) * std::cos(TwoPi * Y / Ly); } - KOKKOS_FUNCTION Real tracerDiff(Real X, Real Y) const { - return -TwoPi * TwoPi * std::sin(TwoPi * Y / Ly) * + KOKKOS_FUNCTION Real tracerDiff(Real X, Real Y, Real EddyDiff2) const { + return -EddyDiff2 * TwoPi * TwoPi * std::sin(TwoPi * Y / Ly) * (2 * (1 / Lx / Lx + 1 / Ly / Ly) * std::cos(TwoPi * X / Lx) + (1 / Ly / Ly + (1 / Lx / Lx + 1 / Ly / Ly) * std::cos(2 * TwoPi * X / Lx)) * @@ -152,8 +152,8 @@ struct TestSetupPlane { std::pow(std::sin(TwoPi * Y / Ly), 2); } - KOKKOS_FUNCTION Real tracerHyperDiff(Real X, Real Y) const { - return -2 * TwoPi * TwoPi * + KOKKOS_FUNCTION Real tracerHyperDiff(Real X, Real Y, Real EddyDiff4) const { + return -EddyDiff4 * 2 * TwoPi * TwoPi * (std::cos(2 * TwoPi * X / Lx) / Lx / Lx + std::cos(2 * TwoPi * Y / Ly) / Ly / Ly); } @@ -282,8 +282,9 @@ struct TestSetupSphere { return 2. + std::cos(Lon) * std::sin(Lat); } - KOKKOS_FUNCTION Real tracerDiff(Real Lon, Real Lat) const { - return (4 * std::pow(std::cos(Lon), 2) - + KOKKOS_FUNCTION Real tracerDiff(Real Lon, Real Lat, Real EddyDiff2) const { + return EddyDiff2 * + (4 * std::pow(std::cos(Lon), 2) - 2 * (1. + 3 * std::cos(2 * Lat)) * std::pow(std::sin(Lon), 2) + 2 * std::pow(std::cos(Lon), 3) * std::sin(Lat) - 8 * std::cos(Lon) * std::pow(std::cos(Lat), 2) * @@ -296,8 +297,10 @@ struct TestSetupSphere { std::cos(Lon); } - KOKKOS_FUNCTION Real tracerHyperDiff(Real Lon, Real Lat) const { - return std::sqrt(3. / 2. / Pi) * std::cos(Lat) * std::cos(Lon) / Radius; + KOKKOS_FUNCTION Real tracerHyperDiff(Real Lon, Real Lat, + Real EddyDiff4) const { + return EddyDiff4 * std::sqrt(3. / 2. / Pi) * std::cos(Lat) * + std::cos(Lon) / Radius; } KOKKOS_FUNCTION Real sfcStressForcingX(Real Lon, Real Lat) const { @@ -418,9 +421,10 @@ int testThickFluxDiv(int NVertLayers, Real RTol) { Array2DReal NumThickFluxDiv("NumThickFluxDiv", Mesh->NCellsOwned, NVertLayers); PseudoThicknessFluxDivOnCell ThickFluxDivOnC(Mesh, VCoord); - parallelFor( - {Mesh->NCellsOwned, NVertLayers}, KOKKOS_LAMBDA(int ICell, int KLayer) { - ThickFluxDivOnC(NumThickFluxDiv, ICell, KLayer, OnesEdge, + parallelForOuter( + LaunchConfig({Mesh->NCellsOwned}, TeamScratch(NVertLayers)), + KOKKOS_LAMBDA(int ICell, const TeamMember &Team) { + ThickFluxDivOnC(Team, NumThickFluxDiv, ICell, OnesEdge, ThickFluxEdge); }); @@ -496,9 +500,10 @@ int testPotVortHAdv(int NVertLayers, Real RTol) { Array2DReal NumPotVortHAdv("NumPotVortHAdv", Mesh->NEdgesOwned, NVertLayers); PotentialVortHAdvOnEdge PotVortHAdvOnE(Mesh, VCoord); - parallelFor( - {Mesh->NEdgesOwned, NVertLayers}, KOKKOS_LAMBDA(int IEdge, int KLayer) { - PotVortHAdvOnE(NumPotVortHAdv, IEdge, KLayer, NormRelVortEdge, + parallelForOuter( + LaunchConfig({Mesh->NEdgesOwned}, TeamScratch(NVertLayers)), + KOKKOS_LAMBDA(int IEdge, const TeamMember &Team) { + PotVortHAdvOnE(Team, NumPotVortHAdv, IEdge, NormRelVortEdge, NormPlanetVortEdge, PseudoThickEdge, NormVelEdge); }); @@ -547,9 +552,9 @@ int testKEGrad(int NVertLayers, Real RTol) { Array2DReal NumKEGrad("NumKEGrad", Mesh->NEdgesOwned, NVertLayers); KEGradOnEdge KEGradOnE(Mesh, VCoord); - parallelFor( - {Mesh->NEdgesOwned, NVertLayers}, KOKKOS_LAMBDA(int IEdge, int KLayer) { - KEGradOnE(NumKEGrad, IEdge, KLayer, KECell); + parallelForOuter( + {Mesh->NEdgesOwned}, KOKKOS_LAMBDA(int IEdge, const TeamMember &Team) { + KEGradOnE(Team, NumKEGrad, IEdge, KECell); }); // Compute errors @@ -595,9 +600,9 @@ int testSSHGrad(int NVertLayers, Real RTol) { Array2DReal NumSSHGrad("NumSSHGrad", Mesh->NEdgesOwned, NVertLayers); SSHGradOnEdge SSHGradOnE(Mesh, VCoord); - parallelFor( - {Mesh->NEdgesOwned, NVertLayers}, KOKKOS_LAMBDA(int IEdge, int KLayer) { - SSHGradOnE(NumSSHGrad, IEdge, KLayer, SSHCell); + parallelForOuter( + {Mesh->NEdgesOwned}, KOKKOS_LAMBDA(int IEdge, const TeamMember &Team) { + SSHGradOnE(Team, NumSSHGrad, IEdge, SSHCell); }); // Compute errors @@ -657,9 +662,9 @@ int testVelDiff(int NVertLayers, Real RTol) { // Compute numerical result Array2DReal NumVelDiff("NumVelDiff", Mesh->NEdgesOwned, NVertLayers); - parallelFor( - {Mesh->NEdgesOwned, NVertLayers}, KOKKOS_LAMBDA(int IEdge, int KLayer) { - VelDiffOnE(NumVelDiff, IEdge, KLayer, DivCell, RVortVertex); + parallelForOuter( + {Mesh->NEdgesOwned}, KOKKOS_LAMBDA(int IEdge, const TeamMember &Team) { + VelDiffOnE(Team, NumVelDiff, IEdge, DivCell, RVortVertex); }); // Compute errors @@ -728,9 +733,9 @@ int testVelHyperDiff(int NVertLayers, Real RTol) { Array2DReal NumVelHyperDiff("NumVelHyperDiff", Mesh->NEdgesOwned, NVertLayers); - parallelFor( - {Mesh->NEdgesOwned, NVertLayers}, KOKKOS_LAMBDA(int IEdge, int KLayer) { - VelHyperDiffOnE(NumVelHyperDiff, IEdge, KLayer, DivCell, RVortVertex); + parallelForOuter( + {Mesh->NEdgesOwned}, KOKKOS_LAMBDA(int IEdge, const TeamMember &Team) { + VelHyperDiffOnE(Team, NumVelHyperDiff, IEdge, DivCell, RVortVertex); }); // Compute errors @@ -798,9 +803,9 @@ int testSfcStressForcing(int NVertLayers) { SfcStressForcingOnEdge SfcStressForcingOnE(Mesh, VCoord); - parallelFor( - {Mesh->NEdgesOwned, NVertLayers}, KOKKOS_LAMBDA(int IEdge, int KLayer) { - SfcStressForcingOnE(NumSfcStressForcing, IEdge, KLayer, + parallelForOuter( + {Mesh->NEdgesOwned}, KOKKOS_LAMBDA(int IEdge, const TeamMember &Team) { + SfcStressForcingOnE(Team, NumSfcStressForcing, IEdge, NormalStressEdge, PseudoThickEdge); }); @@ -1115,16 +1120,18 @@ int testTracerHorzAdvOnCell(int NVertLayers, int NTracers, Real RTol) { TrHorzAdvOnC.ForceLowOrder = true; TrHorzAdvOnC.init(); - parallelFor( - {NTracers, Mesh->NEdgesAll, NVertLayers}, - KOKKOS_LAMBDA(int L, int IEdge, int KLayer) { - TrHorzAdvOnC(L, IEdge, KLayer, TrCell, ThickEdge, NormalVelocity); + parallelForOuter( + LaunchConfig({NTracers, Mesh->NEdgesAll}, + TeamScratch(NVertLayers)), + KOKKOS_LAMBDA(int L, int IEdge, const TeamMember &Team) { + TrHorzAdvOnC(Team, L, IEdge, TrCell, ThickEdge, NormalVelocity); }); - parallelFor( - {NTracers, Mesh->NCellsOwned, NVertLayers}, - KOKKOS_LAMBDA(int L, int ICell, int KLayer) { - TrHorzAdvOnC(NumTrFluxDiv, L, ICell, KLayer); + parallelForOuter( + LaunchConfig({NTracers, Mesh->NCellsOwned}, + TeamScratch(NVertLayers)), + KOKKOS_LAMBDA(int L, int ICell, const TeamMember &Team) { + TrHorzAdvOnC(Team, NumTrFluxDiv, L, ICell); }); ErrorMeasures TrHAdvErrors; @@ -1149,12 +1156,16 @@ int testTracerDiffOnCell(int NVertLayers, int NTracers, Real RTol) { const auto Mesh = HorzMesh::getDefault(); const auto VCoord = VertCoord::getDefault(); + const Real EddyDiff2 = 0.84_Real; + // Compute exact result Array3DReal ExactTracerDiff("ExactTracerDiff", NTracers, Mesh->NCellsOwned, NVertLayers); Err += setScalar( - KOKKOS_LAMBDA(Real X, Real Y) { return Setup.tracerDiff(X, Y); }, + KOKKOS_LAMBDA(Real X, Real Y) { + return Setup.tracerDiff(X, Y, EddyDiff2); + }, ExactTracerDiff, Geom, Mesh, OnCell, ExchangeHalos::No); // Set input arrays @@ -1176,13 +1187,13 @@ int testTracerDiffOnCell(int NVertLayers, int NTracers, Real RTol) { Array3DReal NumTracerDiff("NumTracerDiff", NTracers, Mesh->NCellsOwned, NVertLayers); TracerDiffOnCell TrDiffOnC(Mesh, VCoord); - TrDiffOnC.EddyDiff2 = 1._Real; + TrDiffOnC.EddyDiff2 = EddyDiff2; - parallelFor( - {NTracers, Mesh->NCellsOwned, NVertLayers}, - KOKKOS_LAMBDA(int L, int ICell, int KLayer) { - TrDiffOnC(NumTracerDiff, L, ICell, KLayer, TracerCell, - PseudoThickEdge); + parallelForOuter( + LaunchConfig({NTracers, Mesh->NCellsOwned}, + TeamScratch(NVertLayers)), + KOKKOS_LAMBDA(int L, int ICell, const TeamMember &Team) { + TrDiffOnC(Team, NumTracerDiff, L, ICell, TracerCell, PseudoThickEdge); }); ErrorMeasures TrDiffErrors; @@ -1207,12 +1218,16 @@ int testTracerHyperDiffOnCell(int NVertLayers, int NTracers, Real RTol) { const auto Mesh = HorzMesh::getDefault(); const auto VCoord = VertCoord::getDefault(); + const Real EddyDiff4 = 1.23_Real; + // Compute exact result Array3DReal ExactTracerHyperDiff("ExactTracerHyperDiff", NTracers, Mesh->NCellsOwned, NVertLayers); Err += setScalar( - KOKKOS_LAMBDA(Real X, Real Y) { return -Setup.tracerHyperDiff(X, Y); }, + KOKKOS_LAMBDA(Real X, Real Y) { + return -Setup.tracerHyperDiff(X, Y, EddyDiff4); + }, ExactTracerHyperDiff, Geom, Mesh, OnCell, ExchangeHalos::No); // Set input arrays @@ -1227,11 +1242,12 @@ int testTracerHyperDiffOnCell(int NVertLayers, int NTracers, Real RTol) { Array3DReal NumTracerHyperDiff("NumTracerHyperDiff", NTracers, Mesh->NCellsOwned, NVertLayers); TracerHyperDiffOnCell TrHypDiffOnC(Mesh, VCoord); - TrHypDiffOnC.EddyDiff4 = 1._Real; - parallelFor( - {NTracers, Mesh->NCellsOwned, NVertLayers}, - KOKKOS_LAMBDA(int L, int ICell, int KLayer) { - TrHypDiffOnC(NumTracerHyperDiff, L, ICell, KLayer, TrDel2Cell); + TrHypDiffOnC.EddyDiff4 = EddyDiff4; + parallelForOuter( + LaunchConfig({NTracers, Mesh->NCellsOwned}, + TeamScratch(NVertLayers)), + KOKKOS_LAMBDA(int L, int ICell, const TeamMember &Team) { + TrHypDiffOnC(Team, NumTracerHyperDiff, L, ICell, TrDel2Cell); }); ErrorMeasures TrHyperDiffErrors;