Add Kokkos timers

env/bash

@@ -69,7 +69,7 @@ if [[ $PARTITION == "unknown" ]]; then
 elif [[ $PARTITION == "chicoma-gpu" ]]; then
     module unload cray-libsci
     module load PrgEnv-gnu
-    module load cudatoolkit/24.7_12.5
+    module load cudatoolkit
     export CRAY_ACCEL_TARGET=nvidia80
     export MPICH_GPU_SUPPORT_ENABLED=1
     export MPICH_GPU_MANAGED_MEMORY_SUPPORT_ENABLED=1

src/artemis_driver.cpp

@@ -114,6 +114,7 @@ ArtemisDriver<GEOM>::ArtemisDriver(ParameterInput *pin, ApplicationInput *app_in
 //! \brief Assembles the tasks associated with a step for the ArtemisDriver
 template <Coordinates GEOM>
 TaskListStatus ArtemisDriver<GEOM>::Step() {
+  PARTHENON_INSTRUMENT
   // Prepare registers
   PreStepTasks();
   TaskListStatus status = TaskListStatus::complete;
@@ -158,6 +159,7 @@ TaskListStatus ArtemisDriver<GEOM>::Step() {
 //! \brief Defines the tasks executed prior to the main integrator in the ArtemisDriver
 template <Coordinates GEOM>
 void ArtemisDriver<GEOM>::PreStepTasks() {
+  PARTHENON_INSTRUMENT
   // set the integration timestep
   integrator->dt = tm.dt;
   if (do_nbody) nbody_integrator->dt = tm.dt;
@@ -194,6 +196,7 @@ void ArtemisDriver<GEOM>::PreStepTasks() {
 //! \brief Defines the main integrator's TaskCollection for the ArtemisDriver
 template <Coordinates GEOM>
 TaskCollection ArtemisDriver<GEOM>::StepTasks() {
+  PARTHENON_INSTRUMENT
   using TQ = TaskQualifier;
   using namespace ::parthenon::Update;
   TaskCollection tc;
@@ -338,6 +341,7 @@ TaskCollection ArtemisDriver<GEOM>::StepTasks() {
 //! \brief Defines the TaskCollection for post step tasks in the ArtemisDriver
 template <Coordinates GEOM>
 TaskCollection ArtemisDriver<GEOM>::PostStepTasks() {
+  PARTHENON_INSTRUMENT
   using namespace ::parthenon::Update;
   TaskCollection tc;
   TaskID none(0);

src/derived/fill_derived.cpp

@@ -29,6 +29,7 @@ namespace ArtemisDerived {
 //! NOTE(PDM): Note that this function is not called during remeshing.
 template <Coordinates GEOM>
 TaskStatus SetAuxillaryFields(MeshData<Real> *md) {
+  PARTHENON_INSTRUMENT
   using parthenon::MakePackDescriptor;
   auto pm = md->GetParentPointer();
   auto &resolved_pkgs = pm->resolved_packages;
@@ -94,6 +95,7 @@ TaskStatus SetAuxillaryFields(MeshData<Real> *md) {
 //! or remeshing event in preparation for FillGhost
 template <Coordinates GEOM>
 void ConsToPrim(MeshData<Real> *md) {
+  PARTHENON_INSTRUMENT
   using parthenon::MakePackDescriptor;
   auto pm = md->GetParentPointer();
   auto &resolved_pkgs = pm->resolved_packages;
@@ -224,6 +226,7 @@ void ConsToPrim(MeshData<Real> *md) {
 //! \brief Executes P2C following integrator updates and/or remeshing events
 template <typename T, Coordinates GEOM>
 void PrimToCons(T *md) {
+  PARTHENON_INSTRUMENT
   using parthenon::MakePackDescriptor;
   auto pm = md->GetParentPointer();
   auto &resolved_pkgs = pm->resolved_packages;
@@ -383,6 +386,7 @@ void PrimToCons(T *md) {
 //! but before PreCommFillDerived
 template <Coordinates GEOM>
 void PostInitialization(MeshBlock *pmb, ParameterInput *pin) {
+  PARTHENON_INSTRUMENT
   auto &md = pmb->meshblock_data.Get();
   PrimToCons<MeshBlockData<Real>, GEOM>(md.get());
 }
@@ -392,6 +396,7 @@ void PostInitialization(MeshBlock *pmb, ParameterInput *pin) {
 //! \brief Syncs unsplit fields following an operator split update
 template <Coordinates GEOM>
 TaskCollection SyncFields(Mesh *pmesh, const Real time, const Real dt) {
+  PARTHENON_INSTRUMENT
   using namespace ::parthenon::Update;
   TaskCollection tc;
   TaskID none(0);

src/drag/drag.cpp

@@ -89,6 +89,7 @@ std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin) {
 //! \brief Calls source terms for drag
 template <Coordinates GEOM>
 TaskStatus DragSource(MeshData<Real> *md, const Real time, const Real dt) {
+  PARTHENON_INSTRUMENT
   using parthenon::MakePackDescriptor;
   auto pm = md->GetParentPointer();
 

src/drag/drag.hpp

@@ -164,6 +164,7 @@ template <Diffusion::DiffType DTYP, Coordinates GEOM>
 TaskStatus SelfDragSourceImpl(MeshData<Real> *md, const Real time, const Real dt,
                               const Diffusion::DiffCoeffParams &dp, const EOS &eos_d,
                               const SelfDragParams &gasp, const SelfDragParams &dustp) {
+  PARTHENON_INSTRUMENT
   using TE = parthenon::TopologicalElement;
   auto pm = md->GetParentPointer();
   auto &resolved_pkgs = pm->resolved_packages;
@@ -344,6 +345,7 @@ TaskStatus SimpleDragSourceImpl(MeshData<Real> *md, const Real time, const Real
                                 const Diffusion::DiffCoeffParams &dp, const EOS &eos_d,
                                 const SelfDragParams &gasp, const SelfDragParams &dustp,
                                 const StoppingTimeParams &tp) {
+  PARTHENON_INSTRUMENT
   using TE = parthenon::TopologicalElement;
   auto pm = md->GetParentPointer();
   auto &resolved_pkgs = pm->resolved_packages;

src/dust/dust.cpp

@@ -223,6 +223,7 @@ std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin,
 //! \brief Compute dust hydrodynamics timestep
 template <Coordinates GEOM>
 Real EstimateTimestepMesh(MeshData<Real> *md) {
+  PARTHENON_INSTRUMENT
   using parthenon::MakePackDescriptor;
   using RotatingFrame::BackgroundVelocity;
   auto pm = md->GetParentPointer();
@@ -271,6 +272,7 @@ Real EstimateTimestepMesh(MeshData<Real> *md) {
 //! \fn  TaskStatus Dust::CalculateFluxes
 //! \brief Evaluates advective fluxes for dust evolution
 TaskStatus CalculateFluxes(MeshData<Real> *md, const bool pcm) {
+  PARTHENON_INSTRUMENT
   auto pm = md->GetParentPointer();
   auto &resolved_pkgs = pm->resolved_packages;
 
@@ -300,6 +302,7 @@ TaskStatus CalculateFluxes(MeshData<Real> *md, const bool pcm) {
 //! \fn  TaskStatus Dust::FluxSource
 //! \brief Evaluates coordinate terms from advective fluxes for dust evolution
 TaskStatus FluxSource(MeshData<Real> *md, const Real dt) {
+  PARTHENON_INSTRUMENT
   auto pm = md->GetParentPointer();
   auto &resolved_pkgs = pm->resolved_packages;
 
@@ -336,6 +339,7 @@ TaskStatus FluxSource(MeshData<Real> *md, const Real dt) {
 //! \brief Add history outputs for dust quantities for generic coordinate system
 template <Coordinates GEOM>
 void AddHistoryImpl(Params &params) {
+  PARTHENON_INSTRUMENT
   using namespace ArtemisUtils;
   auto HstSum = parthenon::UserHistoryOperation::sum;
   using parthenon::HistoryOutputVar;
@@ -361,6 +365,7 @@ void AddHistoryImpl(Params &params) {
 //! \fn  void Dust::AddHistory
 //! \brief Add history outputs for dust quantities
 void AddHistory(Coordinates coords, Params &params) {
+  PARTHENON_INSTRUMENT
   if (coords == Coordinates::cartesian) {
     AddHistoryImpl<Coordinates::cartesian>(params);
   } else if (coords == Coordinates::cylindrical) {

src/gas/cooling/beta_cooling.cpp

@@ -39,6 +39,7 @@ namespace Cooling {
 //!     Tp - T = -om dt (T - T0) / (beta + om dt)
 template <Coordinates GEOM, TempRefType TTYP>
 TaskStatus BetaCooling(MeshData<Real> *md, const Real time, const Real dt) {
+  PARTHENON_INSTRUMENT
   using parthenon::MakePackDescriptor;
   using TE = parthenon::TopologicalElement;
   auto pm = md->GetParentPointer();

src/gas/cooling/cooling.cpp

@@ -90,6 +90,7 @@ std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin) {
 //! \brief Wrapper function for external cooling options
 template <Coordinates GEOM>
 TaskStatus CoolingSource(MeshData<Real> *md, const Real time, const Real dt) {
+  PARTHENON_INSTRUMENT
   auto pm = md->GetParentPointer();
   auto &pkg = pm->packages.Get("cooling");
   CoolingType ctype = pkg->template Param<CoolingType>("type");

src/gas/gas.cpp

@@ -439,6 +439,7 @@ std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin,
 //! \brief Compute gas hydrodynamics timestep
 template <Coordinates GEOM>
 Real EstimateTimestepMesh(MeshData<Real> *md) {
+  PARTHENON_INSTRUMENT
   using parthenon::MakePackDescriptor;
   using RotatingFrame::BackgroundVelocity;
   auto pm = md->GetParentPointer();
@@ -527,6 +528,7 @@ Real EstimateTimestepMesh(MeshData<Real> *md) {
 //! \fn  TaskStatus Gas::CalculateFluxes
 //! \brief Evaluates advective fluxes for gas evolution
 TaskStatus CalculateFluxes(MeshData<Real> *md, const bool pcm) {
+  PARTHENON_INSTRUMENT
   auto pm = md->GetParentPointer();
   auto &resolved_pkgs = pm->resolved_packages;
 
@@ -559,6 +561,7 @@ TaskStatus CalculateFluxes(MeshData<Real> *md, const bool pcm) {
 //! \fn  TaskStatus Gas::FluxSource
 //! \brief Evaluates coordinate terms from advective fluxes for gas evolution
 TaskStatus FluxSource(MeshData<Real> *md, const Real dt) {
+  PARTHENON_INSTRUMENT
   auto pm = md->GetParentPointer();
   auto &resolved_pkgs = pm->resolved_packages;
 
@@ -590,6 +593,7 @@ TaskStatus FluxSource(MeshData<Real> *md, const Real dt) {
 //  \brief Evaluates viscous flux
 template <Coordinates GEOM>
 TaskStatus ViscousFlux(MeshData<Real> *md) {
+  PARTHENON_INSTRUMENT
   auto pm = md->GetParentPointer();
   auto &pkg = pm->packages.Get("gas");
 
@@ -629,6 +633,7 @@ TaskStatus ViscousFlux(MeshData<Real> *md) {
 //  \brief Evaluates thermal flux
 template <Coordinates GEOM>
 TaskStatus ThermalFlux(MeshData<Real> *md) {
+  PARTHENON_INSTRUMENT
   auto pm = md->GetParentPointer();
   auto &pkg = pm->packages.Get("gas");
 
@@ -666,6 +671,7 @@ TaskStatus ThermalFlux(MeshData<Real> *md) {
 //! \fn  TaskStatus Gas::ZeroDiffusionFlux
 //  \brief Resets the diffusion flux
 TaskStatus ZeroDiffusionFlux(MeshData<Real> *md) {
+  PARTHENON_INSTRUMENT
   auto pm = md->GetParentPointer();
   auto &pkg = pm->packages.Get("gas");
 
@@ -683,6 +689,7 @@ TaskStatus ZeroDiffusionFlux(MeshData<Real> *md) {
 //  \brief Applies the diffusion fluxes to update the momenta and energy
 template <Coordinates GEOM>
 TaskStatus DiffusionUpdate(MeshData<Real> *md, const Real dt) {
+  PARTHENON_INSTRUMENT
   auto pm = md->GetParentPointer();
   auto &pkg = pm->packages.Get("gas");
 
@@ -709,6 +716,7 @@ TaskStatus DiffusionUpdate(MeshData<Real> *md, const Real dt) {
 }
 
 TaskStatus DepositEnergy(MeshData<Real> *md, const Real dt) {
+  PARTHENON_INSTRUMENT
   using parthenon::MakePackDescriptor;
   auto pm = md->GetParentPointer();
   auto &resolved_pkgs = pm->resolved_packages;
@@ -739,6 +747,7 @@ TaskStatus DepositEnergy(MeshData<Real> *md, const Real dt) {
 //! \brief Add history outputs for gas quantities for generic coordinate system
 template <Coordinates GEOM>
 void AddHistoryImpl(Params &params) {
+  PARTHENON_INSTRUMENT
   using namespace ArtemisUtils;
   auto HstSum = parthenon::UserHistoryOperation::sum;
   using parthenon::HistoryOutputVar;
@@ -772,6 +781,7 @@ void AddHistoryImpl(Params &params) {
 //! \fn  void Gas::AddHistory
 //! \brief Add history outputs for gas quantities
 void AddHistory(Coordinates coords, Params &params) {
+  PARTHENON_INSTRUMENT
   if (coords == Coordinates::cartesian) {
     AddHistoryImpl<Coordinates::cartesian>(params);
   } else if (coords == Coordinates::cylindrical) {

src/geometry/geometry.cpp

@@ -106,6 +106,7 @@ std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin) {
 //! \brief Initializes the geometry meshblock data
 template <Coordinates GEOM>
 void InitBlockGeom(MeshBlock *pmb, ParameterInput *pin) {
+  PARTHENON_INSTRUMENT
   using parthenon::MakePackDescriptor;
   auto pm = pmb->pmy_mesh;
   const int ndim_ = pm->ndim;

src/gravity/binary_mass.cpp

@@ -25,6 +25,7 @@ namespace Gravity {
 //! \brief Applies accelerations due to a binary
 template <Coordinates GEOM>
 TaskStatus BinaryMassGravity(MeshData<Real> *md, const Real time, const Real dt) {
+  PARTHENON_INSTRUMENT
   using parthenon::MakePackDescriptor;
   using TE = parthenon::TopologicalElement;
   auto pm = md->GetParentPointer();

src/gravity/gravity.cpp

@@ -133,6 +133,7 @@ std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin,
 //! \brief Wrapper function for external gravity options
 template <Coordinates GEOM>
 TaskStatus ExternalGravity(MeshData<Real> *md, const Real time, const Real dt) {
+  PARTHENON_INSTRUMENT
   auto pm = md->GetParentPointer();
 
   auto &pkg = pm->packages.Get("gravity");

src/gravity/nbody_gravity.hpp

@@ -162,6 +162,7 @@ NBodyGravityImpl(V1 &vmesh, V2 &vg, const geometry::Coords<GEOM> &coords,
 //! \brief Applies accelerations due to collection of point masses
 template <Coordinates GEOM>
 TaskStatus NBodyGravity(MeshData<Real> *md, const Real time, const Real dt) {
+  PARTHENON_INSTRUMENT
   using parthenon::MakePackDescriptor;
   using TE = parthenon::TopologicalElement;
   auto pm = md->GetParentPointer();

src/gravity/point_mass.cpp

@@ -25,6 +25,7 @@ namespace Gravity {
 //! \brief Applies accelerations due to a point mass gravitational potential
 template <Coordinates GEOM>
 TaskStatus PointMassGravity(MeshData<Real> *md, const Real time, const Real dt) {
+  PARTHENON_INSTRUMENT
   using parthenon::MakePackDescriptor;
   using TE = parthenon::TopologicalElement;
   auto pm = md->GetParentPointer();

src/gravity/uniform.cpp

@@ -26,6 +26,7 @@ namespace Gravity {
 //! \brief Applies accelerations due to a constant g
 template <Coordinates GEOM>
 TaskStatus UniformGravity(MeshData<Real> *md, const Real time, const Real dt) {
+  PARTHENON_INSTRUMENT
   using parthenon::MakePackDescriptor;
   using TE = parthenon::TopologicalElement;
   auto pm = md->GetParentPointer();

src/main.cpp

@@ -24,6 +24,7 @@
 
 parthenon::DriverStatus LaunchWorkFlow(parthenon::ParthenonManager &pman,
                                        parthenon::ParameterInput *pin) {
+  PARTHENON_INSTRUMENT
   // Geometry specific routines
   // (1) Handle ProblemGenerators and associated modifiers
   // (2) Call ParthenonInit to set up the mesh and packages

src/nbody/nbody.cpp

@@ -261,6 +261,7 @@ std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin,
 //! \fn  Real NBody::EstimateTimestepMesh
 //! \brief Compute NBody timestep
 Real EstimateTimestepMesh(MeshData<Real> *md) {
+  PARTHENON_INSTRUMENT
   using parthenon::MakePackDescriptor;
   auto pm = md->GetParentPointer();
 
@@ -274,6 +275,7 @@ Real EstimateTimestepMesh(MeshData<Real> *md) {
 //! \brief Distance based refinement criterion
 template <Coordinates GEOM>
 AmrTag DistanceRefinement(MeshBlockData<Real> *md) {
+  PARTHENON_INSTRUMENT
   auto pmb = md->GetBlockPointer();
   auto pm = pmb->pmy_mesh;
   auto &resolved_pkgs = pm->resolved_packages;
@@ -324,6 +326,7 @@ AmrTag DistanceRefinement(MeshBlockData<Real> *md) {
 //! \brief Create REBOUND restart file and store in Params to reuse as Parthenon restart
 void UserWorkBeforeRestartOutputMesh(Mesh *pmesh, ParameterInput *, SimTime &,
                                      OutputParameters *) {
+  PARTHENON_INSTRUMENT
   auto &artemis_pkg = pmesh->packages.Get("artemis");
   if (!(artemis_pkg->Param<bool>("do_nbody"))) return;
 

src/nbody/nbody_advance.cpp

@@ -84,6 +84,7 @@ namespace NBody {
 //!         stage 2 integrates to end
 TaskStatus Advance(Mesh *pm, const Real time, const int stage,
                    const parthenon::LowStorageIntegrator *nbody_integ) {
+  PARTHENON_INSTRUMENT
   auto &nbody_pkg = pm->packages.Get("nbody");
   auto particle_id = nbody_pkg->Param<std::vector<int>>("particle_id");
   auto particles = nbody_pkg->Param<ParArray1D<Particle>>("particles");

src/nbody/nbody_outputs.cpp

@@ -30,6 +30,7 @@ namespace NBody {
 //! \fn  void NBody::Outputs
 //! \brief
 void Outputs(parthenon::Mesh *pm, const Real time) {
+  PARTHENON_INSTRUMENT
   // Return immediately if nbody outputs are disabled
   auto nbody = pm->packages.Get("nbody").get();
   if (nbody->Param<bool>("disable_outputs")) return;

src/nbody/nbody_utils.hpp

@@ -122,6 +122,7 @@ static void PrintSystem(const int npart, ParArray1D<Particle> particles) {
 //! \brief Copy the positions and velocities from the rebound sim to the Particles list
 static void SyncWithRebound(RebSim &r_sim, std::vector<int> particle_id,
                             ParArray1D<Particle> particles) {
+  PARTHENON_INSTRUMENT
   const auto npart = particle_id.size();
   auto particles_h = particles.GetHostMirrorAndCopy();
 
@@ -211,6 +212,7 @@ static void enable_stderr(int stderr_save_fd) {
 //! \fn  int NBody::write_bytes_to_file
 //! \brief
 inline void write_bytes_to_file(std::string filename, std::vector<BYTE> &bytes) {
+  PARTHENON_INSTRUMENT
   std::ofstream outfile(filename.c_str(), std::ios::binary);
   if (outfile.is_open()) {
     outfile.write(reinterpret_cast<char *>(bytes.data()), bytes.size());
@@ -224,6 +226,7 @@ inline void write_bytes_to_file(std::string filename, std::vector<BYTE> &bytes)
 //! \fn  int NBody::read_bytes_from_file
 //! \brief
 inline std::vector<BYTE> read_bytes_from_file(std::string filename) {
+  PARTHENON_INSTRUMENT
   std::ifstream file(NBody::rebound_filename, std::ios::binary | std::ios::ate);
   if (file.is_open()) {
     auto size = file.tellg();

src/pgen/advection.hpp

@@ -62,6 +62,7 @@ static AdvectionVariables av;
 //! \brief Sets initial conditions for advection tests
 template <Coordinates GEOM>
 inline void ProblemGenerator(MeshBlock *pmb, ParameterInput *pin) {
+  PARTHENON_INSTRUMENT
   using parthenon::MakePackDescriptor;
   const Mesh *pmesh = pmb->pmy_mesh;
   const int ndim = pmesh->ndim;
@@ -240,6 +241,7 @@ inline void ProblemGenerator(MeshBlock *pmb, ParameterInput *pin) {
 //! periods.
 template <Coordinates GEOM>
 inline void UserWorkAfterLoop(Mesh *pmesh, ParameterInput *pin, parthenon::SimTime &tm) {
+  PARTHENON_INSTRUMENT
   using parthenon::MakePackDescriptor;
   const int nhyd_vars = 5;
   const int nspec_vars = 4;