Merge pull request #5449 from ye-luo/pinned-mem

Pin host memory for fast/stable H2D and D2H transfers in non-batched delayed update implementation

Author: prckent

src/Containers/OhmmsPETE/OhmmsMatrix.h

@@ -132,6 +132,9 @@ class Matrix
     X.attachReference(other, ref, n * m);
   }
 
+  /// return true if this container is attached to another container
+  inline bool isAttached() const { return X.isAttached(); }
+
   template<typename Allocator = Alloc, typename = IsHostSafe<Allocator>>
   inline void add(size_type n) // you can add rows: adding columns are forbidden
   {

src/Containers/OhmmsPETE/OhmmsVector.h

@@ -151,6 +151,9 @@ class Vector
     qmc_allocator_traits<Alloc>::attachReference(other.mAllocator, mAllocator, ref - other.data());
   }
 
+  /// return true if this container is attached to another
+  inline bool isAttached() const { return nLocal > 0 && nAllocated == 0; }
+
   //! return the current size
   inline size_type size() const { return nLocal; }
 

src/Containers/OhmmsSoA/VectorSoaContainer.h

@@ -200,6 +200,9 @@ struct VectorSoaContainer
     qmc_allocator_traits<Alloc>::attachReference(other.mAllocator, mAllocator, other.data(), ptr);
   }
 
+  /// return true if this container is attached to another container
+  inline bool isAttached() { return nLocal > 0 && nAllocated == 0; }
+
   ///return the physical size
   inline size_type size() const { return nLocal; }
   ///return the physical size

src/QMCWaveFunctions/Fermion/DelayedUpdate.h

@@ -62,6 +62,10 @@ class DelayedUpdate
     delay_list.resize(delay);
   }
 
+  inline void prepareForDeviceCopy(const Matrix<T>& mat) const {}
+
+  inline void releaseFromDeviceCopy(const Matrix<T>& mat) const {}
+
   /** initialize internal objects when Ainv is refreshed
    * @param Ainv inverse matrix
    */

src/QMCWaveFunctions/Fermion/DelayedUpdateCUDA.h

@@ -85,6 +85,18 @@ class DelayedUpdateCUDA
     Ainv_gpu.resize(norb, norb);
   }
 
+  inline void prepareForDeviceCopy(Matrix<T>& mat) const
+  {
+    if(mat.size())
+      compute::MemManage<PL>::registerHost(mat.data(), mat.size() * sizeof(T));
+  }
+
+  inline void releaseFromDeviceCopy(Matrix<T>& mat) const
+  {
+    if(mat.size())
+      compute::MemManage<PL>::unregisterHost(mat.data());
+  }
+
   /** compute the inverse of the transpose of matrix A and its determinant value in log
    * @tparam TREAL real type
    */

src/QMCWaveFunctions/Fermion/DiracDeterminant.cpp

@@ -47,6 +47,14 @@ DiracDeterminant<PL, VT, FPVT>::DiracDeterminant(std::unique_ptr<SPOSet>&& spos,
     rot_spo->buildOptVariables(NumPtcls);
 }
 
+template<PlatformKind PL, typename VT, typename FPVT>
+DiracDeterminant<PL, VT, FPVT>::~DiracDeterminant()
+{
+  if(!psiM.isAttached())
+    accel_engine_.update_eng_.releaseFromDeviceCopy(psiM);
+  accel_engine_.update_eng_.releaseFromDeviceCopy(psiM_temp);
+}
+
 template<PlatformKind PL, typename VT, typename FPVT>
 void DiracDeterminant<PL, VT, FPVT>::invertPsiM(const ValueMatrix& logdetT, ValueMatrix& invMat)
 {
@@ -115,6 +123,10 @@ void DiracDeterminant<PL, VT, FPVT>::resize(int nel, int morb)
   int norb = morb;
   if (norb <= 0)
     norb = nel; // for morb == -1 (default)
+
+  accel_engine_.update_eng_.releaseFromDeviceCopy(psiM);
+  accel_engine_.update_eng_.releaseFromDeviceCopy(psiM_temp);
+
   accel_engine_.update_eng_.resize(norb, ndelay_);
   psiM.resize(nel, norb);
   dpsiM.resize(nel, norb);
@@ -123,6 +135,9 @@ void DiracDeterminant<PL, VT, FPVT>::resize(int nel, int morb)
   invRow.resize(norb);
   psiM_temp.resize(nel, norb);
 
+  accel_engine_.update_eng_.prepareForDeviceCopy(psiM);
+  accel_engine_.update_eng_.prepareForDeviceCopy(psiM_temp);
+
   dpsiV.resize(NumOrbitals);
   dspin_psiV.resize(NumOrbitals);
   d2psiV.resize(NumOrbitals);
@@ -345,6 +360,7 @@ void DiracDeterminant<PL, VT, FPVT>::registerData(ParticleSet& P, WFBufferType&
     buf.add(d2psiM.first_address(), d2psiM.last_address());
     Bytes_in_WFBuffer = buf.current() - Bytes_in_WFBuffer;
     // free local space
+    accel_engine_.update_eng_.releaseFromDeviceCopy(psiM);
     psiM.free();
     dpsiM.free();
     d2psiM.free();

src/QMCWaveFunctions/Fermion/DiracDeterminant.h

@@ -114,6 +114,8 @@ class DiracDeterminant : public DiracDeterminantBase
                    int ndelay                          = 1,
                    DetMatInvertor matrix_inverter_kind = DetMatInvertor::ACCEL);
 
+  ~DiracDeterminant() override;
+
   // copy constructor and assign operator disabled
   DiracDeterminant(const DiracDeterminant& s)            = delete;
   DiracDeterminant& operator=(const DiracDeterminant& s) = delete;