continue refactoring

lucafedeli88 · Jul 23, 2024 · d82d686 · d82d686
1 parent be96db2
commit d82d686
Show file tree

Hide file tree

Showing 20 changed files with 91 additions and 128 deletions.
diff --git a/Source/FieldSolver/FiniteDifferenceSolver/EvolveB.cpp b/Source/FieldSolver/FiniteDifferenceSolver/EvolveB.cpp
@@ -187,6 +187,8 @@ void FiniteDifferenceSolver::EvolveBCartesian (
                 }
             );
         }
+
+        cost_tracker.add();
     }
 }
 
@@ -340,6 +342,7 @@ void FiniteDifferenceSolver::EvolveBCartesianECT (
 
             });
 
+            cost_tracker.add();
         }
     }
 #else
@@ -449,6 +452,8 @@ void FiniteDifferenceSolver::EvolveBCylindrical (
             }
 
         );
+
+        cost_tracker.add();
     }
 }
 

diff --git a/Source/FieldSolver/FiniteDifferenceSolver/EvolveE.cpp b/Source/FieldSolver/FiniteDifferenceSolver/EvolveE.cpp
@@ -212,6 +212,8 @@ void FiniteDifferenceSolver::EvolveECartesian (
             );
 
         }
+
+        cost_tracker.add();
     }
 
 }
@@ -426,6 +428,8 @@ void FiniteDifferenceSolver::EvolveECylindrical (
 
         } // end of if condition for F
 
+        cost_tracker.add();
+
     } // end of loop over grid/tiles
 
 }

diff --git a/Source/FieldSolver/FiniteDifferenceSolver/EvolveECTRho.cpp b/Source/FieldSolver/FiniteDifferenceSolver/EvolveECTRho.cpp
@@ -146,6 +146,8 @@ void FiniteDifferenceSolver::EvolveRhoCartesianECT (
 #ifdef WARPX_DIM_XZ
         amrex::ignore_unused(Ey, Rhox, Rhoz, ly);
 #endif
+
+        cost_tracker.add();
     }
 #else
     amrex::ignore_unused(Efield, edge_lengths, face_areas, ECTRhofield, lev);

diff --git a/Source/FieldSolver/FiniteDifferenceSolver/HybridPICSolveE.cpp b/Source/FieldSolver/FiniteDifferenceSolver/HybridPICSolveE.cpp
@@ -232,6 +232,8 @@ void FiniteDifferenceSolver::CalculateCurrentAmpereCylindrical (
                 }
             }
         );
+
+        cost_tracker.add();
     }
 }
 
@@ -338,6 +340,8 @@ void FiniteDifferenceSolver::CalculateCurrentAmpereCartesian (
                 );
             }
         );
+
+        cost_tracker.add();
     }
 }
 #endif
@@ -501,6 +505,8 @@ void FiniteDifferenceSolver::HybridPICSolveECylindrical (
                 - (jt_interp - jit_interp - Jextt(i, j, 0)) * Br_interp
             );
         });
+
+        cost_tracker.add();
     }
 
     // Loop through the grids, and over the tiles within each grid again
@@ -670,6 +676,8 @@ void FiniteDifferenceSolver::HybridPICSolveECylindrical (
                 }
             }
         );
+
+        cost_tracker.add();
     }
 }
 
@@ -790,6 +798,8 @@ void FiniteDifferenceSolver::HybridPICSolveECartesian (
                 - (jy_interp - jiy_interp - Jexty(i, j, k)) * Bx_interp
             );
         });
+
+        cost_tracker.add();
     }
 
     // Loop through the grids, and over the tiles within each grid again
@@ -953,6 +963,8 @@ void FiniteDifferenceSolver::HybridPICSolveECartesian (
                 }
             }
         );
+
+        cost_tracker.add();
     }
 }
 #endif
diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp b/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp
@@ -7,6 +7,7 @@
  */
 #include "SpectralFieldData.H"
 
+#include "LoadBalance/LoadBalance.H"
 #include "Utils/WarpXAlgorithmSelection.H"
 #include "Utils/WarpXUtil.H"
 #include "WarpX.H"
