Named SoA Support (#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 #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]>
AMReX-Codes · 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();