Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Enable OpenMP in particle push and coordinate transformation routines. #241

Merged
merged 8 commits into from Sep 22, 2022
Merged

Enable OpenMP in particle push and coordinate transformation routines. #241

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
96c86b6
Enable OpenMP in particle push and coordinate transformation routines.
atmyers Sep 9, 2022
ecccc14
enable tiling by default if not running on GPU
atmyers Sep 20, 2022
479c1e3
use dynamic scheduling by default
atmyers Sep 20, 2022
eea6241
do in nice way with named function
atmyers Sep 20, 2022
6e0c594
add some docstrings
atmyers Sep 21, 2022
91419a6
also enable omp in the new spacecharge particle gather and push
atmyers Sep 21, 2022
49d3f6f
turn on dynamic scheduling by default
atmyers Sep 21, 2022
10d5e87
also accelerate mfiter loop with omp
atmyers Sep 22, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
13 changes: 13 additions & 0 deletions src/initialization/InitParser.cpp
View file
Open in desktop
Expand Up @@ -22,5 +22,18 @@ namespace impactx::initialization
bool abort_on_out_of_gpu_memory = true; // AMReX' default: false
pp_amrex.query("abort_on_out_of_gpu_memory", abort_on_out_of_gpu_memory);
pp_amrex.add("abort_on_out_of_gpu_memory", abort_on_out_of_gpu_memory);

// Here we override the default tiling option for particles, which is always
// "false" in AMReX, to "false" if compiling for GPU execution and "true"
// if compiling for CPU.
{
amrex::ParmParse pp_particles("particles");
#ifdef AMREX_USE_GPU
bool do_tiling = false; // By default, tiling is off on GPU
#else
bool do_tiling = true;
#endif
pp_particles.queryAdd("do_tiling", do_tiling);
}
}
} // namespace impactx
36 changes: 34 additions & 2 deletions src/particles/ImpactXParticleContainer.H
View file
Open in desktop
Expand Up @@ -72,6 +72,38 @@ namespace impactx
};
};

/** AMReX iterator for particle boxes
*
* 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>
{
public:
using amrex::ParIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>::ParIter;

ParIter (ContainerType& pc, int level);

ParIter (ContainerType& pc, int level, amrex::MFItInfo& info);
};

/** Const AMReX iterator for particle boxes - data is read only.
*
* 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>
{
public:
using amrex::ParConstIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>::ParConstIter;

ParConstIter (ContainerType& pc, int level);

ParConstIter (ContainerType& pc, int level, amrex::MFItInfo& info);
};

/** Beam Particles in ImpactX
*
* This class stores particles, distributed over MPI ranks.
Expand All @@ -81,10 +113,10 @@ namespace impactx
{
public:
//! amrex iterator for particle boxes
using iterator = amrex::ParIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>;
using iterator = impactx::ParIter;

//! amrex constant iterator for particle boxes (read-only)
using const_iterator = amrex::ParConstIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>;
using const_iterator = impactx::ParConstIter;

//! Construct a new particle container
ImpactXParticleContainer (amrex::AmrCore* amr_core);
Expand Down
23 changes: 23 additions & 0 deletions src/particles/ImpactXParticleContainer.cpp
View file
Open in desktop
Expand Up @@ -24,6 +24,29 @@

namespace impactx
{
bool do_omp_dynamic () {
bool do_dynamic = true;
amrex::ParmParse pp_impactx("impactx");
pp_impactx.query("do_dynamic_scheduling", do_dynamic);
atmyers marked this conversation as resolved.
Show resolved Hide resolved
return do_dynamic;
}

ParIter::ParIter (ContainerType& pc, int level)
: amrex::ParIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>(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,
info.SetDynamic(do_omp_dynamic())) {}

ParConstIter::ParConstIter (ContainerType& pc, int level)
: amrex::ParConstIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>(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,
info.SetDynamic(do_omp_dynamic())) {}

ImpactXParticleContainer::ImpactXParticleContainer (amrex::AmrCore* amr_core)
: amrex::ParticleContainer<0, 0, RealSoA::nattribs, IntSoA::nattribs>(amr_core->GetParGDB())
{
Expand Down
3 changes: 3 additions & 0 deletions src/particles/Push.cpp
View file
Open in desktop
Expand Up @@ -118,6 +118,9 @@ namespace detail

// loop over all particle boxes
using ParIt = ImpactXParticleContainer::iterator;
#ifdef AMREX_USE_OMP
#pragma omp parallel if (amrex::Gpu::notInLaunchRegion())
#endif
for (ParIt pti(pc, lev); pti.isValid(); ++pti) {
const int np = pti.numParticles();
//const auto t_lev = pti.GetLevel();
Expand Down
3 changes: 3 additions & 0 deletions src/particles/spacecharge/ForceFromSelfFields.cpp
View file
Open in desktop
Expand Up @@ -40,6 +40,9 @@ namespace impactx::spacecharge
space_charge_field.at(lev).at("y").setVal(0.);
space_charge_field.at(lev).at("z").setVal(0.);

#ifdef AMREX_USE_OMP
#pragma omp parallel if (amrex::Gpu::notInLaunchRegion())
#endif
for (amrex::MFIter mfi(phi.at(lev)); mfi.isValid(); ++mfi) {
Copy link
Member

@ax3l ax3l Sep 27, 2022 •
edited

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@WeiqunZhang commented:
we can also do tiling on CPU for MFIter loops (not yet done here).

Let's investigate if this helps us here and potentially make it a user option.

Example:
https://github.com/ECP-WarpX/WarpX/blob/3fe406c9701f61e07b23f7123cf0a7bad492c6dc/Source/Diagnostics/ComputeDiagFunctors/BackTransformFunctor.cpp#L110


amrex::Box bx = mfi.validbox();
Expand Down
3 changes: 3 additions & 0 deletions src/particles/spacecharge/GatherAndPush.cpp
View file
Open in desktop
Expand Up @@ -43,6 +43,9 @@ namespace impactx::spacecharge

// loop over all particle boxes
using ParIt = ImpactXParticleContainer::iterator;
#ifdef AMREX_USE_OMP
#pragma omp parallel if (amrex::Gpu::notInLaunchRegion())
#endif
for (ParIt pti(pc, lev); pti.isValid(); ++pti) {
const int np = pti.numParticles();

Expand Down
3 changes: 3 additions & 0 deletions src/particles/transformation/CoordinateTransformation.cpp
View file
Open in desktop
Expand Up @@ -37,6 +37,9 @@ namespace transformation {
for (int lev = 0; lev <= nLevel; ++lev) {
// loop over all particle boxes
using ParIt = ImpactXParticleContainer::iterator;
#ifdef AMREX_USE_OMP
#pragma omp parallel if (amrex::Gpu::notInLaunchRegion())
#endif
for (ParIt pti(pc, lev); pti.isValid(); ++pti) {
const int np = pti.numParticles();

Expand Down