Enable OpenMP in particle push and coordinate transformation routines. (

#241)

* Enable OpenMP in particle push and coordinate transformation routines.

* enable tiling by default if not running on GPU

* use dynamic scheduling by default

* do in nice way with named function

* add some docstrings

* also enable omp in the new spacecharge particle gather and push

* turn on dynamic scheduling by default

* also accelerate mfiter loop with omp

src/initialization/InitParser.cpp

@@ -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

src/particles/ImpactXParticleContainer.H

@@ -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.
@@ -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);

src/particles/ImpactXParticleContainer.cpp

@@ -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);
+        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())
     {

src/particles/Push.cpp

@@ -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();

src/particles/spacecharge/ForceFromSelfFields.cpp

@@ -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) {
 
                 amrex::Box bx = mfi.validbox();

src/particles/spacecharge/GatherAndPush.cpp

@@ -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();
 

src/particles/transformation/CoordinateTransformation.cpp

@@ -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();