Improve support for Particles with PolymorphicArenaAllocator (#4603)

AlexanderSinn · web-flow · commit 7f21584d8681 · 2025-08-13T11:27:28.000-07:00
## Summary This PR adds a `SetArena(Arena*)` function to ParticleContainerBase that allows setting a memory arena that is used for all the particle vectors if the allocator is PolymorphicArenaAllocator. The function has to be called before particle tiles are defined. This functionality is used to fix a bunch of places where a polymorphic vector would previously not have its arena set properly. Additionally the `RunOnGpu` logic in `AMReX_WriteBinaryParticleData.H` is extended to work with a polymorphic allocator. Uses changes extracted from #4404. ## Additional background Previously all components of all particle tiles needed their arena set individually by the user if PolymorphicArenaAllocator was used. In case this was done this PR is a braking change due to the `AMREX_ALWAYS_ASSERT_WITH_MESSAGE(a_arena != nullptr` assert in `ParticleTile::define()`. ## Checklist The proposed changes: - [ ] 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
diff --git a/Src/Base/AMReX_GpuAllocators.H b/Src/Base/AMReX_GpuAllocators.H
@@ -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 {};
@@ -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
diff --git a/Src/Particle/AMReX_NeighborParticlesI.H b/Src/Particle/AMReX_NeighborParticlesI.H
@@ -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();
         }
diff --git a/Src/Particle/AMReX_ParticleContainer.H b/Src/Particle/AMReX_ParticleContainer.H
@@ -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 ()
@@ -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);
     }
 
@@ -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);
     }
 
diff --git a/Src/Particle/AMReX_ParticleContainerBase.H b/Src/Particle/AMReX_ParticleContainerBase.H
@@ -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;
+    }
+
     static const std::string& CheckpointVersion ();
     static const std::string& PlotfileVersion ();
     static const std::string& DataPrefix ();
@@ -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();
diff --git a/Src/Particle/AMReX_ParticleContainerI.H b/Src/Particle/AMReX_ParticleContainerI.H
@@ -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,
@@ -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();
@@ -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();
@@ -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);
@@ -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());
+                    }
+                }
             }
         }
     }
diff --git a/Src/Particle/AMReX_ParticleIO.H b/Src/Particle/AMReX_ParticleIO.H
@@ -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();
diff --git a/Src/Particle/AMReX_ParticleTile.H b/Src/Particle/AMReX_ParticleTile.H
@@ -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
@@ -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
diff --git a/Src/Particle/AMReX_WriteBinaryParticleData.H b/Src/Particle/AMReX_WriteBinaryParticleData.H
