Initial LCAO vgl batched implementation using GEMM

src/QMCWaveFunctions/BasisSetBase.h

@@ -23,6 +23,8 @@
 
 #include "Particle/ParticleSet.h"
 #include "QMCWaveFunctions/OrbitalSetTraits.h"
+#include "OMPTarget/OffloadAlignedAllocators.hpp"
+
 
 namespace qmcplusplus
 {
@@ -128,10 +130,13 @@ struct BasisSetBase : public OrbitalSetTraits<T>
 template<typename T>
 struct SoaBasisSetBase
 {
-  using value_type = T;
-  using vgl_type   = VectorSoaContainer<T, OHMMS_DIM + 2>;
-  using vgh_type   = VectorSoaContainer<T, 10>;
-  using vghgh_type = VectorSoaContainer<T, 20>;
+  using value_type        = T;
+  using vgl_type          = VectorSoaContainer<T, OHMMS_DIM + 2>;
+  using vgh_type          = VectorSoaContainer<T, 10>;
+  using vghgh_type        = VectorSoaContainer<T, 20>;
+  using ValueType         = QMCTraits::ValueType;
+  using OffloadMWVGLArray = Array<ValueType, 3, OffloadPinnedAllocator<ValueType>>; // [VGL, walker, Orbs]
+
   ///size of the basis set
   int BasisSetSize;
 
@@ -143,6 +148,8 @@ struct SoaBasisSetBase
 
   //Evaluates value, gradient, and laplacian for electron "iat".  Parks them into a temporary data structure "vgl".
   virtual void evaluateVGL(const ParticleSet& P, int iat, vgl_type& vgl) = 0;
+  //Evaluates value, gradient, and laplacian for electron "iat".  places them in a offload array for batched code.
+  virtual void mw_evaluateVGL(const RefVectorWithLeader<ParticleSet>& P_list, int iat, OffloadMWVGLArray& vgl) = 0;
   //Evaluates value, gradient, and Hessian for electron "iat".  Parks them into a temporary data structure "vgh".
   virtual void evaluateVGH(const ParticleSet& P, int iat, vgh_type& vgh) = 0;
   //Evaluates value, gradient, and Hessian, and Gradient Hessian for electron "iat".  Parks them into a temporary data structure "vghgh".

src/QMCWaveFunctions/LCAO/LCAOrbitalSet.cpp

@@ -346,6 +346,67 @@ void LCAOrbitalSet::evaluateVGL(const ParticleSet& P, int iat, ValueVector& psi,
   }
 }
 
+void LCAOrbitalSet::mw_evaluateVGL(const RefVectorWithLeader<SPOSet>& spo_list,
+                                   const RefVectorWithLeader<ParticleSet>& P_list,
+                                   int iat,
+                                   const RefVector<ValueVector>& psi_v_list,
+                                   const RefVector<GradVector>& dpsi_v_list,
+                                   const RefVector<ValueVector>& d2psi_v_list) const
+{
+  OffloadMWVGLArray phi_vgl_v;
+  phi_vgl_v.resize(DIM_VGL, spo_list.size(), OrbitalSetSize);
+  mw_evaluateVGLImplGEMM(spo_list, P_list, iat, phi_vgl_v);
+
+  const size_t output_size = phi_vgl_v.size(2);
+  const size_t nw          = phi_vgl_v.size(1);
+
+  //TODO: make this cleaner?
+  for (int iw = 0; iw < nw; iw++)
+  {
+    std::copy_n(phi_vgl_v.data_at(0, iw, 0), output_size, psi_v_list[iw].get().data());
+    std::copy_n(phi_vgl_v.data_at(4, iw, 0), output_size, d2psi_v_list[iw].get().data());
+    // grads are [dim, walker, orb] in phi_vgl_v
+    //           [walker][orb, dim] in dpsi_v_list
+    for (size_t idim = 0; idim < DIM; idim++)
+      BLAS::copy(output_size, phi_vgl_v.data_at(idim + 1, iw, 0), 1, &dpsi_v_list[iw].get().data()[0][idim], DIM);
+  }
+}
+
+void LCAOrbitalSet::mw_evaluateVGLImplGEMM(const RefVectorWithLeader<SPOSet>& spo_list,
+                                           const RefVectorWithLeader<ParticleSet>& P_list,
+                                           int iat,
+                                           OffloadMWVGLArray& phi_vgl_v) const
+{
+  // [5][NW][NumAO]
+  OffloadMWVGLArray basis_mw;
+  basis_mw.resize(DIM_VGL, spo_list.size(), BasisSetSize);
+
+  if (Identity)
+  {
+    myBasisSet->mw_evaluateVGL(P_list, iat, basis_mw);
+    // output_size can be smaller than BasisSetSize
+    const size_t output_size = phi_vgl_v.size(2);
+    const size_t nw          = phi_vgl_v.size(1);
+
+    for (size_t idim = 0; idim < DIM_VGL; idim++)
+      for (int iw = 0; iw < nw; iw++)
+        std::copy_n(basis_mw.data_at(idim, iw, 0), output_size, phi_vgl_v.data_at(idim, iw, 0));
+  }
+  else
+  {
+    const size_t requested_orb_size = phi_vgl_v.size(2);
+    assert(requested_orb_size <= OrbitalSetSize);
+    ValueMatrix C_partial_view(C->data(), requested_orb_size, BasisSetSize);
+    myBasisSet->mw_evaluateVGL(P_list, iat, basis_mw);
+    BLAS::gemm('T', 'N',
+               requested_orb_size,        // MOs
+               spo_list.size() * DIM_VGL, // walkers * DIM_VGL
+               BasisSetSize,              // AOs
+               1, C_partial_view.data(), BasisSetSize, basis_mw.data(), BasisSetSize, 0, phi_vgl_v.data(),
+               requested_orb_size);
+  }
+}
+
 void LCAOrbitalSet::evaluateDetRatios(const VirtualParticleSet& VP,
                                       ValueVector& psi,
                                       const ValueVector& psiinv,
@@ -365,6 +426,34 @@ void LCAOrbitalSet::evaluateDetRatios(const VirtualParticleSet& VP,
   }
 }
 
+void LCAOrbitalSet::mw_evaluateVGLandDetRatioGrads(const RefVectorWithLeader<SPOSet>& spo_list,
+                                                   const RefVectorWithLeader<ParticleSet>& P_list,
+                                                   int iat,
+                                                   const std::vector<const ValueType*>& invRow_ptr_list,
+                                                   OffloadMWVGLArray& phi_vgl_v,
+                                                   std::vector<ValueType>& ratios,
+                                                   std::vector<GradType>& grads) const
+{
+  assert(this == &spo_list.getLeader());
+  assert(phi_vgl_v.size(0) == DIM_VGL);
+  assert(phi_vgl_v.size(1) == spo_list.size());
+
+  mw_evaluateVGLImplGEMM(spo_list, P_list, iat, phi_vgl_v);
+  // Device data of phi_vgl_v must be up-to-date upon return
+  phi_vgl_v.updateTo();
+
+  const size_t nw             = spo_list.size();
+  const size_t norb_requested = phi_vgl_v.size(2);
+  for (int iw = 0; iw < nw; iw++)
+  {
+    ratios[iw] = simd::dot(invRow_ptr_list[iw], phi_vgl_v.data_at(0, iw, 0), norb_requested);
+    GradType dphi;
+    for (size_t idim = 0; idim < DIM; idim++)
+      dphi[idim] = simd::dot(invRow_ptr_list[iw], phi_vgl_v.data_at(idim + 1, iw, 0), norb_requested) / ratios[iw];
+    grads[iw] = dphi;
+  }
+}
+
 void LCAOrbitalSet::evaluateVGH(const ParticleSet& P, int iat, ValueVector& psi, GradVector& dpsi, HessVector& dhpsi)
 {
   //TAKE CARE OF IDENTITY

src/QMCWaveFunctions/LCAO/LCAOrbitalSet.h

@@ -30,10 +30,11 @@ namespace qmcplusplus
 struct LCAOrbitalSet : public SPOSet
 {
 public:
-  using basis_type = SoaBasisSetBase<ValueType>;
-  using vgl_type   = basis_type::vgl_type;
-  using vgh_type   = basis_type::vgh_type;
-  using vghgh_type = basis_type::vghgh_type;
+  using basis_type        = SoaBasisSetBase<ValueType>;
+  using vgl_type          = basis_type::vgl_type;
+  using vgh_type          = basis_type::vgh_type;
+  using vghgh_type        = basis_type::vghgh_type;
+  using OffloadMWVGLArray = Array<ValueType, 3, OffloadPinnedAllocator<ValueType>>; // [VGL, walker, Orbs]
 
   ///pointer to the basis set
   std::unique_ptr<basis_type> myBasisSet;
@@ -78,11 +79,27 @@ struct LCAOrbitalSet : public SPOSet
 
   void evaluateVGL(const ParticleSet& P, int iat, ValueVector& psi, GradVector& dpsi, ValueVector& d2psi) override;
 
+
+  void mw_evaluateVGL(const RefVectorWithLeader<SPOSet>& spo_list,
+                      const RefVectorWithLeader<ParticleSet>& P_list,
+                      int iat,
+                      const RefVector<ValueVector>& psi_v_list,
+                      const RefVector<GradVector>& dpsi_v_list,
+                      const RefVector<ValueVector>& d2psi_v_list) const override;
+
   void evaluateDetRatios(const VirtualParticleSet& VP,
                          ValueVector& psi,
                          const ValueVector& psiinv,
                          std::vector<ValueType>& ratios) override;
 
+  void mw_evaluateVGLandDetRatioGrads(const RefVectorWithLeader<SPOSet>& spo_list,
+                                      const RefVectorWithLeader<ParticleSet>& P_list,
+                                      int iat,
+                                      const std::vector<const ValueType*>& invRow_ptr_list,
+                                      OffloadMWVGLArray& phi_vgl_v,
+                                      std::vector<ValueType>& ratios,
+                                      std::vector<GradType>& grads) const override;
+
   void evaluateVGH(const ParticleSet& P,
                    int iat,
                    ValueVector& psi,
@@ -219,7 +236,7 @@ struct LCAOrbitalSet : public SPOSet
   vghgh_type Tempghv;
 
 private:
-  ///helper functions to handl Identity
+  ///helper functions to handle Identity
   void evaluate_vgl_impl(const vgl_type& temp, ValueVector& psi, GradVector& dpsi, ValueVector& d2psi) const;
 
   void evaluate_vgl_impl(const vgl_type& temp,
@@ -228,6 +245,8 @@ struct LCAOrbitalSet : public SPOSet
                          GradMatrix& dlogdet,
                          ValueMatrix& d2logdet) const;
   ///These two functions unpack the data in vgh_type temp object into wavefunction friendly data structures.
+
+
   ///This unpacks temp into vectors psi, dpsi, and d2psi.
   void evaluate_vgh_impl(const vgh_type& temp, ValueVector& psi, GradVector& dpsi, HessVector& d2psi) const;
 
@@ -266,6 +285,11 @@ struct LCAOrbitalSet : public SPOSet
 
   ///Unpacks data in vgl object and calculates/places ionic gradient of a single row (phi_j(r)) into dlogdet.
   void evaluate_ionderiv_v_row_impl(const vgl_type& temp, GradVector& dlogdet) const;
+
+  void mw_evaluateVGLImplGEMM(const RefVectorWithLeader<SPOSet>& spo_list,
+                              const RefVectorWithLeader<ParticleSet>& P_list,
+                              int iat,
+                              OffloadMWVGLArray& phi_vgl_v) const;
 };
 } // namespace qmcplusplus
 #endif

src/QMCWaveFunctions/LCAO/SoaLocalizedBasisSet.cpp

@@ -138,6 +138,37 @@ void SoaLocalizedBasisSet<COT, ORBT>::evaluateVGL(const ParticleSet& P, int iat,
   }
 }
 
+template<class COT, typename ORBT>
+void SoaLocalizedBasisSet<COT, ORBT>::mw_evaluateVGL(const RefVectorWithLeader<ParticleSet>& P_list,
+                                                     int iat,
+                                                     OffloadMWVGLArray& vgl_v)
+{
+  for (size_t iw = 0; iw < P_list.size(); iw++)
+  {
+    const auto& IonID(ions_.GroupID);
+    const auto& coordR  = P_list[iw].activeR(iat);
+    const auto& d_table = P_list[iw].getDistTableAB(myTableIndex);
+    const auto& dist    = (P_list[iw].getActivePtcl() == iat) ? d_table.getTempDists() : d_table.getDistRow(iat);
+    const auto& displ   = (P_list[iw].getActivePtcl() == iat) ? d_table.getTempDispls() : d_table.getDisplRow(iat);
+
+    PosType Tv;
+
+    // number of walkers * BasisSetSize
+    auto stride = vgl_v.size(1) * BasisSetSize;
+    assert(BasisSetSize == vgl_v.size(2));
+    vgl_type vgl_iw(vgl_v.data_at(0, iw, 0), BasisSetSize, stride);
+
+    for (int c = 0; c < NumCenters; c++)
+    {
+      Tv[0] = (ions_.R[c][0] - coordR[0]) - displ[c][0];
+      Tv[1] = (ions_.R[c][1] - coordR[1]) - displ[c][1];
+      Tv[2] = (ions_.R[c][2] - coordR[2]) - displ[c][2];
+      LOBasisSet[IonID[c]]->evaluateVGL(P_list[iw].getLattice(), dist[c], displ[c], BasisOffset[c], vgl_iw, Tv);
+    }
+  }
+}
+
+
 template<class COT, typename ORBT>
 void SoaLocalizedBasisSet<COT, ORBT>::evaluateVGH(const ParticleSet& P, int iat, vgh_type& vgh)
 {

src/QMCWaveFunctions/LCAO/SoaLocalizedBasisSet.h

@@ -21,6 +21,7 @@
 
 #include <memory>
 #include "QMCWaveFunctions/BasisSetBase.h"
+#include "OMPTarget/OffloadAlignedAllocators.hpp"
 
 namespace qmcplusplus
 {
@@ -36,12 +37,15 @@ template<class COT, typename ORBT>
 class SoaLocalizedBasisSet : public SoaBasisSetBase<ORBT>
 {
 public:
-  using RealType   = typename COT::RealType;
-  using BaseType   = SoaBasisSetBase<ORBT>;
-  using vgl_type   = typename BaseType::vgl_type;
-  using vgh_type   = typename BaseType::vgh_type;
-  using vghgh_type = typename BaseType::vghgh_type;
-  using PosType    = typename ParticleSet::PosType;
+  using RealType  = typename COT::RealType;
+  using BaseType  = SoaBasisSetBase<ORBT>;
+  using ValueType = QMCTraits::ValueType;
+
+  using vgl_type          = typename BaseType::vgl_type;
+  using vgh_type          = typename BaseType::vgh_type;
+  using vghgh_type        = typename BaseType::vghgh_type;
+  using PosType           = typename ParticleSet::PosType;
+  using OffloadMWVGLArray = Array<ValueType, 3, OffloadPinnedAllocator<ValueType>>; // [VGL, walker, Orbs]
 
   using BaseType::BasisSetSize;
 
@@ -104,6 +108,8 @@ class SoaLocalizedBasisSet : public SoaBasisSetBase<ORBT>
    */
   void evaluateVGL(const ParticleSet& P, int iat, vgl_type& vgl) override;
 
+  void mw_evaluateVGL(const RefVectorWithLeader<ParticleSet>& P_list, int iat, OffloadMWVGLArray& vgl) override;
+
   /** compute VGH 
    * @param P quantum particleset
    * @param iat active particle