Fix TM1/3 when using both batched algorithm and delayed updates

src/QMCWaveFunctions/Fermion/DelayedUpdate.h

@@ -41,20 +41,9 @@ namespace qmcplusplus {
       /// current number of delays, increase one for each acceptance, reset to 0 after updating Ainv
       int delay_count;
 
-      /// pointer to the row of up-to-date Ainv
-      const T* Ainv_row_ptr;
-      /** row id correspond to the up-to-date Ainv_row_ptr. [0 norb), Ainv_row_ptr is valid; -1, Ainv_row_ptr is not valid.
-       *  This id is set by getInvRow indicating Ainv_row_ptr has been prepared for the Ainv_row_id row
-       *  ratioGrad checks if Ainv_row_id is consistent. If not, Ainv_row_ptr needs to be recomputed.
-       *  acceptRow and updateInvMat mark Ainv_row_ptr invalid by setting Ainv_row_id to -1
-       */
-      int Ainv_row_id;
-      /// current determinant ratio
-      T curRatio;
-
     public:
       /// default constructor
-      DelayedUpdate(): delay_count(0), Ainv_row_ptr(nullptr), Ainv_row_id(-1) {}
+      DelayedUpdate(): delay_count(0) {}
 
       ///resize the internal storage
       /** resize the internal storage
@@ -76,72 +65,25 @@ namespace qmcplusplus {
        * @param Ainv inverse matrix
        * @param rowchanged the row id corresponding to the proposed electron
        */
-      inline void getInvRow(const Matrix<T>& Ainv, int rowchanged)
+      template<typename VVT>
+      inline void getInvRow(const Matrix<T>& Ainv, int rowchanged, VVT& invRow)
       {
-        Ainv_row_id = rowchanged;
         if ( delay_count == 0 )
         {
           // Ainv is fresh, directly access Ainv
-          Ainv_row_ptr = Ainv[rowchanged];
+          std::copy_n(Ainv[rowchanged], invRow.size(), invRow.data());
           return;
         }
         const T cone(1);
         const T czero(0);
-        const T* AinvRow = Ainv[rowchanged];
         const int norb = Ainv.rows();
         const int lda_Binv = Binv.cols();
-        // save AinvRow to new_AinvRow
-        std::copy_n(AinvRow, norb, V[delay_count]);
-        // multiply V (NxK) Binv(KxK) U(KxN) AinvRow right to the left
-        BLAS::gemv('T', norb, delay_count, cone, U.data(), norb, AinvRow, 1, czero, p.data(), 1);
+        // save Ainv[rowchanged] to invRow
+        std::copy_n(Ainv[rowchanged], norb, invRow.data());
+        // multiply V (NxK) Binv(KxK) U(KxN) invRow right to the left
+        BLAS::gemv('T', norb, delay_count, cone, U.data(), norb, invRow.data(), 1, czero, p.data(), 1);
         BLAS::gemv('N', delay_count, delay_count, cone, Binv.data(), lda_Binv, p.data(), 1, czero, Binv[delay_count], 1);
-        BLAS::gemv('N', norb, delay_count, -cone, V.data(), norb, Binv[delay_count], 1, cone, V[delay_count], 1);
-        Ainv_row_ptr = V[delay_count];
-      }
-
-      /** compute determinant ratio of new determinant
-       * @param Ainv inverse matrix
-       * @param rowchanged the row id corresponding to the proposed electron
-       * @param psiV new orbital values
-       */
-      template<typename VVT>
-      inline T ratio(const Matrix<T>& Ainv, int rowchanged, const VVT& psiV)
-      {
-        // check Ainv_row_id against rowchanged to see if getInvRow() has been called
-        // This is intended to save redundant compuation in TM1 and TM3
-        if(Ainv_row_id != rowchanged)
-          getInvRow(Ainv, rowchanged);
-        return curRatio = simd::dot(Ainv_row_ptr,psiV.data(),Ainv.cols());
-      }
-
-      /** compute the old gradient
-       * @param Ainv inverse matrix
-       * @param rowchanged the row id corresponding to the proposed electron
-       * @param dpsiV old orbital derivatives
-       */
-      template<typename GT>
-      inline GT evalGrad(const Matrix<T>& Ainv, int rowchanged, const GT* dpsiV)
-      {
-        getInvRow(Ainv, rowchanged);
-        return simd::dot(Ainv_row_ptr,dpsiV,Ainv.cols());
-      }
-
-      /** compute determinant ratio and gradients of new determinant
-       * @param Ainv inverse matrix
-       * @param rowchanged the row id corresponding to the proposed electron
-       * @param psiV new orbital values
-       * @param dpsiV new orbital derivatives
-       * @param g new gradients
-       */
-      template<typename VVT, typename GGT, typename GT>
-      inline T ratioGrad(const Matrix<T>& Ainv, int rowchanged, const VVT& psiV, const GGT& dpsiV, GT& g)
-      {
-        // check Ainv_row_id against rowchanged to ensure getInvRow() called before ratioGrad()
-        // This is not a safety check. Some code paths do call ratioGrad without calling evalGrad first.
-        if(Ainv_row_id != rowchanged)
-          getInvRow(Ainv, rowchanged);
-        g = simd::dot(Ainv_row_ptr,dpsiV.data(),Ainv.cols());
-        return curRatio = simd::dot(Ainv_row_ptr,psiV.data(),Ainv.cols());
+        BLAS::gemv('N', norb, delay_count, -cone, V.data(), norb, Binv[delay_count], 1, cone, invRow.data(), 1);
       }
 
       /** accept a move with the update delayed
@@ -154,9 +96,6 @@ namespace qmcplusplus {
       template<typename VVT>
       inline void acceptRow(Matrix<T>& Ainv, int rowchanged, const VVT& psiV)
       {
-        // Ainv_row_ptr is no more valid by marking Ainv_row_id -1
-        Ainv_row_id = -1;
-
         const T cminusone(-1);
         const T czero(0);
         const int norb = Ainv.rows();
@@ -188,8 +127,6 @@ namespace qmcplusplus {
        */
       inline void updateInvMat(Matrix<T>& Ainv)
       {
-        // Ainv_row_ptr is no more valid by marking Ainv_row_id -1
-        Ainv_row_id = -1;
         if(delay_count==0) return;
         // update the inverse matrix
         const T cone(1);

src/QMCWaveFunctions/Fermion/DiracDeterminant.cpp

@@ -32,7 +32,7 @@ namespace qmcplusplus
  *@param first index of the first particle
  */
 DiracDeterminant::DiracDeterminant(SPOSetPtr const spos, int first):
-  DiracDeterminantBase(spos,first), ndelay(1)
+  DiracDeterminantBase(spos,first), ndelay(1), invRow_id(-1)
 {
   ClassName = "DiracDeterminant";
 }
@@ -87,6 +87,7 @@ void DiracDeterminant::resize(int nel, int morb)
   dpsiM.resize(nel,norb);
   d2psiM.resize(nel,norb);
   psiV.resize(norb);
+  invRow.resize(norb);
   psiM_temp.resize(nel,norb);
   if( typeid(ValueType) != typeid(mValueType) )
     psiM_hp.resize(nel,norb);
@@ -116,7 +117,9 @@ DiracDeterminant::evalGrad(ParticleSet& P, int iat)
 {
   const int WorkingIndex = iat-FirstIndex;
   RatioTimer.start();
-  GradType g = updateEng.evalGrad(psiM, WorkingIndex, dpsiM[WorkingIndex]);
+  invRow_id = WorkingIndex;
+  updateEng.getInvRow(psiM, WorkingIndex, invRow);
+  GradType g = simd::dot(invRow.data(), dpsiM[WorkingIndex], invRow.size());
   RatioTimer.stop();
   return g;
 }
@@ -131,9 +134,17 @@ DiracDeterminant::ratioGrad(ParticleSet& P, int iat, GradType& grad_iat)
   const int WorkingIndex = iat-FirstIndex;
   UpdateMode=ORB_PBYP_PARTIAL;
   GradType rv;
-  curRatio = updateEng.ratioGrad(psiM, WorkingIndex, psiV, dpsiV, rv);
-  grad_iat += ((RealType)1.0/curRatio) * rv;
-  RatioTimer.stop();
+
+  // This is an optimization.
+  // check invRow_id against WorkingIndex to see if getInvRow() has been called already
+  // Some code paths call evalGrad before calling ratioGrad.
+  if(invRow_id != WorkingIndex)
+  {
+    invRow_id = WorkingIndex;
+    updateEng.getInvRow(psiM, WorkingIndex, invRow);
+  }
+  curRatio = simd::dot(invRow.data(), psiV.data(), invRow.size());
+  grad_iat += ((RealType)1.0/curRatio) * simd::dot(invRow.data(), dpsiV.data(), invRow.size());
   return curRatio;
 }
 
@@ -146,6 +157,8 @@ void DiracDeterminant::acceptMove(ParticleSet& P, int iat)
   LogValue +=std::log(std::abs(curRatio));
   UpdateTimer.start();
   updateEng.acceptRow(psiM,WorkingIndex,psiV);
+  // invRow becomes invalid after accepting a move
+  invRow_id = -1;
   if(UpdateMode == ORB_PBYP_PARTIAL)
   {
     simd::copy(dpsiM[WorkingIndex],  dpsiV.data(),  NumOrbitals);
@@ -165,6 +178,8 @@ void DiracDeterminant::restore(int iat)
 void DiracDeterminant::completeUpdates()
 {
   UpdateTimer.start();
+  // invRow becomes invalid after updating the inverse matrix
+  invRow_id = -1;
   updateEng.updateInvMat(psiM);
   UpdateTimer.stop();
 }
@@ -250,6 +265,8 @@ void DiracDeterminant::copyFromBuffer(ParticleSet& P, WFBufferType& buf)
   d2psiM.attachReference(buf.lendReference<ValueType>(d2psiM.size()));
   buf.get(LogValue);
   buf.get(PhaseValue);
+  // start with invRow labelled invalid
+  invRow_id = -1;
   BufferTimer.stop();
 }
 
@@ -265,16 +282,26 @@ DiracDeterminant::ValueType DiracDeterminant::ratio(ParticleSet& P, int iat)
   Phi->evaluate(P, iat, psiV);
   SPOVTimer.stop();
   RatioTimer.start();
-  curRatio = updateEng.ratio(psiM, WorkingIndex, psiV);
+  // This is an optimization.
+  // check invRow_id against WorkingIndex to see if getInvRow() has been called
+  // This is intended to save redundant compuation in TM1 and TM3
+  if(invRow_id != WorkingIndex)
+  {
+    invRow_id = WorkingIndex;
+    updateEng.getInvRow(psiM, WorkingIndex, invRow);
+  }
+  curRatio = simd::dot(invRow.data(), psiV.data(), invRow.size());
   RatioTimer.stop();
   return curRatio;
 }
 
 void DiracDeterminant::evaluateRatios(VirtualParticleSet& VP, std::vector<ValueType>& ratios)
 {
-  ValueVector_t psiM_row(psiM[VP.refPtcl-FirstIndex], psiM.cols());
   SPOVTimer.start();
-  Phi->evaluateDetRatios(VP, psiV, psiM_row, ratios);
+  const int WorkingIndex = VP.refPtcl-FirstIndex;
+  invRow_id = WorkingIndex;
+  updateEng.getInvRow(psiM, WorkingIndex, invRow);
+  Phi->evaluateDetRatios(VP, psiV, invRow, ratios);
   SPOVTimer.stop();
 }
 

