From e37aeefe023a7df9a0e6ad3629b8ded86f17c53c Mon Sep 17 00:00:00 2001
From: AlexanderSinn <alexander.sinn@desy.de>
Date: Sun, 10 Aug 2025 15:22:13 +0200
Subject: [PATCH 1/5] Improve support for Particles with
 PolymorphicArenaAllocator

---
 Src/Base/AMReX_GpuAllocators.H               |   7 +-
 Src/Particle/AMReX_NeighborParticlesI.H      |   3 +-
 Src/Particle/AMReX_ParticleContainer.H       |  16 +-
 Src/Particle/AMReX_ParticleContainerBase.H   |   9 +
 Src/Particle/AMReX_ParticleContainerI.H      |  28 ++-
 Src/Particle/AMReX_ParticleIO.H              |   5 +-
 Src/Particle/AMReX_ParticleTile.H            |  43 ++++-
 Src/Particle/AMReX_WriteBinaryParticleData.H | 185 +++++++++++++------
 8 files changed, 226 insertions(+), 70 deletions(-)

diff --git a/Src/Base/AMReX_GpuAllocators.H b/Src/Base/AMReX_GpuAllocators.H
index de2fc0ad6bc..5accd7ee84a 100644
--- 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
index de5146fbe9c..1065ae09970 100644
--- 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, 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
index 3baafe1d34e..ca709aa5576 100644
--- 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
index 0f573b9d163..7cb4fe62839 100644
--- 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
index 1d90fde2c76..eab4509685a 100644
--- 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,11 @@ 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);
+                if constexpr (has_polymorphic_allocator) {
+                    tmp_local[lev][index][t].setArena(arena());
+                }
+                soa_local[lev][index][t].define(m_num_runtime_real, m_num_runtime_int,
+                    &m_soa_rdata_names, &m_soa_idata_names, arena());
             }
         }
     }
diff --git a/Src/Particle/AMReX_ParticleIO.H b/Src/Particle/AMReX_ParticleIO.H
index 71b893d2af8..05bcde40c0c 100644
--- 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
index d4407942521..2bcc1b569fd 100644
--- 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
index 4a5117a1d29..bac3be3ec4d 100644
--- a/Src/Particle/AMReX_WriteBinaryParticleData.H
+++ b/Src/Particle/AMReX_WriteBinaryParticleData.H
@@ -29,12 +29,11 @@ std::size_t PSizeInFile (const Vector<int>& wrc, const Vector<int>& wic)
     return rsize + isize + AMREX_SPACEDIM*sizeof(ParticleReal) + 2*sizeof(int);
 }
 
-template <template <class, class> class Container,
-          class Allocator,
+template <class Container,
           class PTile,
           class F>
-std::enable_if_t<RunOnGpu<typename Container<int, Allocator>::allocator_type>::value>
-fillFlags (Container<int, Allocator>& pflags, const PTile& ptile, F const& f)
+void
+fillFlagsGpu (Container& pflags, const PTile& ptile, F const& f)
 {
     const auto& ptd = ptile.getConstParticleTileData();
     const auto np = ptile.numParticles();
@@ -57,12 +56,11 @@ fillFlags (Container<int, Allocator>& pflags, const PTile& ptile, F const& f)
         });
 }
 
-template <template <class, class> class Container,
-          class Allocator,
+template <class Container,
           class PTile,
           class F>
-std::enable_if_t<!RunOnGpu<typename Container<int, Allocator>::allocator_type>::value>
-fillFlags (Container<int, Allocator>& pflags, const PTile& ptile, F const& f)
+void
+fillFlagsCpu (Container& pflags, const PTile& ptile, F const& f)
 {
     const auto& ptd = ptile.getConstParticleTileData();
     const auto np = ptile.numParticles();
@@ -82,9 +80,31 @@ fillFlags (Container<int, Allocator>& pflags, const PTile& ptile, F const& f)
     }
 }
 
-template <template <class, class> class Container, class Allocator, class PC>
-std::enable_if_t<RunOnGpu<typename Container<int, Allocator>::allocator_type>::value, amrex::Long>
-countFlags (const Vector<std::map<std::pair<int,int>,Container<int,Allocator>>>& particle_io_flags, const PC& pc)
+template <template <class, class> class Container,
+          class Allocator,
+          class PTile,
+          class F>
+void
+fillFlags (Container<int, Allocator>& pflags, const PTile& ptile, F const& f)
+{
+    if constexpr (IsPolymorphicArenaAllocator<Allocator>::value) {
+        if (pflags.arena()->isManaged() ||  pflags.arena()->isDevice()) {
+            fillFlagsGpu(pflags, ptile, f);
+        } else {
+            fillFlagsCpu(pflags, ptile, f);
+        }
+    } else {
+        if constexpr (RunOnGpu<Allocator>::value) {
+            fillFlagsGpu(pflags, ptile, f);
+        } else {
+            fillFlagsCpu(pflags, ptile, f);
+        }
+    }
+}
+
+template <class Container, class PC>
+amrex::Long
+countFlagsGpu (const Vector<std::map<std::pair<int,int>,Container>>& particle_io_flags, const PC& pc)
 {
     ReduceOps<ReduceOpSum> reduce_op;
     ReduceData<Long> reduce_data(reduce_op);
@@ -108,9 +128,9 @@ countFlags (const Vector<std::map<std::pair<int,int>,Container<int,Allocator>>>&
     return amrex::get<0>(hv);
 }
 
-template <template <class, class> class Container, class Allocator>
-std::enable_if_t<RunOnGpu<typename Container<int, Allocator>::allocator_type>::value, int>
-countFlags (const Container<int,Allocator>& pflags)
+template <class Container>
+amrex::Long
+countFlagsGpu (const Container& pflags)
 {
     ReduceOps<ReduceOpSum> reduce_op;
     ReduceData<Long> reduce_data(reduce_op);
@@ -118,7 +138,7 @@ countFlags (const Container<int,Allocator>& pflags)
 
     const auto flag_ptr = pflags.data();
     reduce_op.eval(pflags.size(), reduce_data,
-        [=] AMREX_GPU_DEVICE (const int i) -> ReduceTuple
+        [=] AMREX_GPU_DEVICE (const amrex::Long i) -> ReduceTuple
         {
             return flag_ptr[i] ? 1 : 0;
         });
@@ -126,9 +146,9 @@ countFlags (const Container<int,Allocator>& pflags)
     return amrex::get<0>(hv);
 }
 
-template <template <class, class> class Container, class Allocator, class PC>
-std::enable_if_t<!RunOnGpu<typename Container<int, Allocator>::allocator_type>::value, amrex::Long>
-countFlags (const Vector<std::map<std::pair<int,int>,Container<int,Allocator>>>& particle_io_flags, const PC& pc)
+template <class Container, class PC>
+amrex::Long
+countFlagsCpu (const Vector<std::map<std::pair<int,int>,Container>>& particle_io_flags, const PC& pc)
 {
     amrex::Long nparticles = 0;
     for (int lev = 0; lev < pc.GetParticles().size();  lev++)
@@ -146,11 +166,11 @@ countFlags (const Vector<std::map<std::pair<int,int>,Container<int,Allocator>>>&
     return nparticles;
 }
 
-template <template <class, class> class Container, class Allocator>
-std::enable_if_t<!RunOnGpu<typename Container<int, Allocator>::allocator_type>::value, int>
-countFlags (const Container<int,Allocator>& pflags)
+template <class Container>
+amrex::Long
+countFlagsCpu (const Container& pflags)
 {
-    int nparticles = 0;
+    amrex::Long nparticles = 0;
     for (std::size_t k = 0; k < pflags.size(); ++k)
     {
         if (pflags[k]) { nparticles++; }
@@ -158,6 +178,44 @@ countFlags (const Container<int,Allocator>& pflags)
     return nparticles;
 }
 
+template <template <class, class> class Container, class Allocator, class PC>
+amrex::Long
+countFlags (const Vector<std::map<std::pair<int,int>,Container<int,Allocator>>>& particle_io_flags, const PC& pc)
+{
+    if constexpr (IsPolymorphicArenaAllocator<Allocator>::value) {
+        if (pc.arena()->isManaged() ||  pc.arena()->isDevice()) {
+            return countFlagsGpu(particle_io_flags, pc);
+        } else {
+            return countFlagsCpu(particle_io_flags, pc);
+        }
+    } else {
+        if constexpr (RunOnGpu<Allocator>::value) {
+            return countFlagsGpu(particle_io_flags, pc);
+        } else {
+            return countFlagsCpu(particle_io_flags, pc);
+        }
+    }
+}
+
+template <template <class, class> class Container, class Allocator>
+amrex::Long
+countFlags (const Container<int,Allocator>& pflags)
+{
+    if constexpr (IsPolymorphicArenaAllocator<Allocator>::value) {
+        if (pflags.arena()->isManaged() || pflags.arena()->isDevice()) {
+            return countFlagsGpu(pflags);
+        } else {
+            return countFlagsCpu(pflags);
+        }
+    } else {
+        if constexpr (RunOnGpu<Allocator>::value) {
+            return countFlagsGpu(pflags);
+        } else {
+            return countFlagsCpu(pflags);
+        }
+    }
+}
+
 template <typename I>
 AMREX_GPU_HOST_DEVICE
 void packParticleIDs (I* idata, const std::uint64_t idcpu, bool is_checkpoint) noexcept
@@ -254,11 +312,11 @@ iPackParticleData (const PTD& ptd, int idx, typename PTD::IntType * idata_ptr,
 }
 
 template <class PC>
-std::enable_if_t<RunOnGpu<typename PC::template AllocatorType<int>>::value, void>
-packIOData (Vector<int>& idata, Vector<ParticleReal>& rdata, const PC& pc, int lev, int grid,
-            const Vector<int>& write_real_comp, const Vector<int>& write_int_comp,
-            const Vector<std::map<std::pair<int, int>, typename PC::IntVector>>& particle_io_flags,
-            const Vector<int>& tiles, int np, bool is_checkpoint)
+void
+packIODataGpu (Vector<int>& idata, Vector<ParticleReal>& rdata, const PC& pc, int lev, int grid,
+               const Vector<int>& write_real_comp, const Vector<int>& write_int_comp,
+               const Vector<std::map<std::pair<int, int>, typename PC::IntVector>>& particle_io_flags,
+               const Vector<int>& tiles, int np, bool is_checkpoint)
 {
     int num_output_int = 0;
     for (int i = 0; i < pc.NumIntComps() + PC::NStructInt; ++i) {
@@ -276,8 +334,8 @@ packIOData (Vector<int>& idata, Vector<ParticleReal>& rdata, const PC& pc, int l
     const Long rChunkSize = AMREX_SPACEDIM + num_output_real;
     rdata.resize(np*rChunkSize);
 
-    typename PC::IntVector write_int_comp_d(write_int_comp.size());
-    typename PC::IntVector write_real_comp_d(write_real_comp.size());
+    Gpu::DeviceVector<int> write_int_comp_d(write_int_comp.size());
+    Gpu::DeviceVector<int> write_real_comp_d(write_real_comp.size());
     Gpu::copyAsync(Gpu::hostToDevice, write_int_comp.begin(), write_int_comp.end(),
                    write_int_comp_d.begin());
     Gpu::copyAsync(Gpu::hostToDevice, write_real_comp.begin(), write_real_comp.end(),
@@ -291,17 +349,17 @@ packIOData (Vector<int>& idata, Vector<ParticleReal>& rdata, const PC& pc, int l
         const auto& ptile = pc.ParticlesAt(lev, grid, tile);
         const auto& pflags = particle_io_flags[lev].at(std::make_pair(grid, tile));
         int np_tile = ptile.numParticles();
-        typename PC::IntVector offsets(np_tile);
+        Gpu::DeviceVector<int> offsets(np_tile);
         int num_copies = Scan::ExclusiveSum(np_tile, pflags.begin(), offsets.begin(), Scan::retSum);
 
-        typename PC::IntVector  idata_d(num_copies*iChunkSize);
-        typename PC::RealVector rdata_d(num_copies*rChunkSize);
+        Gpu::DeviceVector<int> idata_d(num_copies*iChunkSize);
+        Gpu::DeviceVector<ParticleReal> rdata_d(num_copies*rChunkSize);
 
         const auto flag_ptr = pflags.data();
         const auto offset_ptr = offsets.data();
 
-        auto idata_d_ptr = idata_d.data();
-        auto rdata_d_ptr = rdata_d.data();
+        auto* idata_d_ptr = idata_d.data();
+        auto* rdata_d_ptr = rdata_d.data();
 
         const auto& ptd = ptile.getConstParticleTileData();
         amrex::ParallelFor(ptile.numParticles(),
@@ -318,9 +376,9 @@ packIOData (Vector<int>& idata, Vector<ParticleReal>& rdata, const PC& pc, int l
         });
 
         Gpu::copyAsync(Gpu::deviceToHost, idata_d.begin(), idata_d.end(),
-                       idata.begin() + typename PC::IntVector::difference_type(poffset));
+                       idata.begin() + poffset;
         Gpu::copyAsync(Gpu::deviceToHost, rdata_d.begin(), rdata_d.end(),
-                       rdata.begin() + typename PC::RealVector::difference_type(poffset));
+                       rdata.begin() + poffset;
         Gpu::Device::streamSynchronize();
 
         poffset += num_copies;
@@ -328,11 +386,11 @@ packIOData (Vector<int>& idata, Vector<ParticleReal>& rdata, const PC& pc, int l
 }
 
 template <class PC>
-std::enable_if_t<!RunOnGpu<typename PC::template AllocatorType<int>>::value, void>
-packIOData (Vector<int>& idata, Vector<ParticleReal>& rdata, const PC& pc, int lev, int grid,
-            const Vector<int>& write_real_comp, const Vector<int>& write_int_comp,
-            const Vector<std::map<std::pair<int, int>, typename PC::IntVector>>& particle_io_flags,
-            const Vector<int>& tiles, int np, bool is_checkpoint)
+void
+packIODataCpu (Vector<int>& idata, Vector<ParticleReal>& rdata, const PC& pc, int lev, int grid,
+               const Vector<int>& write_real_comp, const Vector<int>& write_int_comp,
+               const Vector<std::map<std::pair<int, int>, typename PC::IntVector>>& particle_io_flags,
+               const Vector<int>& tiles, int np, bool is_checkpoint)
 {
     int num_output_int = 0;
     for (int i = 0; i < pc.NumIntComps() + PC::NStructInt; ++i) {
@@ -370,6 +428,33 @@ packIOData (Vector<int>& idata, Vector<ParticleReal>& rdata, const PC& pc, int l
         }
     }
 }
+
+template <class PC>
+void
+packIOData (Vector<int>& idata, Vector<ParticleReal>& rdata, const PC& pc, int lev, int grid,
+            const Vector<int>& write_real_comp, const Vector<int>& write_int_comp,
+            const Vector<std::map<std::pair<int, int>, typename PC::IntVector>>& particle_io_flags,
+            const Vector<int>& tiles, int np, bool is_checkpoint)
+{
+    if constexpr (IsPolymorphicArenaAllocator<typename PC::IntVector::allocator_type>::value) {
+        if (pc.arena()->isManaged() || pc.arena()->isDevice()) {
+            packIODataGpu(idata, rdata, pc, lev, grid, write_real_comp, write_int_comp,
+                          particle_io_flags, tiles, np, is_checkpoint);
+        } else {
+            packIODataCpu(idata, rdata, pc, lev, grid, write_real_comp, write_int_comp,
+                          particle_io_flags, tiles, np, is_checkpoint);
+        }
+    } else {
+        if constexpr (RunOnGpu<typename PC::IntVector::allocator_type>::value) {
+            packIODataGpu(idata, rdata, pc, lev, grid, write_real_comp, write_int_comp,
+                          particle_io_flags, tiles, np, is_checkpoint);
+        } else {
+            packIODataCpu(idata, rdata, pc, lev, grid, write_real_comp, write_int_comp,
+                          particle_io_flags, tiles, np, is_checkpoint);
+        }
+    }
+}
+
 }
 
 template <class PC, class F, std::enable_if_t<IsParticleContainer<PC>::value, int> foo = 0>
@@ -429,6 +514,9 @@ void WriteBinaryParticleDataSync (PC const& pc,
         for (const auto& kv : pmap)
         {
             auto& flags = particle_io_flags[lev][kv.first];
+            if constexpr (PC::has_polymorphic_allocator) {
+                flags.setArena(pc.arena());
+            }
             particle_detail::fillFlags(flags, kv.second, f);
         }
     }
@@ -811,7 +899,7 @@ void WriteBinaryParticleDataAsync (PC const& pc,
 
     // make tmp particle tiles in pinned memory to write
     using PinnedPTile = ParticleTile<typename PC::ParticleType, NArrayReal, NArrayInt,
-                                     PinnedArenaAllocator>;
+                                     PolymorphicArenaAllocator>;
     auto myptiles = std::make_shared<Vector<std::map<std::pair<int, int>,PinnedPTile> > >();
     myptiles->resize(pc.finestLevel()+1);
     for (int lev = 0; lev <= pc.finestLevel(); lev++)
@@ -827,20 +915,13 @@ void WriteBinaryParticleDataAsync (PC const& pc,
 
                 const auto np = np_per_grid_local[lev][mfi.index()];
 
-                new_ptile.resize(np);
-
                 const auto runtime_real_comps = ptile.NumRuntimeRealComps();
                 const auto runtime_int_comps = ptile.NumRuntimeIntComps();
 
-                new_ptile.define(runtime_real_comps, runtime_int_comps);
-
-                for (auto comp(0); comp < runtime_real_comps; ++comp) {
-                    new_ptile.push_back_real(NArrayReal+comp, np, 0.);
-                }
+                new_ptile.define(runtime_real_comps, runtime_int_comps,
+                    nullptr, nullptr, The_Pinned_Arena());
 
-                for (auto comp(0); comp < runtime_int_comps; ++comp) {
-                    new_ptile.push_back_int(NArrayInt+comp, np, 0);
-                }
+                new_ptile.resize(np);
 
                 amrex::filterParticles(new_ptile, ptile, KeepValidFilter());
             }

From d28a9695db94b7de8b59d2fa1fe43534b0e8af88 Mon Sep 17 00:00:00 2001
From: AlexanderSinn <alexander.sinn@desy.de>
Date: Sun, 10 Aug 2025 15:42:08 +0200
Subject: [PATCH 2/5] fix )

---
 Src/Particle/AMReX_WriteBinaryParticleData.H | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/Src/Particle/AMReX_WriteBinaryParticleData.H b/Src/Particle/AMReX_WriteBinaryParticleData.H
index bac3be3ec4d..873d41c2ab4 100644
--- a/Src/Particle/AMReX_WriteBinaryParticleData.H
+++ b/Src/Particle/AMReX_WriteBinaryParticleData.H
@@ -376,9 +376,9 @@ packIODataGpu (Vector<int>& idata, Vector<ParticleReal>& rdata, const PC& pc, in
         });
 
         Gpu::copyAsync(Gpu::deviceToHost, idata_d.begin(), idata_d.end(),
-                       idata.begin() + poffset;
+                       idata.begin() + poffset);
         Gpu::copyAsync(Gpu::deviceToHost, rdata_d.begin(), rdata_d.end(),
-                       rdata.begin() + poffset;
+                       rdata.begin() + poffset);
         Gpu::Device::streamSynchronize();
 
         poffset += num_copies;

From 63c922339760df31ef1ac2360cd61e9006c7ffb2 Mon Sep 17 00:00:00 2001
From: AlexanderSinn <alexander.sinn@desy.de>
Date: Sun, 10 Aug 2025 16:05:12 +0200
Subject: [PATCH 3/5] fix RedistributeCPU

---
 Src/Particle/AMReX_ParticleContainerI.H | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/Src/Particle/AMReX_ParticleContainerI.H b/Src/Particle/AMReX_ParticleContainerI.H
index eab4509685a..0f908ad8f20 100644
--- a/Src/Particle/AMReX_ParticleContainerI.H
+++ b/Src/Particle/AMReX_ParticleContainerI.H
@@ -1655,11 +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);
                 if constexpr (has_polymorphic_allocator) {
-                    tmp_local[lev][index][t].setArena(arena());
+                    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());
+                    }
                 }
-                soa_local[lev][index][t].define(m_num_runtime_real, m_num_runtime_int,
-                    &m_soa_rdata_names, &m_soa_idata_names, arena());
             }
         }
     }

