Skip to content

Improve support for Particles with PolymorphicArenaAllocator #4603

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.

Sign up for GitHub

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

Merged
merged 5 commits into from
Aug 13, 2025
Merged

Improve support for Particles with PolymorphicArenaAllocator #4603

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
e37aeef
Improve support for Particles with PolymorphicArenaAllocator
AlexanderSinn Aug 10, 2025
d28a969
fix )
AlexanderSinn Aug 10, 2025
63c9223
fix RedistributeCPU
AlexanderSinn Aug 10, 2025
b7cea29
add this to neighbors
AlexanderSinn Aug 10, 2025
7ec3a3b
fix
AlexanderSinn Aug 10, 2025
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
7 changes: 3 additions & 4 deletions Src/Base/AMReX_GpuAllocators.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -171,6 +171,9 @@ namespace amrex {
template <typename T>
struct IsPolymorphicArenaAllocator : std::false_type {};

template <typename T>
struct IsPolymorphicArenaAllocator<PolymorphicArenaAllocator<T> > : std::true_type {};

#ifdef AMREX_USE_GPU
template <typename T>
struct RunOnGpu<ArenaAllocator<T> > : std::true_type {};
Expand All @@ -183,10 +186,6 @@ namespace amrex {

template <typename T>
struct RunOnGpu<AsyncArenaAllocator<T> > : std::true_type {};

template <typename T>
struct IsPolymorphicArenaAllocator<PolymorphicArenaAllocator<T> > : std::true_type {};

#endif // AMREX_USE_GPU

#ifdef AMREX_USE_GPU
Expand Down
3 changes: 2 additions & 1 deletion Src/Particle/AMReX_NeighborParticlesI.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -454,7 +454,8 @@ NeighborParticleContainer<NStructReal, NStructInt, NArrayReal, NArrayInt>
const Vector<NeighborIndexMap>& map = local_map[lev][dst_index];
const int num_ghosts = map.size();
neighbors[lev][dst_index].define(this->NumRuntimeRealComps(),
this->NumRuntimeIntComps());
this->NumRuntimeIntComps(),
nullptr, nullptr, this->arena());
neighbors[lev][dst_index].resize(num_ghosts);
local_neighbor_sizes[lev][dst_index] = neighbors[lev][dst_index].size();
}
Expand Down
16 changes: 13 additions & 3 deletions Src/Particle/AMReX_ParticleContainer.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -197,6 +197,9 @@ public:
using ParIterType = ParIter_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssignor>;
using ParConstIterType = ParConstIter_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssignor>;

static constexpr bool has_polymorphic_allocator =
IsPolymorphicArenaAllocator<Allocator<RealType>>::value;

//! \brief Default constructor - construct an empty particle container that has no concept
//! of a level hierarchy. Must be properly initialized later.
ParticleContainer_impl ()
Expand Down Expand Up @@ -1156,8 +1159,11 @@ public:
*/
ParticleTileType& DefineAndReturnParticleTile (int lev, int grid, int tile)
{
m_particles[lev][std::make_pair(grid, tile)].define(NumRuntimeRealComps(), NumRuntimeIntComps(), &m_soa_rdata_names, &m_soa_idata_names);

m_particles[lev][std::make_pair(grid, tile)].define(
NumRuntimeRealComps(), NumRuntimeIntComps(),
&m_soa_rdata_names, &m_soa_idata_names,
arena()
);
return ParticlesAt(lev, grid, tile);
}

Expand Down Expand Up @@ -1186,7 +1192,11 @@ public:
ParticleTileType& DefineAndReturnParticleTile (int lev, const Iterator& iter)
{
auto index = std::make_pair(iter.index(), iter.LocalTileIndex());
m_particles[lev][index].define(NumRuntimeRealComps(), NumRuntimeIntComps(), &m_soa_rdata_names, &m_soa_idata_names);
m_particles[lev][index].define(
NumRuntimeRealComps(), NumRuntimeIntComps(),
&m_soa_rdata_names, &m_soa_idata_names,
arena()
);
return ParticlesAt(lev, iter);
}

Expand Down
9 changes: 9 additions & 0 deletions Src/Particle/AMReX_ParticleContainerBase.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -245,6 +245,14 @@ public:
template <class MF>
bool OnSameGrids (int level, const MF& mf) const { return m_gdb->OnSameGrids(level, mf); }

[[nodiscard]] Arena* arena () const {
return m_arena;
}

void SetArena (Arena* a) {
m_arena = a;
}
Comment on lines +252 to +254
Copy link
Member

@ax3l ax3l Nov 10, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

@AlexanderSinn @atmyers I am wondering: should we add an assert here that setting the arena can only be done once?

There are already good asserts / exceptions now if a PC with an unset allocator is trying to add particles. Thus, the logic here could be as easy as: assert if m_arena is not nullptr when callint SetArena. What do you think?

I want to avoid that someone calls SetArena for an already defined PC with live particles and gets into an unexpected state by that.

Suggested change
void SetArena (Arena* a) {
m_arena = a;
}
void SetArena (Arena* a) {
AMREX_ALWAYS_ASSERT_WITH_MESSAGE(m_arena == nullptr, "SetArena can only be called once per PC!");
m_arena = a;
}

(Technically, we can make it more complicated, like allowing to switch the arena until any tiles/particles are added. That could be useful in some usage patterns when specializing a user container that might set a default in its constructor.)

Copy link
Member

@ax3l ax3l Nov 10, 2025

Choose a reason for hiding this comment

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

x-ref: AMReX-Codes/pyamrex#460 (comment)

Copy link
Member Author

@AlexanderSinn AlexanderSinn Nov 10, 2025

Choose a reason for hiding this comment

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

There is the a_arena == GetStructOfArrays().GetIdCPUData().arena() check in ParticleTile that should catch this issue, but I guess we could add another check to ParticleContainerBase 


static const std::string& CheckpointVersion ();
static const std::string& PlotfileVersion ();
static const std::string& DataPrefix ();
Expand All @@ -269,6 +277,7 @@ protected:
std::unique_ptr<ParGDB> m_gdb_object = std::make_unique<ParGDB>();
ParGDBBase* m_gdb{nullptr};
Vector<std::unique_ptr<MultiFab> > m_dummy_mf;
Arena* m_arena = nullptr;

