Merge pull request #3846 from ye-luo/skip-donePbyP

Recompute DTAA at donePbyP on demand only.
QMCPACK · Feb 18, 2022 · 4250568 · 4250568
2 parents 3765fcc + a31932b
commit 4250568
Show file tree

Hide file tree

Showing 7 changed files with 21 additions and 9 deletions.
diff --git a/src/Particle/DTModes.h b/src/Particle/DTModes.h
@@ -20,7 +20,7 @@ namespace qmcplusplus
 enum class DTModes : uint_fast8_t
 {
   ALL_OFF = 0x0,
-  /** whether full table needs to be ready at anytime or not
+  /** whether full table needs to be ready at anytime or not during PbyP
    * Optimization can be implemented during forward PbyP move when the full table is not needed all the time.
    * DT consumers should know if full table is needed or not and request via addTable.
    */
@@ -34,7 +34,12 @@ enum class DTModes : uint_fast8_t
   /** skip data transfer back to host after mw_evalaute full distance table.
    * this optimization can be used for distance table consumed directly on the device without copying back to the host.
    */
-  MW_EVALUATE_RESULT_NO_TRANSFER_TO_HOST = 0x4
+  MW_EVALUATE_RESULT_NO_TRANSFER_TO_HOST = 0x4,
+  /** whether full table needs to be ready at anytime or not after donePbyP
+   * Optimization can be implemented during forward PbyP move when the full table is not needed all the time.
+   * DT consumers should know if full table is needed or not and request via addTable.
+   */
+  NEED_FULL_TABLE_ON_HOST_AFTER_DONEPBYP = 0x8,
 };
 
 constexpr bool operator&(DTModes x, DTModes y)

diff --git a/src/Particle/SoaDistanceTableAAOMPTarget.h b/src/Particle/SoaDistanceTableAAOMPTarget.h
@@ -487,8 +487,8 @@ struct SoaDistanceTableAAOMPTarget : public DTD_BConds<T, D, SC>, public Distanc
                        const RefVectorWithLeader<ParticleSet>& p_list) const override
   {
     // if the distance table is not updated by mw_move during p-by-p, needs to recompute the whole table
-    // before being used by Hamiltonian.
-    if (!(modes_ & DTModes::NEED_TEMP_DATA_ON_HOST))
+    // before being used by Hamiltonian if requested
+    if (!(modes_ & DTModes::NEED_TEMP_DATA_ON_HOST) && (modes_ & DTModes::NEED_FULL_TABLE_ON_HOST_AFTER_DONEPBYP))
       mw_evaluate(dt_list, p_list);
   }
 

diff --git a/src/Particle/tests/test_distance_table.cpp b/src/Particle/tests/test_distance_table.cpp
@@ -541,7 +541,7 @@ void test_distance_pbc_z_batched_APIs(DynamicCoordinateKind test_kind)
 
   // calculate particle distances
   ions.update();
-  const int ee_tid = electrons.addTable(electrons);
+  const int ee_tid = electrons.addTable(electrons, DTModes::NEED_FULL_TABLE_ON_HOST_AFTER_DONEPBYP);
   // get target particle set's distance table data
   const auto& ee_dtable = electrons.getDistTableAA(ee_tid);
   CHECK(ee_dtable.getName() == "e_e");

diff --git a/src/QMCHamiltonians/CoulombPBCAA.cpp b/src/QMCHamiltonians/CoulombPBCAA.cpp
@@ -31,7 +31,7 @@ CoulombPBCAA::CoulombPBCAA(ParticleSet& ref, bool active, bool computeForces, bo
       ComputeForces(computeForces),
       Ps(ref),
       use_offload_(active && !computeForces && use_offload),
-      d_aa_ID(ref.addTable(ref)),
+      d_aa_ID(ref.addTable(ref, use_offload_ ? DTModes::ALL_OFF : DTModes::NEED_FULL_TABLE_ON_HOST_AFTER_DONEPBYP)),
       evalLR_timer_(*timer_manager.createTimer("CoulombPBCAA::LongRange", timer_level_fine)),
       evalSR_timer_(*timer_manager.createTimer("CoulombPBCAA::ShortRange", timer_level_fine)),
       offload_timer_(*timer_manager.createTimer("CoulombPBCAA::offload", timer_level_fine))

diff --git a/src/QMCHamiltonians/CoulombPotential.h b/src/QMCHamiltonians/CoulombPotential.h
@@ -66,7 +66,7 @@ struct CoulombPotential : public OperatorBase, public ForceBase
       : ForceBase(s, s),
         Pa(s),
         Pb(s),
-        myTableIndex(s.addTable(s)),
+        myTableIndex(s.addTable(s, DTModes::NEED_FULL_TABLE_ON_HOST_AFTER_DONEPBYP)),
         is_AA(true),
         is_active(active),
         ComputeForces(computeForces)

diff --git a/src/QMCHamiltonians/MPC.cpp b/src/QMCHamiltonians/MPC.cpp
@@ -29,7 +29,11 @@ namespace qmcplusplus
 {
 void MPC::resetTargetParticleSet(ParticleSet& ptcl) {}
 
-MPC::MPC(ParticleSet& ptcl, double cutoff) : Ecut(cutoff), d_aa_ID(ptcl.addTable(ptcl)), PtclRef(&ptcl), FirstTime(true)
+MPC::MPC(ParticleSet& ptcl, double cutoff)
+    : Ecut(cutoff),
+      d_aa_ID(ptcl.addTable(ptcl, DTModes::NEED_FULL_TABLE_ON_HOST_AFTER_DONEPBYP)),
+      PtclRef(&ptcl),
+      FirstTime(true)
 {
   initBreakup();
 }

diff --git a/src/QMCHamiltonians/PairCorrEstimator.cpp b/src/QMCHamiltonians/PairCorrEstimator.cpp
@@ -23,7 +23,10 @@
 namespace qmcplusplus
 {
 PairCorrEstimator::PairCorrEstimator(ParticleSet& elns, std::string& sources)
-    : Dmax(10.), Delta(0.5), num_species(2), d_aa_ID_(elns.addTable(elns))
+    : Dmax(10.),
+      Delta(0.5),
+      num_species(2),
+      d_aa_ID_(elns.addTable(elns, DTModes::NEED_FULL_TABLE_ON_HOST_AFTER_DONEPBYP))
 {
   update_mode_.set(COLLECTABLE, 1);
   num_species = elns.groups();