src/QMCWaveFunctions/Fermion/DiracDeterminant.h

@@ -183,6 +183,16 @@ class DiracDeterminant: public DiracDeterminantBase
   DiracMatrix<mValueType> detEng;
   DelayedUpdate<ValueType> updateEng;
 
+  /// the row of up-to-date inverse matrix
+  ValueVector_t invRow;
+
+  /** row id correspond to the up-to-date invRow. [0 norb), invRow is ready; -1, invRow is not valid.
+   *  This id is set after calling getInvRow indicating invRow has been prepared for the invRow_id row
+   *  ratioGrad checks if invRow_id is consistent. If not, invRow needs to be recomputed.
+   *  acceptMove and completeUpdates mark invRow invalid by setting invRow_id to -1
+   */
+  int invRow_id;
+
   ValueType curRatio,cumRatio;
   ParticleSet::SingleParticleValue_t *FirstAddressOfG;
   ParticleSet::SingleParticleValue_t *LastAddressOfG;

src/QMCWaveFunctions/tests/test_dirac_matrix.cpp

@@ -121,13 +121,13 @@ TEST_CASE("DiracMatrix_update_row", "[wavefunction][fermion]")
   dm.invert(a,false);
 
   // new row
-  Vector<ValueType> v;
-  v.resize(3);
+  Vector<ValueType> v(3), invRow(3);
   v[0] = 1.9;
   v[1] = 2.0;
   v[2] = 3.1;
 
-  ValueType det_ratio1 = updateEng.ratio(a,0,v);
+  updateEng.getInvRow(a,0,invRow);
+  ValueType det_ratio1 = simd::dot(v.data(), invRow.data(), invRow.size());
 
   ValueType det_ratio = 0.178276269185;
   REQUIRE(det_ratio1 == ValueApprox(det_ratio));