Named SoA Support (AMReX-Codes#4163)

## Summary Add support to optionally name SoA Real and Int components, both for compile-time and runtime components. - [x] define and defaults - [x] single tile `Get*Data(std::string)` - [x] refactor out default name generation for reuse, finalize defaults in `AddReal/IntComp` - [x] test ## Additional background Close AMReX-Codes#3614 ## Checklist The proposed changes: - [ ] fix a bug or incorrect behavior in AMReX - [x] add new capabilities to AMReX - [ ] changes answers in the test suite to more than roundoff level - [ ] are likely to significantly affect the results of downstream AMReX users - [ ] include documentation in the code and/or rst files, if appropriate --------- Co-authored-by: Andrew Myers <[email protected]>
WeiqunZhang · Oct 9, 2024 · 8df11b6 · 8df11b6
1 parent bd56c53
commit 8df11b6
Show file tree

Hide file tree

Showing 10 changed files with 367 additions and 60 deletions.
diff --git a/Src/Particle/AMReX_ParticleContainer.H b/Src/Particle/AMReX_ParticleContainer.H
@@ -60,6 +60,7 @@
 #include <memory>
 #include <numeric>
 #include <random>
+#include <string>
 #include <tuple>
 #include <type_traits>
 #include <utility>
@@ -1144,7 +1145,8 @@ public:
  */
  ParticleTileType& DefineAndReturnParticleTile (int lev, int grid, int tile)
  {
- m_particles[lev][std::make_pair(grid, tile)].define(NumRuntimeRealComps(), NumRuntimeIntComps());
+ m_particles[lev][std::make_pair(grid, tile)].define(NumRuntimeRealComps(), NumRuntimeIntComps(), &m_soa_rdata_names, &m_soa_idata_names);
+
  return ParticlesAt(lev, grid, tile);
  }
 
@@ -1247,10 +1249,10 @@ public:
 
  Long superParticleSize() const { return superparticle_size; }
 
- template <typename T,
- std::enable_if_t<std::is_same_v<T,bool>,int> = 0>
- void AddRealComp (T communicate=true)
+ void AddRealComp (std::string const & name, int communicate=1)
  {
+ m_soa_rdata_names.push_back(name);
+
  m_runtime_comps_defined = true;
  m_num_runtime_real++;
  h_redistribute_real_comp.push_back(communicate);
@@ -1270,10 +1272,15 @@ public:
  }
  }
 
- template <typename T,
- std::enable_if_t<std::is_same_v<T,bool>,int> = 0>
- void AddIntComp (T communicate=true)
+ void AddRealComp (int communicate=1)
  {
+ AddRealComp(getDefaultCompNameReal<ParticleType>(NArrayReal+m_num_runtime_real), communicate);
+ }
+
+ void AddIntComp (std::string const & name, int communicate=1)
+ {
+ m_soa_idata_names.push_back(name);
+
  m_runtime_comps_defined = true;
  m_num_runtime_int++;
  h_redistribute_int_comp.push_back(communicate);
@@ -1293,6 +1300,11 @@ public:
  }
  }
 
+ void AddIntComp (int communicate=1)
+ {
+ AddIntComp(getDefaultCompNameInt<ParticleType>(NArrayInt+m_num_runtime_int), communicate);
+ }
+
  int NumRuntimeRealComps () const { return m_num_runtime_real; }
  int NumRuntimeIntComps () const { return m_num_runtime_int; }
 
@@ -1403,6 +1415,15 @@ public:
 #include "AMReX_ParticlesHDF5.H"
 #endif
 
+ /** Overwrite the default names for the compile-time SoA components */
+ void SetSoACompileTimeNames (std::vector<std::string> const & rdata_name, std::vector<std::string> const & idata_name);
+
+ /** Get the names for the real SoA components **/
+ std::vector<std::string> GetRealSoANames () const {return m_soa_rdata_names;}
+
+ /** Get the names for the int SoA components **/
+ std::vector<std::string> GetIntSoANames () const {return m_soa_idata_names;}
+
 protected:
 
  template <class RTYPE>
@@ -1435,6 +1456,10 @@ private:
  size_t particle_size, superparticle_size;
  int num_real_comm_comps, num_int_comm_comps;
  Vector<ParticleLevel> m_particles;
