Merge pull request QMCPACK#4240 from jptowns/orbopt_perf_sandbox

Added new timers to help with optimization scaling.
williamfgc · Oct 24, 2022 · c860467 · c860467
2 parents 3b0e1ea + c55081b
commit c860467
Show file tree

Hide file tree

Showing 6 changed files with 51 additions and 20 deletions.
diff --git a/src/QMCDrivers/WFOpt/QMCCostFunction.cpp b/src/QMCDrivers/WFOpt/QMCCostFunction.cpp
@@ -26,7 +26,9 @@
 namespace qmcplusplus
 {
 QMCCostFunction::QMCCostFunction(MCWalkerConfiguration& w, TrialWaveFunction& psi, QMCHamiltonian& h, Communicate* comm)
-    : QMCCostFunctionBase(w, psi, h, comm)
+    : QMCCostFunctionBase(w, psi, h, comm),
+      fill_timer_(
+          *timer_manager.createTimer("QMCCostFunction::fillOverlapHamiltonianMatrices", timer_level_medium))
 {
   CSWeight = 1.0;
   app_log() << " Using QMCCostFunction::QMCCostFunction" << std::endl;
@@ -635,6 +637,8 @@ QMCCostFunction::EffectiveWeight QMCCostFunction::correlatedSampling(bool needGr
 QMCCostFunction::Return_rt QMCCostFunction::fillOverlapHamiltonianMatrices(Matrix<Return_rt>& Left,
                                                                            Matrix<Return_rt>& Right)
 {
+  ScopedTimer tmp_timer(fill_timer_);
+
   RealType b1, b2;
   if (GEVType == "H2")
   {

diff --git a/src/QMCDrivers/WFOpt/QMCCostFunction.h b/src/QMCDrivers/WFOpt/QMCCostFunction.h
@@ -66,6 +66,8 @@ class QMCCostFunction : public QMCCostFunctionBase, public CloneManager
   int total_samples();
   Return_rt LMYEngineCost_detail(cqmc::engine::LMYEngine<Return_t>* EngineObj) override;
 #endif
+
+  NewTimer& fill_timer_;
 };
 } // namespace qmcplusplus
 #endif
diff --git a/src/QMCDrivers/WFOpt/QMCFixedSampleLinearOptimize.cpp b/src/QMCDrivers/WFOpt/QMCFixedSampleLinearOptimize.cpp
@@ -93,12 +93,12 @@ QMCFixedSampleLinearOptimize::QMCFixedSampleLinearOptimize(const ProjectData& pr
       Max_iterations(1),
       wfNode(NULL),
       param_tol(1e-4),
-      generate_samples_timer_(
-          *timer_manager.createTimer("QMCLinearOptimizeBatched::GenerateSamples", timer_level_medium)),
-      initialize_timer_(*timer_manager.createTimer("QMCLinearOptimizeBatched::Initialize", timer_level_medium)),
-      eigenvalue_timer_(*timer_manager.createTimer("QMCLinearOptimizeBatched::Eigenvalue", timer_level_medium)),
-      line_min_timer_(*timer_manager.createTimer("QMCLinearOptimizeBatched::Line_Minimization", timer_level_medium)),
-      cost_function_timer_(*timer_manager.createTimer("QMCLinearOptimizeBatched::CostFunction", timer_level_medium))
+      generate_samples_timer_(*timer_manager.createTimer("QMCLinearOptimize::GenerateSamples", timer_level_medium)),
+      initialize_timer_(*timer_manager.createTimer("QMCLinearOptimize::Initialize", timer_level_medium)),
+      eigenvalue_timer_(*timer_manager.createTimer("QMCLinearOptimize::EigenvalueSolve", timer_level_medium)),
+      involvmat_timer_(*timer_manager.createTimer("QMCLinearOptimize::invertOverlapMat", timer_level_medium)),
+      line_min_timer_(*timer_manager.createTimer("QMCLinearOptimize::Line_Minimization", timer_level_medium)),
+      cost_function_timer_(*timer_manager.createTimer("QMCLinearOptimize::CostFunction", timer_level_medium))
 {
   IsQMCDriver = false;
   //set the optimization flag
@@ -305,7 +305,10 @@ bool QMCFixedSampleLinearOptimize::run()
     if (apply_inverse)
     {
       Matrix<RealType> RightT(Left);
-      invert_matrix(RightT, false);
+      {
+        ScopedTimer local(involvmat_timer_);
+        invert_matrix(RightT, false);
+      }
       Left = 0;
       product(RightT, Right, Left);
       //       Now the left matrix is the Hamiltonian with the inverse of the overlap applied ot it.
@@ -341,10 +344,12 @@ bool QMCFixedSampleLinearOptimize::run()
       for (int i = 1; i < N; i++)
         Right(i, i) += std::exp(XS);
       app_log() << "  Using XS:" << XS << " " << failedTries << " " << stability << std::endl;
-      eigenvalue_timer_.start();
-      getLowestEigenvector(Right, currentParameterDirections);
-      Lambda = getNonLinearRescale(currentParameterDirections, S, *optTarget);
-      eigenvalue_timer_.stop();
+
+      {
+        ScopedTimer local(eigenvalue_timer_);
+        getLowestEigenvector(Right, currentParameterDirections);
+        Lambda = getNonLinearRescale(currentParameterDirections, S, *optTarget);
+      }
       //       biggest gradient in the parameter direction vector
       RealType bigVec(0);
       for (int i = 0; i < numParams; i++)
@@ -1314,7 +1319,10 @@ bool QMCFixedSampleLinearOptimize::one_shift_run()
   }
 
   // compute the inverse of the overlap matrix
-  invert_matrix(invMat, false);
+  {
+    ScopedTimer local(involvmat_timer_);
+    invert_matrix(invMat, false);
+  }
 
   // apply the overlap shift
   for (int i = 1; i < N; i++)
@@ -1330,7 +1338,10 @@ bool QMCFixedSampleLinearOptimize::one_shift_run()
       std::swap(prdMat(i, j), prdMat(j, i));
 
   // compute the lowest eigenvalue of the product matrix and the corresponding eigenvector
-  getLowestEigenvector(prdMat, parameterDirections);
+  {
+    ScopedTimer local(eigenvalue_timer_);
+    getLowestEigenvector(prdMat, parameterDirections);
+  }
 
   // compute the scaling constant to apply to the update
   Lambda = getNonLinearRescale(parameterDirections, ovlMat, *optTarget);
@@ -1522,8 +1533,11 @@ void QMCFixedSampleLinearOptimize::start()
     Timer t2;
     optTarget->getConfigurations(h5FileRoot);
     optTarget->setRng(vmcEngine->getRngRefs());
+
+    // Compute wfn parameter derivatives
     NullEngineHandle handle;
     optTarget->checkConfigurations(handle);
+
     // check recomputed variance against VMC
     auto sigma2_vmc   = vmcEngine->getBranchEngine()->vParam[SimpleFixedNodeBranch::SBVP::SIGMA2];
     auto sigma2_check = optTarget->getVariance();

diff --git a/src/QMCDrivers/WFOpt/QMCFixedSampleLinearOptimize.h b/src/QMCDrivers/WFOpt/QMCFixedSampleLinearOptimize.h
@@ -243,6 +243,7 @@ class QMCFixedSampleLinearOptimize : public QMCDriver, public LinearMethod, priv
   NewTimer& generate_samples_timer_;
   NewTimer& initialize_timer_;
   NewTimer& eigenvalue_timer_;
+  NewTimer& involvmat_timer_;
   NewTimer& line_min_timer_;
   NewTimer& cost_function_timer_;
   Timer t1;

diff --git a/src/QMCDrivers/WFOpt/QMCFixedSampleLinearOptimizeBatched.cpp b/src/QMCDrivers/WFOpt/QMCFixedSampleLinearOptimizeBatched.cpp
@@ -90,6 +90,7 @@ QMCFixedSampleLinearOptimizeBatched::QMCFixedSampleLinearOptimizeBatched(
           *timer_manager.createTimer("QMCLinearOptimizeBatched::GenerateSamples", timer_level_medium)),
       initialize_timer_(*timer_manager.createTimer("QMCLinearOptimizeBatched::Initialize", timer_level_medium)),
       eigenvalue_timer_(*timer_manager.createTimer("QMCLinearOptimizeBatched::Eigenvalue", timer_level_medium)),
+      involvmat_timer_(*timer_manager.createTimer("QMCLinearOptimizedBatched::invert_matrix", timer_level_medium)),
       line_min_timer_(*timer_manager.createTimer("QMCLinearOptimizeBatched::Line_Minimization", timer_level_medium)),
       cost_function_timer_(*timer_manager.createTimer("QMCLinearOptimizeBatched::CostFunction", timer_level_medium)),
       wfNode(NULL),
@@ -426,10 +427,11 @@ bool QMCFixedSampleLinearOptimizeBatched::previous_linear_methods_run()
       for (int i = 1; i < N; i++)
         Right(i, i) += std::exp(XS);
       app_log() << "  Using XS:" << XS << " " << failedTries << " " << stability << std::endl;
-      eigenvalue_timer_.start();
-      getLowestEigenvector(Right, currentParameterDirections);
-      objFuncWrapper_.Lambda = getNonLinearRescale(currentParameterDirections, S, *optTarget);
-      eigenvalue_timer_.stop();
+      {
+        ScopedTimer local(eigenvalue_timer_);
+        getLowestEigenvector(Right, currentParameterDirections);
+        objFuncWrapper_.Lambda = getNonLinearRescale(currentParameterDirections, S, *optTarget);
+      }
       //       biggest gradient in the parameter direction vector
       RealType bigVec(0);
       for (int i = 0; i < numParams; i++)
@@ -1650,7 +1652,10 @@ bool QMCFixedSampleLinearOptimizeBatched::one_shift_run()
   }
 
   // compute the inverse of the overlap matrix
-  invert_matrix(invMat, false);
+  {
+    ScopedTimer local(involvmat_timer_);
+    invert_matrix(invMat, false);
+  }
 
   // apply the overlap shift
   for (int i = 1; i < N; i++)
@@ -1666,7 +1671,11 @@ bool QMCFixedSampleLinearOptimizeBatched::one_shift_run()
       std::swap(prdMat(i, j), prdMat(j, i));
 
   // compute the lowest eigenvalue of the product matrix and the corresponding eigenvector
-  RealType lowestEV = getLowestEigenvector(prdMat, parameterDirections);
+  RealType lowestEV = 0.;
+  {
+    ScopedTimer local(eigenvalue_timer_);
+    lowestEV = getLowestEigenvector(prdMat, parameterDirections);
+  }
 
   // compute the scaling constant to apply to the update
   objFuncWrapper_.Lambda = getNonLinearRescale(parameterDirections, ovlMat, *optTarget);

diff --git a/src/QMCDrivers/WFOpt/QMCFixedSampleLinearOptimizeBatched.h b/src/QMCDrivers/WFOpt/QMCFixedSampleLinearOptimizeBatched.h
@@ -277,6 +277,7 @@ class QMCFixedSampleLinearOptimizeBatched : public QMCDriverNew, LinearMethod
   NewTimer& generate_samples_timer_;
   NewTimer& initialize_timer_;
   NewTimer& eigenvalue_timer_;
+  NewTimer& involvmat_timer_;
   NewTimer& line_min_timer_;
   NewTimer& cost_function_timer_;
   Timer t1;