WIP: Performance Improvements (for Burgers)

benchmarks/burgers/burgers.pin

@@ -22,7 +22,7 @@ problem_id = burgers
 nghost = 4
 refinement = adaptive
 numlevel = 2
-deref_count = 10
+derefine_count = 10
 
 nx1 = 128
 x1min = -0.5

benchmarks/burgers/burgers_driver.cpp

@@ -10,6 +10,8 @@
 // license in this material to reproduce, prepare derivative works, distribute copies to
 // the public, perform publicly and display publicly, and to permit others to do so.
 //========================================================================================
+// Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+//========================================================================================
 
 #include <memory>
 #include <string>
@@ -120,17 +122,26 @@ TaskCollection BurgersDriver::MakeTaskCollection(BlockList_t &blocks, const int
 
     auto fill_deriv = tl.AddTask(update, FillDerived<MeshData<Real>>, mc1.get());
 
-    auto set_bc = tl.AddTask(update, parthenon::ApplyBoundaryConditionsMD, mc1);
+    //auto set_bc = tl.AddTask(update, parthenon::ApplyBoundaryConditionsMD, mc1);
 
     // estimate next time step
     if (stage == integrator->nstages) {
       auto new_dt = tl.AddTask(update, EstimateTimestep<MeshData<Real>>, mc1.get());
-      if (pmesh->adaptive) {
-        auto tag_refine =
-            tl.AddTask(set_bc, parthenon::Refinement::Tag<MeshData<Real>>, mc1.get());
-      }
     }
-  }
+  } 
+
+  TaskRegion &single_tasklist_per_pack_region3 = tc.AddRegion(num_partitions);
+  for (int i = 0; i < num_partitions; i++) {
+    auto &tl = single_tasklist_per_pack_region3[i];
+    auto &mc1 = pmesh->mesh_data.GetOrAdd(stage_name[stage], i);
+
+    auto set_bc = tl.AddTask(none, parthenon::ApplyBoundaryConditionsOnCoarseOrFineMD, mc1, false);
+
+    if (stage == integrator->nstages && pmesh->adaptive) {
+      auto tag_refine = tl.AddTask(set_bc, parthenon::Refinement::TagFused<MeshData<Real>>, mc1.get());
+    }
+  } 
+
   return tc;
 }
 

benchmarks/burgers/burgers_package.cpp

@@ -10,6 +10,8 @@
 // license in this material to reproduce, prepare derivative works, distribute copies to
 // the public, perform publicly and display publicly, and to permit others to do so.
 //========================================================================================
+// Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+//========================================================================================
 
 #include <algorithm>
 #include <cmath>
