From bd92dee2bea6ccd19998c8019d8418f622336bd0 Mon Sep 17 00:00:00 2001 From: shen guo <1257960165@qq.com> Date: Tue, 1 Sep 2026 14:52:05 +0200 Subject: [PATCH 1/6] support cpu lrf --- src/core/common/include/context.h | 3 +- src/core/common/include/energy.h | 7 + src/core/common/include/md_types.h | 4 + src/core/common/include/nonbonded_force.h | 16 +- src/core/common/src/energy.cpp | 3 +- src/core/common/src/handler.cpp | 2 + src/core/common/src/inp_parser.cpp | 352 +++++++++++---- src/core/common/src/nonbonded_force.cpp | 18 + src/core/common/src/std_output.cpp | 1 + src/core/cpu/include/cpu_nonbonded_force.h | 18 + src/core/cpu/src/cpu_nonbonded_force.cpp | 493 +++++++++++++++++++-- 11 files changed, 795 insertions(+), 122 deletions(-) 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/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..53cc0005 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,7 +96,9 @@ 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 @@ -106,6 +108,15 @@ struct NonbondedData { 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 +138,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..12c558af 100644 --- a/src/core/common/src/nonbonded_force.cpp +++ b/src/core/common/src/nonbonded_force.cpp @@ -11,9 +11,27 @@ 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; 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..0d04235d 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 non_q_slot_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..fa4f2f40 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,468 @@ 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 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; + for (int j = i; j < n_groups; j++) { + if (!group_is_active(j)) continue; + + 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(); + + const auto* atom_to_group = data_.atom_to_group->cpu_data_p; + + // Dense lookup table for exact charge-group pairs. + std::vector is_exact_group_pair(static_cast(n_groups) * n_groups, 0); + + for (const auto& pair : exact_calculation_groups_) { + is_exact_group_pair[static_cast(pair.first) * n_groups + pair.second] = 1; + + is_exact_group_pair[static_cast(pair.second) * n_groups + pair.first] = 1; + } + + exact_atom_pairs_.clear(); + + 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; + + constexpr uint8_t Q = static_cast(AtomCategory::Q); + + const double rcq2 = ctx.md.q_atom * ctx.md.q_atom; + const coord_t& solute_center = ctx.topo.solute_center; + const auto& config = ctx.charge_group_config; + const coord_t* coords = ctx.coords->cpu_data_p; + + auto inside_rcq = [&](int atom, uint8_t category) { + if (category == static_cast(AtomCategory::W)) { + const int water_offset = atom - ctx.n_atoms_solute; + const int water_oxygen = ctx.n_atoms_solute + (water_offset / 3) * 3; + return norm2(coords[water_oxygen] - solute_center) <= rcq2; + } else if (category == static_cast(AtomCategory::P)) { + const int group = atom_to_group[atom]; + if (group < 0) return false; + if (config.iuse_switch_atom == 1) { + const int switch_atom = groups[group].iswitch - 1; + return norm2(coords[switch_atom] - solute_center) <= rcq2; + } + for (int atom_1based : groups[group].atoms) { + const int group_atom = atom_1based - 1; + if (norm2(coords[group_atom] - solute_center) <= rcq2) { + return true; + } + } + return false; + } + return true; + }; + + for (int slot1 = 0; slot1 < data_.n_total; slot1++) { + const int atom1 = atom_indices[slot1]; + if (atom1 < 0) continue; - 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); + const bool atom1_is_q = categories[slot1] == Q; + for (int slot2 = slot1 + 1; slot2 < data_.n_total; slot2++) { + const int atom2 = atom_indices[slot2]; + if (atom2 < 0 || atom1 == atom2) continue; - 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 bool atom2_is_q = categories[slot2] == Q; - auto [vel, dvel] = calc_electrostatic(qij * scaling, ctx.topo.coulomb_constant, inv_dis); - auto [vvdw, dvvdw] = calc_vdw(pair, inv_dis); + bool calculate_directly = false; - real_t lambda = std::min(data_.atom_lambdas->cpu_data_p[i], data_.atom_lambdas->cpu_data_p[j]); + if (atom1_is_q && atom2_is_q) { + calculate_directly = true; + } else if (atom1_is_q || atom2_is_q) { + /* + * LRF excludes Q atoms. Preserve the current QGPU behavior by + * calculating every Q-containing interaction directly. + */ + const int environment_atom = atom1_is_q ? atom2 : atom1; + const uint8_t environment_category = atom1_is_q ? categories[slot2] : categories[slot1]; + calculate_directly = inside_rcq(environment_atom, environment_category); - real_t dva = (dvel + dvvdw) * inv_dis * lambda; + } else { + const int group1 = atom_to_group[atom1]; + const int group2 = atom_to_group[atom2]; - add_force(dvelocities[atom1].x, -dva * dx); - add_force(dvelocities[atom1].y, -dva * dy); - add_force(dvelocities[atom1].z, -dva * dz); + if (group1 < 0 || group2 < 0) { + continue; + } - add_force(dvelocities[atom2].x, dva * dx); - add_force(dvelocities[atom2].y, dva * dy); - add_force(dvelocities[atom2].z, dva * dz); + calculate_directly = is_exact_group_pair[static_cast(group1) * n_groups + group2] != 0; + } - // Accumulate energy - accumulate_energy(ctx, vel, vvdw, atom1_type, atom2_type, atom1_state, atom2_state); + if (!calculate_directly) continue; + + 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) { + non_q_slot_by_atom_.assign(ctx.n_atoms, -1); + + 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); + + for (int slot = 0; slot < data_.n_total; ++slot) { + const int atom = atom_indices[slot]; + if (atom < 0) continue; + + if (categories[slot] == P || categories[slot] == W) { + non_q_slot_by_atom_[atom] = 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 = non_q_slot_by_atom_[atom]; + if (slot < 0) continue; + 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 From 01c791c3392183d9dfffac5e72c98cb31ec6b752 Mon Sep 17 00:00:00 2001 From: shen guo <1257960165@qq.com> Date: Wed, 2 Sep 2026 15:34:37 +0200 Subject: [PATCH 2/6] simplify cpu lrf --- src/core/cpu/include/cpu_nonbonded_force.h | 2 +- src/core/cpu/src/cpu_nonbonded_force.cpp | 148 +++++++++------------ 2 files changed, 64 insertions(+), 86 deletions(-) diff --git a/src/core/cpu/include/cpu_nonbonded_force.h b/src/core/cpu/include/cpu_nonbonded_force.h index 0d04235d..7cd02d28 100644 --- a/src/core/cpu/include/cpu_nonbonded_force.h +++ b/src/core/cpu/include/cpu_nonbonded_force.h @@ -22,5 +22,5 @@ class CpuNonbondedForce final : public NonbondedForce { std::vector> exact_calculation_groups_, lrf_calculation_groups_; std::vector> exact_atom_pairs_; std::vector lrf_coefficients_; - std::vector non_q_slot_by_atom_; + 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 fa4f2f40..6ca5f969 100644 --- a/src/core/cpu/src/cpu_nonbonded_force.cpp +++ b/src/core/cpu/src/cpu_nonbonded_force.cpp @@ -140,6 +140,25 @@ void CpuNonbondedForce::init_calculation_groups(Context& ctx) { 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(); @@ -147,23 +166,42 @@ void CpuNonbondedForce::init_calculation_groups(Context& ctx) { 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); - 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 + if (atom1_iswitch_is_q && atom2_iswitch_is_q) { 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; } + 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); @@ -179,15 +217,6 @@ void CpuNonbondedForce::init_exact_atom_pairs(Context& ctx) { const auto* atom_to_group = data_.atom_to_group->cpu_data_p; - // Dense lookup table for exact charge-group pairs. - std::vector is_exact_group_pair(static_cast(n_groups) * n_groups, 0); - - for (const auto& pair : exact_calculation_groups_) { - is_exact_group_pair[static_cast(pair.first) * n_groups + pair.second] = 1; - - is_exact_group_pair[static_cast(pair.second) * n_groups + pair.first] = 1; - } - exact_atom_pairs_.clear(); const int* atom_indices = data_.atom_idx->cpu_data_p; @@ -196,72 +225,25 @@ void CpuNonbondedForce::init_exact_atom_pairs(Context& ctx) { constexpr uint8_t Q = static_cast(AtomCategory::Q); - const double rcq2 = ctx.md.q_atom * ctx.md.q_atom; - const coord_t& solute_center = ctx.topo.solute_center; const auto& config = ctx.charge_group_config; const coord_t* coords = ctx.coords->cpu_data_p; - auto inside_rcq = [&](int atom, uint8_t category) { - if (category == static_cast(AtomCategory::W)) { - const int water_offset = atom - ctx.n_atoms_solute; - const int water_oxygen = ctx.n_atoms_solute + (water_offset / 3) * 3; - return norm2(coords[water_oxygen] - solute_center) <= rcq2; - } else if (category == static_cast(AtomCategory::P)) { - const int group = atom_to_group[atom]; - if (group < 0) return false; - if (config.iuse_switch_atom == 1) { - const int switch_atom = groups[group].iswitch - 1; - return norm2(coords[switch_atom] - solute_center) <= rcq2; - } - for (int atom_1based : groups[group].atoms) { - const int group_atom = atom_1based - 1; - if (norm2(coords[group_atom] - solute_center) <= rcq2) { - return true; - } - } - return false; - } - return true; - }; - - for (int slot1 = 0; slot1 < data_.n_total; slot1++) { - const int atom1 = atom_indices[slot1]; - if (atom1 < 0) continue; - - const bool atom1_is_q = categories[slot1] == Q; - for (int slot2 = slot1 + 1; slot2 < data_.n_total; slot2++) { - const int atom2 = atom_indices[slot2]; - if (atom2 < 0 || atom1 == atom2) continue; - - const bool atom2_is_q = categories[slot2] == Q; - - bool calculate_directly = false; + for (const auto& pair : exact_calculation_groups_) { + auto [group1, group2] = pair; - if (atom1_is_q && atom2_is_q) { - calculate_directly = true; - } else if (atom1_is_q || atom2_is_q) { - /* - * LRF excludes Q atoms. Preserve the current QGPU behavior by - * calculating every Q-containing interaction directly. - */ - const int environment_atom = atom1_is_q ? atom2 : atom1; - const uint8_t environment_category = atom1_is_q ? categories[slot2] : categories[slot1]; - calculate_directly = inside_rcq(environment_atom, environment_category); + for (int atom1_1based : groups[group1].atoms) { + int atom1 = atom1_1based - 1; - } else { - const int group1 = atom_to_group[atom1]; - const int group2 = atom_to_group[atom2]; + for (int atom2_1based : groups[group2].atoms) { + int atom2 = atom2_1based - 1; + if (group1 == group2 && atom1 >= atom2) continue; - if (group1 < 0 || group2 < 0) { - continue; + for (auto slot1 : slots_by_atom_[atom1]) { + for (auto slot2 : slots_by_atom_[atom2]) { + exact_atom_pairs_.push_back({slot1, slot2}); + } } - - calculate_directly = is_exact_group_pair[static_cast(group1) * n_groups + group2] != 0; } - - if (!calculate_directly) continue; - - exact_atom_pairs_.push_back({slot1, slot2}); } } } @@ -331,7 +313,7 @@ void CpuNonbondedForce::calc_exact_pairs(Context& ctx) { } void CpuNonbondedForce::init_backend(Context& ctx) { - non_q_slot_by_atom_.assign(ctx.n_atoms, -1); + slots_by_atom_.assign(ctx.n_atoms, std::vector{}); const int* atom_indices = data_.atom_idx->cpu_data_p; const uint8_t* categories = data_.category->cpu_data_p; @@ -342,10 +324,7 @@ void CpuNonbondedForce::init_backend(Context& ctx) { for (int slot = 0; slot < data_.n_total; ++slot) { const int atom = atom_indices[slot]; if (atom < 0) continue; - - if (categories[slot] == P || categories[slot] == W) { - non_q_slot_by_atom_[atom] = slot; - } + slots_by_atom_[atom].push_back(slot); } } @@ -379,8 +358,7 @@ void CpuNonbondedForce::init_lrf_coefficients(Context& ctx) { for (int atom_1based : groups[source_group].atoms) { const int atom = atom_1based - 1; - const int slot = non_q_slot_by_atom_[atom]; - if (slot < 0) continue; + 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]); } }; From 3143c343ecb15e9f04fdffcec03a3693364add72 Mon Sep 17 00:00:00 2001 From: shen guo <1257960165@qq.com> Date: Wed, 2 Sep 2026 16:26:30 +0200 Subject: [PATCH 3/6] simplify cpu lrf --- src/core/cpu/src/cpu_nonbonded_force.cpp | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/core/cpu/src/cpu_nonbonded_force.cpp b/src/core/cpu/src/cpu_nonbonded_force.cpp index 6ca5f969..898045f2 100644 --- a/src/core/cpu/src/cpu_nonbonded_force.cpp +++ b/src/core/cpu/src/cpu_nonbonded_force.cpp @@ -215,7 +215,6 @@ void CpuNonbondedForce::init_exact_atom_pairs(Context& ctx) { const auto& groups = ctx.charge_group_config.charge_groups; const int n_groups = groups.size(); - const auto* atom_to_group = data_.atom_to_group->cpu_data_p; exact_atom_pairs_.clear(); @@ -313,7 +312,6 @@ void CpuNonbondedForce::calc_exact_pairs(Context& ctx) { } void CpuNonbondedForce::init_backend(Context& ctx) { - slots_by_atom_.assign(ctx.n_atoms, std::vector{}); const int* atom_indices = data_.atom_idx->cpu_data_p; const uint8_t* categories = data_.category->cpu_data_p; @@ -321,6 +319,7 @@ void CpuNonbondedForce::init_backend(Context& ctx) { 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; From 7853a1db3576f340e9805cf313ee6eac217b16e5 Mon Sep 17 00:00:00 2001 From: shen guo <1257960165@qq.com> Date: Mon, 7 Sep 2026 10:24:00 +0200 Subject: [PATCH 4/6] support cuda lrf, unfinish --- src/core/common/include/nonbonded_force.h | 4 + src/core/common/src/nonbonded_force.cpp | 92 +++++++++-- .../cuda/include/cuda_nonbonded_force.cuh | 15 ++ src/core/cuda/src/cuda_nonbonded_force.cu | 152 +++++++++++++++--- 4 files changed, 232 insertions(+), 31 deletions(-) diff --git a/src/core/common/include/nonbonded_force.h b/src/core/common/include/nonbonded_force.h index 53cc0005..694f8474 100644 --- a/src/core/common/include/nonbonded_force.h +++ b/src/core/common/include/nonbonded_force.h @@ -102,6 +102,10 @@ struct NonbondedData { 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; diff --git a/src/core/common/src/nonbonded_force.cpp b/src/core/common/src/nonbonded_force.cpp index 12c558af..e9080b05 100644 --- a/src/core/common/src/nonbonded_force.cpp +++ b/src/core/common/src/nonbonded_force.cpp @@ -37,6 +37,9 @@ void NonbondedForce::build_combinded_list(Context& ctx) { 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++) { @@ -47,48 +50,117 @@ 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); + } + + // 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/cuda/include/cuda_nonbonded_force.cuh b/src/core/cuda/include/cuda_nonbonded_force.cuh index 58111926..d44b6ebd 100644 --- a/src/core/cuda/include/cuda_nonbonded_force.cuh +++ b/src/core/cuda/include/cuda_nonbonded_force.cuh @@ -2,6 +2,19 @@ #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 x_len; // <= 128 + int y_start; + int y_len; // <= 32 +}; + class CudaNonbondedForce final : public NonbondedForce { public: void calc(Context& ctx) override; @@ -11,4 +24,6 @@ class CudaNonbondedForce final : public NonbondedForce { private: std::unique_ptr> coord_x, coord_y, coord_z; + void calc_all_direct_pairs(Context& ctx); + void init_calculation_groups(Context& ctx); }; diff --git a/src/core/cuda/src/cuda_nonbonded_force.cu b/src/core/cuda/src/cuda_nonbonded_force.cu index 6d99cfa6..1a0b74ae 100644 --- a/src/core/cuda/src/cuda_nonbonded_force.cu +++ b/src/core/cuda/src/cuda_nonbonded_force.cu @@ -1,5 +1,6 @@ #include "cuda_force_accumulation.cuh" #include "cuda_nonbonded_force.cuh" +#include "geometry.h" namespace { @@ -104,16 +105,6 @@ __global__ void update_nonbonded_coords_kernel( 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); -} - __global__ void nonbonded_kernel( // ---- dimensions ---- int sz, // data_.n_total, number of participating atoms @@ -222,18 +213,104 @@ __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 classify_group_pairs_by_switch_kernel( + int n_groups_ranges, - /* - Do calculation - */ + double solute_solute_cutoff2, + double solute_solvent_cutoff2, + double solvent_solvent_cutoff2, + double rcq2, + double lrf_cutoff2, + + const coord_t solute_center, + const int* group_start_idx, + const int* atom_idx, + const uint8_t* category, + const int* q_state, + const coord_t* coords, + + uint8_t* group_pair_modes) { + const int pair_index = blockIdx.x * blockDim.x + threadIdx.x; + const int total_pairs = n_groups_ranges * (n_groups_ranges + 1) / 2; + + if (pair_index >= total_pairs) { + return; + } + + const int2 pair = get_tile_idx(n_groups_ranges, pair_index); + const int group1 = pair.x; + const int group2 = pair.y; + + group_pair_modes[pair_index] = GROUP_PAIR_IGNORE; + const int switch_atom1 = group_start_idx[group1]; + const int switch_atom2 = group_start_idx[group2]; + + const uint8_t category1 = category[switch_atom1]; + const uint8_t category2 = category[switch_atom2]; + + constexpr uint8_t P = static_cast(AtomCategory::P); + constexpr uint8_t Q = static_cast(AtomCategory::Q); + constexpr uint8_t W = static_cast(AtomCategory::W); + + const bool group1_is_q = category1 == Q; + const bool group2_is_q = category2 == Q; + + if (group1_is_q && group2_is_q) { + const int state1 = q_state[switch_atom1]; + const int state2 = q_state[switch_atom2]; + if (state1 == state2) { + group_pair_modes[pair_index] = GROUP_PAIR_EXACT; + } + return; + } + + if (group1_is_q || group2_is_q) { + // Q-P or Q-W + const int environment_switch_atom = group1_is_q ? switch_atom2 : switch_atom1; + const double environment_distance2 = norm2(coords[atom_idx[environment_switch_atom]] - solute_center); + + if (environment_distance2 <= rcq2) { + group_pair_modes[pair_index] = GROUP_PAIR_EXACT; + } + return; + } + + // P-P, P-W, or W-W + const double group_distance2 = norm2(coords[atom_idx[switch_atom1]] - coords[atom_idx[switch_atom2]]); + double 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; + } + + if (group_distance2 <= normal_cutoff2) { + group_pair_modes[pair_index] = GROUP_PAIR_EXACT; + } else if (group_distance2 <= lrf_cutoff2) { + group_pair_modes[pair_index] = GROUP_PAIR_LRF; + } +} + + + + + +__global__ void init_calculation_groups_by_all() { +} + +} // 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); +} + +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; @@ -247,4 +324,37 @@ void CudaNonbondedForce::calc(Context& ctx) { 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 auto& config = ctx.charge_group_config; + if (config.iuse_switch_atom == 1) { + // Use groups.iswitch to check the distance + + } else { + // Should use every atoms to check the distance + } +} + +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); + } } \ No newline at end of file From 3cf2ae91c2334afba0f1af49bc12a2a75be60b54 Mon Sep 17 00:00:00 2001 From: shen Date: Mon, 7 Sep 2026 18:32:25 +0200 Subject: [PATCH 5/6] support cuda lrf first version --- .../cuda/include/cuda_nonbonded_force.cuh | 36 +- src/core/cuda/src/cuda_nonbonded_force.cu | 878 +++++++++++++++++- 2 files changed, 867 insertions(+), 47 deletions(-) diff --git a/src/core/cuda/include/cuda_nonbonded_force.cuh b/src/core/cuda/include/cuda_nonbonded_force.cuh index d44b6ebd..06dcb65b 100644 --- a/src/core/cuda/include/cuda_nonbonded_force.cuh +++ b/src/core/cuda/include/cuda_nonbonded_force.cuh @@ -10,9 +10,18 @@ enum GroupPairMode : uint8_t { struct ExactEntry { int x_start; - int x_len; // <= 128 + int x_len; // <= 32 int y_start; int y_len; // <= 32 + uint8_t diagonal; +}; + +struct LrfPairEntry { + int range1; + int range2; + + int group1; + int group2; }; class CudaNonbondedForce final : public NonbondedForce { @@ -23,7 +32,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_calculation_groups_by_switch(Context& ctx); + void init_calculation_groups_by_all_atoms(Context& ctx); + void init_lrf_coefficients(Context &ctx); + void calc_exact_tiles(Context& ctx); + void calc_lrf(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> lrf_group_pairs_; + + std::unique_ptr> exact_tile_count_; + std::unique_ptr> lrf_pair_count_; + std::unique_ptr> list_overflow_; + + std::unique_ptr> lrf_coefficients_; + + size_t exact_tile_capacity_ = 0; + size_t lrf_pair_capacity_ = 0; + + int n_exact_tiles_ = 0; + int n_lrf_pairs_ = 0; }; diff --git a/src/core/cuda/src/cuda_nonbonded_force.cu b/src/core/cuda/src/cuda_nonbonded_force.cu index 1a0b74ae..1de63b04 100644 --- a/src/core/cuda/src/cuda_nonbonded_force.cu +++ b/src/core/cuda/src/cuda_nonbonded_force.cu @@ -90,24 +90,13 @@ __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); -} +__device__ void nonbonded_force_calculation( + int x_idx, + int y_idx, + bool is_diag, + 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 @@ -132,26 +121,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; + energy_accum_t* e - 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; - - 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]; @@ -161,7 +135,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]; @@ -171,7 +145,6 @@ __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); 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, @@ -213,6 +186,474 @@ __global__ void nonbonded_kernel( } } +__device__ void accumulate_lrf_direction( + int source_range, + int target_group, + + const int* group_start_idx, + const int* group_sizes, + + const real_t* atom_charge, + const real_t* cx, + const real_t* cy, + const real_t* cz, + + LrfCoefficients* coefficients) { + constexpr unsigned FULL_MASK = 0xffffffffu; + + const int lane = threadIdx.x & 31; + + const int source_start = group_start_idx[source_range]; + + const int source_size = group_sizes[source_range]; + + const coord_t target_center = coefficients[target_group].center; + + double local_phi0 = 0.0; + double local_phi1[3] = {}; + double local_phi2[9] = {}; + double local_phi3[27] = {}; + + for (int local_atom = lane; local_atom < source_size; local_atom += 32) { + const int slot = source_start + local_atom; + + const double charge = static_cast(atom_charge[slot]); + + const double rx = static_cast(cx[slot]) - target_center.x; + + const double ry = static_cast(cy[slot]) - target_center.y; + + const double rz = static_cast(cz[slot]) - target_center.z; + + const double r[3] = {rx, ry, rz}; + + const double r2 = rx * rx + ry * ry + rz * rz; + + const double r_length = sqrt(r2); + + const double inv_r = 1.0 / r_length; + + const double inv_r2 = 1.0 / r2; + + const double inv_r3 = inv_r * inv_r2; + + const double inv_r5 = inv_r3 * inv_r2; + + const double inv_r7 = inv_r5 * inv_r2; + + /* + * phi0 += q/r + */ + local_phi0 += charge * inv_r; + + /* + * phi1[a] -= q*r[a]/r^3 + */ + for (int a = 0; a < 3; ++a) { + local_phi1[a] -= charge * r[a] * inv_r3; + } + /* + * phi2[a,b] += q * + * (3*r[a]*r[b]/r^5 - delta[a,b]/r^3) + */ + for (int a = 0; a < 3; ++a) { + for (int b = 0; b < 3; ++b) { + const int index = a * 3 + b; + + const double delta_ab = a == b ? 1.0 : 0.0; + + local_phi2[index] += charge * (3.0 * r[a] * r[b] * inv_r5 - delta_ab * inv_r3); + } + } + + /* + * phi3[a,b,c] += q * ( + * 3*(delta_ab*r[c] + + * delta_ac*r[b] + + * delta_bc*r[a])/r^5 + * - 15*r[a]*r[b]*r[c]/r^7 + * ) + */ + 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; + + const double delta_ab = a == b ? 1.0 : 0.0; + const double delta_ac = a == c ? 1.0 : 0.0; + const double delta_bc = b == c ? 1.0 : 0.0; + + const double v1 = 3.0 * (delta_ab * r[c] + delta_ac * r[b] + delta_bc * r[a]) * inv_r5; + + const double v2 = -15.0 * r[a] * r[b] * r[c] * inv_r7; + + local_phi3[index] += charge * (v1 + v2); + } + } + } + } + + /* + * Reduce all lane-local coefficients to lane 0. + */ + for (int offset = 16; offset > 0; offset >>= 1) { + local_phi0 += __shfl_down_sync(FULL_MASK, local_phi0, offset); + + for (int i = 0; i < 3; ++i) { + local_phi1[i] += __shfl_down_sync(FULL_MASK, local_phi1[i], offset); + } + + for (int i = 0; i < 9; ++i) { + local_phi2[i] += __shfl_down_sync(FULL_MASK, local_phi2[i], offset); + } + + for (int i = 0; i < 27; ++i) { + local_phi3[i] += __shfl_down_sync(FULL_MASK, local_phi3[i], offset); + } + } + + if (lane == 0) { + atomicAdd(&coefficients[target_group].phi0, local_phi0); + + for (int i = 0; i < 3; ++i) { + atomicAdd(&coefficients[target_group].phi1[i], local_phi1[i]); + } + + for (int i = 0; i < 9; ++i) { + atomicAdd(&coefficients[target_group].phi2[i], local_phi2[i]); + } + + for (int i = 0; i < 27; ++i) { + atomicAdd(&coefficients[target_group].phi3[i], local_phi3[i]); + } + } +} + +__global__ void build_lrf_coefficients_kernel( + int n_lrf_pairs, + + const LrfPairEntry* lrf_pairs, + + const int* group_start_idx, + const int* group_sizes, + + const real_t* atom_charge, + + const real_t* cx, + const real_t* cy, + const real_t* cz, + + LrfCoefficients* coefficients) { + const int warp_in_block = threadIdx.x >> 5; + + const int warps_per_block = blockDim.x >> 5; + + const int warp_index = blockIdx.x * warps_per_block + warp_in_block; + + if (warp_index >= n_lrf_pairs) { + return; + } + + const LrfPairEntry pair = lrf_pairs[warp_index]; + + /* + * source group1 contributes to the expansion around + * target group2's center. + */ + accumulate_lrf_direction( + pair.range1, + pair.group2, + + group_start_idx, + group_sizes, + + atom_charge, + cx, + cy, + cz, + + coefficients); + + /* + * source group2 contributes to the expansion around + * target group1's center. + */ + accumulate_lrf_direction( + pair.range2, + pair.group1, + + group_start_idx, + group_sizes, + + atom_charge, + cx, + cy, + cz, + + coefficients); +} + +__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, 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 build_pair_lists_kernel( + int n_groups_ranges, + + const uint8_t* group_pair_modes, + const int* group_indices, + const int* group_start_idx, + const int* group_sizes, + + const int exact_tile_capacity, + int* exact_tile_count, + ExactEntry* exact_tiles, + + const int lrf_pair_capacity, + int* lrf_pair_count, + LrfPairEntry* lrf_pairs, + + int* overflow + +) { + const int pair_index = blockIdx.x * blockDim.x + threadIdx.x; + const int total_pairs = n_groups_ranges * (n_groups_ranges + 1) / 2; + if (pair_index >= total_pairs) { + return; + } + const int2 pair = get_tile_idx(n_groups_ranges, pair_index); + const int group1 = pair.x; + const int group2 = pair.y; + + if (group2 < group1) { + return; + } + + const uint8_t mode = group_pair_modes[pair_index]; + if (mode == GROUP_PAIR_IGNORE) { + return; + } + + if (mode == GROUP_PAIR_LRF) { + const int dst = atomicAdd(lrf_pair_count, 1); + + if (dst >= lrf_pair_capacity) { + atomicExch(overflow, 1); + return; + } + + lrf_pairs[dst] = { + group1, + group2, + group_indices[group1], + group_indices[group2], + }; + return; + } + + const int start1 = group_start_idx[group1]; + const int start2 = group_start_idx[group2]; + const int size1 = group_sizes[group1]; + const int size2 = group_sizes[group2]; + + const int nx = (size1 + 31) / 32; + const int ny = (size2 + 31) / 32; + + int tile_count = 0; + if (group1 == group2) { + tile_count = nx * (nx + 1) / 2; + } else { + tile_count = nx * ny; + } + + const int base = atomicAdd(exact_tile_count, tile_count); + if (base + tile_count > exact_tile_capacity) { + atomicExch(overflow, 1); + return; + } + + int dst = base; + for (int ix = 0; ix < nx; ix++) { + const int x_offset = ix * 32; + const int x_len = min(32, size1 - x_offset); + + for (int iy = 0; iy < ny; iy++) { + if (group1 == group2 && iy < ix) { + continue; + } + + const int y_offset = iy * 32; + const int y_len = min(32, size2 - y_offset); + + ExactEntry tile; + tile.x_start = start1 + x_offset; + tile.y_start = start2 + y_offset; + tile.x_len = x_len; + tile.y_len = y_len; + tile.diagonal = (group1 == group2 && ix == iy); + exact_tiles[dst++] = tile; + } + } +} + __global__ void classify_group_pairs_by_switch_kernel( int n_groups_ranges, @@ -293,11 +734,124 @@ __global__ void classify_group_pairs_by_switch_kernel( } } +__global__ void exact_tiles_nonbonded_force_kernel( + int n_exact_tiles, + const ExactEntry* exact_entries, + 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 -__global__ void init_calculation_groups_by_all() { + // ---- 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]; + + /* + * 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 = lane < tile.x_len ? tile.x_start + lane : -1; + const int y_idx = lane < tile.y_len ? tile.y_start + lane : -1; + + nonbonded_force_calculation(x_idx, y_idx, tile.diagonal, 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}; + } } } // namespace @@ -305,9 +859,45 @@ __global__ void init_calculation_groups_by_all() { 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); + 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; + int max_exact_tiles = 0; + + const int* group_sizes = data_.group_sizes->cpu_data_p; + + for (int group1 = 0; group1 < n_group_ranges; group1++) { + int nx = (group_sizes[group1] + 31) / 32; + for (int group2 = group1; group2 < n_group_ranges; group2++) { + int ny = (group_sizes[group2] + 31) / 32; + if (group1 == group2) { + max_exact_tiles += nx * (nx + 1) / 2; + } else { + max_exact_tiles += nx * ny; + } + } + } + + exact_tile_capacity_ = max_exact_tiles; + lrf_pair_capacity_ = max_group_pairs; + + group_pair_modes_ = std::make_unique>(max_group_pairs, false, true); + + exact_tiles_ = std::make_unique>(exact_tile_capacity_, false, true); + + lrf_group_pairs_ = std::make_unique>(lrf_pair_capacity_, false, true); + + exact_tile_count_ = std::make_unique>(1, true, true); + + lrf_pair_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); } void CudaNonbondedForce::calc_all_direct_pairs(Context& ctx) { @@ -323,19 +913,213 @@ void CudaNonbondedForce::calc_all_direct_pairs(Context& ctx) { 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()); + 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_by_switch(Context& ctx) { + exact_tile_count_->zero(); + lrf_pair_count_->zero(); + list_overflow_->zero(); + + 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 thread_num = 256; + const int n_group_ranges = data_.group_indices->length; + int total_pairs = n_group_ranges * (n_group_ranges + 1) >> 1; + int grid_sz = (total_pairs + thread_num - 1) / thread_num; + + dim3 grid = dim3(grid_sz); + classify_group_pairs_by_switch_kernel<<>>(n_group_ranges, + solute_solute_cutoff2, + solute_solvent_cutoff2, + solvent_solvent_cutoff2, + rcq2, + lrf_cutoff2, + ctx.topo.solute_center, + data_.group_start_idx->gpu_data_p, + data_.atom_idx->gpu_data_p, + data_.category->gpu_data_p, + data_.q_state->gpu_data_p, + ctx.coords->gpu_data_p, + group_pair_modes_->gpu_data_p); + check_cuda(cudaGetLastError()); + build_pair_lists_kernel<<>>(n_group_ranges, + group_pair_modes_->gpu_data_p, + data_.group_indices->gpu_data_p, + data_.group_start_idx->gpu_data_p, + data_.group_sizes->gpu_data_p, + exact_tile_capacity_, + exact_tile_count_->gpu_data_p, + exact_tiles_->gpu_data_p, + lrf_pair_capacity_, + lrf_pair_count_->gpu_data_p, + lrf_group_pairs_->gpu_data_p, + list_overflow_->gpu_data_p); + + check_cuda(cudaGetLastError()); + exact_tile_count_->download(); + lrf_pair_count_->download(); + list_overflow_->download(); + + if (list_overflow_->cpu_data_p[0] != 0) { + throw std::runtime_error("CUDA nonbonded pair-list capacity exceeded"); + } + + n_exact_tiles_ = exact_tile_count_->cpu_data_p[0]; + n_lrf_pairs_ = lrf_pair_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 tile count"); + } + + if (n_lrf_pairs_ < 0 || static_cast(n_lrf_pairs_) > lrf_pair_capacity_) { + throw std::runtime_error("Invalid CUDA LRF pair count"); + } +} + +void CudaNonbondedForce::init_calculation_groups_by_all_atoms(Context& ctx) { + throw std::runtime_error("CUDA LRF with iuse_switch_atom == 0 is not implemented yet"); } void CudaNonbondedForce::init_calculation_groups(Context& ctx) { const auto& config = ctx.charge_group_config; if (config.iuse_switch_atom == 1) { // Use groups.iswitch to check the distance - + init_calculation_groups_by_switch(ctx); } else { // Should use every atoms to check the distance + // init_calculation_groups_by_all_atoms(ctx); + // todo: now alwasys use switch to test + init_calculation_groups_by_switch(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, + + 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(); + constexpr int thread_num = 256; + constexpr int warps_per_block = thread_num / 32; + const int n_group_ranges = static_cast(data_.group_indices->length); + const int grid = (n_group_ranges + warps_per_block - 1) / warps_per_block; + + if (n_group_ranges > 0) { + 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()); + } + if (n_lrf_pairs_ > 0) { + const int coefficient_grid = (n_lrf_pairs_ + warps_per_block - 1) / warps_per_block; + + build_lrf_coefficients_kernel<<>>( + n_lrf_pairs_, + + lrf_group_pairs_->gpu_data_p, + + data_.group_start_idx->gpu_data_p, + data_.group_sizes->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::calc(Context& ctx) { /* Sync the coords to CudaNonbondedForce::coords first. @@ -343,7 +1127,7 @@ void CudaNonbondedForce::calc(Context& ctx) { 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); + 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 @@ -356,5 +1140,9 @@ void CudaNonbondedForce::calc(Context& ctx) { 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 From 4d4b0da29bf5f2e8a2c2839fb794a610d7073ae4 Mon Sep 17 00:00:00 2001 From: shen Date: Mon, 7 Sep 2026 23:51:13 +0200 Subject: [PATCH 6/6] support cuda lrf first version --- src/core/cuda/src/cuda_nonbonded_force.cu | 214 ++++++++++++++++++---- 1 file changed, 174 insertions(+), 40 deletions(-) diff --git a/src/core/cuda/src/cuda_nonbonded_force.cu b/src/core/cuda/src/cuda_nonbonded_force.cu index 1de63b04..585fc1da 100644 --- a/src/core/cuda/src/cuda_nonbonded_force.cu +++ b/src/core/cuda/src/cuda_nonbonded_force.cu @@ -4,6 +4,18 @@ namespace { +__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)); @@ -330,13 +342,14 @@ __device__ void accumulate_lrf_direction( } __global__ void build_lrf_coefficients_kernel( - int n_lrf_pairs, - - const LrfPairEntry* lrf_pairs, + int n_group_ranges, + const uint8_t* group_pair_modes, + const int* group_indices, const int* group_start_idx, const int* group_sizes, + const uint8_t* category, const real_t* atom_charge, const real_t* cx, @@ -344,53 +357,171 @@ __global__ void build_lrf_coefficients_kernel( const real_t* cz, LrfCoefficients* coefficients) { - const int warp_in_block = threadIdx.x >> 5; + constexpr unsigned FULL_MASK = 0xffffffffu; + const int lane = threadIdx.x & 31; + const int warp_in_block = threadIdx.x >> 5; const int warps_per_block = blockDim.x >> 5; + const int target_range = blockIdx.x * warps_per_block + warp_in_block; - const int warp_index = blockIdx.x * warps_per_block + warp_in_block; + if (target_range >= n_group_ranges) { + return; + } + const int target_start = group_start_idx[target_range]; + const uint8_t target_category = category[target_start]; - if (warp_index >= n_lrf_pairs) { + constexpr uint8_t P = static_cast(AtomCategory::P); + constexpr uint8_t W = static_cast(AtomCategory::W); + if (target_category != P && target_category != W) { return; } + const int target_group = group_indices[target_range]; + const coord_t target_center = coefficients[target_group].center; + + double local_phi0 = 0.0; + double local_phi1[3] = {}; + double local_phi2[9] = {}; + double local_phi3[27] = {}; - const LrfPairEntry pair = lrf_pairs[warp_index]; + for (int source_range = 0; source_range < n_group_ranges; ++source_range) { + const int pair_index = get_pair_index(n_group_ranges, target_range, source_range); + if (group_pair_modes[pair_index] != GROUP_PAIR_LRF) { + continue; + } - /* - * source group1 contributes to the expansion around - * target group2's center. - */ - accumulate_lrf_direction( - pair.range1, - pair.group2, + const int source_start = group_start_idx[source_range]; + const int source_size = group_sizes[source_range]; - group_start_idx, - group_sizes, + for (int local_atom = lane; local_atom < source_size; local_atom += 32) { + const int slot = source_start + local_atom; - atom_charge, - cx, - cy, - cz, + const double charge = static_cast(atom_charge[slot]); - coefficients); + const double rx = static_cast(cx[slot]) - target_center.x; + const double ry = static_cast(cy[slot]) - target_center.y; + const double rz = static_cast(cz[slot]) - target_center.z; - /* - * source group2 contributes to the expansion around - * target group1's center. - */ - accumulate_lrf_direction( - pair.range2, - pair.group1, + const double r[3] = {rx, ry, rz}; + + const double r2 = rx * rx + ry * ry + rz * rz; + + /* + * A valid LRF pair should never contain the target group + * itself, but guard against singular input anyway. + */ + if (r2 == 0.0) { + continue; + } + + const double inv_r = rsqrt(r2); + const double inv_r2 = 1.0 / r2; + const double inv_r3 = inv_r * inv_r2; + const double inv_r5 = inv_r3 * inv_r2; + const double inv_r7 = inv_r5 * inv_r2; + + /* + * phi0 = sum(q/r) + */ + local_phi0 += charge * inv_r; + + /* + * phi1[a] = sum(-q*r[a]/r^3) + */ +#pragma unroll + for (int a = 0; a < 3; ++a) { + local_phi1[a] -= charge * r[a] * inv_r3; + } + + /* + * phi2[a,b] = + * q * (3*r[a]*r[b]/r^5 - delta[a,b]/r^3) + */ +#pragma unroll + for (int a = 0; a < 3; ++a) { +#pragma unroll + for (int b = 0; b < 3; ++b) { + const int index = a * 3 + b; + const double delta_ab = a == b ? 1.0 : 0.0; + + local_phi2[index] += charge * (3.0 * r[a] * r[b] * inv_r5 - delta_ab * inv_r3); + } + } + + /* + * phi3[a,b,c] = + * q * ( + * 3*(delta_ab*r[c] + + * delta_ac*r[b] + + * delta_bc*r[a])/r^5 + * - 15*r[a]*r[b]*r[c]/r^7 + * ) + */ +#pragma unroll + for (int a = 0; a < 3; ++a) { +#pragma unroll + for (int b = 0; b < 3; ++b) { +#pragma unroll + for (int c = 0; c < 3; ++c) { + const int index = (a * 3 + b) * 3 + c; + + const double delta_ab = a == b ? 1.0 : 0.0; + const double delta_ac = a == c ? 1.0 : 0.0; + const double delta_bc = b == c ? 1.0 : 0.0; + + const double v1 = 3.0 * (delta_ab * r[c] + delta_ac * r[b] + delta_bc * r[a]) * inv_r5; + + const double v2 = -15.0 * r[a] * r[b] * r[c] * inv_r7; + + local_phi3[index] += charge * (v1 + v2); + } + } + } + } + } + + for (int offset = 16; offset > 0; offset >>= 1) { + local_phi0 += __shfl_down_sync(FULL_MASK, local_phi0, offset); + +#pragma unroll + for (int i = 0; i < 3; ++i) { + local_phi1[i] += __shfl_down_sync(FULL_MASK, local_phi1[i], offset); + } + +#pragma unroll + for (int i = 0; i < 9; ++i) { + local_phi2[i] += __shfl_down_sync(FULL_MASK, local_phi2[i], offset); + } + +#pragma unroll + for (int i = 0; i < 27; ++i) { + local_phi3[i] += __shfl_down_sync(FULL_MASK, local_phi3[i], offset); + } + } + + if (lane == 0) { + LrfCoefficients& output = coefficients[target_group]; + + /* + * Do not overwrite output.center, which was initialized by + * compute_lrf_centers_kernel(). + */ + output.phi0 = local_phi0; - group_start_idx, - group_sizes, +#pragma unroll + for (int i = 0; i < 3; ++i) { + output.phi1[i] = local_phi1[i]; + } - atom_charge, - cx, - cy, - cz, +#pragma unroll + for (int i = 0; i < 9; ++i) { + output.phi2[i] = local_phi2[i]; + } - coefficients); +#pragma unroll + for (int i = 0; i < 27; ++i) { + output.phi3[i] = local_phi3[i]; + } + } } __global__ void calc_lrf_kernel( @@ -1065,23 +1196,26 @@ void CudaNonbondedForce::init_lrf_coefficients(Context& ctx) { check_cuda(cudaGetLastError()); } if (n_lrf_pairs_ > 0) { - const int coefficient_grid = (n_lrf_pairs_ + warps_per_block - 1) / warps_per_block; + const int coefficient_grid = (n_group_ranges + warps_per_block - 1) / warps_per_block; build_lrf_coefficients_kernel<<>>( - n_lrf_pairs_, - - lrf_group_pairs_->gpu_data_p, + n_group_ranges, + group_pair_modes_->gpu_data_p, + data_.group_indices->gpu_data_p, data_.group_start_idx->gpu_data_p, data_.group_sizes->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); + lrf_coefficients_->gpu_data_p + + ); check_cuda(cudaGetLastError()); }