+
+ // names of both compile-time and runtime Real and Int SoA data
+ std::vector<std::string> m_soa_rdata_names;
+ std::vector<std::string> m_soa_idata_names;
 };
 
 template <int T_NStructReal, int T_NStructInt, int T_NArrayReal, int T_NArrayInt, template<class> class Allocator, class CellAssignor>

diff --git a/Src/Particle/AMReX_ParticleContainerI.H b/Src/Particle/AMReX_ParticleContainerI.H
@@ -1,6 +1,10 @@
-#include <type_traits>
 #include <AMReX_MakeParticle.H>
 
+#include <string>
+#include <type_traits>
+#include <vector>
+
+
 template <typename ParticleType, int NArrayReal, int NArrayInt,
  template<class> class Allocator, class CellAssignor>
 void
@@ -60,10 +64,40 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
  pp.query("do_unlink", doUnlink);
  pp.queryAdd("do_mem_efficient_sort", memEfficientSort);
 
+ // add default names for SoA Real and Int compile-time arguments
+ for (int i=0; i<NArrayReal; ++i)
+ {
+ m_soa_rdata_names.push_back(getDefaultCompNameReal<ParticleType>(i));
+ }
+ for (int i=0; i<NArrayInt; ++i)
+ {
+ m_soa_idata_names.push_back(getDefaultCompNameInt<ParticleType>(i));
+ }
+
  initialized = true;
  }
 }
 
+template <typename ParticleType, int NArrayReal, int NArrayInt,
+ template<class> class Allocator, class CellAssignor>
+void
+ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssignor> :: SetSoACompileTimeNames (
+ std::vector<std::string> const & rdata_name, std::vector<std::string> const & idata_name
+)
+{
+ AMREX_ALWAYS_ASSERT_WITH_MESSAGE(rdata_name.size() == NArrayReal, "rdata_name must be equal to NArrayReal");
+ AMREX_ALWAYS_ASSERT_WITH_MESSAGE(idata_name.size() == NArrayInt, "idata_name must be equal to NArrayInt");
+
+ for (int i=0; i<NArrayReal; ++i)
+ {
+ m_soa_rdata_names.at(i) = rdata_name.at(i);
+ }
+ for (int i=0; i<NArrayInt; ++i)
+ {
+ m_soa_idata_names.at(i) = idata_name.at(i);
+ }
+}
+
 template <typename ParticleType, int NArrayReal, int NArrayInt,
  template<class> class Allocator, class CellAssignor>
 template <typename P, typename Assignor>
@@ -1161,7 +1195,7 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
  }
  } else {
  ParticleTileType ptile_tmp;
- ptile_tmp.define(m_num_runtime_real, m_num_runtime_int);
+ ptile_tmp.define(m_num_runtime_real, m_num_runtime_int, &m_soa_rdata_names, &m_soa_idata_names);
  ptile_tmp.resize(np_total);
  // copy re-ordered particles
  gatherParticles(ptile_tmp, ptile, np, permutations);
@@ -1498,7 +1532,7 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
  tmp_local[lev][index].resize(num_threads);
  soa_local[lev][index].resize(num_threads);
  for (int t = 0; t < num_threads; ++t) {
- soa_local[lev][index][t].define(m_num_runtime_real, m_num_runtime_int);
+ soa_local[lev][index][t].define(m_num_runtime_real, m_num_runtime_int, &m_soa_rdata_names, &m_soa_idata_names);
  }
  }
  }

