Fix: Pinned Allocators First

They are used in methods that copy from device to host and thus need to be known first, before other allocator flavors.
AMReX-Codes · Oct 19, 2023 · dbffc34 · dbffc34
1 parent ee6537d
commit dbffc34
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 7 deletions.
diff --git a/src/Particle/ArrayOfStructs.cpp b/src/Particle/ArrayOfStructs.cpp
@@ -137,7 +137,7 @@ void make_ArrayOfStructs(py::module &m, std::string allocstr)
         .def("__getitem__", [](AOSType &aos, int const v){ return aos[v]; }, py::return_value_policy::reference)
 
         .def("to_host", [](AOSType const & aos) {
-            ArrayOfStructs<T_ParticleType, std::allocator> h_data;
+            ArrayOfStructs<T_ParticleType, amrex::PinnedArenaAllocator> h_data;
             h_data.resize(aos.size());
             //py::array_t<T_ParticleType> h_data(aos.size());
             amrex::Gpu::copy(amrex::Gpu::deviceToHost,
@@ -158,6 +158,9 @@ void make_ArrayOfStructs(py::module &m)
     // AMReX legacy AoS position + id/cpu particle ype
     using ParticleType = Particle<NReal, NInt>;
 
+    // first, because used as copy target in methods in containers with other allocators
+    make_ArrayOfStructs<ParticleType, amrex::PinnedArenaAllocator> (m, "pinned");
+
     // see Src/Base/AMReX_GpuContainers.H
     //   !AMREX_USE_GPU: DefaultAllocator = std::allocator
     //    AMREX_USE_GPU: DefaultAllocator = amrex::ArenaAllocator
@@ -173,7 +176,6 @@ void make_ArrayOfStructs(py::module &m)
     make_ArrayOfStructs<ParticleType, amrex::ArenaAllocator> (m, "arena");
 #endif
     //   end work-around
-    make_ArrayOfStructs<ParticleType, amrex::PinnedArenaAllocator> (m, "pinned");
 #ifdef AMREX_USE_GPU
     make_ArrayOfStructs<ParticleType, amrex::DeviceArenaAllocator> (m, "device");
     make_ArrayOfStructs<ParticleType, amrex::ManagedArenaAllocator> (m, "managed");

diff --git a/src/Particle/ParticleContainer.H b/src/Particle/ParticleContainer.H
@@ -394,6 +394,10 @@ void make_ParticleContainer_and_Iterators (py::module &m)
     if constexpr (!T_ParticleType::is_soa_particle)
         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
@@ -415,8 +419,6 @@ void make_ParticleContainer_and_Iterators (py::module &m)
             amrex::ArenaAllocator>(m, "arena");
 #endif
     //   end work-around
-    make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
-            amrex::PinnedArenaAllocator>(m, "pinned");
 #ifdef AMREX_USE_GPU
     make_ParticleContainer_and_Iterators<T_ParticleType, T_NArrayReal, T_NArrayInt,
                                          amrex::DeviceArenaAllocator>(m, "device");

diff --git a/src/Particle/ParticleTile.cpp b/src/Particle/ParticleTile.cpp
@@ -143,6 +143,10 @@ void make_ParticleTile(py::module &m)
         make_ParticleTileData<T_ParticleType, NArrayReal, NArrayInt>(m);
     }
 
+    // first, because used as copy target in methods in containers with other allocators
+    make_ParticleTile<T_ParticleType, NArrayReal, NArrayInt,
+                      amrex::PinnedArenaAllocator>(m, "pinned");
+
     // see Src/Base/AMReX_GpuContainers.H
     //   !AMREX_USE_GPU: DefaultAllocator = std::allocator
     //    AMREX_USE_GPU: DefaultAllocator = amrex::ArenaAllocator
@@ -164,8 +168,6 @@ void make_ParticleTile(py::module &m)
                       amrex::ArenaAllocator>(m, "arena");
 #endif
     //   end work-around
-    make_ParticleTile<T_ParticleType, NArrayReal, NArrayInt,
-                      amrex::PinnedArenaAllocator>(m, "pinned");
 #ifdef AMREX_USE_GPU
     make_ParticleTile<T_ParticleType, NArrayReal, NArrayInt,
                       amrex::DeviceArenaAllocator>(m, "device");

diff --git a/src/Particle/StructOfArrays.cpp b/src/Particle/StructOfArrays.cpp
@@ -62,6 +62,9 @@ void make_StructOfArrays(py::module &m, std::string allocstr)
 template <int NReal, int NInt>
 void make_StructOfArrays(py::module &m)
 {
+    // first, because used as copy target in methods in containers with other allocators
+    make_StructOfArrays<NReal, NInt, amrex::PinnedArenaAllocator>(m, "pinned");
+
     // see Src/Base/AMReX_GpuContainers.H
     //   !AMREX_USE_GPU: DefaultAllocator = std::allocator
     //    AMREX_USE_GPU: DefaultAllocator = amrex::ArenaAllocator
@@ -77,7 +80,6 @@ void make_StructOfArrays(py::module &m)
     make_StructOfArrays<NReal, NInt, amrex::ArenaAllocator>(m, "arena");
 #endif
     //   end work-around
-    make_StructOfArrays<NReal, NInt, amrex::PinnedArenaAllocator>(m, "pinned");
 #ifdef AMREX_USE_GPU
     make_StructOfArrays<NReal, NInt, amrex::DeviceArenaAllocator>(m, "device");
     make_StructOfArrays<NReal, NInt, amrex::ManagedArenaAllocator>(m, "managed");