@@ -125,9 +126,6 @@ SpectralFieldData::SpectralFieldData( const int lev,
                                       const bool periodic_single_box):
     m_periodic_single_box{periodic_single_box}
 {
-    amrex::LayoutData<amrex::Real>* cost = WarpX::getCosts(lev);
-    const bool do_costs = WarpXUtilLoadBalance::doCosts(cost, realspace_ba, dm);
-
     const BoxArray& spectralspace_ba = k_space.spectralspace_ba;
 
     // Allocate the arrays that contain the fields in spectral space
@@ -168,11 +166,8 @@ SpectralFieldData::SpectralFieldData( const int lev,
     // Loop over boxes and allocate the corresponding plan
     // for each box owned by the local MPI proc
     for ( MFIter mfi(spectralspace_ba, dm); mfi.isValid(); ++mfi ){
-        if (do_costs)
-        {
-            amrex::Gpu::synchronize();
-        }
-        auto wt = static_cast<amrex::Real>(amrex::second());
+
+        const auto cost_tracker = warpx::LoadBalance::CostTracker(lev, mfi);
 
         // Note: the size of the real-space box and spectral-space box
         // differ when using real-to-complex FFT. When initializing
@@ -189,12 +184,7 @@ SpectralFieldData::SpectralFieldData( const int lev,
             reinterpret_cast<ablastr::math::anyfft::Complex*>( tmpSpectralField[mfi].dataPtr()),
             ablastr::math::anyfft::direction::C2R, AMREX_SPACEDIM);
 
-        if (do_costs)
-        {
-            amrex::Gpu::synchronize();
-            wt = static_cast<amrex::Real>(amrex::second()) - wt;
-            amrex::HostDevice::Atomic::Add( &(*cost)[mfi.index()], wt);
-        }
+        cost_tracker.add();
     }
 }
 
@@ -217,9 +207,6 @@ SpectralFieldData::ForwardTransform (const int lev,
                                      const MultiFab& mf, const int field_index,
                                      const int i_comp)
 {
-    amrex::LayoutData<amrex::Real>* cost = WarpX::getCosts(lev);
-    const bool do_costs = WarpXUtilLoadBalance::doCosts(cost, mf.boxArray(), mf.DistributionMap());
-
     // Check field index type, in order to apply proper shift in spectral space
 #if (AMREX_SPACEDIM >= 2)
     const bool is_nodal_x = mf.is_nodal(0);
@@ -237,11 +224,8 @@ SpectralFieldData::ForwardTransform (const int lev,
     // Note: we do NOT OpenMP parallelize here, since we use OpenMP threads for
     //       the FFTs on each box!
     for ( MFIter mfi(mf); mfi.isValid(); ++mfi ){
-        if (do_costs)
-        {
-            amrex::Gpu::synchronize();
-        }
-        auto wt = static_cast<amrex::Real>(amrex::second());
+
+        const auto cost_tracker = warpx::LoadBalance::CostTracker(lev, mfi);
 
         // Copy the real-space field `mf` to the temporary field `tmpRealField`
         // This ensures that all fields have the same number of points
@@ -305,12 +289,7 @@ SpectralFieldData::ForwardTransform (const int lev,
             });
         }
 
-        if (do_costs)
-        {
-            amrex::Gpu::synchronize();
-            wt = static_cast<amrex::Real>(amrex::second()) - wt;
-            amrex::HostDevice::Atomic::Add( &(*cost)[mfi.index()], wt);
-        }
+        cost_tracker.add();
     }
 }
 
@@ -324,9 +303,6 @@ SpectralFieldData::BackwardTransform (const int lev,
                                       const amrex::IntVect& fill_guards,
                                       const int i_comp)
 {
-    amrex::LayoutData<amrex::Real>* cost = WarpX::getCosts(lev);
-    const bool do_costs = WarpXUtilLoadBalance::doCosts(cost, mf.boxArray(), mf.DistributionMap());
-
     // Check field index type, in order to apply proper shift in spectral space
 #if (AMREX_SPACEDIM >= 2)
     const bool is_nodal_x = mf.is_nodal(0);
@@ -362,11 +338,8 @@ SpectralFieldData::BackwardTransform (const int lev,
     // Note: we do NOT OpenMP parallelize here, since we use OpenMP threads for
     //       the iFFTs on each box!
     for ( MFIter mfi(mf); mfi.isValid(); ++mfi ){
-        if (do_costs)
-        {
-            amrex::Gpu::synchronize();
-        }
-        auto wt = static_cast<amrex::Real>(amrex::second());
+
+        const auto cost_tracker = warpx::LoadBalance::CostTracker(lev, mfi);
 
         // Copy the spectral-space field `tmpSpectralField` to the appropriate
         // field (specified by the input argument field_index)
@@ -466,12 +439,8 @@ SpectralFieldData::BackwardTransform (const int lev,
             });
         }
 
-        if (do_costs)
-        {
-            amrex::Gpu::synchronize();
-            wt = static_cast<amrex::Real>(amrex::second()) - wt;
-            amrex::HostDevice::Atomic::Add( &(*cost)[mfi.index()], wt);
-        }
+        cost_tracker.add();
+
     }
 }
-Original file line number
+Diff line change
@@ Expand Up / @@ -187,6 +187,8 @@ void FiniteDifferenceSolver::EvolveBCartesian ( @@
                     }
                 );
             }
+            cost_tracker.add();
         }
     }
@@ Expand Down Expand Up / @@ -340,6 +342,7 @@ void FiniteDifferenceSolver::EvolveBCartesianECT ( @@
                 });
+                cost_tracker.add();
             }
         }
     #else
@@ Expand Down Expand Up / @@ -449,6 +452,8 @@ void FiniteDifferenceSolver::EvolveBCylindrical ( @@
                 }
             );
+            cost_tracker.add();
         }
     }
@@ Expand Down @@