diff --git a/Src/Particle/AMReX_ParticleIO.H b/Src/Particle/AMReX_ParticleIO.H
@@ -51,20 +51,18 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
 {
  Vector<int> write_real_comp;
  Vector<std::string> tmp_real_comp_names;
- int nrc = ParticleType::is_soa_particle ? NStructReal + NumRealComps() - AMREX_SPACEDIM : NStructReal + NumRealComps();
 
- for (int i = 0; i < nrc; ++i )
+ int first_rcomp = ParticleType::is_soa_particle ? AMREX_SPACEDIM : 0;
+ for (int i = first_rcomp; i < NStructReal + NumRealComps(); ++i )
  {
  write_real_comp.push_back(1);
  if (real_comp_names.empty())
  {
- std::stringstream ss;
- ss << "real_comp" << i;
- tmp_real_comp_names.push_back(ss.str());
+ tmp_real_comp_names.push_back(getDefaultCompNameReal<ParticleType>(i));
  }
  else
  {
- tmp_real_comp_names.push_back(real_comp_names[i]);
+ tmp_real_comp_names.push_back(real_comp_names[i-first_rcomp]);
  }
  }
 
@@ -75,9 +73,7 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
  write_int_comp.push_back(1);
  if (int_comp_names.empty())
  {
- std::stringstream ss;
- ss << "int_comp" << i;
- tmp_int_comp_names.push_back(ss.str());
+ tmp_int_comp_names.push_back(getDefaultCompNameInt<ParticleType>(i));
  }
  else
  {
@@ -98,24 +94,20 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
 {
  Vector<int> write_real_comp;
  Vector<std::string> real_comp_names;
- int nrc = ParticleType::is_soa_particle ? NStructReal + NumRealComps() - AMREX_SPACEDIM : NStructReal + NumRealComps();
 
- for (int i = 0; i < nrc; ++i )
+ int first_rcomp = ParticleType::is_soa_particle ? AMREX_SPACEDIM : 0;
+ for (int i = first_rcomp; i < NStructReal + NumRealComps(); ++i )
  {
  write_real_comp.push_back(1);
- std::stringstream ss;
- ss << "real_comp" << i;
- real_comp_names.push_back(ss.str());
+ real_comp_names.push_back(getDefaultCompNameReal<ParticleType>(i));
  }
 
  Vector<int> write_int_comp;
  Vector<std::string> int_comp_names;
  for (int i = 0; i < NStructInt + NumIntComps(); ++i )
  {
  write_int_comp.push_back(1);
- std::stringstream ss;
- ss << "int_comp" << i;
- int_comp_names.push_back(ss.str());
+ int_comp_names.push_back(getDefaultCompNameInt<ParticleType>(i));
  }
 
  WriteBinaryParticleData(dir, name, write_real_comp, write_int_comp,
@@ -182,9 +174,7 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
  Vector<std::string> int_comp_names;
  for (int i = 0; i < NStructInt + NumIntComps(); ++i )
  {
- std::stringstream ss;
- ss << "int_comp" << i;
- int_comp_names.push_back(ss.str());
+ int_comp_names.push_back(getDefaultCompNameInt<ParticleType>(i));
  }
 
  WriteBinaryParticleData(dir, name,
@@ -211,20 +201,16 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
  AMREX_ASSERT(write_int_comp.size() == NStructInt + NArrayInt );
 
  Vector<std::string> real_comp_names;
- int nrc = ParticleType::is_soa_particle ? NStructReal + NumRealComps() - AMREX_SPACEDIM : NStructReal + NumRealComps();
- for (int i = 0; i < nrc; ++i )
+ int first_rcomp = ParticleType::is_soa_particle ? AMREX_SPACEDIM : 0;
+ for (int i = first_rcomp; i < NStructReal + NumRealComps(); ++i )
  {
- std::stringstream ss;
- ss << "real_comp" << i;
- real_comp_names.push_back(ss.str());
+ real_comp_names.push_back(getDefaultCompNameReal<ParticleType>(i));
  }
 
  Vector<std::string> int_comp_names;
  for (int i = 0; i < NStructInt + NumIntComps(); ++i )
  {
- std::stringstream ss;
- ss << "int_comp" << i;
- int_comp_names.push_back(ss.str());
+ int_comp_names.push_back(getDefaultCompNameInt<ParticleType>(i));
  }
 
  WriteBinaryParticleData(dir, name, write_real_comp, write_int_comp,
@@ -259,24 +245,20 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
 {
  Vector<int> write_real_comp;
  Vector<std::string> real_comp_names;
- int nrc = ParticleType::is_soa_particle ? NStructReal + NumRealComps() - AMREX_SPACEDIM : NStructReal + NumRealComps();
 
- for (int i = 0; i < nrc; ++i )
+ int first_rcomp = ParticleType::is_soa_particle ? AMREX_SPACEDIM : 0;
+ for (int i = first_rcomp; i < NStructReal + NumRealComps(); ++i )
  {
  write_real_comp.push_back(1);
- std::stringstream ss;
- ss << "real_comp" << i;
- real_comp_names.push_back(ss.str());
+ real_comp_names.push_back(getDefaultCompNameReal<ParticleType>(i));
  }
 
  Vector<int> write_int_comp;
  Vector<std::string> int_comp_names;
  for (int i = 0; i < NStructInt + NumIntComps(); ++i )
  {
  write_int_comp.push_back(1);
- std::stringstream ss;
- ss << "int_comp" << i;
- int_comp_names.push_back(ss.str());
+ int_comp_names.push_back(getDefaultCompNameInt<ParticleType>(i));
  }
 
  WriteBinaryParticleData(dir, name, write_real_comp, write_int_comp,
@@ -345,9 +327,7 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
  Vector<std::string> int_comp_names;
  for (int i = 0; i < NStructInt + NumIntComps(); ++i )
  {
- std::stringstream ss;
- ss << "int_comp" << i;
- int_comp_names.push_back(ss.str());
+ int_comp_names.push_back(getDefaultCompNameInt<ParticleType>(i));
  }
 
  WriteBinaryParticleData(dir, name,
@@ -374,20 +354,16 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
  AMREX_ASSERT(write_int_comp.size() == NStructInt + NumIntComps() );
 
  Vector<std::string> real_comp_names;
- int nrc = ParticleType::is_soa_particle ? NStructReal + NumRealComps() - AMREX_SPACEDIM : NStructReal + NumRealComps();
- for (int i = 0; i < nrc; ++i )
+ int first_rcomp = ParticleType::is_soa_particle ? AMREX_SPACEDIM : 0;
+ for (int i = first_rcomp; i < NStructReal + NumRealComps(); ++i )
  {
- std::stringstream ss;
- ss << "real_comp" << i;
- real_comp_names.push_back(ss.str());
+ real_comp_names.push_back(getDefaultCompNameReal<ParticleType>(i));
  }
 
  Vector<std::string> int_comp_names;
  for (int i = 0; i < NStructInt + NumIntComps(); ++i )
  {
- std::stringstream ss;
- ss << "int_comp" << i;
- int_comp_names.push_back(ss.str());
+ int_comp_names.push_back(getDefaultCompNameInt<ParticleType>(i));
  }
 
  WriteBinaryParticleData(dir, name, write_real_comp, write_int_comp,

diff --git a/Src/Particle/AMReX_ParticleTile.H b/Src/Particle/AMReX_ParticleTile.H
@@ -11,7 +11,10 @@
 #include <AMReX_RealVect.H>
 
 #include <array>
+#include <string>
 #include <type_traits>
+#include <vector>
+
 
 namespace amrex {
 
@@ -730,10 +733,15 @@ struct ParticleTile
  ParticleTile& operator= (ParticleTile &&) noexcept = default;
 #endif
 
- void define (int a_num_runtime_real, int a_num_runtime_int)
+ void define (
+ int a_num_runtime_real,
+ int a_num_runtime_int,
+ std::vector<std::string>* soa_rdata_names=nullptr,
+ std::vector<std::string>* soa_idata_names=nullptr
+ )
  {
  m_defined = true;
- GetStructOfArrays().define(a_num_runtime_real, a_num_runtime_int);
+ GetStructOfArrays().define(a_num_runtime_real, a_num_runtime_int, soa_rdata_names, soa_idata_names);
  m_runtime_r_ptrs.resize(a_num_runtime_real);
  m_runtime_i_ptrs.resize(a_num_runtime_int);
  m_runtime_r_cptrs.resize(a_num_runtime_real);

diff --git a/Src/Particle/AMReX_ParticleUtil.H b/Src/Particle/AMReX_ParticleUtil.H
@@ -883,6 +883,26 @@ void PermutationForDeposition (Gpu::DeviceVector<index_type>& perm, index_type n
  });
 }
 
+template <typename P>
+std::string getDefaultCompNameReal (const int i) {
+ int first_r_name = 0;
+ if constexpr (P::is_soa_particle) {
+ if (i < AMREX_SPACEDIM) {
+ constexpr int x_in_ascii = 120;
+ std::string const name{char(x_in_ascii+i)};
+ return name;
+ }
+ first_r_name = AMREX_SPACEDIM;
+ }
+ std::string const name{("real_comp" + std::to_string(i-first_r_name))};
+ return name;
+}
+
+template <typename P>
+std::string getDefaultCompNameInt (const int i) {
+ std::string const name{("int_comp" + std::to_string(i))};
+ return name;
+}
 
 #ifdef AMREX_USE_HDF5_ASYNC
 void async_vol_es_wait_particle();