Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Named SoA Support #4163

Merged
merged 10 commits into from
Oct 9, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 32 additions & 7 deletions Src/Particle/AMReX_ParticleContainer.H
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,7 @@
#include <memory>
#include <numeric>
#include <random>
#include <string>
#include <tuple>
#include <type_traits>
#include <utility>
Expand Down Expand Up @@ -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);
}

Expand Down Expand Up @@ -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);
Expand All @@ -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);
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

zero indexed and used before m_num_runtime_real++, looks correct.

}

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);
Expand All @@ -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; }

Expand Down Expand Up @@ -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>
Expand Down Expand Up @@ -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>
Expand Down
40 changes: 37 additions & 3 deletions Src/Particle/AMReX_ParticleContainerI.H
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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>
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
}
}
}
Expand Down
70 changes: 23 additions & 47 deletions Src/Particle/AMReX_ParticleIO.H
Original file line number Diff line number Diff line change
Expand Up @@ -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]);
}
}

Expand All @@ -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
{
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down
12 changes: 10 additions & 2 deletions Src/Particle/AMReX_ParticleTile.H
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,10 @@
#include <AMReX_RealVect.H>

#include <array>
#include <string>
#include <type_traits>
#include <vector>


namespace amrex {

Expand Down Expand Up @@ -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);
Expand Down
20 changes: 20 additions & 0 deletions Src/Particle/AMReX_ParticleUtil.H
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
Loading
Loading