Skip to content

Commit

Permalink
New ParticleTileData Interfaces
Browse files Browse the repository at this point in the history
  • Loading branch information
@ax3l
ax3l committed Apr 9, 2023
1 parent 456456b commit 1e6f721
Show file tree
Hide file tree
Showing 2 changed files with 71 additions and 52 deletions.
47 changes: 25 additions & 22 deletions src/Particle/ArrayOfStructs.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,12 +21,12 @@ namespace
*
* https://numpy.org/doc/stable/reference/arrays.interface.html
*/
template <int NReal, int NInt,
template <typename T_ParticleType,
template<class> class Allocator=DefaultAllocator>
py::dict
array_interface(ArrayOfStructs<NReal, NInt, Allocator> const & aos)
array_interface(ArrayOfStructs<T_ParticleType, Allocator> const & aos)
{
using ParticleType = Particle<NReal, NInt>;
using ParticleType = T_ParticleType;
using RealType = typename ParticleType::RealType;

auto d = py::dict();
Expand All @@ -43,14 +43,14 @@ namespace
#if (AMREX_SPACEDIM >= 3)
descr.append(py::make_tuple("z", py::format_descriptor<RealType>::format()));
#endif
if (NReal > 0) {
for(int ii=0; ii < NReal; ii++) {
if constexpr (ParticleType::NReal > 0) {
for(int ii=0; ii < ParticleType::NReal; ii++) {
descr.append(py::make_tuple("rdata_"+std::to_string(ii),py::format_descriptor<RealType>::format()));
}
}
descr.append(py::make_tuple("cpuid", py::format_descriptor<uint64_t>::format()) );
if (NInt > 0) {
for(int ii=0; ii < NInt; ++ii) {
if constexpr (ParticleType::NInt > 0) {
for(int ii=0; ii < ParticleType::NInt; ++ii) {
descr.append(py::make_tuple("idata_"+std::to_string(ii),py::format_descriptor<int>::format()));
}
}
Expand All @@ -61,16 +61,16 @@ namespace
}
}

template <int NReal, int NInt,
template <typename T_ParticleType,
template<class> class Allocator=DefaultAllocator>
void make_ArrayOfStructs(py::module &m, std::string allocstr)
{
using AOSType = ArrayOfStructs<NReal, NInt, Allocator>;
using ParticleType = Particle<NReal, NInt>;
using AOSType = ArrayOfStructs<T_ParticleType, Allocator>;
using ParticleType = T_ParticleType;

auto const aos_name = std::string("ArrayOfStructs_")
.append(std::to_string(NReal)).append("_")
.append(std::to_string(NInt)).append("_")
.append(std::to_string(ParticleType::NReal)).append("_")
.append(std::to_string(ParticleType::NInt)).append("_")
.append(allocstr);
py::class_<AOSType>(m, aos_name.c_str())
.def(py::init())
Expand Down Expand Up @@ -124,26 +124,29 @@ void make_ArrayOfStructs(py::module &m, std::string allocstr)
template <int NReal, int NInt>
void make_ArrayOfStructs(py::module &m)
{
// AMReX legacy AoS position + id/cpu particle ype
using ParticleType = Particle<NReal, NInt>;

// 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_ArrayOfStructs<NReal, NInt, std::allocator> (m, "std");
//make_ArrayOfStructs<NReal, NInt, amrex::ArenaAllocator> (m, "arena");
//make_ArrayOfStructs<ParticleType, std::allocator> (m, "std");
//make_ArrayOfStructs<ParticleType, amrex::ArenaAllocator> (m, "arena");
#ifdef AMREX_USE_GPU
make_ArrayOfStructs<NReal, NInt, std::allocator> (m, "std");
make_ArrayOfStructs<NReal, NInt, amrex::DefaultAllocator> (m, "default"); // amrex::ArenaAllocator
make_ArrayOfStructs<ParticleType, std::allocator> (m, "std");
make_ArrayOfStructs<ParticleType, amrex::DefaultAllocator> (m, "default"); // amrex::ArenaAllocator
#else
make_ArrayOfStructs<NReal, NInt, amrex::DefaultAllocator> (m, "default"); // std::allocator
make_ArrayOfStructs<NReal, NInt, amrex::ArenaAllocator> (m, "arena");
make_ArrayOfStructs<ParticleType, amrex::DefaultAllocator> (m, "default"); // std::allocator
make_ArrayOfStructs<ParticleType, amrex::ArenaAllocator> (m, "arena");
#endif
// end work-around
make_ArrayOfStructs<NReal, NInt, amrex::PinnedArenaAllocator> (m, "pinned");
make_ArrayOfStructs<ParticleType, amrex::PinnedArenaAllocator> (m, "pinned");
#ifdef AMREX_USE_GPU
make_ArrayOfStructs<NReal, NInt, amrex::DeviceArenaAllocator> (m, "device");
make_ArrayOfStructs<NReal, NInt, amrex::ManagedArenaAllocator> (m, "managed");
make_ArrayOfStructs<NReal, NInt, amrex::AsyncArenaAllocator> (m, "async");
make_ArrayOfStructs<ParticleType, amrex::DeviceArenaAllocator> (m, "device");
make_ArrayOfStructs<ParticleType, amrex::ManagedArenaAllocator> (m, "managed");
make_ArrayOfStructs<ParticleType, amrex::AsyncArenaAllocator> (m, "async");
#endif
}

Expand Down
76 changes: 46 additions & 30 deletions src/Particle/ParticleTile.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,16 +14,21 @@

#include <sstream>


namespace py = pybind11;
using namespace amrex;

//Forward declaration
template <int T_NReal, int T_NInt=0>
void make_Particle(py::module &m);

template <int NStructReal, int NStructInt, int NArrayReal, int NArrayInt>
template <typename T_ParticleType, int NArrayReal, int NArrayInt>
void make_ParticleTileData(py::module &m) {
using ParticleTileDataType = ParticleTileData<NStructReal, NStructInt, NArrayReal, NArrayInt>;
using ParticleType = T_ParticleType;
constexpr int NStructReal = ParticleType::NReal;
constexpr int NStructInt = ParticleType::NInt;

using ParticleTileDataType = ParticleTileData<T_ParticleType, NArrayReal, NArrayInt>;
using SuperParticleType = Particle<NStructReal + NArrayReal, NStructInt + NArrayInt>;

auto const particle_tile_data_type =
Expand All @@ -35,21 +40,26 @@ void make_ParticleTileData(py::module &m) {
.def_readonly("m_size", &ParticleTileDataType::m_size)
.def_readonly("m_num_runtime_real", &ParticleTileDataType::m_num_runtime_real)
.def_readonly("m_num_runtime_int", &ParticleTileDataType::m_num_runtime_int)
.def("getSuperParticle", &ParticleTileDataType::getSuperParticle)
//.def("getSuperParticle", &ParticleTileDataType::getSuperParticle) // TODO
.def("setSuperParticle", &ParticleTileDataType::setSuperParticle)
// setter & getter
// setter & getter
.def("__setitem__", [](ParticleTileDataType &pdt, int const v,
SuperParticleType const value) { pdt.setSuperParticle(value, v); })
.def("__getitem__",
[](ParticleTileDataType &pdt, int const v) { return pdt.getSuperParticle(v); });
[](ParticleTileDataType &pdt, int const v) { return pdt.getSuperParticle(v); })

;
}

template <int NStructReal, int NStructInt, int NArrayReal, int NArrayInt,
template <typename T_ParticleType, int NArrayReal, int NArrayInt,
template<class> class Allocator=DefaultAllocator>
void make_ParticleTile(py::module &m, std::string allocstr)
{
using ParticleTileType = ParticleTile<NStructReal, NStructInt, NArrayReal, NArrayInt, Allocator>;
using ParticleType = Particle<NStructReal, NStructInt>;
using ParticleType = T_ParticleType;
constexpr int NStructReal = ParticleType::NReal;
constexpr int NStructInt = ParticleType::NInt;

using ParticleTileType = ParticleTile<T_ParticleType, NArrayReal, NArrayInt, Allocator>;
using SuperParticleType = Particle<NStructReal + NArrayReal, NStructInt + NArrayInt>;

auto const particle_tile_type = std::string("ParticleTile_") + std::to_string(NStructReal) + "_" +
Expand All @@ -66,17 +76,19 @@ void make_ParticleTile(py::module &m, std::string allocstr)
.def("GetStructOfArrays", py::overload_cast<>(&ParticleTileType::GetStructOfArrays),
py::return_value_policy::reference_internal)
.def("empty", &ParticleTileType::empty)
.def("size", &ParticleTileType::size)
.def("numParticles", &ParticleTileType::numParticles)
.def("numRealParticles", &ParticleTileType::numRealParticles)
.def("numNeighborParticles", &ParticleTileType::numNeighborParticles)
.def("numTotalParticles", &ParticleTileType::numTotalParticles)
//.def("size", &ParticleTileType::size) // TODO
//.def("numParticles", &ParticleTileType::numParticles) // TODO
//.def("numRealParticles", &ParticleTileType::numRealParticles) // TODO
//.def("numNeighborParticles", &ParticleTileType::numNeighborParticles) // TODO
//.def("numTotalParticles", &ParticleTileType::numTotalParticles) // TODO
.def("setNumNeighbors", &ParticleTileType::setNumNeighbors)
.def("getNumNeighbors", &ParticleTileType::getNumNeighbors)
.def("resize", &ParticleTileType::resize)

.def("push_back", [](ParticleTileType& ptile, const ParticleType &p){ ptile.push_back(p);})
// .def("push_back", py::overload_cast<const ParticleType&>(&ParticleTileType::push_back), "Add one particle to this tile.")
// .def("push_back", py::overload_cast<const SuperParticleType&>(&ParticleTileType::push_back), "Add one particle to this tile.")

.def("push_back", [](ParticleTileType& ptile, const SuperParticleType &p) {ptile.push_back(p);})
.def("push_back_real", [](ParticleTileType& ptile, int comp, ParticleReal v) {ptile.push_back_real(comp, v);})
.def("push_back_real", [](ParticleTileType& ptile,
Expand Down Expand Up @@ -107,50 +119,54 @@ void make_ParticleTile(py::module &m, std::string allocstr)
;
}

template <int NStructReal, int NStructInt, int NArrayReal, int NArrayInt>
template <typename T_ParticleType, int NArrayReal, int NArrayInt>
void make_ParticleTile(py::module &m)
{
make_ParticleTileData<NStructReal, NStructInt, NArrayReal, NArrayInt>(m);
make_ParticleTileData<T_ParticleType, NArrayReal, NArrayInt>(m);

// 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_ParticleTile<NStructReal, NStructInt, NArrayReal, NArrayInt,
//make_ParticleTile<T_ParticleType, NArrayReal, NArrayInt,
// std::allocator>(m, "std");
//make_ParticleTile<NStructReal, NStructInt, NArrayReal, NArrayInt,
//make_ParticleTile<T_ParticleType, NArrayReal, NArrayInt,
// amrex::ArenaAllocator>(m, "arena");
#ifdef AMREX_USE_GPU
make_ParticleTile<NStructReal, NStructInt, NArrayReal, NArrayInt,
make_ParticleTile<T_ParticleType, NArrayReal, NArrayInt,
std::allocator>(m, "std");
make_ParticleTile<NStructReal, NStructInt, NArrayReal, NArrayInt,
make_ParticleTile<T_ParticleType, NArrayReal, NArrayInt,
amrex::DefaultAllocator>(m, "default"); // amrex::ArenaAllocator
#else
make_ParticleTile<NStructReal, NStructInt, NArrayReal, NArrayInt,
make_ParticleTile<T_ParticleType, NArrayReal, NArrayInt,
amrex::DefaultAllocator>(m, "default"); // std::allocator
make_ParticleTile<NStructReal, NStructInt, NArrayReal, NArrayInt,
make_ParticleTile<T_ParticleType, NArrayReal, NArrayInt,
amrex::ArenaAllocator>(m, "arena");
#endif
// end work-around
make_ParticleTile<NStructReal, NStructInt, NArrayReal, NArrayInt,
make_ParticleTile<T_ParticleType, NArrayReal, NArrayInt,
amrex::PinnedArenaAllocator>(m, "pinned");
#ifdef AMREX_USE_GPU
make_ParticleTile<NStructReal, NStructInt, NArrayReal, NArrayInt,
make_ParticleTile<T_ParticleType, NArrayReal, NArrayInt,
amrex::DeviceArenaAllocator>(m, "device");
make_ParticleTile<NStructReal, NStructInt, NArrayReal, NArrayInt,
make_ParticleTile<T_ParticleType, NArrayReal, NArrayInt,
amrex::ManagedArenaAllocator>(m, "managed");
make_ParticleTile<NStructReal, NStructInt, NArrayReal, NArrayInt,
make_ParticleTile<T_ParticleType, NArrayReal, NArrayInt,
amrex::AsyncArenaAllocator>(m, "async");
#endif
}

void init_ParticleTile(py::module& m) {
// AMReX legacy AoS position + id/cpu particle ype
using ParticleType_0_0 = Particle<0, 0>;
using ParticleType_1_1 = Particle<1, 1>;

// 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_ParticleTile< 1, 1, 2, 1> (m);
make_ParticleTile< 0, 0, 4, 0> (m); // HiPACE++ 22.07
make_ParticleTile< 0, 0, 5, 0> (m); // ImpactX 22.07
make_ParticleTile< 0, 0, 7, 0> (m);
make_ParticleTile< 0, 0, 37, 1> (m); // HiPACE++ 22.07
make_ParticleTile<ParticleType_1_1, 2, 1> (m);
make_ParticleTile<ParticleType_0_0, 4, 0> (m); // HiPACE++ 22.07
make_ParticleTile<ParticleType_0_0, 5, 0> (m); // ImpactX 22.07
make_ParticleTile<ParticleType_0_0, 7, 0> (m);
make_ParticleTile<ParticleType_0_0, 37, 1> (m); // HiPACE++ 22.07
}

0 comments on commit 1e6f721

Please sign in to comment.