Skip to content

Commit f1ef81e

Browse files
atmyersWeiqunZhang
andauthored
Fix GPU restart for pure SoA particles (#3783)
The proposed changes: - [x] fix a bug or incorrect behavior in AMReX - [ ] 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: Weiqun Zhang <[email protected]>
1 parent 3525b4a commit f1ef81e

File tree

2 files changed

+45
-18
lines changed

2 files changed

+45
-18
lines changed

Src/Particle/AMReX_ParticleIO.H

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -999,7 +999,7 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
999999
for (int i = 0; i < cnt; i++) {
10001000
// note: for pure SoA particle layouts, we do write the id, cpu and positions as a struct
10011001
// for backwards compatibility with readers
1002-
if (!ParticleType::is_soa_particle && convert_ids) {
1002+
if (convert_ids) {
10031003
std::int32_t xi, yi;
10041004
std::uint32_t xu, yu;
10051005
xi = iptr[0];

Src/Particle/AMReX_WriteBinaryParticleData.H

Lines changed: 44 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -183,7 +183,7 @@ packIOData (Vector<int>& idata, Vector<ParticleReal>& rdata, const PC& pc, int l
183183
idata.resize(np*iChunkSize);
184184

185185
int num_output_real = 0;
186-
for (int i = 0; i < pc.NumRealComps() + PC::NStructReal; ++i) {
186+
for (int i = 0; i < (int) write_real_comp.size(); ++i) {
187187
if (write_real_comp[i]) { ++num_output_real; }
188188
}
189189

@@ -249,15 +249,18 @@ packIOData (Vector<int>& idata, Vector<ParticleReal>& rdata, const PC& pc, int l
249249
}
250250
}
251251

252-
for (int j = 0; j < PC::SuperParticleType::NReal; j++) {
253-
if (write_real_comp_d_ptr[j]) {
252+
// extra SoA Real components
253+
const int real_start_offset = PC::ParticleType::is_soa_particle ? AMREX_SPACEDIM : 0; // pure SoA: skip positions
254+
for (int j = real_start_offset; j < PC::SuperParticleType::NReal; j++) {
255+
const int write_comp_index = j-real_start_offset;
256+
if (write_real_comp_d_ptr[write_comp_index]) {
254257
rdata_d_ptr[rout_index] = p.rdata(j);
255258
rout_index++;
256259
}
257260
}
258261

259262
for (int j = 0; j < ptd.m_num_runtime_real; j++) {
260-
if (write_real_comp_d_ptr[PC::SuperParticleType::NReal + j]) {
263+
if (write_real_comp_d_ptr[PC::SuperParticleType::NReal+j-real_start_offset]) {
261264
rdata_d_ptr[rout_index] = ptd.m_runtime_rdata[j][pindex];
262265
rout_index++;
263266
}
@@ -336,13 +339,25 @@ packIOData (Vector<int>& idata, Vector<ParticleReal>& rdata, const PC& pc, int l
336339
}
337340
}
338341
else {
339-
amrex::ignore_unused(is_checkpoint);
340-
// Int: id, cpu
341342
uint64_t idcpu = soa.GetIdCPUData()[pindex];
342-
*iptr = (int) ParticleIDWrapper(idcpu);
343-
iptr += 1;
344-
*iptr = (int) ParticleCPUWrapper(idcpu);
345-
iptr += 1;
343+
if (is_checkpoint) {
344+
std::int32_t xi, yi;
345+
std::uint32_t xu, yu;
346+
xu = (std::uint32_t)((idcpu & 0xFFFFFFFF00000000LL) >> 32);
347+
yu = (std::uint32_t)( idcpu & 0xFFFFFFFFLL);
348+
std::memcpy(&xi, &xu, sizeof(xu));
349+
std::memcpy(&yi, &yu, sizeof(yu));
350+
*iptr = xi;
351+
iptr += 1;
352+
*iptr = yi;
353+
iptr += 1;
354+
} else {
355+
// Int: id, cpu
356+
*iptr = (int) ParticleIDWrapper(idcpu);
357+
iptr += 1;
358+
*iptr = (int) ParticleCPUWrapper(idcpu);
359+
iptr += 1;
360+
}
346361

347362
// Real: position
348363
for (int j = 0; j < AMREX_SPACEDIM; j++) { rptr[j] = soa.GetRealData(j)[pindex]; }
@@ -361,8 +376,7 @@ packIOData (Vector<int>& idata, Vector<ParticleReal>& rdata, const PC& pc, int l
361376
// extra SoA Real components
362377
const int real_start_offset = PC::ParticleType::is_soa_particle ? AMREX_SPACEDIM : 0; // pure SoA: skip positions
363378
for (int j = real_start_offset; j < pc.NumRealComps(); j++) {
364-
const int write_comp_offset = PC::ParticleType::is_soa_particle ? AMREX_SPACEDIM : 0; // pure SoA: skip positions
365-
const int write_comp_index = PC::NStructReal+j-write_comp_offset;
379+
const int write_comp_index = PC::NStructReal+j-real_start_offset;
366380
if (write_real_comp[write_comp_index]) {
367381
*rptr = (ParticleReal) soa.GetRealData(j)[pindex];
368382
++rptr;
@@ -1032,12 +1046,25 @@ void WriteBinaryParticleDataAsync (PC const& pc,
10321046
}
10331047
}
10341048
else {
1035-
// Ints: id, cpu
10361049
uint64_t idcpu = soa.GetIdCPUData()[pindex];
1037-
*iptr = (int) ParticleIDWrapper(idcpu);
1038-
iptr += 1;
1039-
*iptr = (int) ParticleCPUWrapper(idcpu);
1040-
iptr += 1;
1050+
if (is_checkpoint) {
1051+
std::int32_t xi, yi;
1052+
std::uint32_t xu, yu;
1053+
xu = (std::uint32_t)((idcpu & 0xFFFFFFFF00000000LL) >> 32);
1054+
yu = (std::uint32_t)( idcpu & 0xFFFFFFFFLL);
1055+
std::memcpy(&xi, &xu, sizeof(xu));
1056+
std::memcpy(&yi, &yu, sizeof(yu));
1057+
*iptr = xi;
1058+
iptr += 1;
1059+
*iptr = yi;
1060+
iptr += 1;
1061+
} else {
1062+
// Int: id, cpu
1063+
*iptr = (int) ParticleIDWrapper(idcpu);
1064+
iptr += 1;
1065+
*iptr = (int) ParticleCPUWrapper(idcpu);
1066+
iptr += 1;
1067+
}
10411068
}
10421069

10431070
// extra SoA Ints

0 commit comments

Comments
 (0)