From 56a47f835eb2f38069394aae14713c46247f7549 Mon Sep 17 00:00:00 2001 From: Thierry Antoun Date: Tue, 13 Sep 2022 09:51:12 -0700 Subject: [PATCH] First change to impactx Adapting to SoA structures Link the CMake file to Axel Warpx branch Avoid branch conflicts Correction of ParIter and ConstParIter Adding corrections to make_alike function call Update Diagnostics Correction of the ParIter loop Still adaptating ImpactX to SoAParticle Update Init Note: https://github.com/AMReX-Codes/amrex/pull/2878#issuecomment-1261115487 Adapted the PushSingleParticle function to SoAParticle --- cmake/dependencies/ABLASTR.cmake | 8 +-- src/particles/ImpactXParticleContainer.H | 53 +++++++++++-------- src/particles/ImpactXParticleContainer.cpp | 40 ++++++-------- src/particles/Push.cpp | 22 +++++--- .../diagnostics/DiagnosticOutput.cpp | 35 ++++++------ .../CoordinateTransformation.cpp | 12 ++--- 6 files changed, 88 insertions(+), 82 deletions(-) diff --git a/cmake/dependencies/ABLASTR.cmake b/cmake/dependencies/ABLASTR.cmake index cda09a677..3d5d3cc26 100644 --- a/cmake/dependencies/ABLASTR.cmake +++ b/cmake/dependencies/ABLASTR.cmake @@ -131,18 +131,18 @@ set(ImpactX_ablastr_src "" "Local path to ABLASTR source directory (preferred if set)") # Git fetcher -set(ImpactX_ablastr_repo "https://github.com/ECP-WarpX/WarpX.git" +set(ImpactX_ablastr_repo "https://github.com/Thierry992/WarpX.git" CACHE STRING "Repository URI to pull and build ABLASTR from if(ImpactX_ablastr_internal)") -set(ImpactX_ablastr_branch "22.10" +set(ImpactX_ablastr_branch "soa-particle" CACHE STRING "Repository branch for ImpactX_ablastr_repo if(ImpactX_ablastr_internal)") # AMReX is transitively pulled through ABLASTR -set(ImpactX_amrex_repo "https://github.com/AMReX-Codes/amrex.git" +set(ImpactX_amrex_repo "https://github.com/Thierry992/amrex.git" CACHE STRING "Repository URI to pull and build AMReX from if(ImpactX_amrex_internal)") -set(ImpactX_amrex_branch "ed1ecd62acb3fd7d39b8a23aa4e9ad09669741bb" +set(ImpactX_amrex_branch "particle_soa_refactor" CACHE STRING "Repository branch for ImpactX_amrex_repo if(ImpactX_amrex_internal)") diff --git a/src/particles/ImpactXParticleContainer.H b/src/particles/ImpactXParticleContainer.H index 6744cf2fd..5dfbbee43 100644 --- a/src/particles/ImpactXParticleContainer.H +++ b/src/particles/ImpactXParticleContainer.H @@ -17,6 +17,7 @@ #include #include #include +#include #include #include @@ -35,15 +36,16 @@ namespace impactx * because we change the meaning of these "positions" depending on the * coordinate system we are currently in. */ - struct RealAoS - { - enum - { - x, ///< position in x [m] (at fixed t OR fixed s) - y, ///< position in y [m] (at fixed t OR fixed s) - z, ///< position in z [m] (at fixed t) OR time-of-flight ct [m] (at fixed s) - }; - }; + + //struct RealAoS + //{ + // enum + // { + // x, ///< position in x [m] (at fixed t OR fixed s) + // y, ///< position in y [m] (at fixed t OR fixed s) + // z, ///< position in z [m] (at fixed t) OR time-of-flight ct [m] (at fixed s) + // }; + //}; /** This struct indexes the additional Real attributes * stored in an SoA in ImpactXParticleContainer @@ -52,6 +54,9 @@ namespace impactx { enum { + x, + y, + z, ux, ///< momentum in x, scaled by the magnitude of the reference momentum [unitless] (at fixed t or s) uy, ///< momentum in y, scaled by the magnitude of the reference momentum [unitless] (at fixed t or s) pt, ///< momentum in z, scaled by the magnitude of the reference momentum [unitless] (at fixed t) OR energy deviation, scaled by speed of light * the magnitude of the reference momentum [unitless] (at fixed s) @@ -68,7 +73,9 @@ namespace impactx { enum { - nattribs ///< the number of particles above (always last) + id, + cpu, + nattribs ///< the number of attributes above (always last) }; }; @@ -77,15 +84,15 @@ namespace impactx * We subclass here to change the default threading strategy, which is * `static` in AMReX, to `dynamic` in ImpactX. */ - class ParIter - : public amrex::ParIter<0, 0, RealSoA::nattribs, IntSoA::nattribs> + class ParIterSoA + : public amrex::ParIterSoA { public: - using amrex::ParIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>::ParIter; + using amrex::ParIterSoA::ParIterSoA; - ParIter (ContainerType& pc, int level); + ParIterSoA (ContainerType& pc, int level); - ParIter (ContainerType& pc, int level, amrex::MFItInfo& info); + ParIterSoA (ContainerType& pc, int level, amrex::MFItInfo& info); }; /** Const AMReX iterator for particle boxes - data is read only. @@ -93,15 +100,15 @@ namespace impactx * We subclass here to change the default threading strategy, which is * `static` in AMReX, to `dynamic` in ImpactX. */ - class ParConstIter - : public amrex::ParConstIter<0, 0, RealSoA::nattribs, IntSoA::nattribs> + class ParConstIterSoA + : public amrex::ParConstIterSoA { public: - using amrex::ParConstIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>::ParConstIter; + using amrex::ParConstIterSoA::ParConstIterSoA; - ParConstIter (ContainerType& pc, int level); + ParConstIterSoA (ContainerType& pc, int level); - ParConstIter (ContainerType& pc, int level, amrex::MFItInfo& info); + ParConstIterSoA (ContainerType& pc, int level, amrex::MFItInfo& info); }; /** Beam Particles in ImpactX @@ -109,14 +116,14 @@ namespace impactx * This class stores particles, distributed over MPI ranks. */ class ImpactXParticleContainer - : public amrex::ParticleContainer<0, 0, RealSoA::nattribs, IntSoA::nattribs> + : public amrex::ParticleContainerPureSoA { public: //! amrex iterator for particle boxes - using iterator = impactx::ParIter; + using iterator = impactx::ParIterSoA; //! amrex constant iterator for particle boxes (read-only) - using const_iterator = impactx::ParConstIter; + using const_iterator = impactx::ParConstIterSoA; //! Construct a new particle container ImpactXParticleContainer (amrex::AmrCore* amr_core); diff --git a/src/particles/ImpactXParticleContainer.cpp b/src/particles/ImpactXParticleContainer.cpp index 9f80c3955..e75223e3f 100644 --- a/src/particles/ImpactXParticleContainer.cpp +++ b/src/particles/ImpactXParticleContainer.cpp @@ -17,7 +17,6 @@ #include #include #include -#include #include @@ -31,24 +30,24 @@ namespace impactx return do_dynamic; } - ParIter::ParIter (ContainerType& pc, int level) - : amrex::ParIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>(pc, level, + ParIterSoA::ParIterSoA (ContainerType& pc, int level) + : amrex::ParIterSoA(pc, level, amrex::MFItInfo().SetDynamic(do_omp_dynamic())) {} - ParIter::ParIter (ContainerType& pc, int level, amrex::MFItInfo& info) - : amrex::ParIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>(pc, level, + ParIterSoA::ParIterSoA (ContainerType& pc, int level, amrex::MFItInfo& info) + : amrex::ParIterSoA(pc, level, info.SetDynamic(do_omp_dynamic())) {} - ParConstIter::ParConstIter (ContainerType& pc, int level) - : amrex::ParConstIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>(pc, level, + ParConstIterSoA::ParConstIterSoA (ContainerType& pc, int level) + : amrex::ParConstIterSoA(pc, level, amrex::MFItInfo().SetDynamic(do_omp_dynamic())) {} - ParConstIter::ParConstIter (ContainerType& pc, int level, amrex::MFItInfo& info) - : amrex::ParConstIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>(pc, level, + ParConstIterSoA::ParConstIterSoA (ContainerType& pc, int level, amrex::MFItInfo& info) + : amrex::ParConstIterSoA(pc, level, info.SetDynamic(do_omp_dynamic())) {} ImpactXParticleContainer::ImpactXParticleContainer (amrex::AmrCore* amr_core) - : amrex::ParticleContainer<0, 0, RealSoA::nattribs, IntSoA::nattribs>(amr_core->GetParGDB()) + : amrex::ParticleContainerPureSoA(amr_core->GetParGDB()) { SetParticleSize(); } @@ -105,32 +104,26 @@ namespace impactx reserveData(); resizeData(); + // write Real attributes (SoA) to particle initialized zero auto& particle_tile = DefineAndReturnParticleTile(0, 0, 0); /* Create a temporary tile to obtain data from simulation. This data * is then copied to the permanent tile which is stored on the particle * (particle_tile). */ - using PinnedTile = amrex::ParticleTile; + using PinnedTile = typename ContainerLike::ParticleTileType; PinnedTile pinned_tile; pinned_tile.define(NumRuntimeRealComps(), NumRuntimeIntComps()); for (int i = 0; i < np; i++) { - ParticleType p; - p.id() = ParticleType::NextID(); - p.cpu() = amrex::ParallelDescriptor::MyProc(); - p.pos(0) = x[i]; - p.pos(1) = y[i]; - p.pos(2) = z[i]; - // write position, creating cpu id, and particle id - pinned_tile.push_back(p); + pinned_tile.push_back_int(IntSoA::id, ParticleType::NextID()); + pinned_tile.push_back_int(IntSoA::cpu, amrex::ParallelDescriptor::MyProc()); } - // write Real attributes (SoA) to particle initialized zero - DefineAndReturnParticleTile(0, 0, 0); - + pinned_tile.push_back_real(RealSoA::x, x); + pinned_tile.push_back_real(RealSoA::y, y); + pinned_tile.push_back_real(RealSoA::z, z); pinned_tile.push_back_real(RealSoA::ux, px); pinned_tile.push_back_real(RealSoA::uy, py); pinned_tile.push_back_real(RealSoA::pt, pz); @@ -142,6 +135,7 @@ namespace impactx */ auto old_np = particle_tile.numParticles(); auto new_np = old_np + pinned_tile.numParticles(); + amrex::AllPrint() << "old_np=" << old_np << " new_np=" << new_np << "\n"; particle_tile.resize(new_np); amrex::copyParticles( particle_tile, pinned_tile, 0, old_np, pinned_tile.numParticles()); diff --git a/src/particles/Push.cpp b/src/particles/Push.cpp index 10c17d571..e8d05b1a7 100644 --- a/src/particles/Push.cpp +++ b/src/particles/Push.cpp @@ -35,6 +35,8 @@ namespace detail struct PushSingleParticle { using PType = ImpactXParticleContainer::ParticleType; + using ParticleTileType = amrex::ParticleTile,RealSoA::nattribs, IntSoA::nattribs>; + using ParticleTileDataType = amrex::ParticleTileData; /** Constructor taking in pointers to particle data * @@ -46,12 +48,12 @@ namespace detail * @param ref_part the struct containing the reference particle */ PushSingleParticle (T_Element element, - PType* AMREX_RESTRICT aos_ptr, + ParticleTileDataType ptd, amrex::ParticleReal* AMREX_RESTRICT part_px, amrex::ParticleReal* AMREX_RESTRICT part_py, amrex::ParticleReal* AMREX_RESTRICT part_pt, RefPart ref_part) - : m_element(element), m_aos_ptr(aos_ptr), + : m_element(element), m_ptd(ptd), m_part_px(part_px), m_part_py(part_py), m_part_pt(part_pt), m_ref_part(ref_part) { @@ -70,8 +72,8 @@ namespace detail void operator() (long i) const { - // access AoS data such as positions and cpu/id - PType& AMREX_RESTRICT p = m_aos_ptr[i]; + + PType p(m_ptd,i); // access SoA Real data amrex::ParticleReal & AMREX_RESTRICT px = m_part_px[i]; @@ -85,7 +87,7 @@ namespace detail private: T_Element const m_element; - PType* const AMREX_RESTRICT m_aos_ptr; + ParticleTileDataType const m_ptd; amrex::ParticleReal* const AMREX_RESTRICT m_part_px; amrex::ParticleReal* const AMREX_RESTRICT m_part_py; amrex::ParticleReal* const AMREX_RESTRICT m_part_pt; @@ -129,8 +131,12 @@ namespace detail // preparing access to particle data: AoS using PType = ImpactXParticleContainer::ParticleType; - auto& aos = pti.GetArrayOfStructs(); - PType* AMREX_RESTRICT aos_ptr = aos().dataPtr(); + using ParticleTile = amrex::ParticleTile,RealSoA::nattribs, IntSoA::nattribs>; + using ParticleTileDataType = amrex::ParticleTileData; + + ParticleTileDataType tile_data = pti.GetParticleTile().getParticleTileData(); + //auto& aos = pti.GetArrayOfStructs(); + //PType* AMREX_RESTRICT aos_ptr = aos().dataPtr(); // preparing access to particle data: SoA of Reals auto& soa_real = pti.GetStructOfArrays().GetRealData(); @@ -143,7 +149,7 @@ namespace detail [=, &ref_part](auto element) { // push beam particles relative to reference particle detail::PushSingleParticle const pushSingleParticle( - element, aos_ptr, part_px, part_py, part_pt, ref_part); + element, tile_data, part_px, part_py, part_pt, ref_part); // loop over beam particles in the box amrex::ParallelFor(np, pushSingleParticle); diff --git a/src/particles/diagnostics/DiagnosticOutput.cpp b/src/particles/diagnostics/DiagnosticOutput.cpp index a9aae3f8e..a6854faf4 100644 --- a/src/particles/diagnostics/DiagnosticOutput.cpp +++ b/src/particles/diagnostics/DiagnosticOutput.cpp @@ -17,6 +17,7 @@ #include // for ParmParse #include // for ParticleReal #include // for PrintToFile +#include // for constructor of SoAParticle namespace impactx::diagnostics @@ -46,7 +47,8 @@ namespace impactx::diagnostics } // create a host-side particle buffer - auto tmp = pc.make_alike(); + using ParticleType = amrex::SoAParticle; + auto tmp = pc.template make_alike(); // copy device-to-host bool const local = true; @@ -56,34 +58,34 @@ namespace impactx::diagnostics int const nLevel = tmp.finestLevel(); for (int lev = 0; lev <= nLevel; ++lev) { // loop over all particle boxes - using ParIt = typename decltype(tmp)::ParConstIterType; - for (ParIt pti(tmp, lev); pti.isValid(); ++pti) { + using MyPinnedParIter = amrex::ParIterSoA; + for (MyPinnedParIter pti(tmp, lev); pti.isValid(); ++pti) { const int np = pti.numParticles(); // preparing access to particle data: AoS - using PType = ImpactXParticleContainer::ParticleType; - auto const &aos = pti.GetArrayOfStructs(); - PType const *const AMREX_RESTRICT aos_ptr = aos().dataPtr(); + //using PType = ImpactXParticleContainer::ParticleType; + //auto const &aos = pti.GetArrayOfStructs(); + //PType const *const AMREX_RESTRICT aos_ptr = aos().dataPtr(); // preparing access to particle data: SoA of Reals - auto &soa_real = pti.GetStructOfArrays().GetRealData(); - amrex::ParticleReal const *const AMREX_RESTRICT part_px = soa_real[RealSoA::ux].dataPtr(); - amrex::ParticleReal const *const AMREX_RESTRICT part_py = soa_real[RealSoA::uy].dataPtr(); - amrex::ParticleReal const *const AMREX_RESTRICT part_pt = soa_real[RealSoA::pt].dataPtr(); + auto const& particle_attributes = pti.GetStructOfArrays(); + auto const& part_px = particle_attributes.GetRealData(RealSoA::ux); + auto const& part_py = particle_attributes.GetRealData(RealSoA::uy); + auto const& part_pt = particle_attributes.GetRealData(RealSoA::pt); + + // Preparing the constructor for the new SoA Particle Type + auto ptd = pti.GetParticleTile().getParticleTileData(); if (otype == OutputType::PrintParticles) { // print out particles (this hack works only on CPU and on GPUs with // unified memory access) for (int i = 0; i < np; ++i) { - - // access AoS data such as positions and cpu/id - PType const &p = aos_ptr[i]; + ParticleType p(ptd, i); amrex::ParticleReal const x = p.pos(0); amrex::ParticleReal const y = p.pos(1); amrex::ParticleReal const t = p.pos(2); uint64_t const global_id = ablastr::particles::localIDtoGlobal(p.id(), p.cpu()); - // access SoA Real data amrex::ParticleReal const px = part_px[i]; amrex::ParticleReal const py = part_py[i]; amrex::ParticleReal const pt = part_pt[i]; @@ -119,14 +121,11 @@ namespace impactx::diagnostics // print out particles (this hack works only on CPU and on GPUs with // unified memory access) for (int i = 0; i < np; ++i) { - - // access AoS data such as positions and cpu/id - PType const &p = aos_ptr[i]; + ParticleType p(ptd, i); amrex::ParticleReal const x = p.pos(0); amrex::ParticleReal const y = p.pos(1); uint64_t const global_id = ablastr::particles::localIDtoGlobal(p.id(), p.cpu()); - // access SoA Real data amrex::ParticleReal const px = part_px[i]; amrex::ParticleReal const py = part_py[i]; diff --git a/src/particles/transformation/CoordinateTransformation.cpp b/src/particles/transformation/CoordinateTransformation.cpp index 71d4741c3..433738995 100644 --- a/src/particles/transformation/CoordinateTransformation.cpp +++ b/src/particles/transformation/CoordinateTransformation.cpp @@ -43,10 +43,7 @@ namespace transformation { for (ParIt pti(pc, lev); pti.isValid(); ++pti) { const int np = pti.numParticles(); - // preparing access to particle data: AoS using PType = ImpactXParticleContainer::ParticleType; - auto &aos = pti.GetArrayOfStructs(); - PType *AMREX_RESTRICT aos_ptr = aos().dataPtr(); // preparing access to particle data: SoA of Reals auto &soa_real = pti.GetStructOfArrays().GetRealData(); @@ -60,9 +57,11 @@ namespace transformation { amrex::ParticleReal const pzd = sqrt(pow(pd, 2) - 1.0); ToFixedS const to_s(pzd); + auto tile_data = pti.GetParticleTile().getParticleTileData(); + amrex::ParallelFor(np, [=] AMREX_GPU_DEVICE(long i) { - // access AoS data such as positions and cpu/id - PType &p = aos_ptr[i]; + + PType p(tile_data,i); // access SoA Real data amrex::ParticleReal &px = part_px[i]; @@ -75,9 +74,10 @@ namespace transformation { BL_PROFILE("impactx::transformation::CoordinateTransformation::to_fixed_t"); amrex::ParticleReal const ptd = pd; // Design value of pt/mc2 = -gamma. ToFixedT const to_t(ptd); + auto tile_data = pti.GetParticleTile().getParticleTileData(); amrex::ParallelFor(np, [=] AMREX_GPU_DEVICE(long i) { // access AoS data such as positions and cpu/id - PType &p = aos_ptr[i]; + PType p(tile_data,i); // access SoA Real data amrex::ParticleReal &px = part_px[i];