Add WARPX_PROFILE calls to each python callback. (#2573)

When python callbacks take some time, this is useful as otherwise many
callbacks are lumped together in WarpX::Evolve::step.

1 files changed, 40 insertions, 10 deletions

diff --git a/Source/Evolve/WarpXEvolve.cpp b/Source/Evolve/WarpXEvolve.cpp
index 02d9fcf4c..d04a54ae3 100644
--- a/Source/Evolve/WarpXEvolve.cpp
+++ b/Source/Evolve/WarpXEvolve.cpp
@@ -82,7 +82,10 @@ WarpX::Evolve (int numsteps)
         if (verbose) {
             amrex::Print() << "\nSTEP " << step+1 << " starts ...\n";
         }
-        if (warpx_py_beforestep) warpx_py_beforestep();
+        if (warpx_py_beforestep) {
+            WARPX_PROFILE("warpx_py_beforestep");
+            warpx_py_beforestep();
+        }
 
         amrex::LayoutData<amrex::Real>* cost = WarpX::getCosts(0);
         if (cost) {
@@ -164,7 +167,10 @@ WarpX::Evolve (int numsteps)
         // Main PIC operation:
         // gather fields, push particles, deposit sources, update fields
 
-        if (warpx_py_particleinjection) warpx_py_particleinjection();
+        if (warpx_py_particleinjection) {
+            WARPX_PROFILE("warpx_py_particleinjection");
+            warpx_py_particleinjection();
+        }
         // Electrostatic case: only gather fields and push particles,
         // deposition and calculation of fields done further below
         if (do_electrostatic != ElectrostaticSolverAlgo::None)
@@ -294,7 +300,10 @@ WarpX::Evolve (int numsteps)
         }
 
         if( do_electrostatic != ElectrostaticSolverAlgo::None ) {
-            if (warpx_py_beforeEsolve) warpx_py_beforeEsolve();
+            if (warpx_py_beforeEsolve) {
+                WARPX_PROFILE("warpx_py_beforeEsolve");
+                warpx_py_beforeEsolve();
+            }
             // Electrostatic solver:
             // For each species: deposit charge and add the associated space-charge
             // E and B field to the grid ; this is done at the end of the PIC
@@ -303,7 +312,10 @@ WarpX::Evolve (int numsteps)
             // and so that the fields are at the correct time in the output.
             bool const reset_fields = true;
             ComputeSpaceChargeField( reset_fields );
-            if (warpx_py_afterEsolve) warpx_py_afterEsolve();
+            if (warpx_py_afterEsolve) {
+                WARPX_PROFILE("warpx_py_afterEsolve");
+                warpx_py_afterEsolve();
+            }
         }
 
         // sync up time
@@ -313,7 +325,10 @@ WarpX::Evolve (int numsteps)
 
         // warpx_py_afterstep runs with the updated global time. It is included
         // in the evolve timing.
-        if (warpx_py_afterstep) warpx_py_afterstep();
+        if (warpx_py_afterstep) {
+            WARPX_PROFILE("warpx_py_afterstep");
+            warpx_py_afterstep();
+        }
 
         /// reduced diags
         if (reduced_diags->m_plot_rd != 0)
@@ -370,11 +385,20 @@ WarpX::OneStep_nosub (Real cur_time)
     //               from p^{n-1/2} to p^{n+1/2}
     // Deposit current j^{n+1/2}
     // Deposit charge density rho^{n}
-    if (warpx_py_particlescraper) warpx_py_particlescraper();
-    if (warpx_py_beforedeposition) warpx_py_beforedeposition();
+    if (warpx_py_particlescraper) {
+        WARPX_PROFILE("warpx_py_particlescraper");
+        warpx_py_particlescraper();
+    }
+    if (warpx_py_beforedeposition) {
+        WARPX_PROFILE("warpx_py_beforedeposition");
+        warpx_py_beforedeposition();
+    }
     PushParticlesandDepose(cur_time);
 
-    if (warpx_py_afterdeposition) warpx_py_afterdeposition();
+    if (warpx_py_afterdeposition) {
+        WARPX_PROFILE("warpx_py_afterdeposition");
+        warpx_py_afterdeposition();
+    }
 
     // Synchronize J and rho
     SyncCurrent();
@@ -396,7 +420,10 @@ WarpX::OneStep_nosub (Real cur_time)
     if (do_pml && pml_has_particles) CopyJPML();
     if (do_pml && do_pml_j_damping) DampJPML();
 
-    if (warpx_py_beforeEsolve) warpx_py_beforeEsolve();
+    if (warpx_py_beforeEsolve) {
+        WARPX_PROFILE("warpx_py_beforeEsolve");
+        warpx_py_beforeEsolve();
+    }
 
     // Push E and B from {n} to {n+1}
     // (And update guard cells immediately afterwards)
@@ -476,7 +503,10 @@ WarpX::OneStep_nosub (Real cur_time)
             FillBoundaryB(guard_cells.ng_alloc_EB);
     } // !PSATD
 
-    if (warpx_py_afterEsolve) warpx_py_afterEsolve();
+    if (warpx_py_afterEsolve) {
+        WARPX_PROFILE("warpx_py_afterEsolve");
+        warpx_py_afterEsolve();
+    }
 }
 
 void