From 731efdc2d06a95df69583df5f850cf1497980bf8 Mon Sep 17 00:00:00 2001
From: Andrew Myers <atmyers2@gmail.com>
Date: Thu, 9 Jan 2025 09:17:04 -0800
Subject: [PATCH] Make using The_Comms_Arena() for particle communication
 opt-in (#4286)

For mesh data it is better to use a small, separate Arena for
communication. But, we found that for particle communication this
approach uses too much memory (see the ImpactX Issue here:
https://github.com/ECP-WarpX/impactx/issues/769)

This restores the old default behavior (prior to PR #4175), but allows
users to opt-in to using The_Comms_Arena if desired.

The proposed changes:
- [ ] 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
---
 Src/Particle/AMReX_NeighborParticlesGPUImpl.H | 2 +-
 Src/Particle/AMReX_ParticleContainerBase.H    | 1 +
 Src/Particle/AMReX_ParticleContainerBase.cpp  | 1 +
 Src/Particle/AMReX_ParticleContainerI.H       | 3 ++-
 4 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/Src/Particle/AMReX_NeighborParticlesGPUImpl.H b/Src/Particle/AMReX_NeighborParticlesGPUImpl.H
index b4e50bef7b8..ce0444ac44d 100644
--- a/Src/Particle/AMReX_NeighborParticlesGPUImpl.H
+++ b/Src/Particle/AMReX_NeighborParticlesGPUImpl.H
@@ -256,7 +256,7 @@ updateNeighborsGPU (bool boundary_neighbors_only)
 
     clearNeighbors();
 
-    if (ParallelDescriptor::UseGpuAwareMpi()) {
+    if (this->use_comms_arena) {
         snd_buffer.setArena(The_Comms_Arena());
         rcv_buffer.setArena(The_Comms_Arena());
     }
diff --git a/Src/Particle/AMReX_ParticleContainerBase.H b/Src/Particle/AMReX_ParticleContainerBase.H
index 433890b1579..0f573b9d163 100644
--- a/Src/Particle/AMReX_ParticleContainerBase.H
+++ b/Src/Particle/AMReX_ParticleContainerBase.H
@@ -256,6 +256,7 @@ public:
     static AMREX_EXPORT bool do_tiling;
     static AMREX_EXPORT IntVect tile_size;
     static AMREX_EXPORT bool memEfficientSort;
+    static AMREX_EXPORT bool use_comms_arena;
     mutable AmrParticleLocator<DenseBins<Box> > m_particle_locator;
 
 protected:
diff --git a/Src/Particle/AMReX_ParticleContainerBase.cpp b/Src/Particle/AMReX_ParticleContainerBase.cpp
index 4f0ec83f4cb..b8bd739e43c 100644
--- a/Src/Particle/AMReX_ParticleContainerBase.cpp
+++ b/Src/Particle/AMReX_ParticleContainerBase.cpp
@@ -10,6 +10,7 @@ using namespace amrex;
 bool    ParticleContainerBase::do_tiling = false;
 IntVect ParticleContainerBase::tile_size { AMREX_D_DECL(1024000,8,8) };
 bool    ParticleContainerBase::memEfficientSort = true;
+bool    ParticleContainerBase::use_comms_arena = false;
 
 void ParticleContainerBase::Define (const Geometry            & geom,
                                     const DistributionMapping & dmap,
diff --git a/Src/Particle/AMReX_ParticleContainerI.H b/Src/Particle/AMReX_ParticleContainerI.H
index 3bd292973f3..9d749bb3b8f 100644
--- a/Src/Particle/AMReX_ParticleContainerI.H
+++ b/Src/Particle/AMReX_ParticleContainerI.H
@@ -63,6 +63,7 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
         pp.query("use_prepost", usePrePost);
         pp.query("do_unlink", doUnlink);
         pp.queryAdd("do_mem_efficient_sort", memEfficientSort);
+        pp.queryAdd("use_comms_arena", use_comms_arena);
 
         // add default names for SoA Real and Int compile-time arguments
         for (int i=0; i<NArrayReal; ++i)
@@ -1392,7 +1393,7 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
     amrex::PODVector<char, PolymorphicArenaAllocator<char> > snd_buffer;
     amrex::PODVector<char, PolymorphicArenaAllocator<char> > rcv_buffer;
 
-    if (ParallelDescriptor::UseGpuAwareMpi()) {
+    if (use_comms_arena) {
         snd_buffer.setArena(The_Comms_Arena());
         rcv_buffer.setArena(The_Comms_Arena());
     }