Feature/qgpu support lrf

src/core/common/include/context.h

@@ -22,14 +22,15 @@ class Context {
 
     CommandInfo command_info;
     bool fresh_start = false;
+    int step = 0;            // the current step
     int n_atoms = 0;         // the total number of atoms
     int n_atoms_solute = 0;  // the total number of solute number, in our system [0, n_atoms_solute) are solute, [n_atoms_solute, n_atoms) are water atoms
     double dt = 0.0;
     int n_excluded = 0;
     md_t md;
     topo_t topo;
     NativeOutputConfig native_output;
-    charge_group_config_t charge_group_config;
+    charge_group_config_t charge_group_config;  // todo: when applying LRF, we should change it to HostDeviceBuffer
     std::unique_ptr<HostDeviceBuffer<coord_t>> coords;
     std::unique_ptr<HostDeviceBuffer<vel_t>> velocities;
     std::unique_ptr<HostDeviceBuffer<dvel_t>> dvelocities;

src/core/common/include/energy.h

@@ -7,21 +7,27 @@ enum EnergySlot : int {  // fixed, non-FEP slots
     E_BOND_P_ANGLE,
     E_BOND_P_TOR,
     E_BOND_P_IMP,
+
     E_BOND_W_BOND,
     E_BOND_W_ANGLE,
     E_BOND_W_TOR,
     E_BOND_W_IMP,
+
     E_NB_PP_COUL,
     E_NB_PP_VDW,
     E_NB_PW_COUL,
     E_NB_PW_VDW,
     E_NB_WW_COUL,
     E_NB_WW_VDW,
+
+    E_LRF,
+
     E_RESTR_RADX,
     E_RESTR_POLX,
     E_RESTR_FIX,
     E_RESTR_SHELL,
     E_RESTR_PRES,
+
     ENERGY_FIXED_COUNT
 };
 
@@ -46,6 +52,7 @@ struct EnergyData {
     E_bonded_t bond_p, bond_w;
     E_nonbonded_t nb_pp, nb_pw, nb_ww;
     E_restraint_t restraint;  // Uradx, Upolx, Ufix, Ushell, Upres, Urestr
+    double lrf = 0;
 
     // per-state (FEP)
     std::vector<E_bonded_t> eq_bond;

src/core/common/include/geometry.h

@@ -1,9 +1,8 @@
 #pragma once
 #include <math.h>
 
-#include "md_types.h"
 #include "cuda_runtime_utility.h"
-
+#include "md_types.h"
 
 template <class T>
 HD inline Real3<T> operator-(const Real3<T>& a, const Real3<T>& b) {
@@ -58,3 +57,39 @@ HD inline Real3<To> real3_cast(const Real3<From>& a) {
         static_cast<To>(a.y),
         static_cast<To>(a.z)};
 }
+
+inline uint32_t expand_morton_bits(uint32_t value) {
+    value &= 0x000003ffu;
+    value = (value | (value << 16)) & 0x030000ffu;
+
+    value = (value | (value << 8)) & 0x0300f00fu;
+
+    value = (value | (value << 4)) & 0x030c30c3u;
+
+    value = (value | (value << 2)) & 0x09249249u;
+
+    return value;
+}
+
+inline uint32_t morton_code(uint32_t x, uint32_t y, uint32_t z) {
+    return expand_morton_bits(x) | (expand_morton_bits(y) << 1) | (expand_morton_bits(z) << 2);
+}
+inline uint32_t quantize_morton_coordinate(double value, double minimum, double maximum) {
+    const double extent = maximum - minimum;
+
+    if (!(extent > 0.0)) {
+        return 0;
+    }
+
+    double normalized = (value - minimum) / extent;
+    normalized = std::max(0.0, std::min(1.0, normalized));
+
+    return static_cast<uint32_t>(normalized * 1023.0 + 0.5);
+}
+
+inline uint32_t get_morton_code(const coord_t& p, double x_min, double x_max, double y_min, double y_max, double z_min, double z_max) {
+    auto x = quantize_morton_coordinate(p.x, x_min, x_max);
+    auto y = quantize_morton_coordinate(p.y, y_min, y_max);
+    auto z = quantize_morton_coordinate(p.z, z_min, z_max);
+    return morton_code(x, y, z);
+}

src/core/common/include/md_types.h

@@ -33,6 +33,10 @@ struct md_t {
     double solvent_solvent;
     double solute_solvent;
     double q_atom;
+    double lrf_cutoff;
+
+
+
     // [sphere]
     double shell_radius;  // Note: this is for the pshell
     double shell_force;   // Note: this is for the pshell

src/core/common/include/nonbonded_force.h

@@ -29,7 +29,7 @@ HD inline int nb_coul_slot(uint8_t t1, uint8_t t2, int s1, int s2, int n_states)
     return EnergyBuffer::eq_index(ENERGY_FIXED_COUNT, state, EQ_NB_QP_COUL);  // else Q-P
 }
 
-// The vdw slot is always coul+1 
+// The vdw slot is always coul+1
 HD inline int nb_vdw_slot(uint8_t t1, uint8_t t2, int s1, int s2, int n_states) {
     return nb_coul_slot(t1, t2, s1, s2, n_states) + 1;
 }
@@ -96,16 +96,31 @@ HD inline real_t2 combine_vdw(int vdw_rule, real_t aii_i, real_t bii_i, real_t a
 
 struct NonbondedData {
     int n_total = 0;
-    std::unique_ptr<HostDeviceBuffer<int>> atom_idx;         // global atom index
+    std::unique_ptr<HostDeviceBuffer<int>> atom_idx;       // global atom index
+    std::unique_ptr<HostDeviceBuffer<int>> atom_to_group;  // global atom index
+
     std::unique_ptr<HostDeviceBuffer<uint8_t>> category;     // Atom Category
     std::unique_ptr<HostDeviceBuffer<int>> q_state;          // segment idx; -1 for P/W
     std::unique_ptr<HostDeviceBuffer<real_t>> atom_lambdas;  // lambdas[state]; 1.0 for P/W
+    std::unique_ptr<HostDeviceBuffer<int>> group_indices;    // group idx;
+    std::unique_ptr<HostDeviceBuffer<int>> group_start_idx;  // In the atom_idx, the first index of the atom that belongs to the group
+    std::unique_ptr<HostDeviceBuffer<int>> group_sizes;       // The group size
+
     std::unique_ptr<HostDeviceBuffer<real_t>> atom_charge;
     std::unique_ptr<HostDeviceBuffer<vdw_atom_param_t>> atom_vdw;
 
     bool enabled() const { return n_total > 0; }
 };
 
+struct LrfCoefficients {
+    coord_t center{};
+
+    double phi0 = 0;
+    double phi1[3]{};
+    double phi2[9]{};
+    double phi3[27]{};
+};
+
 class NonbondedForce {
    public:
     virtual ~NonbondedForce() = default;
@@ -127,4 +142,5 @@ class NonbondedForce {
     void build_combinded_list(Context& ctx);  // atom_idx, category, q_state, atom_lambdas
     void build_charge_table(Context& ctx);    // charge_pair_products + charge_types + counts
     void build_catype_table(Context& ctx);    // catype_pair_params + catype_types + counts
+    void build_atom_to_group(Context& ctx);   // atom_to_group
 };

src/core/common/src/energy.cpp

@@ -40,6 +40,7 @@ void EnergyBuffer::unpack() {
     d.nb_pp = {E(E_NB_PP_COUL), E(E_NB_PP_VDW)};
     d.nb_pw = {E(E_NB_PW_COUL), E(E_NB_PW_VDW)};
     d.nb_ww = {E(E_NB_WW_COUL), E(E_NB_WW_VDW)};
+    d.lrf = E(E_LRF);
 
     // fixed restraint components (raw). Upres here is only the direct kernel
     // contribution; the lambda-weighted per-state term is added in combine.
@@ -95,6 +96,6 @@ void EnergyBuffer::combine(const double* lambdas) {
              d.nb_pp.Ucoul + d.nb_pp.Uvdw + d.nb_pw.Ucoul + d.nb_pw.Uvdw +
              d.nb_ww.Ucoul + d.nb_ww.Uvdw +
              d.bond_q.Ubond + d.bond_q.Uangle + d.bond_q.Utor + d.bond_q.Uimp +
-             d.nb_qx.Ucoul + d.nb_qx.Uvdw + d.restraint.Urestr;
+             d.nb_qx.Ucoul + d.nb_qx.Uvdw + d.restraint.Urestr + d.lrf;
     d.Utot = d.Upot + d.Ukin;
 }

src/core/common/src/handler.cpp

@@ -83,6 +83,7 @@ void Handler::stop_cm_translation() {
 }
 
 void Handler::calc_final_potential(int iteration) {
+    ctx.step = iteration;
     reset_energies();
     calc_nonbonded_forces();
     calc_internal_forces(iteration);
@@ -96,6 +97,7 @@ void Handler::run() {
     auto t0 = std::chrono::steady_clock::now();
 
     for (int i = 0; i < num_iterations; i++) {
+        ctx.step = i;
         run_iteration(i);
     }
     auto t1 = std::chrono::steady_clock::now();