Add MemoryLocation {Host, Device} member in bunch_particles to keep track of where particles are

src/synergia/bunch/bunch.h

@@ -658,7 +658,7 @@ class bunch_t {
 
  // checkpoint partiles
  void
- save_checkpoint_particles(Hdf5_file& file, int idx) const
+ save_checkpoint_particles(Hdf5_file& file, int idx)
  {
  get_bunch_particles(PG::regular).save_checkpoint_particles(file, idx);
  get_bunch_particles(PG::spectator).save_checkpoint_particles(file, idx);

src/synergia/bunch/bunch_particles.cc

@@ -718,8 +718,7 @@ bunch_particles_t<double>::print_particle(size_t idx, Logger& logger) const
 
 template <>
 void
-bunch_particles_t<double>::save_checkpoint_particles(Hdf5_file& file,
- int idx) const
+bunch_particles_t<double>::save_checkpoint_particles(Hdf5_file& file, int idx)
 {
  checkout_particles();
 

src/synergia/bunch/bunch_particles.h

@@ -16,6 +16,9 @@
 #include <openPMD/openPMD.hpp>
 #endif
 
+// To keep track of the memory location for the particles
+enum class MemoryLocation { Host, Device };
+
 enum class ParticleGroup { regular = 0, spectator = 1 };
 
 using Particles = Kokkos::View<double* [7],
@@ -195,6 +198,9 @@ class bunch_particles_t {
  * ones. This one should be used when looping through the
  * particles array
  *
+ * 'memory_location' indicates whether the particles are currently
+ * on the host or on the device.
+ *
  */
 
  // particle group (regular or spectator)
@@ -205,6 +211,7 @@ class bunch_particles_t {
  int n_valid;
  int n_active;
  int n_reserved;
+ MemoryLocation memory_location;
 
  // sum of n_valid on all ranks
  int n_total;
@@ -259,6 +266,12 @@ class bunch_particles_t {
  return n_last_discarded;
  }
 
+ MemoryLocation
+ get_memory_location() const
+ {
+ return memory_location;
+ }
+
  std::pair<size_t, size_t>
  get_local_particle_count_in_range(int num_part, int offset) const
  {
@@ -303,17 +316,19 @@ class bunch_particles_t {
 
  // copy particles/masks between host and device memories
  void
- checkin_particles() const
+ checkin_particles()
  {
  Kokkos::deep_copy(parts, hparts);
  Kokkos::deep_copy(masks, hmasks);
+ memory_location = MemoryLocation::Device;
  }
 
  void
- checkout_particles() const
+ checkout_particles()
  {
  Kokkos::deep_copy(hparts, parts);
  Kokkos::deep_copy(hmasks, masks);
+ memory_location = MemoryLocation::Host;
  }
 
  // change capacity (can only increase)
@@ -395,7 +410,7 @@ class bunch_particles_t {
  Commxx const& comm) const;
 
  // checkpoint save/load
- void save_checkpoint_particles(Hdf5_file& file, int idx) const;
+ void save_checkpoint_particles(Hdf5_file& file, int idx);
  void load_checkpoint_particles(Hdf5_file& file, int idx);
 
  // assign ids cooperatively
@@ -427,6 +442,7 @@ class bunch_particles_t {
  save(AR& ar) const
  {
  ar(CEREAL_NVP(label));
+ ar(CEREAL_NVP(memory_location));
  ar(CEREAL_NVP(n_valid));
  ar(CEREAL_NVP(n_active));
  ar(CEREAL_NVP(n_reserved));
@@ -442,6 +458,7 @@ class bunch_particles_t {
  load(AR& ar)
  {
  ar(CEREAL_NVP(label));
+ ar(CEREAL_NVP(memory_location));
  ar(CEREAL_NVP(n_valid));
  ar(CEREAL_NVP(n_active));
  ar(CEREAL_NVP(n_reserved));
@@ -580,6 +597,7 @@ inline bunch_particles_t<PART>::bunch_particles_t(ParticleGroup pg,
  Commxx const& comm)
  : group(pg)
  , label(pg == PG::regular ? "particles" : "spectators")
+ , memory_location(MemoryLocation::Host)
  , n_valid(0)
  , n_active(0)
  , n_reserved(0)

src/synergia/bunch/core_diagnostics.cc

@@ -1,10 +1,6 @@
 #include <cmath>
-#include <functional>
 #include <stdexcept>
 
-#include "synergia/foundation/physical_constants.h"
-#include "synergia/utils/logger.h"
-
 #include "core_diagnostics.h"
 
 namespace core_diagnostics_impl {
@@ -210,8 +206,13 @@ Core_diagnostics::calculate_mean(Bunch const& bunch)
 
  karray1d mean("mean", 6);
 
- auto particles = bunch.get_local_particles();
- auto masks = bunch.get_local_particle_masks();
+ auto bparts = bunch.get_bunch_particles();
+ if (bparts.get_memory_location() == MemoryLocation::Host) {
+ std::runtime_error("Bunch particles are active on Host memory space!");
+ }
+
+ auto particles = bparts.parts;
+ auto masks = bparts.masks;
  const int npart = bunch.size();
 
  particle_reducer<mean_tag> pr(particles, masks);
@@ -235,8 +236,12 @@ Core_diagnostics::calculate_z_mean(Bunch const& bunch)
 
  double mean = 0;
 
- auto particles = bunch.get_local_particles();
- auto masks = bunch.get_local_particle_masks();
+ auto bparts = bunch.get_bunch_particles();
+ if (bparts.get_memory_location() == MemoryLocation::Host) {
+ std::runtime_error("Bunch particles are active on Host memory space!");
+ }
+ auto particles = bparts.parts;
+ auto masks = bparts.masks;
  const int npart = bunch.size();
 
  particle_reducer<z_mean_tag> pr(particles, masks);
@@ -259,8 +264,12 @@ Core_diagnostics::calculate_abs_mean(Bunch const& bunch)
 
  karray1d abs_mean("abs_mean", 6);
 
- auto particles = bunch.get_local_particles();
- auto masks = bunch.get_local_particle_masks();
+ auto bparts = bunch.get_bunch_particles();
+ if (bparts.get_memory_location() == MemoryLocation::Host) {
+ std::runtime_error("Bunch particles are active on Host memory space!");
+ }
+ auto particles = bparts.parts;
+ auto masks = bparts.masks;
  const int npart = bunch.size();
 
  particle_reducer<abs_mean_tag> pr(particles, masks);
@@ -288,8 +297,12 @@ Core_diagnostics::calculate_std(Bunch const& bunch, karray1d const& mean)
 
  karray1d std("std", 6);
 
- auto particles = bunch.get_local_particles();
- auto masks = bunch.get_local_particle_masks();
+ auto bparts = bunch.get_bunch_particles();
+ if (bparts.get_memory_location() == MemoryLocation::Host) {
+ std::runtime_error("Bunch particles are active on Host memory space!");
+ }
+ auto particles = bparts.parts;
+ auto masks = bparts.masks;
  const int npart = bunch.size();
 
  particle_reducer<std_tag> pr(particles, masks, mean);
@@ -312,8 +325,12 @@ Core_diagnostics::calculate_sum2(Bunch const& bunch, karray1d const& mean)
 
  karray2d_row sum2("sum2", 6, 6);
 
- auto particles = bunch.get_local_particles();
- auto masks = bunch.get_local_particle_masks();
+ auto bparts = bunch.get_bunch_particles();
+ if (bparts.get_memory_location() == MemoryLocation::Host) {
+ std::runtime_error("Bunch particles are active on Host memory space!");
+ }
+ auto particles = bparts.parts;
+ auto masks = bparts.masks;
  auto npart = bunch.size();
 
  particle_reducer<mom2_tag> pr(particles, masks, mean);
@@ -354,8 +371,12 @@ Core_diagnostics::calculate_min(Bunch const& bunch)
  min(1) = 1e100;
  min(2) = 1e100;
 
- auto particles = bunch.get_local_particles();
- auto masks = bunch.get_local_particle_masks();
+ auto bparts = bunch.get_bunch_particles();
+ if (bparts.get_memory_location() == MemoryLocation::Host) {
+ std::runtime_error("Bunch particles are active on Host memory space!");
+ }
+ auto particles = bparts.parts;
+ auto masks = bparts.masks;
  const int npart = bunch.size();
 
  particle_reducer<min_tag> pr(particles, masks);
@@ -379,8 +400,12 @@ Core_diagnostics::calculate_max(Bunch const& bunch)
  max(1) = -1e100;
  max(2) = -1e100;
 
- auto particles = bunch.get_local_particles();
- auto masks = bunch.get_local_particle_masks();
+ auto bparts = bunch.get_bunch_particles();
+ if (bparts.get_memory_location() == MemoryLocation::Host) {
+ std::runtime_error("Bunch particles are active on Host memory space!");
+ }
+ auto particles = bparts.parts;
+ auto masks = bparts.masks;
  const int npart = bunch.size();
 
  particle_reducer<max_tag> pr(particles, masks);
@@ -399,8 +424,12 @@ Core_diagnostics::calculate_spatial_mean_stddev(Bunch const& bunch)
  using core_diagnostics_impl::particle_reducer;
  using core_diagnostics_impl::spatial_mean_stddev_tag;
 
- auto particles = bunch.get_local_particles();
- auto masks = bunch.get_local_particle_masks();
+ auto bparts = bunch.get_bunch_particles();
+ if (bparts.get_memory_location() == MemoryLocation::Host) {
+ std::runtime_error("Bunch particles are active on Host memory space!");
+ }
+ auto particles = bparts.parts;
+ auto masks = bparts.masks;
  const int npart = bunch.size();
 
  const auto total_bunch_particles = bunch.get_total_num();

src/synergia/bunch/tests/test_bunch.cc

@@ -41,17 +41,16 @@ TEST_CASE("Bunch", "[Bunch]")
  CHECK(p2(0, 6) == 123);
  CHECK(p2(1, 6) == 124);
  CHECK(p2(4, 6) == 127);
-
- CHECK(bunch.get_real_num() == Approx(1e13) );
- bunch.set_real_num(1.2e13);
- CHECK(bunch.get_real_num() == Approx(1.2e13) );
 
+ CHECK(bunch.get_real_num() == Approx(1e13));
+ bunch.set_real_num(1.2e13);
+ CHECK(bunch.get_real_num() == Approx(1.2e13));
 }
 
 #if defined SYNERGIA_HAVE_OPENPMD
 
 void
-check_particle_values(BunchParticles const& bp1, BunchParticles const& bp2)
+check_particle_values(BunchParticles& bp1, BunchParticles& bp2)
 {
  bp1.checkout_particles();
  bp2.checkout_particles();

src/synergia/bunch/tests/test_bunch_particles.cc

@@ -25,7 +25,7 @@ init_particle_values(BunchParticles& bp)
 }
 
 void
-check_particle_values(BunchParticles const& bp)
+check_particle_values(BunchParticles& bp)
 {
  bp.checkout_particles();
 

src/synergia/simulation/bunch_simulator.cc

@@ -502,7 +502,7 @@ namespace {
 }
 
 void
-Bunch_simulator::save_checkpoint_particles(std::string const& fname) const
+Bunch_simulator::save_checkpoint_particles(std::string const& fname)
 {
  Hdf5_file file(fname, Hdf5_file::Flag::truncate, *comm);
  auto bunches = get_bunch_ptrs(trains);

src/synergia/simulation/bunch_simulator.h

@@ -430,7 +430,7 @@ class Bunch_simulator {
  const_karray1d limits);
 
  // serialization helper
- void save_checkpoint_particles(std::string const& fname) const;
+ void save_checkpoint_particles(std::string const& fname);
  void load_checkpoint_particles(std::string const& fname);
 
  std::string

src/synergia/simulation/propagator.cc

@@ -46,6 +46,21 @@ Propagator::do_step(Bunch_simulator& simulator,
  // lattice elements has been updated
  lattice.update();
 
+ // ensure that particles are on the device in case any of the
+ // following might have transferred them onto the host
+ // custom diagnostics routines, turn_end_action, etc
+ Kokkos::Profiling::pushRegion("memory-location check before step-apply");
+ for (auto& train : simulator.get_trains()) {
+ for (auto& bunch : train.get_bunches()) {
+ auto bparts = bunch.get_bunch_particles();
+ if (bparts.get_memory_location() == MemoryLocation::Host) {
+ std::runtime_error(
+ "Bunch particles are active on Host memory space!");
+ }
+ }
+ }
+ Kokkos::Profiling::popRegion();
+
  // propagate through the step
  step.apply(simulator, logger);