diff --git a/src/core/common/include/context.h b/src/core/common/include/context.h index 37ad57b3..f44a8685 100644 --- a/src/core/common/include/context.h +++ b/src/core/common/include/context.h @@ -22,6 +22,7 @@ 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; @@ -29,7 +30,7 @@ class Context { 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> coords; std::unique_ptr> velocities; std::unique_ptr> dvelocities; diff --git a/src/core/common/include/energy.h b/src/core/common/include/energy.h index 677a13a2..2ad47079 100644 --- a/src/core/common/include/energy.h +++ b/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 eq_bond; diff --git a/src/core/common/include/geometry.h b/src/core/common/include/geometry.h index 3c4d0a5c..981a3c62 100644 --- a/src/core/common/include/geometry.h +++ b/src/core/common/include/geometry.h @@ -1,9 +1,8 @@ #pragma once #include -#include "md_types.h" #include "cuda_runtime_utility.h" - +#include "md_types.h" template HD inline Real3 operator-(const Real3& a, const Real3& b) { @@ -58,3 +57,39 @@ HD inline Real3 real3_cast(const Real3& a) { static_cast(a.y), static_cast(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(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); +} \ No newline at end of file diff --git a/src/core/common/include/md_types.h b/src/core/common/include/md_types.h index bc7cad52..8f225f50 100644 --- a/src/core/common/include/md_types.h +++ b/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 diff --git a/src/core/common/include/nonbonded_force.h b/src/core/common/include/nonbonded_force.h index f4f284f7..694f8474 100644 --- a/src/core/common/include/nonbonded_force.h +++ b/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> atom_idx; // global atom index + std::unique_ptr> atom_idx; // global atom index + std::unique_ptr> atom_to_group; // global atom index + std::unique_ptr> category; // Atom Category std::unique_ptr> q_state; // segment idx; -1 for P/W std::unique_ptr> atom_lambdas; // lambdas[state]; 1.0 for P/W + std::unique_ptr> group_indices; // group idx; + std::unique_ptr> group_start_idx; // In the atom_idx, the first index of the atom that belongs to the group + std::unique_ptr> group_sizes; // The group size + std::unique_ptr> atom_charge; std::unique_ptr> 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 }; diff --git a/src/core/common/src/energy.cpp b/src/core/common/src/energy.cpp index dfe3d9ec..cc33d7f2 100644 --- a/src/core/common/src/energy.cpp +++ b/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; } \ No newline at end of file diff --git a/src/core/common/src/handler.cpp b/src/core/common/src/handler.cpp index 10119846..693764ef 100644 --- a/src/core/common/src/handler.cpp +++ b/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(); diff --git a/src/core/common/src/inp_parser.cpp b/src/core/common/src/inp_parser.cpp index 6a135000..cf6edeed 100644 --- a/src/core/common/src/inp_parser.cpp +++ b/src/core/common/src/inp_parser.cpp @@ -320,6 +320,8 @@ struct InpParser::TopData { std::string vdw_rule = "1"; std::vector solucenter; std::vector solvcenter; + + std::string use_switch_atom = "1"; }; struct InpParser::FepData { @@ -422,42 +424,61 @@ void InpParser::ensure_topology() { block = 1; continue; } - if (line.find("No. of integer atom codes") != std::string::npos) { block = 2; continue; } + if (line.find("No. of integer atom codes") != std::string::npos) { + block = 2; + continue; + } if (line.find("No. of bonds") != std::string::npos) { std::vector f = split_ws(line); if (f.size() > 1) top_->nbonds_solute = f[1]; block = 3; continue; } - if (line.find("No. of bond codes") != std::string::npos) { block = 4; continue; } + if (line.find("No. of bond codes") != std::string::npos) { + block = 4; + continue; + } if (line.find("No. of angles") != std::string::npos) { std::vector f = split_ws(line); if (f.size() > 1) top_->nangles_solute = f[1]; block = 5; continue; } - if (line.find("No. of angle codes") != std::string::npos) { block = 6; continue; } + if (line.find("No. of angle codes") != std::string::npos) { + block = 6; + continue; + } if (line.find("No. of torsions") != std::string::npos) { std::vector f = split_ws(line); if (f.size() > 1) top_->ntorsions_solute = f[1]; block = 7; continue; } - if (line.find("No. of torsion codes") != std::string::npos) { block = 8; continue; } + if (line.find("No. of torsion codes") != std::string::npos) { + block = 8; + continue; + } if (line.find("No. of impropers") != std::string::npos) { std::vector f = split_ws(line); if (f.size() > 1) top_->nimpropers_solute = f[1]; block = 9; continue; } - if (line.find("No. of improper codes") != std::string::npos) { block = 10; continue; } - if (line.find("No. of atomic charges") != std::string::npos) { block = 11; continue; } + if (line.find("No. of improper codes") != std::string::npos) { + block = 10; + continue; + } + if (line.find("No. of atomic charges") != std::string::npos) { + block = 11; + continue; + } if (line.find("No. of charge groups") != std::string::npos) { std::vector f = split_ws(line); int total = f.empty() ? 0 : parse_int(f[0]); int solute = f.size() > 1 ? parse_int(f[1]) : 0; top_->solute_cgps = std::to_string(solute); top_->solvent_cgps = std::to_string(total - solute); + top_->use_switch_atom = f.size() > 2 ? f[2] : "1"; block = 12; charge_group_switch = 1; continue; @@ -468,24 +489,77 @@ void InpParser::ensure_topology() { block = 13; continue; } - if (line.find("Electrostatic 1-4 scaling factor") != std::string::npos) { block = 14; } - if (line.find("Masses") != std::string::npos) { block = 15; continue; } - if (line.find("sqrt (Aii) normal") != std::string::npos || line.find("R* normal:") != std::string::npos) { block = 16; continue; } - if (line.find("sqrt (Bii) normal") != std::string::npos || line.find("epsilon normal:") != std::string::npos) { block = 17; continue; } - if (line.find("sqrt (Aii) polar") != std::string::npos || line.find("R* polar:") != std::string::npos) { block = 18; continue; } - if (line.find("sqrt (Bii) polar") != std::string::npos || line.find("epsilon polar:") != std::string::npos) { block = 19; continue; } - if (line.find("sqrt (Aii) 1-4") != std::string::npos || line.find("R* 1-4:") != std::string::npos) { block = 20; continue; } - if (line.find("sqrt (Bii) 1-4") != std::string::npos || line.find("epsilon 1-4:") != std::string::npos) { block = 21; continue; } - if (line.find("No. of type-2 vdW interactions") != std::string::npos) { block = 22; continue; } - if (line.find("No. of 1-4 neighbours") != std::string::npos) { block = 23; continue; } - if (line.find("No. of long 1-4 nbrs") != std::string::npos) { block = 24; continue; } - if (line.find("No. of exclusions") != std::string::npos) { block = 25; continue; } - if (line.find("No. of long exclusions") != std::string::npos) { block = 26; continue; } - if (line.find("No. of residues") != std::string::npos) { block = 27; continue; } - if (line.find("Sequence") != std::string::npos) { block = 28; continue; } - if (line.find("No. of separate molecules") != std::string::npos) { block = 29; continue; } - if (line.find("No. of atom types") != std::string::npos) { block = 30; continue; } - if (line.find("No. of SYBYL atom types") != std::string::npos) { block = 31; continue; } + if (line.find("Electrostatic 1-4 scaling factor") != std::string::npos) { + block = 14; + } + if (line.find("Masses") != std::string::npos) { + block = 15; + continue; + } + if (line.find("sqrt (Aii) normal") != std::string::npos || line.find("R* normal:") != std::string::npos) { + block = 16; + continue; + } + if (line.find("sqrt (Bii) normal") != std::string::npos || line.find("epsilon normal:") != std::string::npos) { + block = 17; + continue; + } + if (line.find("sqrt (Aii) polar") != std::string::npos || line.find("R* polar:") != std::string::npos) { + block = 18; + continue; + } + if (line.find("sqrt (Bii) polar") != std::string::npos || line.find("epsilon polar:") != std::string::npos) { + block = 19; + continue; + } + if (line.find("sqrt (Aii) 1-4") != std::string::npos || line.find("R* 1-4:") != std::string::npos) { + block = 20; + continue; + } + if (line.find("sqrt (Bii) 1-4") != std::string::npos || line.find("epsilon 1-4:") != std::string::npos) { + block = 21; + continue; + } + if (line.find("No. of type-2 vdW interactions") != std::string::npos) { + block = 22; + continue; + } + if (line.find("No. of 1-4 neighbours") != std::string::npos) { + block = 23; + continue; + } + if (line.find("No. of long 1-4 nbrs") != std::string::npos) { + block = 24; + continue; + } + if (line.find("No. of exclusions") != std::string::npos) { + block = 25; + continue; + } + if (line.find("No. of long exclusions") != std::string::npos) { + block = 26; + continue; + } + if (line.find("No. of residues") != std::string::npos) { + block = 27; + continue; + } + if (line.find("Sequence") != std::string::npos) { + block = 28; + continue; + } + if (line.find("No. of separate molecules") != std::string::npos) { + block = 29; + continue; + } + if (line.find("No. of atom types") != std::string::npos) { + block = 30; + continue; + } + if (line.find("No. of SYBYL atom types") != std::string::npos) { + block = 31; + continue; + } if (line.find("solvent type (0=SPC,1=3-atom,2=general)") != std::string::npos) { std::vector f = split_ws(line); if (!f.empty()) top_->solvtype = f[0]; @@ -493,31 +567,46 @@ void InpParser::ensure_topology() { block = 32; continue; } - if (line.find("No. of excluded atoms") != std::string::npos) { block = 33; continue; } + if (line.find("No. of excluded atoms") != std::string::npos) { + block = 33; + continue; + } std::vector f = split_ws(line); switch (block) { - case 1: coord_flat.insert(coord_flat.end(), f.begin(), f.end()); break; + case 1: + coord_flat.insert(coord_flat.end(), f.begin(), f.end()); + break; case 2: for (const std::string& value : f) top_->atypes.push_back({++atype_count, parse_int(value)}); break; - case 3: bond_flat.insert(bond_flat.end(), f.begin(), f.end()); break; + case 3: + bond_flat.insert(bond_flat.end(), f.begin(), f.end()); + break; case 4: if (f.size() >= 3) top_->cbonds[parse_int(f[0])] = std::vector(f.begin() + 1, f.begin() + 3); break; - case 5: angle_flat.insert(angle_flat.end(), f.begin(), f.end()); break; + case 5: + angle_flat.insert(angle_flat.end(), f.begin(), f.end()); + break; case 6: if (f.size() >= 3) top_->cangles[parse_int(f[0])] = std::vector(f.begin() + 1, f.begin() + 3); break; - case 7: torsion_flat.insert(torsion_flat.end(), f.begin(), f.end()); break; + case 7: + torsion_flat.insert(torsion_flat.end(), f.begin(), f.end()); + break; case 8: if (f.size() >= 5) top_->ctorsions[parse_int(f[0])] = std::vector(f.begin() + 1, f.begin() + 5); break; - case 9: improper_flat.insert(improper_flat.end(), f.begin(), f.end()); break; + case 9: + improper_flat.insert(improper_flat.end(), f.begin(), f.end()); + break; case 10: if (f.size() >= 3) top_->cimpropers[parse_int(f[0])] = std::vector(f.begin() + 1, f.begin() + 3); break; - case 11: charges_tmp.insert(charges_tmp.end(), f.begin(), f.end()); break; + case 11: + charges_tmp.insert(charges_tmp.end(), f.begin(), f.end()); + break; case 12: if (charge_group_switch == 1 && f.size() >= 2) { current_group_header = f; @@ -539,26 +628,46 @@ void InpParser::ensure_topology() { top_->coulomb = f[1]; } break; - case 15: masses.insert(masses.end(), f.begin(), f.end()); break; - case 16: aii_normal.insert(aii_normal.end(), f.begin(), f.end()); break; - case 17: bii_normal.insert(bii_normal.end(), f.begin(), f.end()); break; - case 18: aii_polar.insert(aii_polar.end(), f.begin(), f.end()); break; - case 19: bii_polar.insert(bii_polar.end(), f.begin(), f.end()); break; - case 20: aii14.insert(aii14.end(), f.begin(), f.end()); break; - case 21: bii14.insert(bii14.end(), f.begin(), f.end()); break; - case 23: ngbr14_flat += trim(line); break; + case 15: + masses.insert(masses.end(), f.begin(), f.end()); + break; + case 16: + aii_normal.insert(aii_normal.end(), f.begin(), f.end()); + break; + case 17: + bii_normal.insert(bii_normal.end(), f.begin(), f.end()); + break; + case 18: + aii_polar.insert(aii_polar.end(), f.begin(), f.end()); + break; + case 19: + bii_polar.insert(bii_polar.end(), f.begin(), f.end()); + break; + case 20: + aii14.insert(aii14.end(), f.begin(), f.end()); + break; + case 21: + bii14.insert(bii14.end(), f.begin(), f.end()); + break; + case 23: + ngbr14_flat += trim(line); + break; case 24: { auto groups = checked_split_groups(f, 2); top_->ngbr14long.insert(top_->ngbr14long.end(), groups.begin(), groups.end()); break; } - case 25: ngbr23_flat += trim(line); break; + case 25: + ngbr23_flat += trim(line); + break; case 26: { auto groups = checked_split_groups(f, 2); top_->ngbr23long.insert(top_->ngbr23long.end(), groups.begin(), groups.end()); break; } - case 29: top_->molecules.insert(top_->molecules.end(), f.begin(), f.end()); break; + case 29: + top_->molecules.insert(top_->molecules.end(), f.begin(), f.end()); + break; case 32: if (line.find("Exclusion") != std::string::npos && f.size() >= 2) { if (parse_double(f[0]) > 30.0) throw parse_error("Sphere sizes exceeding 30A are currently not supported"); @@ -577,7 +686,8 @@ void InpParser::ensure_topology() { if (!std::isspace(static_cast(ch))) top_->excluded.push_back(ch == 'F' ? "0" : "1"); } break; - default: break; + default: + break; } } @@ -648,28 +758,99 @@ void InpParser::ensure_fep() { while (std::getline(in, raw)) { std::string line = strip_comment(raw); if (line.empty()) continue; - if (line.find("[atoms]") != std::string::npos) { block = 1; continue; } - if (line.find("[FEP]") != std::string::npos) { block = 2; continue; } - if (line.find("[change_charges]") != std::string::npos) { block = 3; continue; } - if (line.find("[atom_types]") != std::string::npos) { block = 4; atype_index = 0; continue; } - if (line.find("[change_atoms]") != std::string::npos) { block = 5; continue; } - if (line.find("[soft_pairs]") != std::string::npos) { block = 6; continue; } - if (line.find("[excluded_pairs]") != std::string::npos) { block = 7; continue; } - if (line.find("[el_scale]") != std::string::npos) { block = 8; continue; } - if (line.find("[softcore]") != std::string::npos) { block = 9; continue; } - if (line.find("[bond_types]") != std::string::npos) { block = 12; fep_->q_cbonds.push_back({"0", "0.0", "0.0"}); continue; } - if (line.find("[change_bonds]") != std::string::npos) { block = 13; continue; } - if (line.find("[angle_types]") != std::string::npos) { block = 14; fep_->q_cangles.push_back({"0", "0.0", "0.0"}); continue; } - if (line.find("[change_types]") != std::string::npos) { block = 15; continue; } - if (line.find("[torsion_types]") != std::string::npos) { block = 16; fep_->q_ctorsions.push_back({"0", "0.0", "0.0", "0.0"}); continue; } - if (line.find("[change_torsions]") != std::string::npos) { block = 17; continue; } - if (line.find("[improper_types]") != std::string::npos) { block = 18; fep_->q_cimpropers.push_back({"0", "0.0", "0.0"}); continue; } - if (line.find("[change_impropers]") != std::string::npos) { block = 19; continue; } - if (line.find("[angle_couplings]") != std::string::npos) { block = 20; continue; } - if (line.find("[torsion_couplings]") != std::string::npos) { block = 21; continue; } - if (line.find("[improper_couplings]") != std::string::npos) { block = 22; continue; } - if (line.find("[shake_constraints]") != std::string::npos) { block = 23; continue; } - if (line.find("[off-diagonals]") != std::string::npos) { block = 24; continue; } + if (line.find("[atoms]") != std::string::npos) { + block = 1; + continue; + } + if (line.find("[FEP]") != std::string::npos) { + block = 2; + continue; + } + if (line.find("[change_charges]") != std::string::npos) { + block = 3; + continue; + } + if (line.find("[atom_types]") != std::string::npos) { + block = 4; + atype_index = 0; + continue; + } + if (line.find("[change_atoms]") != std::string::npos) { + block = 5; + continue; + } + if (line.find("[soft_pairs]") != std::string::npos) { + block = 6; + continue; + } + if (line.find("[excluded_pairs]") != std::string::npos) { + block = 7; + continue; + } + if (line.find("[el_scale]") != std::string::npos) { + block = 8; + continue; + } + if (line.find("[softcore]") != std::string::npos) { + block = 9; + continue; + } + if (line.find("[bond_types]") != std::string::npos) { + block = 12; + fep_->q_cbonds.push_back({"0", "0.0", "0.0"}); + continue; + } + if (line.find("[change_bonds]") != std::string::npos) { + block = 13; + continue; + } + if (line.find("[angle_types]") != std::string::npos) { + block = 14; + fep_->q_cangles.push_back({"0", "0.0", "0.0"}); + continue; + } + if (line.find("[change_types]") != std::string::npos) { + block = 15; + continue; + } + if (line.find("[torsion_types]") != std::string::npos) { + block = 16; + fep_->q_ctorsions.push_back({"0", "0.0", "0.0", "0.0"}); + continue; + } + if (line.find("[change_torsions]") != std::string::npos) { + block = 17; + continue; + } + if (line.find("[improper_types]") != std::string::npos) { + block = 18; + fep_->q_cimpropers.push_back({"0", "0.0", "0.0"}); + continue; + } + if (line.find("[change_impropers]") != std::string::npos) { + block = 19; + continue; + } + if (line.find("[angle_couplings]") != std::string::npos) { + block = 20; + continue; + } + if (line.find("[torsion_couplings]") != std::string::npos) { + block = 21; + continue; + } + if (line.find("[improper_couplings]") != std::string::npos) { + block = 22; + continue; + } + if (line.find("[shake_constraints]") != std::string::npos) { + block = 23; + continue; + } + if (line.find("[off-diagonals]") != std::string::npos) { + block = 24; + continue; + } std::vector f = split_ws(line); if (f.empty()) continue; @@ -697,34 +878,53 @@ void InpParser::ensure_fep() { case 7: for (int s = 0; s < fep_->states && static_cast(s + 1) < f.size(); s++) fep_->q_exclpairs[s].push_back(f[s + 1]); break; - case 8: fep_->q_elscales.push_back(f); break; + case 8: + fep_->q_elscales.push_back(f); + break; case 9: for (int s = 0; s < fep_->states && static_cast(s + 1) < f.size(); s++) fep_->q_softcores[s].push_back(f[s + 1]); break; - case 12: fep_->q_cbonds.push_back(f.size() > 1 ? std::vector(f.begin() + 1, f.end()) : f); break; + case 12: + fep_->q_cbonds.push_back(f.size() > 1 ? std::vector(f.begin() + 1, f.end()) : f); + break; case 13: for (int s = 0; s < fep_->states && static_cast(s + 1) < f.size(); s++) fep_->q_bonds[s].push_back(f[s + 1]); break; - case 14: fep_->q_cangles.push_back(f.size() > 1 ? std::vector(f.begin() + 1, f.end()) : f); break; + case 14: + fep_->q_cangles.push_back(f.size() > 1 ? std::vector(f.begin() + 1, f.end()) : f); + break; case 15: for (int s = 0; s < fep_->states && static_cast(s + 1) < f.size(); s++) fep_->q_angles[s].push_back(f[s + 1]); break; - case 16: fep_->q_ctorsions.push_back(f.size() > 1 ? std::vector(f.begin() + 1, f.end()) : f); break; + case 16: + fep_->q_ctorsions.push_back(f.size() > 1 ? std::vector(f.begin() + 1, f.end()) : f); + break; case 17: for (int s = 0; s < fep_->states && static_cast(s + 1) < f.size(); s++) fep_->q_torsions[s].push_back(f[s + 1]); break; - case 18: fep_->q_cimpropers.push_back(f.size() > 1 ? std::vector(f.begin() + 1, f.end()) : f); break; + case 18: + fep_->q_cimpropers.push_back(f.size() > 1 ? std::vector(f.begin() + 1, f.end()) : f); + break; case 19: for (int s = 0; s < fep_->states && static_cast(s + 1) < f.size(); s++) fep_->q_impropers[s].push_back(f[s + 1]); break; - case 20: fep_->q_angcouples.push_back(f); break; - case 21: fep_->q_torcouples.push_back(f); break; - case 22: fep_->q_imprcouples.push_back(f); break; + case 20: + fep_->q_angcouples.push_back(f); + break; + case 21: + fep_->q_torcouples.push_back(f); + break; + case 22: + fep_->q_imprcouples.push_back(f); + break; case 23: for (int s = 0; s < fep_->states && static_cast(s + 1) < f.size(); s++) fep_->q_shakes[s].push_back(f[s + 1]); break; - case 24: fep_->q_offdiags.push_back(f); break; - default: break; + case 24: + fep_->q_offdiags.push_back(f); + break; + default: + break; } } @@ -875,6 +1075,8 @@ void InpParser::parse_md() { md.solute_solute = parse_double(value_or(cut, "solute-solute", value_or(cut, "solute_solute", "10"))); md.solvent_solvent = parse_double(value_or(cut, "solvent-solvent", value_or(cut, "solvent_solvent", "10"))); md.solute_solvent = parse_double(value_or(cut, "solute-solvent", value_or(cut, "solute_solvent", "10"))); + md.lrf_cutoff = parse_double(value_or(cut, "lrf", "99")); + md.q_atom = parse_double(value_or(cut, "q-atom", value_or(cut, "q_atom", "99"))); md.shell_radius = parse_double(value_or(sphere, "shell-radius", value_or(sphere, "shell_radius", "0"))); md.shell_force = parse_double(value_or(sphere, "shell-force", value_or(sphere, "shell_force", "10.0"))); @@ -1112,7 +1314,7 @@ void InpParser::parse_charge_groups() { charge_group_config_t config; config.n_cgrps_solute = parse_int(top_->solute_cgps); config.n_cgrps_solvent = parse_int(top_->solvent_cgps); - config.iuse_switch_atom = 0; + config.iuse_switch_atom = parse_int(top_->use_switch_atom); config.charge_groups.resize(top_->charge_group_headers.size()); for (size_t i = 0; i < top_->charge_group_headers.size(); i++) { config.charge_groups[i].iswitch = row_int(top_->charge_group_headers[i], 1); diff --git a/src/core/common/src/nonbonded_force.cpp b/src/core/common/src/nonbonded_force.cpp index 9cd242bc..6e0ab2f2 100644 --- a/src/core/common/src/nonbonded_force.cpp +++ b/src/core/common/src/nonbonded_force.cpp @@ -5,20 +5,147 @@ #include #include "constants.h" +#include "geometry.h" #include "vdw_rules.h" +namespace { + +struct SpatialGroupEntry { + uint32_t morton; + int group; +}; + +coord_t group_spatial_position(const Context& ctx, int group_index, uint8_t group_category) { + const auto& config = ctx.charge_group_config; + + const auto& group = config.charge_groups[group_index]; + + const coord_t* coords = ctx.coords->cpu_data_p; + + constexpr uint8_t W = static_cast(AtomCategory::W); + + /* + * Switch mode always sorts using the switch atom. + * + * Water also uses its switch atom when + * iuse_switch_atom == 0, matching the CPU + * W-W and P-W distance definitions. + */ + if (config.iuse_switch_atom == 1 || group_category == W) { + const int switch_atom = group.iswitch - 1; + + return coords[switch_atom]; + } + + /* + * In all-atom mode, use the centroid for P and Q + * groups. This is only a spatial ordering key; it + * does not change group-pair classification. + */ + coord_t center{0.0, 0.0, 0.0}; + + for (int atom_1based : group.atoms) { + const int atom = atom_1based - 1; + + center.x += coords[atom].x; + center.y += coords[atom].y; + center.z += coords[atom].z; + } + + const double inverse_size = 1.0 / static_cast(group.atoms.size()); + + center.x *= inverse_size; + center.y *= inverse_size; + center.z *= inverse_size; + + return center; +} + +void spatially_sort_group_indices(const Context& ctx, uint8_t group_category, std::vector& group_indices) { + if (group_indices.size() < 2) { + return; + } + + std::vector positions; + positions.reserve(group_indices.size()); + + for (int group : group_indices) { + positions.push_back(group_spatial_position(ctx, group, group_category)); + } + + double minimum_x = positions[0].x; + double minimum_y = positions[0].y; + double minimum_z = positions[0].z; + + double maximum_x = positions[0].x; + double maximum_y = positions[0].y; + double maximum_z = positions[0].z; + + for (const coord_t& position : positions) { + minimum_x = std::min(minimum_x, position.x); + minimum_y = std::min(minimum_y, position.y); + minimum_z = std::min(minimum_z, position.z); + + maximum_x = std::max(maximum_x, position.x); + maximum_y = std::max(maximum_y, position.y); + maximum_z = std::max(maximum_z, position.z); + } + + std::vector entries; + entries.reserve(group_indices.size()); + + for (size_t i = 0; i < group_indices.size(); ++i) { + const coord_t& position = positions[i]; + entries.push_back({get_morton_code(position, minimum_x, maximum_x, minimum_y, maximum_y, minimum_z, maximum_z), group_indices[i]}); + } + + /* + * stable_sort keeps the original topology ordering + * for groups with identical Morton keys. + */ + std::stable_sort(entries.begin(), entries.end(), [](const SpatialGroupEntry& lhs, const SpatialGroupEntry& rhs) { + return lhs.morton < rhs.morton; + }); + + for (size_t i = 0; i < entries.size(); ++i) { + group_indices[i] = entries[i].group; + } +} +} // namespace + void NonbondedForce::init(Context& ctx) { build_combinded_list(ctx); build_charge_table(ctx); build_catype_table(ctx); + build_atom_to_group(ctx); init_backend(ctx); } +void NonbondedForce::build_atom_to_group(Context& ctx) { + const auto& groups = ctx.charge_group_config.charge_groups; + const int n_groups = groups.size(); + + std::vector atom_to_group(ctx.n_atoms, -1); + for (int group = 0; group < n_groups; group++) { + for (int atom_1based : groups[group].atoms) { + const int atom = atom_1based - 1; + + if (atom >= 0 && atom < ctx.n_atoms) { + atom_to_group[atom] = group; + } + } + } + data_.atom_to_group = HostDeviceBuffer::from_vector(atom_to_group, ctx.command_info.requested_gpu); +} + void NonbondedForce::build_combinded_list(Context& ctx) { std::vector atom_idx; std::vector category; std::vector q_state; std::vector atom_lambdas; + std::vector group_indices; + std::vector group_start_idx; + std::vector group_sizes; auto push_dummy = [&](int count) { for (int i = 0; i < count; i++) { @@ -29,48 +156,127 @@ void NonbondedForce::build_combinded_list(Context& ctx) { } }; + std::vector atom_type(ctx.n_atoms); for (int i = 0; i < ctx.n_patoms(); i++) { int idx = ctx.p_atoms[i]; - if (ctx.excluded->cpu_data_p[idx]) continue; - atom_idx.push_back(idx); + atom_type[idx] = static_cast(AtomCategory::P); + } + for (int i = 0; i < ctx.n_qatoms(); i++) { + int idx = ctx.q_atoms[i]; + atom_type[idx] = static_cast(AtomCategory::Q); + } + for (int i = ctx.n_atoms_solute; i < ctx.n_atoms; i++) { + atom_type[i] = static_cast(AtomCategory::W); + } + + const auto& groups = ctx.charge_group_config.charge_groups; + int group_size = groups.size(); + + std::vector> category_groups(3); + for (int i = 0; i < group_size; i++) { + int atom = groups[i].iswitch - 1; + if (ctx.excluded->cpu_data_p[atom]) continue; + category_groups[atom_type[atom]].push_back(i); + } + + if (ctx.command_info.requested_gpu) { + constexpr uint8_t P = static_cast(AtomCategory::P); + constexpr uint8_t Q = static_cast(AtomCategory::Q); + constexpr uint8_t W = static_cast(AtomCategory::W); + + spatially_sort_group_indices(ctx, P, category_groups[P]); + spatially_sort_group_indices(ctx, Q, category_groups[Q]); + spatially_sort_group_indices(ctx, W, category_groups[W]); + } + + // P + for (int i = 0; i < category_groups[0].size(); i++) { + int group_idx = category_groups[0][i]; + group_indices.push_back(group_idx); + group_start_idx.push_back(atom_idx.size()); + group_sizes.push_back(groups[group_idx].atoms.size()); + + int switch_atom = groups[group_idx].iswitch - 1; + atom_idx.push_back(switch_atom); category.push_back(static_cast(AtomCategory::P)); q_state.push_back(-1); atom_lambdas.push_back(1.0); + + for (int j = 0; j < groups[group_idx].atoms.size(); j++) { + int atom = groups[group_idx].atoms[j] - 1; + if (atom == switch_atom) continue; + atom_idx.push_back(atom); + category.push_back(static_cast(AtomCategory::P)); + q_state.push_back(-1); + atom_lambdas.push_back(1.0); + } } int sz = atom_idx.size(); push_dummy((32 - (sz % 32)) % 32); + // Q for (int state = 0; state < ctx.n_lambdas(); state++) { - for (int i = 0; i < ctx.n_qatoms(); i++) { - int idx = ctx.q_atoms[i]; - if (ctx.excluded->cpu_data_p[idx]) continue; - atom_idx.push_back(idx); + for (int i = 0; i < category_groups[1].size(); i++) { + int group_idx = category_groups[1][i]; + group_indices.push_back(group_idx); + group_start_idx.push_back(atom_idx.size()); + group_sizes.push_back(groups[group_idx].atoms.size()); + + int switch_atom = groups[group_idx].iswitch - 1; + atom_idx.push_back(switch_atom); category.push_back(static_cast(AtomCategory::Q)); q_state.push_back(state); atom_lambdas.push_back(ctx.lambdas->cpu_data_p[state]); + + for (int j = 0; j < groups[group_idx].atoms.size(); j++) { + int atom = groups[group_idx].atoms[j] - 1; + if (atom == switch_atom) continue; + atom_idx.push_back(atom); + category.push_back(static_cast(AtomCategory::Q)); + q_state.push_back(state); + atom_lambdas.push_back(ctx.lambdas->cpu_data_p[state]); + } } sz = atom_idx.size(); push_dummy((32 - (sz % 32)) % 32); } - for (int i = ctx.n_atoms_solute; i < ctx.n_atoms; i++) { - if (ctx.excluded->cpu_data_p[i]) continue; - atom_idx.push_back(i); + // W + for (int i = 0; i < category_groups[2].size(); i++) { + int group_idx = category_groups[2][i]; + group_indices.push_back(group_idx); + group_start_idx.push_back(atom_idx.size()); + group_sizes.push_back(groups[group_idx].atoms.size()); + + int switch_atom = groups[group_idx].iswitch - 1; + atom_idx.push_back(switch_atom); category.push_back(static_cast(AtomCategory::W)); q_state.push_back(-1); atom_lambdas.push_back(1.0); - } + for (int j = 0; j < groups[group_idx].atoms.size(); j++) { + int atom = groups[group_idx].atoms[j] - 1; + if (atom == switch_atom) continue; + atom_idx.push_back(atom); + category.push_back(static_cast(AtomCategory::W)); + q_state.push_back(-1); + atom_lambdas.push_back(1.0); + } + } sz = atom_idx.size(); push_dummy((32 - (sz % 32)) % 32); sz = atom_idx.size(); data_.n_total = sz; + data_.atom_idx = HostDeviceBuffer::from_vector(atom_idx, ctx.command_info.requested_gpu); data_.category = HostDeviceBuffer::from_vector(category, ctx.command_info.requested_gpu); data_.q_state = HostDeviceBuffer::from_vector(q_state, ctx.command_info.requested_gpu); data_.atom_lambdas = HostDeviceBuffer::from_vector(atom_lambdas, ctx.command_info.requested_gpu); + data_.group_indices = HostDeviceBuffer::from_vector(group_indices, ctx.command_info.requested_gpu); + data_.group_start_idx = HostDeviceBuffer::from_vector(group_start_idx, ctx.command_info.requested_gpu); + data_.group_sizes = HostDeviceBuffer::from_vector(group_sizes, ctx.command_info.requested_gpu); } void NonbondedForce::build_charge_table(Context& ctx) { diff --git a/src/core/common/src/std_output.cpp b/src/core/common/src/std_output.cpp index a40f422f..f3d03629 100644 --- a/src/core/common/src/std_output.cpp +++ b/src/core/common/src/std_output.cpp @@ -37,6 +37,7 @@ void StdOutput::output_energy(Context& ctx, int iteration) { std::printf("pw\t%f\t%f\n", energy.nb_pw.Ucoul, energy.nb_pw.Uvdw); std::printf("ww\t%f\t%f\n", energy.nb_ww.Ucoul, energy.nb_ww.Uvdw); std::printf("qx\t%f\t%f\n", energy.nb_qx.Ucoul, energy.nb_qx.Uvdw); + std::printf("LRF\t%f\n", energy.lrf); std::printf("\n"); std::printf("[restraint]\n"); diff --git a/src/core/cpu/include/cpu_nonbonded_force.h b/src/core/cpu/include/cpu_nonbonded_force.h index 71a2ca66..7cd02d28 100644 --- a/src/core/cpu/include/cpu_nonbonded_force.h +++ b/src/core/cpu/include/cpu_nonbonded_force.h @@ -1,8 +1,26 @@ #pragma once +#include + #include "nonbonded_force.h" class CpuNonbondedForce final : public NonbondedForce { public: void calc(Context& ctx) override; + protected: + void init_backend(Context& ctx) override; + + private: + void calc_all_direct_pairs(Context& ctx); + void init_calculation_groups(Context& ctx); + void init_exact_atom_pairs(Context& ctx); + void calc_direct_pair(Context& ctx, int slot1, int slot2); + void calc_exact_pairs(Context& ctx); + void init_lrf_coefficients(Context& ctx); + void calc_lrf(Context& ctx); + + std::vector> exact_calculation_groups_, lrf_calculation_groups_; + std::vector> exact_atom_pairs_; + std::vector lrf_coefficients_; + std::vector> slots_by_atom_; }; \ No newline at end of file diff --git a/src/core/cpu/src/cpu_nonbonded_force.cpp b/src/core/cpu/src/cpu_nonbonded_force.cpp index 4d962662..898045f2 100644 --- a/src/core/cpu/src/cpu_nonbonded_force.cpp +++ b/src/core/cpu/src/cpu_nonbonded_force.cpp @@ -1,7 +1,10 @@ #include "cpu_nonbonded_force.h" +#include + #include "constants.h" #include "cpu_force_accumulation.h" +#include "geometry.h" namespace { void accumulate_energy(Context& ctx, real_t vel, real_t vvdw, @@ -12,64 +15,445 @@ void accumulate_energy(Context& ctx, real_t vel, real_t vvdw, add_energy(e[coul + 1], vvdw); // vdw slot is adjacent (same invariant as GPU) } +void accumulate_lrf_source(LrfCoefficients& target, const coord_t& source_coord, double source_charge) { + coord_t r = source_coord - target.center; + double r_len2 = norm2(r); + double r_len = std::sqrt(r_len2); + double inv_r_len = 1.0 / r_len; + double inv_r_len2 = 1.0 / r_len2; + double inv_r_len3 = inv_r_len * inv_r_len2; + double inv_r_len5 = inv_r_len3 * inv_r_len2; + double inv_r_len7 = inv_r_len5 * inv_r_len2; + + target.phi0 += source_charge * inv_r_len; // q / r + + double r_array[3] = {r.x, r.y, r.z}; + for (int i = 0; i < 3; i++) { + target.phi1[i] -= source_charge * r_array[i] * inv_r_len3; + } + + for (int i = 0; i < 3; i++) { + for (int j = 0; j < 3; j++) { + double delta = i == j ? inv_r_len3 : 0; + target.phi2[i * 3 + j] += source_charge * ((3.0 * r_array[i] * r_array[j] * inv_r_len5) - delta); + } + } + + for (int a = 0; a < 3; a++) { + for (int b = 0; b < 3; b++) { + for (int c = 0; c < 3; c++) { + int idx = (a * 3 + b) * 3 + c; + + double delta_ab = a == b; + double delta_ac = a == c; + double delta_bc = b == c; + + double v1 = 3.0 * (delta_ab * r_array[c] + delta_ac * r_array[b] + delta_bc * r_array[a]) * inv_r_len5; + double v2 = -15.0 * r_array[a] * r_array[b] * r_array[c] * inv_r_len7; + + target.phi3[idx] += source_charge * (v1 + v2); + } + } + } +} + } // namespace -void CpuNonbondedForce::calc(Context& ctx) { - const auto& atom_idxs = data_.atom_idx->cpu_data_p; - const auto& coords = ctx.coords->cpu_data_p; - auto& dvelocities = ctx.dvelocities->cpu_data_p; +void CpuNonbondedForce::calc_all_direct_pairs(Context& ctx) { int sz = data_.n_total; - for (int i = 0; i < sz; i++) { - const int atom1 = atom_idxs[i]; - if (atom1 == -1) continue; - const auto& atom1_type = data_.category->cpu_data_p[i]; - const int atom1_state = data_.q_state->cpu_data_p[i]; - const real_t atom1_charge = data_.atom_charge->cpu_data_p[i]; - const vdw_atom_param_t& atom1_vdw = data_.atom_vdw->cpu_data_p[i]; for (int j = i + 1; j < sz; j++) { - const int atom2 = atom_idxs[j]; - if (atom2 == -1) continue; - const auto& atom2_type = data_.category->cpu_data_p[j]; - const int atom2_state = data_.q_state->cpu_data_p[j]; - const auto& bond_type = get_bond_type(ctx.n_atoms_solute, ctx.LJ_matrix->cpu_data_p, atom1, atom1_type, atom2, atom2_type); - const real_t atom2_charge = data_.atom_charge->cpu_data_p[j]; - const vdw_atom_param_t atom2_vdw = data_.atom_vdw->cpu_data_p[j]; - - if (bond_type == BondType::Bond23) continue; - if (atom1_type == static_cast(AtomCategory::Q) && atom2_type == static_cast(AtomCategory::Q) && atom1_state != atom2_state) { + calc_direct_pair(ctx, i, j); + } + } +} + +void CpuNonbondedForce::init_calculation_groups(Context& ctx) { + const auto& config = ctx.charge_group_config; + const auto& groups = config.charge_groups; + + const int n_groups = groups.size(); + const int n_solute_groups = config.n_cgrps_solute; + + const coord_t* coords = ctx.coords->cpu_data_p; + const bool* excluded = ctx.excluded->cpu_data_p; + + auto normal_cutoff = [&](int group1, int group2) { + const bool solute1 = group1 < n_solute_groups; + const bool solute2 = group2 < n_solute_groups; + if (solute1 && solute2) { + return ctx.md.solute_solute; + } else if (!solute1 && !solute2) { + return ctx.md.solvent_solvent; + } else { + return ctx.md.solute_solvent; + } + }; + + auto group_distance2 = [&](int group1, int group2) { + if (group1 == group2) return 0.0; + const bool solute1 = group1 < n_solute_groups; + const bool solute2 = group2 < n_solute_groups; + + if (config.iuse_switch_atom == 1) { + const int atom1 = groups[group1].iswitch - 1; + const int atom2 = groups[group2].iswitch - 1; + return norm2(coords[atom1] - coords[atom2]); + } + + if (!solute1 && !solute2) { + const int atom1 = groups[group1].iswitch - 1; + const int atom2 = groups[group2].iswitch - 1; + return norm2(coords[atom1] - coords[atom2]); + } + + if (solute1 != solute2) { + const int solute_group = solute1 ? group1 : group2; + const int water_group = solute1 ? group2 : group1; + + const int water_switch = groups[water_group].iswitch - 1; + + double mi = std::numeric_limits::infinity(); + + for (int atom : groups[solute_group].atoms) { + const int atom_idx = atom - 1; + mi = std::min(mi, norm2(coords[atom_idx] - coords[water_switch])); + } + return mi; + } + + double mi = std::numeric_limits::infinity(); + + for (int atom1 : groups[group1].atoms) { + const int atom1_idx = atom1 - 1; + + for (int atom2 : groups[group2].atoms) { + const int atom2_idx = atom2 - 1; + mi = std::min(mi, norm2(coords[atom1_idx] - coords[atom2_idx])); + } + } + return mi; + }; + + auto group_is_active = [&](int group) { + const int switch_atom = groups[group].iswitch - 1; + return switch_atom >= 0 && switch_atom < ctx.n_atoms && !excluded[switch_atom]; + }; + + const double rcq2 = ctx.md.q_atom * ctx.md.q_atom; + const coord_t& solute_center = ctx.topo.solute_center; + auto inside_rcq = [&](int group) { + const bool solvent_group = group >= n_solute_groups; + + if (solvent_group || config.iuse_switch_atom == 1) { + const int switch_atom = groups[group].iswitch - 1; + return norm2(coords[switch_atom] - solute_center) <= rcq2; + } else { + for (int atom_1based : groups[group].atoms) { + const int atom = atom_1based - 1; + if (norm2(coords[atom] - solute_center) <= rcq2) { + return true; + } + } + return false; + } + }; + + const double lrf_cutoff2 = ctx.md.lrf_cutoff * ctx.md.lrf_cutoff; + + exact_calculation_groups_.clear(); + lrf_calculation_groups_.clear(); + + for (int i = 0; i < n_groups; i++) { + if (!group_is_active(i)) continue; + const int atom1_iswitch = groups[i].iswitch - 1; + const int atom1_iswitch_type = data_.category->cpu_data_p[slots_by_atom_[atom1_iswitch][0]]; + const bool atom1_iswitch_is_q = atom1_iswitch_type == static_cast(AtomCategory::Q); + + for (int j = i; j < n_groups; j++) { + if (!group_is_active(j)) continue; + const int atom2_iswitch = groups[j].iswitch - 1; + const int atom2_iswitch_type = data_.category->cpu_data_p[slots_by_atom_[atom2_iswitch][0]]; + const bool atom2_iswitch_is_q = atom2_iswitch_type == static_cast(AtomCategory::Q); + + if (atom1_iswitch_is_q && atom2_iswitch_is_q) { + exact_calculation_groups_.push_back({i, j}); continue; } + if (atom1_iswitch_is_q || atom2_iswitch_is_q) { + const int environment_group = atom1_iswitch_is_q ? j : i; + bool inside = inside_rcq(environment_group); + if (inside) { + exact_calculation_groups_.push_back({i, j}); + } + } else { + const double distance2 = group_distance2(i, j); + const double cutoff = normal_cutoff(i, j); + const double cutoff2 = cutoff * cutoff; + + if (distance2 <= cutoff2) { + // need to calculate each pair + exact_calculation_groups_.push_back({i, j}); + } else if (distance2 <= lrf_cutoff2) { + // need to use lrf + lrf_calculation_groups_.push_back({i, j}); + } else { + // ignore + continue; + } + } + } + } + init_exact_atom_pairs(ctx); + + printf( + "[Non Bonded Force] there are %d exact atom pairs. Total number of atoms: %d. Exact atom pairs ratio is %f\n", + (int)exact_atom_pairs_.size(), ctx.n_atoms, 1.0 * exact_atom_pairs_.size() / ctx.n_atoms / ctx.n_atoms); +} + +void CpuNonbondedForce::init_exact_atom_pairs(Context& ctx) { + const auto& groups = ctx.charge_group_config.charge_groups; + const int n_groups = groups.size(); + - real_t dx = coords[atom2].x - coords[atom1].x; - real_t dy = coords[atom2].y - coords[atom1].y; - real_t dz = coords[atom2].z - coords[atom1].z; - real_t dis2 = dx * dx + dy * dy + dz * dz; - real_t inv_dis2 = static_cast(1.0) / dis2; - real_t inv_dis = sqrt(inv_dis2); + exact_atom_pairs_.clear(); - real_t qij = atom1_charge * atom2_charge; - bool is_14 = (bond_type == BondType::Bond14); - real_t scaling = is_14 ? ctx.topo.el14_scale : 1; - real_t2 pair = is_14 ? combine_vdw(ctx.topo.vdw_rule, atom1_vdw.aii_14, atom1_vdw.bii_14, atom2_vdw.aii_14, atom2_vdw.bii_14) : combine_vdw(ctx.topo.vdw_rule, atom1_vdw.aii_normal, atom1_vdw.bii_normal, atom2_vdw.aii_normal, atom2_vdw.bii_normal); + const int* atom_indices = data_.atom_idx->cpu_data_p; + const uint8_t* categories = data_.category->cpu_data_p; + const int* q_states = data_.q_state->cpu_data_p; - auto [vel, dvel] = calc_electrostatic(qij * scaling, ctx.topo.coulomb_constant, inv_dis); - auto [vvdw, dvvdw] = calc_vdw(pair, inv_dis); + constexpr uint8_t Q = static_cast(AtomCategory::Q); - real_t lambda = std::min(data_.atom_lambdas->cpu_data_p[i], data_.atom_lambdas->cpu_data_p[j]); + const auto& config = ctx.charge_group_config; + const coord_t* coords = ctx.coords->cpu_data_p; - real_t dva = (dvel + dvvdw) * inv_dis * lambda; + for (const auto& pair : exact_calculation_groups_) { + auto [group1, group2] = pair; - add_force(dvelocities[atom1].x, -dva * dx); - add_force(dvelocities[atom1].y, -dva * dy); - add_force(dvelocities[atom1].z, -dva * dz); + for (int atom1_1based : groups[group1].atoms) { + int atom1 = atom1_1based - 1; - add_force(dvelocities[atom2].x, dva * dx); - add_force(dvelocities[atom2].y, dva * dy); - add_force(dvelocities[atom2].z, dva * dz); + for (int atom2_1based : groups[group2].atoms) { + int atom2 = atom2_1based - 1; + if (group1 == group2 && atom1 >= atom2) continue; - // Accumulate energy - accumulate_energy(ctx, vel, vvdw, atom1_type, atom2_type, atom1_state, atom2_state); + for (auto slot1 : slots_by_atom_[atom1]) { + for (auto slot2 : slots_by_atom_[atom2]) { + exact_atom_pairs_.push_back({slot1, slot2}); + } + } + } } } } + +void CpuNonbondedForce::calc_direct_pair(Context& ctx, int slot1, int slot2) { + const auto& atom_idxs = data_.atom_idx->cpu_data_p; + const auto& coords = ctx.coords->cpu_data_p; + auto& dvelocities = ctx.dvelocities->cpu_data_p; + int sz = data_.n_total; + + const int atom1 = atom_idxs[slot1]; + const int atom2 = atom_idxs[slot2]; + + if (atom1 == -1 || atom2 == -1) return; + + const auto& atom1_type = data_.category->cpu_data_p[slot1]; + const int atom1_state = data_.q_state->cpu_data_p[slot1]; + const real_t atom1_charge = data_.atom_charge->cpu_data_p[slot1]; + const vdw_atom_param_t& atom1_vdw = data_.atom_vdw->cpu_data_p[slot1]; + + const auto& atom2_type = data_.category->cpu_data_p[slot2]; + const int atom2_state = data_.q_state->cpu_data_p[slot2]; + const auto& bond_type = get_bond_type(ctx.n_atoms_solute, ctx.LJ_matrix->cpu_data_p, atom1, atom1_type, atom2, atom2_type); + const real_t atom2_charge = data_.atom_charge->cpu_data_p[slot2]; + const vdw_atom_param_t atom2_vdw = data_.atom_vdw->cpu_data_p[slot2]; + + if (bond_type == BondType::Bond23) return; + if (atom1_type == static_cast(AtomCategory::Q) && atom2_type == static_cast(AtomCategory::Q) && atom1_state != atom2_state) { + return; + } + + real_t dx = coords[atom2].x - coords[atom1].x; + real_t dy = coords[atom2].y - coords[atom1].y; + real_t dz = coords[atom2].z - coords[atom1].z; + real_t dis2 = dx * dx + dy * dy + dz * dz; + real_t inv_dis2 = static_cast(1.0) / dis2; + real_t inv_dis = sqrt(inv_dis2); + + real_t qij = atom1_charge * atom2_charge; + bool is_14 = (bond_type == BondType::Bond14); + real_t scaling = is_14 ? ctx.topo.el14_scale : 1; + real_t2 pair = is_14 ? combine_vdw(ctx.topo.vdw_rule, atom1_vdw.aii_14, atom1_vdw.bii_14, atom2_vdw.aii_14, atom2_vdw.bii_14) : combine_vdw(ctx.topo.vdw_rule, atom1_vdw.aii_normal, atom1_vdw.bii_normal, atom2_vdw.aii_normal, atom2_vdw.bii_normal); + + auto [vel, dvel] = calc_electrostatic(qij * scaling, ctx.topo.coulomb_constant, inv_dis); + auto [vvdw, dvvdw] = calc_vdw(pair, inv_dis); + + real_t lambda = std::min(data_.atom_lambdas->cpu_data_p[slot1], data_.atom_lambdas->cpu_data_p[slot2]); + + real_t dva = (dvel + dvvdw) * inv_dis * lambda; + + add_force(dvelocities[atom1].x, -dva * dx); + add_force(dvelocities[atom1].y, -dva * dy); + add_force(dvelocities[atom1].z, -dva * dz); + + add_force(dvelocities[atom2].x, dva * dx); + add_force(dvelocities[atom2].y, dva * dy); + add_force(dvelocities[atom2].z, dva * dz); + + // Accumulate energy + accumulate_energy(ctx, vel, vvdw, atom1_type, atom2_type, atom1_state, atom2_state); +} + +void CpuNonbondedForce::calc_exact_pairs(Context& ctx) { + for (const auto& pair : exact_atom_pairs_) { + calc_direct_pair(ctx, pair.first, pair.second); + } +} + +void CpuNonbondedForce::init_backend(Context& ctx) { + + const int* atom_indices = data_.atom_idx->cpu_data_p; + const uint8_t* categories = data_.category->cpu_data_p; + + constexpr uint8_t P = static_cast(AtomCategory::P); + constexpr uint8_t W = static_cast(AtomCategory::W); + + slots_by_atom_.assign(ctx.n_atoms, std::vector{}); + for (int slot = 0; slot < data_.n_total; ++slot) { + const int atom = atom_indices[slot]; + if (atom < 0) continue; + slots_by_atom_[atom].push_back(slot); + } +} + +void CpuNonbondedForce::init_lrf_coefficients(Context& ctx) { + const auto& groups = ctx.charge_group_config.charge_groups; + + const int n_groups = static_cast(groups.size()); + + const coord_t* coords = ctx.coords->cpu_data_p; + + lrf_coefficients_.assign(n_groups, LrfCoefficients{}); + + for (int group = 0; group < n_groups; group++) { + coord_t center = {}; + const auto& group_atoms = groups[group].atoms; + for (int atom_1based : group_atoms) { + const int atom = atom_1based - 1; + + center = center + coords[atom]; + } + const double inverse_count = 1.0 / group_atoms.size(); + center = center * inverse_count; + + lrf_coefficients_[group].center = center; + } + + const real_t* charges = data_.atom_charge->cpu_data_p; + + auto accumulate_group_into_target = [&](int source_group, int target_group) { + LrfCoefficients& target = lrf_coefficients_[target_group]; + + for (int atom_1based : groups[source_group].atoms) { + const int atom = atom_1based - 1; + const int slot = slots_by_atom_[atom][0]; // Only P or W atoms will be in lrf. + accumulate_lrf_source(target, coords[atom], charges[slot]); + } + }; + + for (const auto& pair : lrf_calculation_groups_) { + const int group1 = pair.first; + const int group2 = pair.second; + accumulate_group_into_target(group1, group2); + accumulate_group_into_target(group2, group1); + } +} + +void CpuNonbondedForce::calc_lrf(Context& ctx) { + const int* atom_indices = data_.atom_idx->cpu_data_p; + const uint8_t* categories = data_.category->cpu_data_p; + const real_t* charges = data_.atom_charge->cpu_data_p; + auto* atom_to_group = data_.atom_to_group->cpu_data_p; + + const coord_t* coords = ctx.coords->cpu_data_p; + auto* dvelocities = ctx.dvelocities->cpu_data_p; + + constexpr uint8_t P = static_cast(AtomCategory::P); + constexpr uint8_t W = static_cast(AtomCategory::W); + + const double coulomb_constant = ctx.topo.coulomb_constant; + double lrf_energy = 0.0; + + for (int slot = 0; slot < data_.n_total; slot++) { + const int atom = atom_indices[slot]; + if (atom < 0) continue; + + // LRF is applied only to non-Q atoms + if (categories[slot] != P && categories[slot] != W) { + continue; + } + + const int group = atom_to_group[atom]; + if (group < 0) continue; + + const LrfCoefficients& lrf = lrf_coefficients_[group]; + + coord_t d = lrf.center - coords[atom]; + double d_array[3] = {d.x, d.y, d.z}; + + double potential = lrf.phi0; + + for (int a = 0; a < 3; a++) { + potential += lrf.phi1[a] * d_array[a]; + } + + for (int a = 0; a < 3; a++) { + for (int b = 0; b < 3; b++) { + potential += 0.5 * lrf.phi2[a * 3 + b] * d_array[a] * d_array[b]; + } + } + + double df[3] = {lrf.phi1[0], lrf.phi1[1], lrf.phi1[2]}; + for (int a = 0; a < 3; a++) { + for (int b = 0; b < 3; b++) { + df[a] += lrf.phi2[a * 3 + b] * d_array[b]; + } + } + + for (int a = 0; a < 3; a++) { + for (int b = 0; b < 3; b++) { + for (int c = 0; c < 3; c++) { + df[a] += 0.5 * lrf.phi3[(a * 3 + b) * 3 + c] * d_array[b] * d_array[c]; + } + } + } + + const double charge = static_cast(charges[slot]); + + lrf_energy += 0.5 * coulomb_constant * charge * potential; + + const double gradient_scale = -coulomb_constant * charge; + + add_force(dvelocities[atom].x, gradient_scale * df[0]); + add_force(dvelocities[atom].y, gradient_scale * df[1]); + add_force(dvelocities[atom].z, gradient_scale * df[2]); + } + + add_energy(ctx.energy.host()[E_LRF], lrf_energy); +} + +void CpuNonbondedForce::calc(Context& ctx) { + if (!ctx.md.lrf || ctx.md.non_bond == 0) { + calc_all_direct_pairs(ctx); + return; + } + + // 1. Examine each pair of charge groups + if (ctx.step == ctx.md.steps || ctx.step % ctx.md.non_bond == 0) { + init_calculation_groups(ctx); + init_lrf_coefficients(ctx); + } + calc_exact_pairs(ctx); + calc_lrf(ctx); +} \ No newline at end of file diff --git a/src/core/cuda/include/cuda_nonbonded_force.cuh b/src/core/cuda/include/cuda_nonbonded_force.cuh index 58111926..a528de64 100644 --- a/src/core/cuda/include/cuda_nonbonded_force.cuh +++ b/src/core/cuda/include/cuda_nonbonded_force.cuh @@ -2,6 +2,17 @@ #include "nonbonded_force.h" #include "precision.h" +enum GroupPairMode : uint8_t { + GROUP_PAIR_IGNORE = 0, + GROUP_PAIR_EXACT = 1, + GROUP_PAIR_LRF = 2 +}; + +struct ExactEntry { + int x_start; + int y_start; +}; + class CudaNonbondedForce final : public NonbondedForce { public: void calc(Context& ctx) override; @@ -10,5 +21,30 @@ class CudaNonbondedForce final : public NonbondedForce { void init_backend(Context& ctx) override; private: - std::unique_ptr> coord_x, coord_y, coord_z; + void calc_all_direct_pairs(Context& ctx); + void init_calculation_groups(Context& ctx); + void init_lrf_coefficients(Context& ctx); + void calc_exact_tiles(Context& ctx); + void calc_lrf(Context& ctx); + void build_exact_atom_tiles(Context& ctx); + + std::unique_ptr> coord_x_, coord_y_, coord_z_; + + std::unique_ptr> group_pair_modes_; + + std::unique_ptr> exact_tiles_; + + std::unique_ptr> exact_tile_count_; + std::unique_ptr> list_overflow_; + + std::unique_ptr> lrf_coefficients_; + + size_t exact_tile_capacity_ = 0; + + int n_exact_tiles_ = 0; + + std::unique_ptr> lrf_slot_to_group_range_; + + std::unique_ptr> exact_pair_masks_; + std::unique_ptr> exact_pair_14_masks_; }; diff --git a/src/core/cuda/src/cuda_nonbonded_force.cu b/src/core/cuda/src/cuda_nonbonded_force.cu index 6d99cfa6..b0dd2c0c 100644 --- a/src/core/cuda/src/cuda_nonbonded_force.cu +++ b/src/core/cuda/src/cuda_nonbonded_force.cu @@ -1,8 +1,182 @@ +#include + #include "cuda_force_accumulation.cuh" #include "cuda_nonbonded_force.cuh" +#include "geometry.h" namespace { +__device__ __forceinline__ real_t combined_slot_distance2( + int slot1, + int slot2, + + const real_t* __restrict__ cx, + const real_t* __restrict__ cy, + const real_t* __restrict__ cz) { + const real_t dx = cx[slot1] - cx[slot2]; + const real_t dy = cy[slot1] - cy[slot2]; + const real_t dz = cz[slot1] - cz[slot2]; + + return dx * dx + dy * dy + dz * dz; +} + +__device__ __forceinline__ real_t combined_slot_center_distance2( + int slot, + const coord_t center, + + const real_t* __restrict__ cx, + const real_t* __restrict__ cy, + const real_t* __restrict__ cz) { + const real_t dx = cx[slot] - center.x; + const real_t dy = cy[slot] - center.y; + const real_t dz = cz[slot] - center.z; + + return dx * dx + dy * dy + dz * dz; +} + +__device__ __forceinline__ void write_unique_lrf_component(LrfCoefficients& output, int phi, double value) { + switch (phi) { + case 0: + output.phi0 = value; + break; + + case 1: + output.phi1[0] = value; + break; + + case 2: + output.phi1[1] = value; + break; + + case 3: + output.phi1[2] = value; + break; + + case 4: + output.phi2[0] = value; + break; + + case 5: + output.phi2[4] = value; + break; + + case 6: + output.phi2[8] = value; + break; + + case 7: + output.phi2[1] = value; + output.phi2[3] = value; + break; + + case 8: + output.phi2[2] = value; + output.phi2[6] = value; + break; + + case 9: + output.phi2[5] = value; + output.phi2[7] = value; + break; + + case 10: + output.phi3[0] = value; + break; + + case 11: + output.phi3[13] = value; + break; + + case 12: + output.phi3[26] = value; + break; + + /* + * xxy: + * (x,x,y), (x,y,x), (y,x,x) + */ + case 13: + output.phi3[1] = value; + output.phi3[3] = value; + output.phi3[9] = value; + break; + + /* + * xxz: + * (x,x,z), (x,z,x), (z,x,x) + */ + case 14: + output.phi3[2] = value; + output.phi3[6] = value; + output.phi3[18] = value; + break; + + /* + * xyy: + * (x,y,y), (y,x,y), (y,y,x) + */ + case 15: + output.phi3[4] = value; + output.phi3[10] = value; + output.phi3[12] = value; + break; + + /* + * yyz: + * (y,y,z), (y,z,y), (z,y,y) + */ + case 16: + output.phi3[14] = value; + output.phi3[16] = value; + output.phi3[22] = value; + break; + + /* + * xzz: + * (x,z,z), (z,x,z), (z,z,x) + */ + case 17: + output.phi3[8] = value; + output.phi3[20] = value; + output.phi3[24] = value; + break; + + /* + * yzz: + * (y,z,z), (z,y,z), (z,z,y) + */ + case 18: + output.phi3[17] = value; + output.phi3[23] = value; + output.phi3[25] = value; + break; + + /* + * xyz: all six permutations. + */ + case 19: + output.phi3[5] = value; + output.phi3[7] = value; + output.phi3[11] = value; + output.phi3[15] = value; + output.phi3[19] = value; + output.phi3[21] = value; + break; + } +} + +__device__ int get_pair_index( + int n, + int group1, + int group2) { + const int x = min(group1, group2); + const int y = max(group1, group2); + + // Row x starts after: + // n + (n-1) + ... + (n-x+1) + return x * n - (x * (x + 1)) / 2 + y; +} + __device__ int2 get_tile_idx(int n, int t) { int x = (int)floorf((2 * n + 1 - sqrtf((2 * n + 1) * (2 * n + 1) - 8 * t)) * 0.5f); int y = t - (x * n - (x * (x - 1) >> 1)); @@ -14,6 +188,7 @@ __device__ int2 get_tile_idx(int n, int t) { return {x, y}; } +template __device__ void compute_pair( // atom1 int atom1, uint8_t atom1_type, int atom1_state, @@ -27,14 +202,16 @@ __device__ void compute_pair( real_t el14_scale, real_t coulomb_constant, int vdw_rule, // output real_t3& atom1_force, real_t3& atom2_force, - real_t& e_coul, real_t& e_vdw) { - if (atom1 == -1 || atom2 == -1 || atom1 == atom2) return; - auto bond_type = get_bond_type(n_atoms_solute, LJ_matrix, atom1, atom1_type, atom2, atom2_type); - if (bond_type == BondType::Bond23) return; - - constexpr uint8_t Q = static_cast(AtomCategory::Q); - if (atom1_type == Q && atom2_type == Q && atom1_state != atom2_state) return; + real_t& e_coul, real_t& e_vdw, bool is_14 = false) { + if constexpr (!BONDS_PRECLASSIFIED) { + if (atom1 == -1 || atom2 == -1 || atom1 == atom2) return; + auto bond_type = get_bond_type(n_atoms_solute, LJ_matrix, atom1, atom1_type, atom2, atom2_type); + if (bond_type == BondType::Bond23) return; + constexpr uint8_t Q = static_cast(AtomCategory::Q); + if (atom1_type == Q && atom2_type == Q && atom1_state != atom2_state) return; + is_14 = bond_type == BondType::Bond14; + } real_t dx = atom2_coord.x - atom1_coord.x; real_t dy = atom2_coord.y - atom1_coord.y; real_t dz = atom2_coord.z - atom1_coord.z; @@ -42,7 +219,6 @@ __device__ void compute_pair( real_t inv_dis2 = static_cast(1.0) / dis2; real_t inv_dis = sqrt(inv_dis2); - bool is_14 = (bond_type == BondType::Bond14); real_t qij = atom1_charge * atom2_charge; real_t scaling = is_14 ? el14_scale : 1; real_t2 pair = (is_14) ? combine_vdw(vdw_rule, atom1_vdw.aii_14, atom1_vdw.bii_14, atom2_vdw.aii_14, atom2_vdw.bii_14) : combine_vdw(vdw_rule, atom1_vdw.aii_normal, atom1_vdw.bii_normal, atom2_vdw.aii_normal, atom2_vdw.bii_normal); @@ -89,34 +265,16 @@ __device__ void shuffle(int& atom, uint8_t& atom_type, int& atom_state, real_t& atom_coord.z = __shfl_sync(FULL_MASK, atom_coord.z, src); } -__global__ void update_nonbonded_coords_kernel( - const coord_t* coords, const int* atom_idx, - real_t* cx, real_t* cy, real_t* cz, int sz) { - const int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i >= sz) return; - const int idx = atom_idx[i]; - if (idx < 0) { // padding slot (atom_idx == -1); main kernel treats these as empty - cx[i] = cy[i] = cz[i] = 0; - return; - } - cx[i] = static_cast(coords[idx].x); - cy[i] = static_cast(coords[idx].y); - cz[i] = static_cast(coords[idx].z); -} - -} // namespace - -void CudaNonbondedForce::init_backend(Context& ctx) { - // Buffers are indexed by combined-list position [0, n_total), which exceeds - // n_atoms because Q atoms are duplicated per FEP state and the list is padded. - coord_x = std::make_unique>(data_.n_total); - coord_y = std::make_unique>(data_.n_total); - coord_z = std::make_unique>(data_.n_total); -} +template +__device__ void nonbonded_force_calculation( + int x_idx, + int y_idx, + bool is_diag, + uint32_t pair_mask, + uint32_t pair_14_mask, + int base_x, + int base_y, -__global__ void nonbonded_kernel( - // ---- dimensions ---- - int sz, // data_.n_total, number of participating atoms int n_states, // ctx.n_lambdas, used by nb_coul_slot int n_atoms_solute, // ctx.n_atoms_solute, water grouping + LJ_matrix row stride @@ -141,26 +299,11 @@ __global__ void nonbonded_kernel( dvel_t* dvelocities, // ctx.dvelocities->gpu_data_p (fixed-point, atomic_add_force) // ---- energy accumulators ---- - energy_accum_t* e) { - const int block_num = (sz + 31) >> 5; - const int total_tiles = (block_num * (block_num + 1)) >> 1; - const int warps_per_block = blockDim.x >> 5; - const int tid = threadIdx.x; - const int lane = tid & 31; - const int warp_in_block = tid >> 5; - - const int tile = blockIdx.x * warps_per_block + warp_in_block; - if (tile >= total_tiles) return; - - auto [tile_x, tile_y] = get_tile_idx(block_num, tile); - - const int base_x = tile_x << 5; - const int base_y = tile_y << 5; + energy_accum_t* e - int x_idx = base_x + lane; - int y_idx = base_y + lane; - - const int atom1 = x_idx < sz ? atom_idx[x_idx] : -1; +) { + int lane = threadIdx.x & 31; + const int atom1 = x_idx == -1 ? -1 : atom_idx[x_idx]; const auto& atom1_type = atom1 == -1 ? static_cast(AtomCategory::INVALID) : category[x_idx]; const int atom1_state = atom1 == -1 ? -1 : q_state[x_idx]; @@ -170,7 +313,7 @@ __global__ void nonbonded_kernel( real_t3 atom1_coord = atom1 == -1 ? real_t3{0, 0, 0} : real_t3{cx[x_idx], cy[x_idx], cz[x_idx]}; real_t3 atom1_force = {0, 0, 0}; - int atom2 = y_idx < sz ? atom_idx[y_idx] : -1; + int atom2 = y_idx == -1 ? -1 : atom_idx[y_idx]; uint8_t atom2_type = atom2 == -1 ? static_cast(AtomCategory::INVALID) : category[y_idx]; int atom2_state = atom2 == -1 ? -1 : q_state[y_idx]; real_t atom2_charge = atom2 == -1 ? 0 : atom_charge[y_idx]; @@ -180,17 +323,21 @@ __global__ void nonbonded_kernel( real_t3 atom2_force = {0, 0, 0}; real_t local_e_coul = 0, local_e_vdw = 0; - bool is_diag = (tile_x == tile_y); + int y_lane = lane; for (int i = 0; i < 32; i++) { - if (!is_diag || atom1 < atom2) { - compute_pair(atom1, atom1_type, atom1_state, atom1_charge, atom1_vdw, atom1_lambda, atom1_coord, - atom2, atom2_type, atom2_state, atom2_charge, atom2_vdw, atom2_lambda, atom2_coord, - n_atoms_solute, LJ_matrix, - el14_scale, coulomb_constant, vdw_rule, - atom1_force, atom2_force, - local_e_coul, local_e_vdw); + const uint32_t bit = uint32_t{1} << y_lane; + const bool mask_enabled = (pair_mask & bit) != 0; + if (mask_enabled && (!is_diag || atom1 < atom2)) { + const bool is_14 = (pair_14_mask & bit) != 0; + compute_pair(atom1, atom1_type, atom1_state, atom1_charge, atom1_vdw, atom1_lambda, atom1_coord, + atom2, atom2_type, atom2_state, atom2_charge, atom2_vdw, atom2_lambda, atom2_coord, + n_atoms_solute, LJ_matrix, + el14_scale, coulomb_constant, vdw_rule, + atom1_force, atom2_force, + local_e_coul, local_e_vdw, is_14); } shuffle(atom2, atom2_type, atom2_state, atom2_charge, atom2_vdw, atom2_lambda, atom2_force, atom2_coord); + y_lane = (y_lane + 1) & 31; } if (atom1 >= 0) { @@ -222,29 +369,1253 @@ __global__ void nonbonded_kernel( } } -void CudaNonbondedForce::calc(Context& ctx) { - /* - Sync the coords to CudaNonbondedForce::coords first. - */ - int sz = data_.n_total; - int sync_block = 256; - int sync_grid = (sz + sync_block - 1) / sync_block; - update_nonbonded_coords_kernel<<>>(ctx.coords->gpu_data_p, data_.atom_idx->gpu_data_p, coord_x->gpu_data_p, coord_y->gpu_data_p, coord_z->gpu_data_p, sz); +__global__ void build_exact_cluster_tiles_kernel( + int n_clusters, + int n_group_ranges, + int max_exact_tiles, - /* - Do calculation - */ - const int thread_num = 256; - int tile_num_per_block = thread_num >> 5; - int n_atom = data_.n_total; - int block_num = (n_atom + 31) >> 5; - int total_tiles = block_num * (block_num + 1) >> 1; - int grid_sz = (total_tiles + tile_num_per_block - 1) / tile_num_per_block; + const uint8_t* group_pair_modes, + const int* slot_to_group_range, - dim3 grid = dim3(grid_sz); - nonbonded_kernel<<>>(n_atom, ctx.n_lambdas(), ctx.n_atoms_solute, - data_.atom_idx->gpu_data_p, data_.category->gpu_data_p, data_.q_state->gpu_data_p, - data_.atom_lambdas->gpu_data_p, data_.atom_charge->gpu_data_p, data_.atom_vdw->gpu_data_p, - ctx.LJ_matrix->gpu_data_p, ctx.topo.el14_scale, ctx.topo.coulomb_constant, ctx.topo.vdw_rule, - coord_x->gpu_data_p, coord_y->gpu_data_p, coord_z->gpu_data_p, ctx.dvelocities->gpu_data_p, ctx.energy.device()); + int n_atoms_solute, + const int* atom_idx, + const uint8_t* category, + const int* q_state, + const int* LJ_matrix, + + int* exact_tile_count, + int* list_overflow, + + ExactEntry* exact_entries, + uint32_t* exact_pair_masks, + uint32_t* exact_pair_14_masks) { + constexpr unsigned FULL_MASK = 0xffffffffu; + constexpr uint8_t Q = static_cast(AtomCategory::Q); + + const int lane = threadIdx.x & 31; + const int warp_in_block = threadIdx.x >> 5; + const int warps_per_block = blockDim.x >> 5; + + const int candidate_index = blockIdx.x * warps_per_block + warp_in_block; + + const int candidate_count = n_clusters * (n_clusters + 1) / 2; + + // Uniform return for the entire warp. + if (candidate_index >= candidate_count) { + return; + } + + const int2 cluster_pair = get_tile_idx(n_clusters, candidate_index); + + const int x_cluster = cluster_pair.x; + const int y_cluster = cluster_pair.y; + + const int x_start = x_cluster * 32; + const int y_start = y_cluster * 32; + + const int x_slot = x_start + lane; + const int x_range = slot_to_group_range[x_slot]; + const int x_atom = atom_idx[x_slot]; + + uint32_t pair_mask = 0; + uint32_t pair_14_mask = 0; + + if (x_range >= 0 && x_atom >= 0) { + const uint8_t x_category = category[x_slot]; + const int x_state = q_state[x_slot]; + +#pragma unroll + for (int y_lane = 0; y_lane < 32; ++y_lane) { + // Unique ownership within a diagonal tile. + if (x_cluster == y_cluster && y_lane <= lane) { + continue; + } + + const int y_slot = y_start + y_lane; + const int y_range = slot_to_group_range[y_slot]; + + if (y_range < 0) { + continue; + } + + const int group_pair_index = get_pair_index(n_group_ranges, x_range, y_range); + + if (group_pair_modes[group_pair_index] != GROUP_PAIR_EXACT) { + continue; + } + + const int y_atom = atom_idx[y_slot]; + + if (y_atom < 0 || x_atom == y_atom) { + continue; + } + + const uint8_t y_category = category[y_slot]; + + if (x_category == Q && y_category == Q && x_state != q_state[y_slot]) { + continue; + } + + const BondType bond_type = get_bond_type(n_atoms_solute, LJ_matrix, x_atom, x_category, y_atom, y_category); + + if (bond_type == BondType::Bond23) { + continue; + } + + const uint32_t bit = uint32_t{1} << y_lane; + + pair_mask |= bit; + + if (bond_type == BondType::Bond14) { + pair_14_mask |= bit; + } + } + } + + // Test activity after removing exclusions. + if (!__any_sync(FULL_MASK, pair_mask != 0)) { + return; + } + + int output_index = -1; + + if (lane == 0) { + output_index = atomicAdd(exact_tile_count, 1); + + if (output_index >= max_exact_tiles) { + atomicExch(list_overflow, 1); + } else { + exact_entries[output_index] = { + x_start, y_start}; + } + } + + output_index = __shfl_sync(FULL_MASK, output_index, 0); + + if (output_index < max_exact_tiles) { + const size_t mask_index = static_cast(output_index) * 32 + lane; + + exact_pair_masks[mask_index] = pair_mask; + exact_pair_14_masks[mask_index] = pair_14_mask; + } +} + +__global__ void calc_lrf_kernel( + int n_slots, + + const int* atom_idx, + const int* atom_to_group, + const uint8_t* category, + const real_t* atom_charge, + + const real_t* cx, + const real_t* cy, + const real_t* cz, + + const LrfCoefficients* coefficients, + double coulomb_constant, + + dvel_t* dvelocities, + energy_accum_t* energy) { + const int slot = blockIdx.x * blockDim.x + threadIdx.x; + + if (slot >= n_slots) { + return; + } + + constexpr uint8_t P = static_cast(AtomCategory::P); + + constexpr uint8_t W = static_cast(AtomCategory::W); + + const uint8_t atom_category = category[slot]; + + if (atom_category != P && atom_category != W) { + return; + } + + const int atom = atom_idx[slot]; + + if (atom < 0) { + return; + } + + const int group = atom_to_group[atom]; + + if (group < 0) { + return; + } + + const LrfCoefficients& lrf = coefficients[group]; + const double d[3] = { + lrf.center.x - static_cast(cx[slot]), + lrf.center.y - static_cast(cy[slot]), + lrf.center.z - static_cast(cz[slot])}; + + double potential = lrf.phi0; + + for (int a = 0; a < 3; ++a) { + potential += lrf.phi1[a] * d[a]; + } + + for (int a = 0; a < 3; ++a) { + for (int b = 0; b < 3; ++b) { + potential += 0.5 * lrf.phi2[a * 3 + b] * d[a] * d[b]; + } + } + + double df[3] = {lrf.phi1[0], lrf.phi1[1], lrf.phi1[2]}; + + for (int a = 0; a < 3; ++a) { + for (int b = 0; b < 3; ++b) { + df[a] += lrf.phi2[a * 3 + b] * d[b]; + } + } + + for (int a = 0; a < 3; ++a) { + for (int b = 0; b < 3; ++b) { + for (int c = 0; c < 3; ++c) { + const int index = (a * 3 + b) * 3 + c; + df[a] += 0.5 * lrf.phi3[index] * d[b] * d[c]; + } + } + } + + const double charge = static_cast(atom_charge[slot]); + + const double energy_value = 0.5 * coulomb_constant * charge * potential; + + const double force_scale = -coulomb_constant * charge; + + atomic_add_force(&dvelocities[atom].x, force_scale * df[0]); + + atomic_add_force(&dvelocities[atom].y, force_scale * df[1]); + + atomic_add_force(&dvelocities[atom].z, force_scale * df[2]); + + atomic_add_energy(&energy[E_LRF], energy_value); +} + +__global__ void update_nonbonded_coords_kernel( + const coord_t* coords, const int* atom_idx, + real_t* cx, real_t* cy, real_t* cz, int sz) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= sz) return; + const int idx = atom_idx[i]; + if (idx < 0) { // padding slot (atom_idx == -1); main kernel treats these as empty + cx[i] = cy[i] = cz[i] = 0; + return; + } + cx[i] = static_cast(coords[idx].x); + cy[i] = static_cast(coords[idx].y); + cz[i] = static_cast(coords[idx].z); +} + +__global__ void nonbonded_kernel( + // ---- dimensions ---- + int sz, // data_.n_total, number of participating atoms + int n_states, // ctx.n_lambdas, used by nb_coul_slot + int n_atoms_solute, // ctx.n_atoms_solute, water grouping + LJ_matrix row stride + + // ---- per-atom arrays (length sz, parallel to atom_idx) ---- + const int* atom_idx, // data_.atom_idx, local i -> global atom index + const uint8_t* category, // data_.category, P/Q/W + const int* q_state, // data_.q_state, Q state; -1 for P/W + const real_t* atom_lambdas, // data_.atom_lambdas + const real_t* atom_charge, // data_.atom_charge + const vdw_atom_param_t* atom_vdw, // data_.atom_vdw + + // ---- exclusion data ---- + const int* LJ_matrix, // ctx.LJ_matrix->gpu_data_p + + // ---- topology scalars (passed by value) ---- + real_t el14_scale, // ctx.topo.el14_scale + real_t coulomb_constant, // ctx.topo.coulomb_constant + int vdw_rule, + + // ---- coordinates / outputs ---- + const real_t* cx, const real_t* cy, const real_t* cz, + dvel_t* dvelocities, // ctx.dvelocities->gpu_data_p (fixed-point, atomic_add_force) + + // ---- energy accumulators ---- + energy_accum_t* e) { + const int block_num = (sz + 31) >> 5; + const int total_tiles = (block_num * (block_num + 1)) >> 1; + const int warps_per_block = blockDim.x >> 5; + const int tid = threadIdx.x; + const int lane = tid & 31; + const int warp_in_block = tid >> 5; + + const int tile = blockIdx.x * warps_per_block + warp_in_block; + if (tile >= total_tiles) return; + + auto [tile_x, tile_y] = get_tile_idx(block_num, tile); + + const int base_x = tile_x << 5; + const int base_y = tile_y << 5; + + int x_idx = base_x + lane; + int y_idx = base_y + lane; + x_idx = x_idx < sz ? x_idx : -1; + y_idx = y_idx < sz ? y_idx : -1; + + nonbonded_force_calculation(x_idx, y_idx, tile_x == tile_y, 0xffffffffu, 0, base_x, base_y, n_states, n_atoms_solute, + atom_idx, category, q_state, atom_lambdas, atom_charge, atom_vdw, LJ_matrix, el14_scale, coulomb_constant, vdw_rule, cx, cy, cz, dvelocities, e); +} + +__global__ void classify_group_pairs_kernel( + int n_group_ranges, + int use_switch_atom, + + real_t solute_solute_cutoff2, + real_t solute_solvent_cutoff2, + real_t solvent_solvent_cutoff2, + real_t rcq2, + real_t lrf_cutoff2, + + const coord_t solute_center, + + const int* __restrict__ group_start_idx, + const int* __restrict__ group_sizes, + + const uint8_t* __restrict__ category, + const int* __restrict__ q_state, + + const real_t* __restrict__ cx, + const real_t* __restrict__ cy, + const real_t* __restrict__ cz, + + uint8_t* __restrict__ group_pair_modes) { + const int pair_index = blockIdx.x * blockDim.x + threadIdx.x; + + const int total_pairs = n_group_ranges * (n_group_ranges + 1) / 2; + + if (pair_index >= total_pairs) { + return; + } + + const int2 pair = get_tile_idx(n_group_ranges, pair_index); + + const int range1 = pair.x; + const int range2 = pair.y; + + const int start1 = group_start_idx[range1]; + const int start2 = group_start_idx[range2]; + + const int size1 = group_sizes[range1]; + const int size2 = group_sizes[range2]; + + const uint8_t category1 = category[start1]; + const uint8_t category2 = category[start2]; + + constexpr uint8_t P = static_cast(AtomCategory::P); + + constexpr uint8_t Q = static_cast(AtomCategory::Q); + + constexpr uint8_t W = static_cast(AtomCategory::W); + + /* + * Default classification. + */ + group_pair_modes[pair_index] = GROUP_PAIR_IGNORE; + + /* + * A group interacts exactly with itself. + */ + if (range1 == range2) { + group_pair_modes[pair_index] = GROUP_PAIR_EXACT; + return; + } + + const bool range1_is_q = category1 == Q; + const bool range2_is_q = category2 == Q; + + /* + * Q-Q is exact only for the same duplicated Q state. + */ + if (range1_is_q && range2_is_q) { + if (q_state[start1] == q_state[start2]) { + group_pair_modes[pair_index] = GROUP_PAIR_EXACT; + } + + return; + } + + /* + * Q-P and Q-W use the rcq rule. + * + * With switch atoms: + * test the environment switch atom. + * + * Without switch atoms: + * W still uses its switch atom; + * P is exact when any atom is inside rcq. + */ + if (range1_is_q || range2_is_q) { + const int environment_range = range1_is_q ? range2 : range1; + + const int environment_start = group_start_idx[environment_range]; + + const int environment_size = group_sizes[environment_range]; + + const uint8_t environment_category = category[environment_start]; + + bool inside_rcq = false; + + if (use_switch_atom != 0 || environment_category == W) { + inside_rcq = combined_slot_center_distance2(environment_start, solute_center, cx, cy, cz) <= rcq2; + } else { + /* + * The environment is a P group. Check every atom, + * matching CpuNonbondedForce::inside_rcq(). + */ + for (int local_atom = 0; local_atom < environment_size; ++local_atom) { + const int slot = environment_start + local_atom; + + if (combined_slot_center_distance2(slot, solute_center, cx, cy, cz) <= rcq2) { + inside_rcq = true; + break; + } + } + } + + if (inside_rcq) { + group_pair_modes[pair_index] = GROUP_PAIR_EXACT; + } + + return; + } + + /* + * At this point the pair is P-P, P-W, or W-W. + */ + real_t normal_cutoff2; + + if (category1 == P && category2 == P) { + normal_cutoff2 = solute_solute_cutoff2; + } else if (category1 == W && category2 == W) { + normal_cutoff2 = solvent_solvent_cutoff2; + } else { + normal_cutoff2 = solute_solvent_cutoff2; + } + + /* + * Switch-atom mode uses the first slot because the combined + * list always places the switch atom first in each group. + */ + if (use_switch_atom != 0) { + const real_t distance2 = combined_slot_distance2(start1, start2, cx, cy, cz); + + if (distance2 <= normal_cutoff2) { + group_pair_modes[pair_index] = GROUP_PAIR_EXACT; + } else if (distance2 <= lrf_cutoff2) { + group_pair_modes[pair_index] = GROUP_PAIR_LRF; + } + + return; + } + + /* + * Without switch atoms, W-W still uses switch atoms. + */ + if (category1 == W && category2 == W) { + const real_t distance2 = combined_slot_distance2(start1, start2, cx, cy, cz); + + if (distance2 <= normal_cutoff2) { + group_pair_modes[pair_index] = GROUP_PAIR_EXACT; + } else if (distance2 <= lrf_cutoff2) { + group_pair_modes[pair_index] = GROUP_PAIR_LRF; + } + + return; + } + + /* + * Without switch atoms, P-W uses the minimum distance from + * every atom in the P group to the W switch atom. + */ + if (category1 == W || category2 == W) { + const int solute_start = category1 == P ? start1 : start2; + + const int solute_size = category1 == P ? size1 : size2; + + const int water_switch = category1 == W ? start1 : start2; + + real_t minimum_distance2 = lrf_cutoff2 + 1.0; + + for (int local_atom = 0; local_atom < solute_size; ++local_atom) { + const real_t distance2 = combined_slot_distance2(solute_start + local_atom, water_switch, cx, cy, cz); + + /* + * Exact is the highest-priority classification, so + * stop as soon as any pair enters the exact cutoff. + */ + if (distance2 <= normal_cutoff2) { + group_pair_modes[pair_index] = GROUP_PAIR_EXACT; + return; + } + + if (distance2 < minimum_distance2) { + minimum_distance2 = distance2; + } + } + + if (minimum_distance2 <= lrf_cutoff2) { + group_pair_modes[pair_index] = GROUP_PAIR_LRF; + } + + return; + } + + /* + * P-P without switch atoms: + * minimum distance over every atom pair. + */ + real_t minimum_distance2 = lrf_cutoff2 + 1.0; + + for (int local1 = 0; local1 < size1; ++local1) { + const int slot1 = start1 + local1; + + for (int local2 = 0; local2 < size2; ++local2) { + const int slot2 = start2 + local2; + + const real_t distance2 = combined_slot_distance2(slot1, slot2, cx, cy, cz); + + if (distance2 <= normal_cutoff2) { + group_pair_modes[pair_index] = GROUP_PAIR_EXACT; + return; + } + + if (distance2 < minimum_distance2) { + minimum_distance2 = distance2; + } + } + } + + if (minimum_distance2 <= lrf_cutoff2) { + group_pair_modes[pair_index] = GROUP_PAIR_LRF; + } +} + +__global__ void exact_tiles_nonbonded_force_kernel( + int n_exact_tiles, + const ExactEntry* exact_entries, + const uint32_t* exact_pair_masks, + const uint32_t* exact_pair_14_masks, + + int n_states, // ctx.n_lambdas, used by nb_coul_slot + int n_atoms_solute, // ctx.n_atoms_solute, water grouping + LJ_matrix row stride + + // ---- per-atom arrays (length sz, parallel to atom_idx) ---- + const int* atom_idx, // data_.atom_idx, local i -> global atom index + const uint8_t* category, // data_.category, P/Q/W + const int* q_state, // data_.q_state, Q state; -1 for P/W + const real_t* atom_lambdas, // data_.atom_lambdas + const real_t* atom_charge, // data_.atom_charge + const vdw_atom_param_t* atom_vdw, // data_.atom_vdw + + // ---- exclusion data ---- + const int* LJ_matrix, // ctx.LJ_matrix->gpu_data_p + + // ---- topology scalars (passed by value) ---- + real_t el14_scale, // ctx.topo.el14_scale + real_t coulomb_constant, // ctx.topo.coulomb_constant + int vdw_rule, + + // ---- coordinates / outputs ---- + const real_t* cx, const real_t* cy, const real_t* cz, + dvel_t* dvelocities, // ctx.dvelocities->gpu_data_p (fixed-point, atomic_add_force) + + // ---- energy accumulators ---- + energy_accum_t* e) { + const int lane = threadIdx.x & 31; + const int warp_in_block = threadIdx.x >> 5; + const int warps_per_block = blockDim.x >> 5; + + const int tile_index = blockIdx.x * warps_per_block + warp_in_block; + + if (tile_index >= n_exact_tiles) { + return; + } + + const ExactEntry tile = exact_entries[tile_index]; + const uint32_t pair_mask = exact_pair_masks[static_cast(tile_index) * 32 + lane]; + const uint32_t pair_14_mask = exact_pair_14_masks[static_cast(tile_index) * 32 + lane]; + + /* + * Each lane loads one X atom and one Y atom. + * Invalid lanes use atom == -1 and still participate in all shuffles. + */ + const int x_idx = tile.x_start + lane; + const int y_idx = tile.y_start + lane; + + nonbonded_force_calculation(x_idx, y_idx, false, pair_mask, pair_14_mask, tile.x_start, tile.y_start, n_states, n_atoms_solute, atom_idx, + category, q_state, atom_lambdas, atom_charge, atom_vdw, LJ_matrix, el14_scale, coulomb_constant, vdw_rule, cx, cy, cz, dvelocities, e); +} + +__global__ void compute_lrf_centers_kernel( + int n_group_ranges, + + const int* group_indices, + const int* group_start_idx, + const int* group_sizes, + + const uint8_t* category, + + const real_t* cx, + const real_t* cy, + const real_t* cz, + + LrfCoefficients* coefficients) { + const int lane = threadIdx.x & 31; + const int warp_in_block = threadIdx.x >> 5; + const int warps_per_block = blockDim.x >> 5; + + const int range = blockIdx.x * warps_per_block + warp_in_block; + + if (range >= n_group_ranges) { + return; + } + + const int start = group_start_idx[range]; + const int size = group_sizes[range]; + + constexpr uint8_t P = static_cast(AtomCategory::P); + constexpr uint8_t W = static_cast(AtomCategory::W); + + const uint8_t type = category[start]; + + if (type != P && type != W) { + return; + } + + double sum_x = 0; + double sum_y = 0; + double sum_z = 0; + + for (int i = lane; i < size; i += 32) { + const int slot = start + i; + + sum_x += static_cast(cx[slot]); + sum_y += static_cast(cy[slot]); + sum_z += static_cast(cz[slot]); + } + + constexpr unsigned MASK = 0xffffffffu; + + for (int offset = 16; offset > 0; offset >>= 1) { + sum_x += __shfl_down_sync(MASK, sum_x, offset); + sum_y += __shfl_down_sync(MASK, sum_y, offset); + sum_z += __shfl_down_sync(MASK, sum_z, offset); + } + + if (lane == 0) { + const int original_group = group_indices[range]; + + const double inv_size = 1.0 / static_cast(size); + + coefficients[original_group].center = { + sum_x * inv_size, + sum_y * inv_size, + sum_z * inv_size}; + } +} + +constexpr int LRF_UNIQUE_COMPONENTS = 20; +constexpr int LRF_COEFFICIENT_THREADS = 256; + +__global__ void build_lrf_coefficients_dense_kernel( + int n_group_ranges, + int n_slots, + + const int* __restrict__ group_indices, + const int* __restrict__ group_start_idx, + const uint8_t* __restrict__ category, + const uint8_t* __restrict__ group_pair_modes, + const int* __restrict__ slot_to_group_range, + + const real_t* __restrict__ atom_charge, + const real_t* __restrict__ cx, + const real_t* __restrict__ cy, + const real_t* __restrict__ cz, + + LrfCoefficients* __restrict__ coefficients) { + constexpr int WARPS_PER_BLOCK = LRF_COEFFICIENT_THREADS / 32; + + constexpr unsigned FULL_MASK = 0xffffffffu; + + /* + * Exactly one block processes one target range. + */ + const int target_range = blockIdx.x; + + if (target_range >= n_group_ranges) { + return; + } + + const int target_start = group_start_idx[target_range]; + + const uint8_t target_category = category[target_start]; + + constexpr uint8_t P = static_cast(AtomCategory::P); + + constexpr uint8_t W = static_cast(AtomCategory::W); + + /* + * LRF is only evaluated for P and W atoms. + * This condition is uniform across the whole block. + */ + if (target_category != P && target_category != W) { + return; + } + + const int target_group = group_indices[target_range]; + + const coord_t center = coefficients[target_group].center; + + const int thread = threadIdx.x; + const int lane = thread & 31; + const int warp = thread >> 5; + + /* + * The 20 independent components are: + * + * 0: phi0 + * + * 1: x + * 2: y + * 3: z + * + * 4: xx + * 5: yy + * 6: zz + * 7: xy + * 8: xz + * 9: yz + * + * 10: xxx + * 11: yyy + * 12: zzz + * 13: xxy + * 14: xxz + * 15: xyy + * 16: yyz + * 17: xzz + * 18: yzz + * 19: xyz + */ + real_t sums[LRF_UNIQUE_COMPONENTS]; + +#pragma unroll + for (int component = 0; component < LRF_UNIQUE_COMPONENTS; ++component) { + sums[component] = 0.0; + } + + /* + * Threads collectively scan the complete contiguous source-slot + * array. Padding slots have source_range == -1. + */ + for (int source_slot = thread; source_slot < n_slots; source_slot += blockDim.x) { + const int source_range = slot_to_group_range[source_slot]; + + if (source_range < 0) { + continue; + } + + const int pair_index = get_pair_index(n_group_ranges, target_range, source_range); + + /* + * This preserves the original CSR semantics exactly: + * + * - GROUP_PAIR_LRF is included. + * - GROUP_PAIR_EXACT is excluded. + * - GROUP_PAIR_IGNORE is excluded. + */ + if (group_pair_modes[pair_index] != GROUP_PAIR_LRF) { + continue; + } + + const real_t charge = atom_charge[source_slot]; + + const real_t x = cx[source_slot] - center.x; + + const real_t y = cy[source_slot] - center.y; + + const real_t z = cz[source_slot] - center.z; + + const real_t r2 = x * x + y * y + z * z; + + /* + * A self relationship should not be LRF, but guard against + * singular values in case classification changes later. + */ + if (r2 == 0.0) { + continue; + } + + const real_t inv_r = rsqrt(r2); + const real_t inv_r2 = 1.0 / r2; + + const real_t inv_r3 = inv_r * inv_r2; + + const real_t inv_r5 = inv_r3 * inv_r2; + + const real_t inv_r7 = inv_r5 * inv_r2; + + const real_t q_r1 = charge * inv_r; + + const real_t q_r3 = charge * inv_r3; + + const real_t q_r5 = charge * inv_r5; + + const real_t q_r7 = charge * inv_r7; + + const real_t three_q_r5 = 3.0 * q_r5; + + const real_t nine_q_r5 = 9.0 * q_r5; + + const real_t fifteen_q_r7 = 15.0 * q_r7; + + const real_t xx = x * x; + const real_t yy = y * y; + const real_t zz = z * z; + + /* + * phi0 + */ + sums[0] += q_r1; + + /* + * phi1 + */ + sums[1] -= x * q_r3; + sums[2] -= y * q_r3; + sums[3] -= z * q_r3; + + /* + * Unique phi2 components. + */ + sums[4] += xx * three_q_r5 - q_r3; + + sums[5] += yy * three_q_r5 - q_r3; + + sums[6] += zz * three_q_r5 - q_r3; + + sums[7] += x * y * three_q_r5; + + sums[8] += x * z * three_q_r5; + + sums[9] += y * z * three_q_r5; + + /* + * Unique phi3 components. + */ + sums[10] += x * nine_q_r5 - x * xx * fifteen_q_r7; + + sums[11] += y * nine_q_r5 - y * yy * fifteen_q_r7; + + sums[12] += z * nine_q_r5 - z * zz * fifteen_q_r7; + + sums[13] += y * three_q_r5 - xx * y * fifteen_q_r7; + + sums[14] += z * three_q_r5 - xx * z * fifteen_q_r7; + + sums[15] += x * three_q_r5 - x * yy * fifteen_q_r7; + + sums[16] += z * three_q_r5 - yy * z * fifteen_q_r7; + + sums[17] += x * three_q_r5 - x * zz * fifteen_q_r7; + + sums[18] += y * three_q_r5 - y * zz * fifteen_q_r7; + + sums[19] -= x * y * z * fifteen_q_r7; + } + + /* + * Reduce each component within every warp. + * + * This is performed once after all source atoms have been + * processed, instead of once per source chunk. + */ +#pragma unroll + for (int component = 0; component < LRF_UNIQUE_COMPONENTS; ++component) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + sums[component] += __shfl_down_sync(FULL_MASK, sums[component], offset); + } + } + + __shared__ double warp_sums[LRF_UNIQUE_COMPONENTS][WARPS_PER_BLOCK]; + + if (lane == 0) { +#pragma unroll + for (int component = 0; component < LRF_UNIQUE_COMPONENTS; ++component) { + warp_sums[component][warp] = sums[component]; + } + } + + __syncthreads(); + + /* + * Warp zero reduces the eight warp results. + */ + if (warp == 0) { +#pragma unroll + for (int component = 0; component < LRF_UNIQUE_COMPONENTS; ++component) { + double value = lane < WARPS_PER_BLOCK ? warp_sums[component][lane] : 0.0; + +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value += __shfl_down_sync(FULL_MASK, value, offset); + } + + if (lane == 0) { + write_unique_lrf_component(coefficients[target_group], component, value); + } + } + } +} + +} // namespace + +void CudaNonbondedForce::init_backend(Context& ctx) { + // Buffers are indexed by combined-list position [0, n_total), which exceeds + // n_atoms because Q atoms are duplicated per FEP state and the list is padded. + coord_x_ = std::make_unique>(data_.n_total); + coord_y_ = std::make_unique>(data_.n_total); + coord_z_ = std::make_unique>(data_.n_total); + + const int n_group_ranges = data_.group_indices->length; + + const int max_group_pairs = n_group_ranges * (n_group_ranges + 1) / 2; + + if ((data_.n_total & 31) != 0) { + throw std::runtime_error("CUDA exact cluster list requires n_total to be padded to 32"); + } + + const int n_exact_clusters = data_.n_total / 32; + + const size_t max_exact_tiles = static_cast(n_exact_clusters) * static_cast(n_exact_clusters + 1) / 2; + + exact_tile_capacity_ = max_exact_tiles; + + group_pair_modes_ = std::make_unique>(max_group_pairs, false, true); + + exact_tiles_ = std::make_unique>(exact_tile_capacity_, false, true); + + exact_pair_masks_ = std::make_unique>(exact_tile_capacity_ * 32, false, true); + + exact_pair_14_masks_ = std::make_unique>(exact_tile_capacity_ * 32, false, true); + + exact_tile_count_ = std::make_unique>(1, true, true); + + list_overflow_ = std::make_unique>(1, true, true); + + lrf_coefficients_ = std::make_unique>(ctx.charge_group_config.charge_groups.size(), false, true); + + std::vector slot_to_group_range(data_.n_total, -1); + + for (int range = 0; range < n_group_ranges; ++range) { + const int begin = data_.group_start_idx->cpu_data_p[range]; + + const int size = data_.group_sizes->cpu_data_p[range]; + + for (int local_atom = 0; local_atom < size; ++local_atom) { + slot_to_group_range[begin + local_atom] = range; + } + } + + lrf_slot_to_group_range_ = HostDeviceBuffer::from_vector(slot_to_group_range, ctx.command_info.requested_gpu); +} + +void CudaNonbondedForce::calc_all_direct_pairs(Context& ctx) { + const int thread_num = 256; + int tile_num_per_block = thread_num >> 5; + int n_atom = data_.n_total; + int block_num = (n_atom + 31) >> 5; + int total_tiles = block_num * (block_num + 1) >> 1; + int grid_sz = (total_tiles + tile_num_per_block - 1) / tile_num_per_block; + + dim3 grid = dim3(grid_sz); + nonbonded_kernel<<>>(n_atom, ctx.n_lambdas(), ctx.n_atoms_solute, + data_.atom_idx->gpu_data_p, data_.category->gpu_data_p, data_.q_state->gpu_data_p, + data_.atom_lambdas->gpu_data_p, data_.atom_charge->gpu_data_p, data_.atom_vdw->gpu_data_p, + ctx.LJ_matrix->gpu_data_p, ctx.topo.el14_scale, ctx.topo.coulomb_constant, ctx.topo.vdw_rule, + coord_x_->gpu_data_p, coord_y_->gpu_data_p, coord_z_->gpu_data_p, ctx.dvelocities->gpu_data_p, ctx.energy.device()); +} + +void CudaNonbondedForce::init_calculation_groups(Context& ctx) { + const double solute_solute_cutoff2 = ctx.md.solute_solute * ctx.md.solute_solute; + + const double solute_solvent_cutoff2 = ctx.md.solute_solvent * ctx.md.solute_solvent; + + const double solvent_solvent_cutoff2 = ctx.md.solvent_solvent * ctx.md.solvent_solvent; + + const double rcq2 = ctx.md.q_atom * ctx.md.q_atom; + + const double lrf_cutoff2 = ctx.md.lrf_cutoff * ctx.md.lrf_cutoff; + + const int n_group_ranges = static_cast(data_.group_indices->length); + + if (n_group_ranges == 0) { + n_exact_tiles_ = 0; + return; + } + + constexpr int threads = 256; + + const int total_pairs = n_group_ranges * (n_group_ranges + 1) / 2; + + const int blocks = (total_pairs + threads - 1) / threads; + + int use_switch_atom = ctx.charge_group_config.iuse_switch_atom == 1 ? 1 : 0; + + classify_group_pairs_kernel<<>>( + n_group_ranges, + use_switch_atom, + + solute_solute_cutoff2, + solute_solvent_cutoff2, + solvent_solvent_cutoff2, + rcq2, + lrf_cutoff2, + + ctx.topo.solute_center, + + data_.group_start_idx->gpu_data_p, + data_.group_sizes->gpu_data_p, + + data_.category->gpu_data_p, + data_.q_state->gpu_data_p, + + coord_x_->gpu_data_p, + coord_y_->gpu_data_p, + coord_z_->gpu_data_p, + + group_pair_modes_->gpu_data_p); + + check_cuda(cudaGetLastError()); + + build_exact_atom_tiles(ctx); +} + +void CudaNonbondedForce::calc_exact_tiles(Context& ctx) { + if (n_exact_tiles_ <= 0) { + return; + } + + constexpr int thread_num = 256; + constexpr int warps_per_block = thread_num / 32; + + const int grid_sz = (n_exact_tiles_ + warps_per_block - 1) / warps_per_block; + + exact_tiles_nonbonded_force_kernel<<>>( + n_exact_tiles_, + exact_tiles_->gpu_data_p, + exact_pair_masks_->gpu_data_p, + exact_pair_14_masks_->gpu_data_p, + + ctx.n_lambdas(), + ctx.n_atoms_solute, + + data_.atom_idx->gpu_data_p, + data_.category->gpu_data_p, + data_.q_state->gpu_data_p, + data_.atom_lambdas->gpu_data_p, + data_.atom_charge->gpu_data_p, + data_.atom_vdw->gpu_data_p, + + ctx.LJ_matrix->gpu_data_p, + + static_cast(ctx.topo.el14_scale), + + static_cast(ctx.topo.coulomb_constant), + + ctx.topo.vdw_rule, + + coord_x_->gpu_data_p, + coord_y_->gpu_data_p, + coord_z_->gpu_data_p, + + ctx.dvelocities->gpu_data_p, + ctx.energy.device()); + + check_cuda(cudaGetLastError()); +} + +void CudaNonbondedForce::init_lrf_coefficients(Context& ctx) { + lrf_coefficients_->zero(); + + const int n_group_ranges = static_cast(data_.group_indices->length); + + if (n_group_ranges <= 0) { + return; + } + + /* + * Calculate target-group centers first. + */ + constexpr int center_threads = 256; + constexpr int center_warps_per_block = center_threads / 32; + + const int center_grid = (n_group_ranges + center_warps_per_block - 1) / center_warps_per_block; + + compute_lrf_centers_kernel<<>>( + n_group_ranges, + + data_.group_indices->gpu_data_p, + data_.group_start_idx->gpu_data_p, + data_.group_sizes->gpu_data_p, + data_.category->gpu_data_p, + + coord_x_->gpu_data_p, + coord_y_->gpu_data_p, + coord_z_->gpu_data_p, + + lrf_coefficients_->gpu_data_p); + + check_cuda(cudaGetLastError()); + + /* + * Exactly one block per target range. + */ + build_lrf_coefficients_dense_kernel<<>>( + n_group_ranges, + data_.n_total, + + data_.group_indices->gpu_data_p, + data_.group_start_idx->gpu_data_p, + data_.category->gpu_data_p, + + group_pair_modes_->gpu_data_p, + + lrf_slot_to_group_range_ + ->gpu_data_p, + + data_.atom_charge->gpu_data_p, + + coord_x_->gpu_data_p, + coord_y_->gpu_data_p, + coord_z_->gpu_data_p, + + lrf_coefficients_->gpu_data_p); + + check_cuda(cudaGetLastError()); +} + +void CudaNonbondedForce::calc_lrf(Context& ctx) { + const int n_slots = data_.n_total; + + if (n_slots <= 0) { + return; + } + + constexpr int thread_num = 256; + + const int grid_sz = (n_slots + thread_num - 1) / thread_num; + + calc_lrf_kernel<<>>( + n_slots, + + data_.atom_idx->gpu_data_p, + data_.atom_to_group->gpu_data_p, + data_.category->gpu_data_p, + data_.atom_charge->gpu_data_p, + + coord_x_->gpu_data_p, + coord_y_->gpu_data_p, + coord_z_->gpu_data_p, + + lrf_coefficients_->gpu_data_p, + + ctx.topo.coulomb_constant, + + ctx.dvelocities->gpu_data_p, + ctx.energy.device()); + + check_cuda(cudaGetLastError()); +} + +void CudaNonbondedForce::build_exact_atom_tiles(Context& ctx) { + n_exact_tiles_ = 0; + + const int n_group_ranges = static_cast(data_.group_indices->length); + + const int n_clusters = data_.n_total / 32; + + if (n_group_ranges == 0 || n_clusters == 0) { + return; + } + + exact_tile_count_->zero(); + list_overflow_->zero(); + + constexpr int threads = 256; + constexpr int warps_per_block = threads / 32; + + const int candidate_count = n_clusters * (n_clusters + 1) / 2; + + const int blocks = (candidate_count + warps_per_block - 1) / warps_per_block; + + build_exact_cluster_tiles_kernel<<>>( + n_clusters, + n_group_ranges, + static_cast(exact_tile_capacity_), + + group_pair_modes_->gpu_data_p, + lrf_slot_to_group_range_->gpu_data_p, + + ctx.n_atoms_solute, + data_.atom_idx->gpu_data_p, + data_.category->gpu_data_p, + data_.q_state->gpu_data_p, + ctx.LJ_matrix->gpu_data_p, + + exact_tile_count_->gpu_data_p, + list_overflow_->gpu_data_p, + + exact_tiles_->gpu_data_p, + exact_pair_masks_->gpu_data_p, + exact_pair_14_masks_->gpu_data_p); + + check_cuda(cudaGetLastError()); + + exact_tile_count_->download(); + list_overflow_->download(); + + if (list_overflow_->cpu_data_p[0] != 0) { + throw std::runtime_error("CUDA exact cluster tile capacity exceeded"); + } + + n_exact_tiles_ = exact_tile_count_->cpu_data_p[0]; + + if (n_exact_tiles_ < 0 || static_cast(n_exact_tiles_) > exact_tile_capacity_) { + throw std::runtime_error("Invalid CUDA exact cluster tile count"); + } +} + +void CudaNonbondedForce::calc(Context& ctx) { + /* + Sync the coords to CudaNonbondedForce::coords first. + */ + int sz = data_.n_total; + int sync_block = 256; + int sync_grid = (sz + sync_block - 1) / sync_block; + update_nonbonded_coords_kernel<<>>(ctx.coords->gpu_data_p, data_.atom_idx->gpu_data_p, coord_x_->gpu_data_p, coord_y_->gpu_data_p, coord_z_->gpu_data_p, sz); + + /* + Do calculation + */ + + if (!ctx.md.lrf || ctx.md.non_bond == 0) { + calc_all_direct_pairs(ctx); + return; + } + + if (ctx.step == ctx.md.steps || ctx.step % ctx.md.non_bond == 0) { + init_calculation_groups(ctx); + init_lrf_coefficients(ctx); + } + + calc_exact_tiles(ctx); + calc_lrf(ctx); } \ No newline at end of file