Skip to content

Commit 2e1106e

Browse files
atmyersax3l
atmyers
and
ax3l
authored
Implement Checkpoint, Restart, and Plotfile for pure SOA (#3332)
Follow-up to #2878. The proposed changes: - [x] 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: Axel Huebl <[email protected]>
1 parent 07f87b0 commit 2e1106e

File tree

15 files changed

+544
-88
lines changed

15 files changed

+544
-88
lines changed

‎Src/Particle/AMReX_ParticleContainer.H‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1346,7 +1346,8 @@ protected:
13461346
* \param lev_max
13471347
* \param local_grid
13481348
*/
1349-
bool EnforcePeriodicWhere (ParticleType& prt, ParticleLocData& pld,
1349+
template <typename P>
1350+
bool EnforcePeriodicWhere (P& prt, ParticleLocData& pld,
13501351
int lev_min = 0, int lev_max = -1, int local_grid=-1) const;
13511352

13521353
public:
@@ -1378,7 +1379,8 @@ private:
13781379
void RedistributeMPI (std::map<int, Vector<char> >& not_ours,
13791380
int lev_min = 0, int lev_max = 0, int nGrow = 0, int local=0);
13801381

1381-
void locateParticle (ParticleType& p, ParticleLocData& pld,
1382+
template <typename P>
1383+
void locateParticle (P& p, ParticleLocData& pld,
13821384
int lev_min, int lev_max, int nGrow, int local_grid=-1) const;
13831385

13841386
void Initialize ();

‎Src/Particle/AMReX_ParticleContainerI.H‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -164,9 +164,10 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator>
164164

165165
template <typename ParticleType, int NArrayReal, int NArrayInt,
166166
template<class> class Allocator>
167+
template <typename P>
167168
bool
168169
ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator>
169-
::EnforcePeriodicWhere (ParticleType& p,
170+
::EnforcePeriodicWhere (P& p,
170171
ParticleLocData& pld,
171172
int lev_min,
172173
int lev_max,
@@ -312,8 +313,9 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator>::resizeDa
312313

313314
template <typename ParticleType, int NArrayReal, int NArrayInt,
314315
template<class> class Allocator>
316+
template <typename P>
315317
void
316-
ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator>::locateParticle (ParticleType& p, ParticleLocData& pld,
318+
ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator>::locateParticle (P& p, ParticleLocData& pld,
317319
int lev_min, int lev_max, int nGrow, int local_grid) const
318320
{
319321
bool success;

‎Src/Particle/AMReX_ParticleIO.H‎

Lines changed: 63 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -919,7 +919,7 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator>
919919
int* iptr = istuff.dataPtr();
920920
RTYPE* rptr = rstuff.dataPtr();
921921

922-
ParticleType p;
922+
Particle<NStructReal, NStructInt> ptemp;
923923
ParticleLocData pld;
924924

925925
Vector<std::map<std::pair<int, int>, Gpu::HostVector<ParticleType> > > host_particles;
@@ -937,77 +937,113 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator>
937937
host_int_attribs.resize(finest_level_in_file+1);
938938

939939
for (int i = 0; i < cnt; i++) {
940-
if (convert_ids) {
940+
// note: for pure SoA particle layouts, we do write the id, cpu and positions as a struct
941+
// for backwards compatibility with readers
942+
if (!ParticleType::is_soa_particle && convert_ids) {
941943
std::int32_t xi, yi;
942944
std::uint32_t xu, yu;
943945
xi = iptr[0];
944946
yi = iptr[1];
945947
std::memcpy(&xu, &xi, sizeof(xi));
946948
std::memcpy(&yu, &yi, sizeof(yi));
947-
p.m_idcpu = ((std::uint64_t)xu) << 32 | yu;
949+
ptemp.m_idcpu = ((std::uint64_t)xu) << 32 | yu;
948950
} else {
949-
p.id() = iptr[0];
950-
p.cpu() = iptr[1];
951+
ptemp.id() = iptr[0];
952+
ptemp.cpu() = iptr[1];
951953
}
952954
iptr += 2;
953955

954956
for (int j = 0; j < NStructInt; j++)
955957
{
956-
p.idata(j) = *iptr;
958+
ptemp.idata(j) = *iptr;
957959
++iptr;
958960
}
959961

960-
AMREX_ASSERT(p.id() > 0);
962+
AMREX_ASSERT(ptemp.id() > 0);
961963

962-
AMREX_D_TERM(p.pos(0) = ParticleReal(rptr[0]);,
963-
p.pos(1) = ParticleReal(rptr[1]);,
964-
p.pos(2) = ParticleReal(rptr[2]););
964+
AMREX_D_TERM(ptemp.pos(0) = ParticleReal(rptr[0]);,
965+
ptemp.pos(1) = ParticleReal(rptr[1]);,
966+
ptemp.pos(2) = ParticleReal(rptr[2]););
965967

966968
rptr += AMREX_SPACEDIM;
967969

968970
for (int j = 0; j < NStructReal; j++)
969971
{
970-
p.rdata(j) = ParticleReal(*rptr);
972+
ptemp.rdata(j) = ParticleReal(*rptr);
971973
++rptr;
972974
}
973975

974-
locateParticle(p, pld, 0, finestLevel(), 0);
976+
locateParticle(ptemp, pld, 0, finestLevel(), 0);
975977

976978
std::pair<int, int> ind(grd, pld.m_tile);
977979

978980
host_real_attribs[lev][ind].resize(NumRealComps());
979981
host_int_attribs[lev][ind].resize(NumIntComps());
980982

981983
// add the struct
982-
host_particles[lev][ind].push_back(p);
984+
if constexpr(!ParticleType::is_soa_particle)
985+
{
986+
host_particles[lev][ind].push_back(ptemp);
983987

984-
// add the real...
985-
for (int icomp = 0; icomp < NumRealComps(); icomp++) {
986-
host_real_attribs[lev][ind][icomp].push_back(*rptr);
987-
++rptr;
988-
}
988+
// add the real...
989+
for (int icomp = 0; icomp < NumRealComps(); icomp++) {
990+
host_real_attribs[lev][ind][icomp].push_back(*rptr);
991+
++rptr;
992+
}
989993

990-
// ... and int array data
991-
for (int icomp = 0; icomp < NumIntComps(); icomp++) {
992-
host_int_attribs[lev][ind][icomp].push_back(*iptr);
993-
++iptr;
994+
// ... and int array data
995+
for (int icomp = 0; icomp < NumIntComps(); icomp++) {
996+
host_int_attribs[lev][ind][icomp].push_back(*iptr);
997+
++iptr;
998+
}
999+
} else {
1000+
host_particles[lev][ind];
1001+
1002+
for (int j = 0; j < AMREX_SPACEDIM; j++) {
1003+
host_real_attribs[pld.m_lev][ind][j].push_back(ptemp.pos(j));
1004+
}
1005+
1006+
host_int_attribs[pld.m_lev][ind][0].push_back(ptemp.id());
1007+
host_int_attribs[pld.m_lev][ind][1].push_back(ptemp.cpu());
1008+
1009+
// read all other SoA
1010+
// add the real...
1011+
for (int icomp = AMREX_SPACEDIM; icomp < NumRealComps(); icomp++) {
1012+
host_real_attribs[lev][ind][icomp].push_back(*rptr);
1013+
++rptr;
1014+
}
1015+
1016+
// ... and int array data
1017+
for (int icomp = 2; icomp < NumIntComps(); icomp++) {
1018+
host_int_attribs[lev][ind][icomp].push_back(*iptr);
1019+
++iptr;
1020+
}
9941021
}
9951022
}
9961023

9971024
for (int host_lev = 0; host_lev < static_cast<int>(host_particles.size()); ++host_lev)
998-
{
1025+
{
9991026
for (auto& kv : host_particles[host_lev]) {
10001027
auto grid = kv.first.first;
10011028
auto tile = kv.first.second;
10021029
const auto& src_tile = kv.second;
10031030

10041031
auto& dst_tile = DefineAndReturnParticleTile(host_lev, grid, tile);
1005-
auto old_size = dst_tile.GetArrayOfStructs().size();
1006-
auto new_size = old_size + src_tile.size();
1032+
auto old_size = dst_tile.size();
1033+
auto new_size = old_size;
1034+
if constexpr(!ParticleType::is_soa_particle)
1035+
{
1036+
new_size += src_tile.size();
1037+
} else {
1038+
new_size += host_real_attribs[host_lev][std::make_pair(grid,tile)][0].size();
1039+
}
10071040
dst_tile.resize(new_size);
10081041

1009-
Gpu::copyAsync(Gpu::hostToDevice, src_tile.begin(), src_tile.end(),
1010-
dst_tile.GetArrayOfStructs().begin() + old_size);
1042+
if constexpr(!ParticleType::is_soa_particle)
1043+
{
1044+
Gpu::copyAsync(Gpu::hostToDevice, src_tile.begin(), src_tile.end(),
1045+
dst_tile.GetArrayOfStructs().begin() + old_size);
1046+
}
10111047

10121048
for (int i = 0; i < NumRealComps(); ++i) {
10131049
Gpu::copyAsync(Gpu::hostToDevice,

‎Src/Particle/AMReX_ParticleTile.H‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -226,6 +226,9 @@ struct ParticleTileData
226226
{
227227
AMREX_ASSERT(index < m_size);
228228
SuperParticleType sp;
229+
for (int i = 0; i < AMREX_SPACEDIM; ++i) {sp.pos(i) = m_rdata[i][index];}
230+
sp.id() = m_idata[0][index];
231+
sp.cpu() = m_idata[1][index];
229232
for (int i = 0; i < NAR; ++i)
230233
sp.rdata(i) = m_rdata[i][index];
231234
for (int i = 0; i < NAI; ++i)
@@ -595,6 +598,9 @@ struct ConstParticleTileData
595598
{
596599
AMREX_ASSERT(index < m_size);
597600
SuperParticleType sp;
601+
for (int i = 0; i < AMREX_SPACEDIM; ++i) {sp.pos(i) = m_rdata[i][index];}
602+
sp.id() = m_idata[0][index];
603+
sp.cpu() = m_idata[1][index];
598604
for (int i = 0; i < NAR; ++i)
599605
sp.rdata(i) = m_rdata[i][index];
600606
for (int i = 0; i < NAI; ++i)

0 commit comments

Comments
 (0)