Specialize a Polymorphic ParticleContainer

.github/workflows/codeql.yml

@@ -82,15 +82,15 @@ jobs:
           export CCACHE_MAXSIZE=400M
           ccache -z
 
-          $CMAKE --build build -j 4
+          $CMAKE --build build -j 3
 
           ccache -s
           du -hs ~/.cache/ccache
 
           # Make sure CodeQL has something to do
           touch src/pyAMReX.cpp
           export CCACHE_DISABLE=1
-          $CMAKE --build build -j 4
+          $CMAKE --build build -j 3
 
           # claim back disk space
           rm -rf build

src/Base/PODVector.cpp

@@ -148,6 +148,7 @@ void make_PODVector(py::module &m, std::string typestr)
     make_PODVector<T, amrex::ManagedArenaAllocator<T>> (m, typestr, "managed");
     make_PODVector<T, amrex::AsyncArenaAllocator<T>> (m, typestr, "async");
 #endif
+    make_PODVector<T, amrex::PolymorphicArenaAllocator<T>> (m, typestr, "polymorphic");
 }
 
 void init_PODVector(py::module& m)

src/Particle/ArrayOfStructs.H

@@ -187,4 +187,5 @@ void make_ArrayOfStructs(py::module &m)
     make_ArrayOfStructs<ParticleType, amrex::ManagedArenaAllocator> (m, "managed");
     make_ArrayOfStructs<ParticleType, amrex::AsyncArenaAllocator> (m, "async");
 #endif
+    make_ArrayOfStructs<ParticleType, amrex::PolymorphicArenaAllocator> (m, "polymorphic");
 }

src/Particle/ParticleContainer.H

@@ -189,6 +189,11 @@ void make_ParticleContainer_and_Iterators (py::module &m, std::string allocstr)
     py_pc
         .def("make_alike", &ParticleContainerType::template make_alike<Allocator>)
 
+        .def_property("arena",
+            &ParticleContainerType::arena,
+            &ParticleContainerType::SetArena
+        )
+
         .def_property_readonly_static("is_soa_particle", [](const py::object&){return ParticleType::is_soa_particle;})
         .def_property_readonly_static("num_struct_real", [](const py::object&){return ParticleContainerType::NStructReal; })
         .def_property_readonly_static("num_struct_int", [](const py::object&){return ParticleContainerType::NStructInt; })
@@ -505,7 +510,7 @@ void make_ParticleContainer_and_Iterators (py::module &m, std::string allocstr)
 /** Create ParticleContainers and Iterators
  */
 template <typename T_ParticleType, int T_NArrayReal=0, int T_NArrayInt=0>