@@ -146,7 +148,7 @@ void CalculateDerived(MeshData<Real> *md) {
   IndexRange kb = md->GetBoundsK(IndexDomain::interior);
 
   std::vector<std::string> vars({"derived", "U"});
-  auto &v = md->PackVariables(vars);
+  auto &v = md->PackVariablesByNames(vars);  
   const int nblocks = md->NumBlocks();
   std::size_t scratch_size = 0;
   constexpr int scratch_level = 0;
@@ -178,7 +180,7 @@ Real EstimateTimestepMesh(MeshData<Real> *md) {
   const auto &cfl = params.Get<Real>("cfl");
 
   std::vector<std::string> vars({"U"});
-  auto &v = md->PackVariables(vars);
+  auto &v = md->PackVariablesByNames(vars);
   const int ndim = pm->ndim;
 
   Real min_dt;
@@ -213,7 +215,7 @@ TaskStatus CalculateFluxes(MeshData<Real> *md) {
   std::vector<std::string> vars({"U", "Ulx", "Urx", "Uly", "Ury", "Ulz", "Urz"});
   std::vector<std::string> flxs({"U"});
   PackIndexMap imap;
-  auto v = md->PackVariablesAndFluxes(vars, flxs, imap);
+  auto v = md->PackVariablesAndFluxesByNames(vars, flxs, imap); 
   const int iu_lo = imap["U"].first;
   const int iu_hi = imap["U"].second;
   const int iulx_lo = imap["Ulx"].first;
@@ -233,9 +235,16 @@ TaskStatus CalculateFluxes(MeshData<Real> *md) {
   // first we'll reconstruct the state to faces
   std::size_t scratch_size = 0;
   constexpr int scratch_level = 0;
+  const int team_size = ib.e - ib.s + 1;
+  #ifdef KOKKOS_ENABLE_CUDA
+  auto perf_opts = parthenon::OuterLoopPerfOpts{team_size};
+  #else 
+  auto perf_opts = parthenon::OuterLoopPerfOpts{};
+  #endif
   parthenon::par_for_outer(
-      DEFAULT_OUTER_LOOP_PATTERN, PARTHENON_AUTO_LABEL, DevExecSpace(), scratch_size,
-      scratch_level, 0, nblocks - 1, kb.s - dk, kb.e + dk, jb.s - dj, jb.e + dj,
+      DEFAULT_OUTER_LOOP_PATTERN, "burgers::Reconstruction", DevExecSpace(), scratch_size,
+      scratch_level, perf_opts,
+      0, nblocks - 1, kb.s - dk, kb.e + dk, jb.s - dj, jb.e + dj,
       KOKKOS_LAMBDA(team_mbr_t member, const int b, const int k, const int j) {
         bool xrec = (k >= kb.s && k <= kb.e) && (j >= jb.s && j <= jb.e);
         bool yrec = (k >= kb.s && k <= kb.e) && (ndim > 1);
@@ -305,8 +314,9 @@ TaskStatus CalculateFluxes(MeshData<Real> *md) {
   // now we'll solve the Riemann problems to get fluxes
   scratch_size = 2 * ScratchPad1D<Real>::shmem_size(ib.e + 1);
   parthenon::par_for_outer(
-      DEFAULT_OUTER_LOOP_PATTERN, PARTHENON_AUTO_LABEL, DevExecSpace(), scratch_size,
-      scratch_level, 0, nblocks - 1, kb.s, kb.e + dk, jb.s, jb.e + dj,
+      DEFAULT_OUTER_LOOP_PATTERN, "burgers::RiemannSolve", DevExecSpace(), scratch_size,
+      scratch_level, perf_opts,
+      0, nblocks - 1, kb.s, kb.e + dk, jb.s, jb.e + dj,
       KOKKOS_LAMBDA(team_mbr_t member, const int b, const int k, const int j) {
         bool xflux = (k <= kb.e && j <= jb.e);
         bool yflux = (ndim > 1 && k <= kb.e);
@@ -414,7 +424,7 @@ Real MassHistory(MeshData<Real> *md, const Real x1min, const Real x1max, const R
   const auto &mesh_vol = params.Get<Real>("mesh_volume");
 
   std::vector<std::string> vars = {"U"};
-  const auto pack = md->PackVariables(vars);
+  const auto pack = md->PackVariablesByNames(vars);
 
   Real result = 0.0;
   parthenon::par_reduce(

src/CMakeLists.txt

@@ -10,6 +10,8 @@
 # license in this material to reproduce, prepare derivative works, distribute copies to
 # the public, perform publicly and display publicly, and to permit others to do so.
 #=========================================================================================
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#=========================================================================================
 
 # This file was made in part with generative AI.
 
@@ -379,7 +381,8 @@ endif()
 # Therefore, we don't use the Kokkos_ENABLE_CUDA_CONSTEXPR option add the flag manually.
 # Also, not checking for NVIDIA as nvcc_wrapper is identified as GNU so we just make sure
 # the flag is not added when compiling with Clang for Cuda.
-if (Kokkos_ENABLE_CUDA AND NOT CMAKE_CXX_COMPILER_ID STREQUAL "Clang")
+if (Kokkos_ENABLE_CUDA AND NOT CMAKE_CXX_COMPILER_ID STREQUAL "Clang"
+   AND NOT CMAKE_CXX_COMPILER_ID STREQUAL "NVHPC" )
    target_compile_options(parthenon PUBLIC $<$<COMPILE_LANGUAGE:CXX>:--expt-relaxed-constexpr>)
 endif()
 

src/amr_criteria/refinement_package.cpp

@@ -10,6 +10,8 @@
 // license in this material to reproduce, prepare derivative works, distribute copies to
 // the public, perform publicly and display publicly, and to permit others to do so.
 //========================================================================================
+// Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+//========================================================================================
 
 #include "amr_criteria/refinement_package.hpp"
 
@@ -141,6 +143,7 @@ void FirstDerivative(const AMRBounds &bnds, MeshData<Real> *md, const std::strin
                      const int &idx, ParArray1D<AmrTag> &amr_tags,
                      const Real refine_criteria_, const Real derefine_criteria_,
                      const int max_level_) {
+  PARTHENON_INSTRUMENT
   CheckRefinementLoop(
       bnds, md, field, idx, amr_tags, refine_criteria_, derefine_criteria_, max_level_,
       KOKKOS_LAMBDA(SparsePack<> pack, const int ndim, const int b, const int var,
@@ -166,6 +169,7 @@ void SecondDerivative(const AMRBounds &bnds, MeshData<Real> *md, const std::stri
                       const int &idx, ParArray1D<AmrTag> &amr_tags,
                       const Real refine_criteria_, const Real derefine_criteria_,
                       const int max_level_) {
+  PARTHENON_INSTRUMENT
   CheckRefinementLoop(
       bnds, md, field, idx, amr_tags, refine_criteria_, derefine_criteria_, max_level_,
       KOKKOS_LAMBDA(SparsePack<> pack, const int ndim, const int b, const int var,
@@ -217,14 +221,107 @@ TaskStatus Tag(MeshBlockData<Real> *rc) {
 template <>
 TaskStatus Tag(MeshData<Real> *md) {
   PARTHENON_INSTRUMENT
+  Kokkos::Profiling::pushRegion("CheckAllRefinement");
   ParArray1D<AmrTag> amr_tags = CheckAllRefinement(md);
   auto amr_tags_h = amr_tags.GetHostMirrorAndCopy();
+  Kokkos::Profiling::popRegion();
 
+  Kokkos::Profiling::pushRegion("SetRefinement1");
   for (int i = 0; i < md->NumBlocks(); i++) {
     SetRefinement_(md->GetBlockData(i).get(), amr_tags_h(i));
   }
+  Kokkos::Profiling::popRegion();
+
+  Kokkos::Profiling::pushRegion("SetRefinement2");
+  for (int i = 0; i < md->NumBlocks(); i++) {
+    auto pmb = md->GetBlockData(i).get()->GetBlockPointer();
+    pmb->pmr->SetRefinement(amr_tags_h(i));
+  }
+  Kokkos::Profiling::popRegion();
+  return TaskStatus::complete;
+}
+
+template <>
+TaskStatus TagFused(MeshData<Real> *md) {
+  PARTHENON_INSTRUMENT
+
+  Kokkos::Profiling::pushRegion("FusedFirstDerivative");
+  std::vector<std::string> vars = {"U"};
+  auto &v = md->PackVariables(vars);
+  IndexRange ib = md->GetBoundsI(IndexDomain::interior);
+  IndexRange jb = md->GetBoundsJ(IndexDomain::interior);
+  IndexRange kb = md->GetBoundsK(IndexDomain::interior);
+  // printf("Tag: %d,%d,%d,%d,%d,%d\n", ib.s, ib.e, jb.s, jb.e, kb.s, kb.e);
+  // printf("v ndim = %d\n", v.GetNdim());
+  // printf("v dim: %d, %d, %d, %d, %d\n", v.GetDim(1), v.GetDim(2), v.GetDim(3), v.GetDim(4), v.GetDim(5));
+  const int nblocks = md->NumBlocks();
+  Kokkos::View<double*> d_maxd("d_maxd", nblocks);
+  AMRBounds bnds(ib,jb,kb);
+  const int ndim = 1 + (bnds.je > bnds.js) + (bnds.ke > bnds.ks);  
+  parthenon::par_for_outer
+    (DEFAULT_OUTER_LOOP_PATTERN, "FusedFirstDerivative", DevExecSpace(), 0, 0,
+     0, nblocks - 1, 
+     KOKKOS_LAMBDA(parthenon::team_mbr_t member, const int b) {
+      Real t_maxd = 0;
+      const int ksize = bnds.ke - bnds.ks + 1;
+      const int jsize = bnds.je - bnds.js + 1;
+      const int isize = bnds.ie - bnds.is + 1;
+      const int nsize = isize*jsize*ksize;
+      Kokkos::parallel_reduce
+        (Kokkos::TeamThreadRange(member, nsize),
+          [=](const int ii, Real &maxd) {
+          int k = ii / (isize*jsize);
+          int j = (ii - k*isize*jsize)/isize;
+          int i = ii - j*isize - k*isize*jsize;
+          k += bnds.ks;
+          j += bnds.js;
+          i += bnds.is;
+          Real scale = std::abs(v(b, 3, k, j, i));
+          Real d =
+            0.5 * std::abs((v(b, 3, k, j, i + 1) - v(b, 3, k, j, i - 1))) / (scale + TINY_NUMBER);
+          maxd = (d > maxd ? d : maxd);
+          if (ndim > 1) {
+            d = 0.5 * std::abs((v(b, 3, k, j + 1, i) - v(b, 3, k, j - 1, i))) / (scale + TINY_NUMBER);
+            maxd = (d > maxd ? d : maxd);
+          }
+          if (ndim > 2) {
+            d = 0.5 * std::abs((v(b, 3, k + 1, j, i) - v(b, 3, k - 1, j, i))) / (scale + TINY_NUMBER);
+            maxd = (d > maxd ? d : maxd);
+          }
+        }, Kokkos::Max<Real>(t_maxd));
+      if (member.team_rank() == 0) d_maxd(b) = t_maxd;
+    });
+  auto h_maxd = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), d_maxd);
+  Kokkos::Profiling::popRegion();
+
+  Kokkos::Profiling::pushRegion("SetRefinement1");
+  std::vector<AmrTag> flags(nblocks);
+  for (int i = 0; i < nblocks; i++) {
+    AmrTag t;    
+    for (auto &pkg : md->GetBlockData(i).get()->GetBlockPointer()->packages.AllPackages()) {
+      for (auto &amr: pkg.second->amr_criteria) {
+        if (h_maxd(i) > amr->refine_criteria) {
+          t = AmrTag::refine;
+        } else if (h_maxd(i) < amr->derefine_criteria) {
+          t = AmrTag::derefine;
+        } else {
+          t = AmrTag::same;
+        }      
+      }
+    }
+    flags[i] = t;
+  }
+  Kokkos::Profiling::popRegion();
+
+  Kokkos::Profiling::pushRegion("SetRefinement2");
+  for (int i = 0; i < md->NumBlocks(); i++) {
+    auto pmb = md->GetBlockData(i).get()->GetBlockPointer();
+    pmb->pmr->SetRefinement(flags[i]);
+  }
+  Kokkos::Profiling::popRegion();
   return TaskStatus::complete;
 }
 
+
 } // namespace Refinement
 } // namespace parthenon

src/amr_criteria/refinement_package.hpp

@@ -10,6 +10,8 @@
 // license in this material to reproduce, prepare derivative works, distribute copies to
 // the public, perform publicly and display publicly, and to permit others to do so.
 //========================================================================================
+// Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+//========================================================================================
 
 #ifndef AMR_CRITERIA_REFINEMENT_PACKAGE_HPP_
 #define AMR_CRITERIA_REFINEMENT_PACKAGE_HPP_
@@ -36,6 +38,8 @@ namespace Refinement {
 std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin);
 template <typename T>
 TaskStatus Tag(T *rc);
+template <typename T>
+TaskStatus TagFused(T *rc);
 
 AmrTag CheckAllRefinement(MeshBlockData<Real> *rc, const AmrTag &level);
 ParArray1D<AmrTag> CheckAllRefinement(MeshData<Real> *md);

src/bvals/comms/boundary_communication.cpp

@@ -14,6 +14,8 @@
 // license in this material to reproduce, prepare derivative works, distribute copies to
 // the public, perform publicly and display publicly, and to permit others to do so.
 //========================================================================================
+// Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+//========================================================================================
 
 #include <algorithm>
 #include <cstdio>
@@ -163,6 +165,12 @@ TaskStatus SendBoundBufsWithRestrictOption(std::shared_ptr<MeshData<Real>> &md,
     else
       buf.SendNull(coal_comm);
   }
+//#ifdef MPI_PARALLEL
+//  //WIP: This optional MPI_Barrier forces many MPI libraries to start the MPI_Isend before continuing. Can improve performance
+//  // If the initial number of meshblocks is less than the MPI ranks, this causes a hang
+//  MPI_Barrier(MPI_COMM_WORLD);
+//  //END OF WIP
+//#endif
   if (pmesh->do_coalesced_comms)
     pmesh->pcoalesced_comms->PackAndSend(md.get(), bound_type);