From b7cea2995ef41dc10f830a9146a09f64c5cb272b Mon Sep 17 00:00:00 2001
From: AlexanderSinn <alexander.sinn@desy.de>
Date: Sun, 10 Aug 2025 16:16:53 +0200
Subject: [PATCH 4/5] add this to neighbors

---
 Src/Particle/AMReX_NeighborParticlesI.H | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Src/Particle/AMReX_NeighborParticlesI.H b/Src/Particle/AMReX_NeighborParticlesI.H
index 1065ae09970..aaf3bb54381 100644
--- a/Src/Particle/AMReX_NeighborParticlesI.H
+++ b/Src/Particle/AMReX_NeighborParticlesI.H
@@ -455,7 +455,7 @@ NeighborParticleContainer<NStructReal, NStructInt, NArrayReal, NArrayInt>
             const int num_ghosts = map.size();
             neighbors[lev][dst_index].define(this->NumRuntimeRealComps(),
                                              this->NumRuntimeIntComps(),
-                                             nullptr, nullptr, arena());
+                                             nullptr, nullptr, this->arena());
             neighbors[lev][dst_index].resize(num_ghosts);
             local_neighbor_sizes[lev][dst_index] = neighbors[lev][dst_index].size();
         }

From 7ec3a3bd5c9d1e68461c4e2588fc2ba5ded6f154 Mon Sep 17 00:00:00 2001
From: AlexanderSinn <alexander.sinn@desy.de>
Date: Sun, 10 Aug 2025 16:36:50 +0200
Subject: [PATCH 5/5] fix

---
 Src/Particle/AMReX_WriteBinaryParticleData.H | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/Src/Particle/AMReX_WriteBinaryParticleData.H b/Src/Particle/AMReX_WriteBinaryParticleData.H
index 873d41c2ab4..a22d05fabc4 100644
--- a/Src/Particle/AMReX_WriteBinaryParticleData.H
+++ b/Src/Particle/AMReX_WriteBinaryParticleData.H
@@ -341,8 +341,8 @@ packIODataGpu (Vector<int>& idata, Vector<ParticleReal>& rdata, const PC& pc, in
     Gpu::copyAsync(Gpu::hostToDevice, write_real_comp.begin(), write_real_comp.end(),
                    write_real_comp_d.begin());
 
-    const auto write_int_comp_d_ptr = write_int_comp_d.data();
-    const auto write_real_comp_d_ptr = write_real_comp_d.data();
+    const auto* write_int_comp_d_ptr = write_int_comp_d.data();
+    const auto* write_real_comp_d_ptr = write_real_comp_d.data();
 
     std::size_t poffset = 0;
     for (int tile : tiles) {
@@ -355,8 +355,8 @@ packIODataGpu (Vector<int>& idata, Vector<ParticleReal>& rdata, const PC& pc, in
         Gpu::DeviceVector<int> idata_d(num_copies*iChunkSize);
         Gpu::DeviceVector<ParticleReal> rdata_d(num_copies*rChunkSize);
 
-        const auto flag_ptr = pflags.data();
-        const auto offset_ptr = offsets.data();
+        const auto* flag_ptr = pflags.data();
+        const auto* offset_ptr = offsets.data();
 
         auto* idata_d_ptr = idata_d.data();
         auto* rdata_d_ptr = rdata_d.data();
@@ -376,9 +376,9 @@ packIODataGpu (Vector<int>& idata, Vector<ParticleReal>& rdata, const PC& pc, in
         });
 
         Gpu::copyAsync(Gpu::deviceToHost, idata_d.begin(), idata_d.end(),
-                       idata.begin() + poffset);
+                       idata.begin() + static_cast<Long>(poffset));
         Gpu::copyAsync(Gpu::deviceToHost, rdata_d.begin(), rdata_d.end(),
-                       rdata.begin() + poffset);
+                       rdata.begin() + static_cast<Long>(poffset));
         Gpu::Device::streamSynchronize();
 
         poffset += num_copies;