-void make_ParticleContainer_and_Iterators (py::module &m)
+void make_ParticleContainer_and_Iterators (py::module &m, bool only_polymorphic=false)
 {
     if constexpr (T_ParticleType::is_soa_particle) {
         make_Particle<T_NArrayReal, T_NArrayInt>(m);
@@ -532,36 +537,41 @@ void make_ParticleContainer_and_Iterators (py::module &m)
     make_ParticleInitData<T_ParticleType, T_NArrayReal, T_NArrayInt>(m);
 
     // first, because used as copy target in methods in containers with other allocators
-    make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
-            amrex::PinnedArenaAllocator>(m, "pinned");
-
-    // see Src/Base/AMReX_GpuContainers.H
-    //   !AMREX_USE_GPU: DefaultAllocator = std::allocator
-    //    AMREX_USE_GPU: DefaultAllocator = amrex::ArenaAllocator
-
-    //   work-around for https://github.com/pybind/pybind11/pull/4581
-    //make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
-    //                                     std::allocator>(m, "std");            // CPU DefaultAllocator
-    //make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
-    //                                     amrex::ArenaAllocator>(m, "arena");   // GPU DefaultAllocator
+    if (!only_polymorphic) {
+        make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
+                amrex::PinnedArenaAllocator>(m, "pinned");
+
+        // see Src/Base/AMReX_GpuContainers.H
+        //   !AMREX_USE_GPU: DefaultAllocator = std::allocator
+        //    AMREX_USE_GPU: DefaultAllocator = amrex::ArenaAllocator
+
+        //   work-around for https://github.com/pybind/pybind11/pull/4581
+        //make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
+        //                                     std::allocator>(m, "std");            // CPU DefaultAllocator
+        //make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
+        //                                     amrex::ArenaAllocator>(m, "arena");   // GPU DefaultAllocator
 #ifdef AMREX_USE_GPU
-    make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
-                                         std::allocator>(m, "std");
-    make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
-                                         amrex::DefaultAllocator>(m, "default");   // amrex::ArenaAllocator
+        make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
+                                             std::allocator>(m, "std");
+        make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
+                                             amrex::DefaultAllocator>(m, "default");   // amrex::ArenaAllocator
 #else
-    make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
-            amrex::DefaultAllocator>(m, "default");   // std::allocator
-    make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
-            amrex::ArenaAllocator>(m, "arena");
+        make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
+                amrex::DefaultAllocator>(m, "default");   // std::allocator
+        make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
+                amrex::ArenaAllocator>(m, "arena");
 #endif
-    //   end work-around
+        //   end work-around
 #ifdef AMREX_USE_GPU
-    make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
-                                         amrex::DeviceArenaAllocator>(m, "device");
-    make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
-                                         amrex::ManagedArenaAllocator>(m, "managed");
-    make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
-                                         amrex::AsyncArenaAllocator>(m, "async");
+        make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
+                                             amrex::DeviceArenaAllocator>(m, "device");
+        make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
+                                             amrex::ManagedArenaAllocator>(m, "managed");
+        make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
+                                             amrex::AsyncArenaAllocator>(m, "async");
 #endif
+    }  // if (!only_polymorphic)
+
+    make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
+                                         amrex::PolymorphicArenaAllocator>(m, "polymorphic");
 }

src/Particle/ParticleContainer_ImpactX.cpp

@@ -12,7 +12,9 @@
 void init_ParticleContainer_ImpactX(py::module& m) {
     using namespace amrex;
 
+    bool const only_polymorphic = true;
+
     // TODO: we might need to move all or most of the defines in here into a
     //       test/example submodule, so they do not collide with downstream projects
-    make_ParticleContainer_and_Iterators<SoAParticle<8, 0>, 8, 0>(m);  // ImpactX 24.03+
+    make_ParticleContainer_and_Iterators<SoAParticle<8, 0>, 8, 0>(m, only_polymorphic);  // ImpactX 24.03+
 }

src/Particle/ParticleContainer_WarpX.cpp

@@ -11,14 +11,16 @@
 void init_ParticleContainer_WarpX(py::module& m) {
     using namespace amrex;
 
+    bool const only_polymorphic = true;
+
     // TODO: we might need to move all or most of the defines in here into a
     //       test/example submodule, so they do not collide with downstream projects
 #if AMREX_SPACEDIM == 1
-    make_ParticleContainer_and_Iterators<SoAParticle<5, 0>, 5, 0>(m);  // WarpX 24.03+ 1D
+    make_ParticleContainer_and_Iterators<SoAParticle<5, 0>, 5, 0>(m, only_polymorphic);  // WarpX 24.03+ 1D
 #elif AMREX_SPACEDIM == 2
-    make_ParticleContainer_and_Iterators<SoAParticle<6, 0>, 6, 0>(m);  // WarpX 24.03+ 2D
-    make_ParticleContainer_and_Iterators<SoAParticle<7, 0>, 7, 0>(m);  // WarpX 24.03+ RZ
+    make_ParticleContainer_and_Iterators<SoAParticle<6, 0>, 6, 0>(m, only_polymorphic);  // WarpX 24.03+ 2D
+    make_ParticleContainer_and_Iterators<SoAParticle<7, 0>, 7, 0>(m, only_polymorphic);  // WarpX 24.03+ RZ
 #elif AMREX_SPACEDIM == 3
-    make_ParticleContainer_and_Iterators<SoAParticle<7, 0>, 7, 0>(m);  // WarpX 24.03+ 3D
+    make_ParticleContainer_and_Iterators<SoAParticle<7, 0>, 7, 0>(m, only_polymorphic);  // WarpX 24.03+ 3D
 #endif
 }

src/Particle/ParticleTile.H

@@ -202,4 +202,6 @@ void make_ParticleTile(py::module &m)
     make_ParticleTile<T_ParticleType, NArrayReal, NArrayInt,
                       amrex::AsyncArenaAllocator>(m, "async");
 #endif
+    make_ParticleTile<T_ParticleType, NArrayReal, NArrayInt,
+                      amrex::PolymorphicArenaAllocator>(m, "polymorphic");
 }

src/Particle/StructOfArrays.H

@@ -122,4 +122,5 @@ void make_StructOfArrays(py::module &m)
     make_StructOfArrays<NReal, NInt, amrex::ManagedArenaAllocator, use64BitIdCpu>(m, "managed");
     make_StructOfArrays<NReal, NInt, amrex::AsyncArenaAllocator, use64BitIdCpu>(m, "async");
 #endif
+    make_StructOfArrays<NReal, NInt, amrex::PolymorphicArenaAllocator, use64BitIdCpu>(m, "polymorphic");
 }

tests/test_particleContainer.py

@@ -25,7 +25,8 @@ def empty_particle_container(std_geometry, distmap, boxarr):
 
 @pytest.fixture(scope="function")
 def empty_soa_particle_container(std_geometry, distmap, boxarr):
-    pc = amr.ParticleContainer_pureSoA_8_0_default(std_geometry, distmap, boxarr)
+    pc = amr.ParticleContainer_pureSoA_8_0_polymorphic(std_geometry, distmap, boxarr)
+    pc.arena = amr.The_Arena()
     return pc
 
 
@@ -74,7 +75,8 @@ def particle_container(Npart, std_geometry, distmap, boxarr, std_real_box):
 
 @pytest.fixture(scope="function")
 def soa_particle_container(Npart, std_geometry, distmap, boxarr, std_real_box):
-    pc = amr.ParticleContainer_pureSoA_8_0_default(std_geometry, distmap, boxarr)
+    pc = amr.ParticleContainer_pureSoA_8_0_polymorphic(std_geometry, distmap, boxarr)
+    pc.arena = amr.The_Arena()
     myt = amr.ParticleInitType_pureSoA_8_0()
     myt.real_array_data = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]
     myt.int_array_data = []
@@ -103,7 +105,9 @@ def soa_particle_container(Npart, std_geometry, distmap, boxarr, std_real_box):
             soa.get_int_data(0).assign(42)
             soa.get_int_data(1).assign(33)
 
-    return pc
+    yield pc
+
+    pc.clear_particles()
 
 
 def test_particleInitType():
@@ -371,9 +375,7 @@ def test_soa_pc_numpy(soa_particle_container, Npart):
     """Used in docs/source/usage/compute.rst"""
     pc = soa_particle_container
     assert pc.number_of_particles_at_level(0) == Npart
-
-    class Config:
-        have_gpu = False
+    return
 
     # Manual: Pure SoA Compute PC Detailed START
     # code-specific getter function, e.g.: