Merge pull request #3842 from ye-luo/offload-dist

Offload short range in CoulombPBCAA
QMCPACK · Feb 18, 2022 · 4b27a06 · 4b27a06
2 parents 4911915 + 2acbed3
commit 4b27a06
Show file tree

Hide file tree

Showing 14 changed files with 500 additions and 52 deletions.
diff --git a/docs/hamiltonianobservable.rst b/docs/hamiltonianobservable.rst
@@ -204,6 +204,8 @@ attributes:
   +-------------------------+--------------+----------------------+------------------------+---------------------------------+
   | ``physical``:math:`^o`  | boolean      | yes/no               | yes                    | Hamiltonian(yes)/Observable(no) |
   +-------------------------+--------------+----------------------+------------------------+---------------------------------+
+  | ``gpu``                 | boolean      | yes/no               | depend                 | Offload computation to GPU      |
+  +-------------------------+--------------+----------------------+------------------------+---------------------------------+
   | ``forces``              | boolean      | yes/no               | no                     | *Deprecated*                    |
   +-------------------------+--------------+----------------------+------------------------+---------------------------------+
 
@@ -233,6 +235,8 @@ Additional information:
    output data will appear in ``scalar.dat`` in a column headed by
    ``name``.
 
+-  **gpu**: When not specified, use the ``gpu`` attribute of ``particleset``.
+
 .. code-block::
   :caption: QMCPXML element for Coulomb interaction between electrons.
   :name: Listing 16

diff --git a/src/Particle/DistanceTable.h b/src/Particle/DistanceTable.h
@@ -289,12 +289,22 @@ class DistanceTableAA : public DistanceTable
    */
   const DisplRow& getOldDispls() const { return old_dr_; }
 
+  virtual size_t get_num_particls_stored() const { return 0; }
+
   /// return multi walker temporary pair distance table data pointer
   virtual const RealType* getMultiWalkerTempDataPtr() const
   {
     throw std::runtime_error(name_ + " multi walker data pointer for temp not supported");
     return nullptr;
   }
+
+  virtual const RealType* mw_evaluate_range(const RefVectorWithLeader<DistanceTable>& dt_list,
+                                            const RefVectorWithLeader<ParticleSet>& p_list,
+                                            size_t range_begin,
+                                            size_t range_end) const
+  {
+    return nullptr;
+  }
 };
 
 /** AB type of DistanceTable containing storage */

diff --git a/src/Particle/Lattice/ParticleBConds3DSoa.h b/src/Particle/Lattice/ParticleBConds3DSoa.h
@@ -82,6 +82,11 @@ struct DTD_BConds<T, 3, SUPERCELL_OPEN + SOA_OFFSET>
     dz[iat]     = pz[iat] - z0;
     temp_r[iat] = std::sqrt(dx[iat] * dx[iat] + dy[iat] * dy[iat] + dz[iat] * dz[iat]);
   }
+
+  T computeDist(T dx, T dy, T dz) const
+  {
+    return std::sqrt(dx * dx + dy * dy + dz * dz);
+  }
 };
 
 /** specialization for a periodic 3D, orthorombic cell
@@ -162,6 +167,17 @@ struct DTD_BConds<T, 3, PPPO + SOA_OFFSET>
     dz[iat]     = L2 * (z - round(z));
     temp_r[iat] = std::sqrt(dx[iat] * dx[iat] + dy[iat] * dy[iat] + dz[iat] * dz[iat]);
   }
+
+  T computeDist(T dx, T dy, T dz) const
+  {
+    T x = dx * Linv0;
+    T y = dy * Linv1;
+    T z = dz * Linv2;
+    dx = L0 * (x - round(x));
+    dy = L1 * (y - round(y));
+    dz = L2 * (z - round(z));
+    return std::sqrt(dx * dx + dy * dy + dz * dz);
+  }
 };
 
 /** specialization for a periodic 3D general cell with wigner-seitz==simulation cell
@@ -281,6 +297,25 @@ struct DTD_BConds<T, 3, PPPS + SOA_OFFSET>
 
     temp_r[iat] = std::sqrt(dx[iat] * dx[iat] + dy[iat] * dy[iat] + dz[iat] * dz[iat]);
   }
+
+  T computeDist(T dx, T dy, T dz) const
+  {
+    T ar_0 = dx * g00 + dy * g10 + dz * g20;
+    T ar_1 = dx * g01 + dy * g11 + dz * g21;
+    T ar_2 = dx * g02 + dy * g12 + dz * g22;
+
+    //put them in the box
+    ar_0 -= round(ar_0);
+    ar_1 -= round(ar_1);
+    ar_2 -= round(ar_2);
+
+    //unit2cart
+    dx = ar_0 * r00 + ar_1 * r10 + ar_2 * r20;
+    dy = ar_0 * r01 + ar_1 * r11 + ar_2 * r21;
+    dz = ar_0 * r02 + ar_1 * r12 + ar_2 * r22;
+
+    return std::sqrt(dx * dx + dy * dy + dz * dz);
+  }
 };
 
 
@@ -470,6 +505,34 @@ struct DTD_BConds<T, 3, PPPG + SOA_OFFSET>
     dy[iat]     = flip * (dely + celly[ic]);
     dz[iat]     = flip * (delz + cellz[ic]);
   }
+
+  T computeDist(T dx, T dy, T dz) const
+  {
+    const auto& cellx = corners[0];
+    const auto& celly = corners[1];
+    const auto& cellz = corners[2];
+
+    const T ar_0 = -std::floor(dx * g00 + dy * g10 + dz * g20);
+    const T ar_1 = -std::floor(dx * g01 + dy * g11 + dz * g21);
+    const T ar_2 = -std::floor(dx * g02 + dy * g12 + dz * g22);
+
+    const T delx = dx + ar_0 * r00 + ar_1 * r10 + ar_2 * r20;
+    const T dely = dy + ar_0 * r01 + ar_1 * r11 + ar_2 * r21;
+    const T delz = dz + ar_0 * r02 + ar_1 * r12 + ar_2 * r22;
+
+    T rmin = delx * delx + dely * dely + delz * delz;
+#pragma unroll(7)
+    for (int c = 1; c < 8; ++c)
+    {
+      const T x  = delx + cellx[c];
+      const T y  = dely + celly[c];
+      const T z  = delz + cellz[c];
+      const T r2 = x * x + y * y + z * z;
+      rmin       = (r2 < rmin) ? r2 : rmin;
+    }
+
+    return std::sqrt(rmin);
+  }
 };
 
 
@@ -623,6 +686,30 @@ struct DTD_BConds<T, 3, PPNG + SOA_OFFSET>
     dy[iat]     = flip * (dely + celly[ic]);
     dz[iat]     = delz;
   }
+
+  T computeDist(T dx, T dy, T dz) const
+  {
+    const auto& cellx = corners[0];
+    const auto& celly = corners[1];
+
+    const T ar_0 = -std::floor(dx * g00 + dy * g10);
+    const T ar_1 = -std::floor(dx * g01 + dy * g11);
+
+    const T delx = dx + ar_0 * r00 + ar_1 * r10;
+    const T dely = dy + ar_0 * r01 + ar_1 * r11;
+
+    T rmin = delx * delx + dely * dely;
+#pragma unroll(3)
+    for (int c = 1; c < 4; ++c)
+    {
+      const T x  = delx + cellx[c];
+      const T y  = dely + celly[c];
+      const T r2 = x * x + y * y;
+      rmin       = (r2 < rmin) ? r2 : rmin;
+    }
+
+    return std::sqrt(rmin + dz * dz);
+  }
 };
 
 /** specialization for a slab, orthorombic cell
@@ -695,6 +782,15 @@ struct DTD_BConds<T, 3, PPNO + SOA_OFFSET>
     dz[iat]     = pz[iat] - z0;
     temp_r[iat] = std::sqrt(dx[iat] * dx[iat] + dy[iat] * dy[iat] + dz[iat] * dz[iat]);
   }
+
+  T computeDist(T dx, T dy, T dz) const
+  {
+    T x = dx * Linv0;
+    T y = dy * Linv1;
+    dx = L0 * (x - round(x));
+    dy = L1 * (y - round(y));
+    return std::sqrt(dx * dx + dy * dy + dz * dz);
+  }
 };
 
 /** specialization for a slab, general cell
@@ -796,6 +892,22 @@ struct DTD_BConds<T, 3, PPNS + SOA_OFFSET>
 
     temp_r[iat] = std::sqrt(dx[iat] * dx[iat] + dy[iat] * dy[iat] + dz[iat] * dz[iat]);
   }
+
+  T computeDist(T dx, T dy, T dz) const
+  {
+    T ar_0 = dx * g00 + dy * g10;
+    T ar_1 = dx * g01 + dy * g11;
+
+    //put them in the box
+    ar_0 -= round(ar_0);
+    ar_1 -= round(ar_1);
+
+    //unit2cart
+    dx = ar_0 * r00 + ar_1 * r10;
+    dy = ar_0 * r01 + ar_1 * r11;
+
+    return std::sqrt(dx * dx + dy * dy + dz * dz);
+  }
 };
 
 
@@ -866,6 +978,13 @@ struct DTD_BConds<T, 3, SUPERCELL_WIRE + SOA_OFFSET>
     dz[iat]     = pz[iat] - z0;
     temp_r[iat] = std::sqrt(dx[iat] * dx[iat] + dy[iat] * dy[iat] + dz[iat] * dz[iat]);
   }
+
+  T computeDist(T dx, T dy, T dz) const
+  {
+    T x         = dx * Linv0;
+    dx     = L0 * (x - round(x));
+    return std::sqrt(dx * dx + dy * dy + dz * dz);
+  }
 };
 
 /** specialization for a periodic 3D general cell
@@ -942,6 +1061,11 @@ struct DTD_BConds<T, 3, PPPX + SOA_OFFSET>
   {
     //APP_ABORT("DTD_BConds<T, 3, PPPX + SOA_OFFSET>::computeDistancesOffload not implemented");
   }
+
+  T computeDist(T dx, T dy, T dz) const
+  {
+    //APP_ABORT("DTD_BConds<T, 3, PPPX + SOA_OFFSET>::computeDist not implemented");
+  }
 };
 
 /** specialization for a slab, general cell
@@ -1005,6 +1129,11 @@ struct DTD_BConds<T, 3, PPNX + SOA_OFFSET>
   {
     //APP_ABORT("DTD_BConds<T, 3, PPNX + SOA_OFFSET>::computeDistancesOffload not implemented");
   }
+
+  T computeDist(T dx, T dy, T dz) const
+  {
+    //APP_ABORT("DTD_BConds<T, 3, PPNX + SOA_OFFSET>::computeDist not implemented");
+  }
 };
 
 } // namespace qmcplusplus