mutable std::unique_ptr<iMultiFab> redistribute_mask_ptr;
mutable int redistribute_mask_nghost = std::numeric_limits<int>::min();
Expand Down
38 changes: 32 additions & 6 deletions Src/Particle/AMReX_ParticleContainerI.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1255,8 +1255,11 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
}
Gpu::streamSynchronize();
} else {
typename SoA::IdCPU tmp_idcpu(np_total);

typename SoA::IdCPU tmp_idcpu;
if constexpr (has_polymorphic_allocator) {
tmp_idcpu.setArena(arena());
}
tmp_idcpu.resize(np_total);
auto src = ptile.GetStructOfArrays().GetIdCPUData().data();
uint64_t* dst = tmp_idcpu.data();
AMREX_HOST_DEVICE_FOR_1D( np_total, i,
Expand All @@ -1270,7 +1273,11 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
}

{ // Create a scope for the temporary vector below
RealVector tmp_real(np_total);
RealVector tmp_real;
if constexpr (has_polymorphic_allocator) {
tmp_real.setArena(arena());
}
tmp_real.resize(np_total);
for (int comp = 0; comp < NArrayReal + m_num_runtime_real; ++comp) {
auto src = ptile.GetStructOfArrays().GetRealData(comp).data();
ParticleReal* dst = tmp_real.data();
Expand All @@ -1285,7 +1292,11 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
}
}

IntVector tmp_int(np_total);
IntVector tmp_int;
if constexpr (has_polymorphic_allocator) {
tmp_int.setArena(arena());
}
tmp_int.resize(np_total);
for (int comp = 0; comp < NArrayInt + m_num_runtime_int; ++comp) {
auto src = ptile.GetStructOfArrays().GetIntData(comp).data();
int* dst = tmp_int.data();
Expand All @@ -1300,7 +1311,8 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
}
} else {
ParticleTileType ptile_tmp;
ptile_tmp.define(m_num_runtime_real, m_num_runtime_int, &m_soa_rdata_names, &m_soa_idata_names);
ptile_tmp.define(m_num_runtime_real, m_num_runtime_int,
&m_soa_rdata_names, &m_soa_idata_names, arena());
ptile_tmp.resize(np_total);
// copy re-ordered particles
gatherParticles(ptile_tmp, ptile, np, permutations);
Expand Down Expand Up @@ -1643,7 +1655,21 @@ 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, &m_soa_rdata_names, &m_soa_idata_names);
soa_local[lev][index][t].define(m_num_runtime_real, m_num_runtime_int,
&m_soa_rdata_names, &m_soa_idata_names);
if constexpr (has_polymorphic_allocator) {
if constexpr (ParticleType::is_soa_particle) {
soa_local[lev][index][t].GetIdCPUData().setArena(arena());
} else {
tmp_local[lev][index][t].setArena(arena());
}
for (int j = 0; j < soa_local[lev][index][t].NumRealComps(); ++j) {
soa_local[lev][index][t].GetRealData(j).setArena(arena());
}
for (int j = 0; j < soa_local[lev][index][t].NumIntComps(); ++j) {
soa_local[lev][index][t].GetIntData(j).setArena(arena());
}
}
}
}
}
Expand Down
5 changes: 3 additions & 2 deletions Src/Particle/AMReX_ParticleIO.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1215,8 +1215,9 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
auto& pmap = m_particles[lev];
for (const auto& kv : pmap) {
ParticleTile<ParticleType, NArrayReal, NArrayInt,
amrex::PinnedArenaAllocator> pinned_ptile;
pinned_ptile.define(NumRuntimeRealComps(), NumRuntimeIntComps());
amrex::PolymorphicArenaAllocator> pinned_ptile;
pinned_ptile.define(NumRuntimeRealComps(), NumRuntimeIntComps(),
nullptr, nullptr, The_Pinned_Arena());
pinned_ptile.resize(kv.second.numParticles());
amrex::copyParticles(pinned_ptile, kv.second);
const auto& host_aos = pinned_ptile.GetArrayOfStructs();
Expand Down
43 changes: 41 additions & 2 deletions Src/Particle/AMReX_ParticleTile.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -748,6 +748,9 @@ struct ParticleTile
using ParticleTileDataType = ParticleTileData<StorageParticleType, NArrayReal, NArrayInt>;
using ConstParticleTileDataType = ConstParticleTileData<StorageParticleType, NArrayReal, NArrayInt>;

static constexpr bool has_polymorphic_allocator =
IsPolymorphicArenaAllocator<Allocator<RealType>>::value;

ParticleTile () = default;

#ifndef _WIN32 // workaround windows compiler bug
Expand All @@ -764,15 +767,51 @@ struct ParticleTile
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
std::vector<std::string>* soa_idata_names=nullptr,
Arena* a_arena=nullptr
)
{
m_defined = true;
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);
m_runtime_i_cptrs.resize(a_num_runtime_int);

if constexpr (has_polymorphic_allocator) {
if (m_defined) {
// it is not allowed to change the arena after the tile has been defined
if constexpr (ParticleType::is_soa_particle) {
AMREX_ALWAYS_ASSERT_WITH_MESSAGE(
a_arena == GetStructOfArrays().GetIdCPUData().arena(),
"ParticleTile with PolymorphicArenaAllocator redefined with "
"different memory arena");
} else {
AMREX_ALWAYS_ASSERT_WITH_MESSAGE(
a_arena == m_aos_tile().arena(),
"ParticleTile with PolymorphicArenaAllocator redefined with "
"different memory arena");
}
}

AMREX_ALWAYS_ASSERT_WITH_MESSAGE(a_arena != nullptr,
"ParticleTile with PolymorphicArenaAllocator defined with no memory arena! "
"Make sure to call setArena() on the ParticleContainer before initialization or "
"to pass an Arena to ParticleTile::define()");

if constexpr (ParticleType::is_soa_particle) {
GetStructOfArrays().GetIdCPUData().setArena(a_arena);
} else {
m_aos_tile().setArena(a_arena);
}
for (int j = 0; j < NumRealComps(); ++j) {
GetStructOfArrays().GetRealData(j).setArena(a_arena);
}
for (int j = 0; j < NumIntComps(); ++j) {
GetStructOfArrays().GetIntData(j).setArena(a_arena);
}
}

m_defined = true;
}

// Get id data
Expand Down
Loading
Loading