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 01/12] 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 02/12] 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 03/12] 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 04/12] 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 05/12] 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 06/12] 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()); } From 62b2c0439db14e72186dca3e23442cdfbef32649 Mon Sep 17 00:00:00 2001 From: shen guo <1257960165@qq.com> Date: Tue, 8 Sep 2026 15:26:44 +0200 Subject: [PATCH 07/12] optimize init lrf part --- .../cuda/include/cuda_nonbonded_force.cuh | 17 +- src/core/cuda/src/cuda_nonbonded_force.cu | 500 +++++++++++++++++- 2 files changed, 506 insertions(+), 11 deletions(-) diff --git a/src/core/cuda/include/cuda_nonbonded_force.cuh b/src/core/cuda/include/cuda_nonbonded_force.cuh index 06dcb65b..c73bcea2 100644 --- a/src/core/cuda/include/cuda_nonbonded_force.cuh +++ b/src/core/cuda/include/cuda_nonbonded_force.cuh @@ -36,9 +36,10 @@ class CudaNonbondedForce final : public NonbondedForce { 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 init_lrf_coefficients(Context& ctx); void calc_exact_tiles(Context& ctx); - void calc_lrf(Context &ctx); + void calc_lrf(Context& ctx); + void build_lrf_atom_csr(Context &ctx); std::unique_ptr> coord_x_, coord_y_, coord_z_; @@ -58,4 +59,16 @@ class CudaNonbondedForce final : public NonbondedForce { int n_exact_tiles_ = 0; int n_lrf_pairs_ = 0; + + std::unique_ptr> lrf_atom_degrees_; + + std::unique_ptr> lrf_atom_offsets_; + + std::unique_ptr> lrf_source_atom_slots_; + + std::unique_ptr> lrf_scan_temp_; + + size_t lrf_source_atom_capacity_ = 0; + size_t lrf_scan_temp_bytes_ = 0; + int n_lrf_source_atom_entries_ = 0; }; diff --git a/src/core/cuda/src/cuda_nonbonded_force.cu b/src/core/cuda/src/cuda_nonbonded_force.cu index 585fc1da..02dfa9fa 100644 --- a/src/core/cuda/src/cuda_nonbonded_force.cu +++ b/src/core/cuda/src/cuda_nonbonded_force.cu @@ -1,9 +1,287 @@ +#include + #include "cuda_force_accumulation.cuh" #include "cuda_nonbonded_force.cuh" #include "geometry.h" namespace { +/* + * Unique component layout: + * + * 0 phi0 + * + * 1 phi1.x + * 2 phi1.y + * 3 phi1.z + * + * 4 phi2.xxx -> xx + * 5 phi2.yy + * 6 phi2.zz + * 7 phi2.xy + * 8 phi2.xz + * 9 phi2.yz + * + * 10 phi3.xxx + * 11 phi3.yyy + * 12 phi3.zzz + * 13 phi3.xxy + * 14 phi3.xxz + * 15 phi3.xyy + * 16 phi3.yyz + * 17 phi3.xzz + * 18 phi3.yzz + * 19 phi3.xyz + */ +__device__ __forceinline__ double +calculate_unique_lrf_component( + int phi, + double charge, + double x, + double y, + double z) { + const double r2 = + x * x + y * y + z * z; + + if (r2 == 0.0) { + return 0.0; + } + + 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; + + switch (phi) { + case 0: + return charge * inv_r; + + case 1: + return -charge * x * inv_r3; + + case 2: + return -charge * y * inv_r3; + + case 3: + return -charge * z * inv_r3; + + case 4: + return charge * + (3.0 * x * x * inv_r5 - + inv_r3); + + case 5: + return charge * + (3.0 * y * y * inv_r5 - + inv_r3); + + case 6: + return charge * + (3.0 * z * z * inv_r5 - + inv_r3); + + case 7: + return charge * + (3.0 * x * y * inv_r5); + + case 8: + return charge * + (3.0 * x * z * inv_r5); + + case 9: + return charge * + (3.0 * y * z * inv_r5); + + case 10: + return charge * + (9.0 * x * inv_r5 - + 15.0 * x * x * x * inv_r7); + + case 11: + return charge * + (9.0 * y * inv_r5 - + 15.0 * y * y * y * inv_r7); + + case 12: + return charge * + (9.0 * z * inv_r5 - + 15.0 * z * z * z * inv_r7); + + case 13: + return charge * + (3.0 * y * inv_r5 - + 15.0 * x * x * y * inv_r7); + + case 14: + return charge * + (3.0 * z * inv_r5 - + 15.0 * x * x * z * inv_r7); + + case 15: + return charge * + (3.0 * x * inv_r5 - + 15.0 * x * y * y * inv_r7); + + case 16: + return charge * + (3.0 * z * inv_r5 - + 15.0 * y * y * z * inv_r7); + + case 17: + return charge * + (3.0 * x * inv_r5 - + 15.0 * x * z * z * inv_r7); + + case 18: + return charge * + (3.0 * y * inv_r5 - + 15.0 * y * z * z * inv_r7); + + case 19: + return charge * + (-15.0 * x * y * z * inv_r7); + + default: + return 0.0; + } +} + +__device__ __forceinline__ void +write_unique_lrf_component( + LrfCoefficients& output, + int phi, + double value) { + switch (phi) { + case 0: + output.phi0 = value; + break; + + case 1: + output.phi1[0] = value; + break; + + case 2: + output.phi1[1] = value; + break; + + case 3: + output.phi1[2] = value; + break; + + case 4: + output.phi2[0] = value; + break; + + case 5: + output.phi2[4] = value; + break; + + case 6: + output.phi2[8] = value; + break; + + case 7: + output.phi2[1] = value; + output.phi2[3] = value; + break; + + case 8: + output.phi2[2] = value; + output.phi2[6] = value; + break; + + case 9: + output.phi2[5] = value; + output.phi2[7] = value; + break; + + case 10: + output.phi3[0] = value; + break; + + case 11: + output.phi3[13] = value; + break; + + case 12: + output.phi3[26] = value; + break; + + /* + * xxy: + * (x,x,y), (x,y,x), (y,x,x) + */ + case 13: + output.phi3[1] = value; + output.phi3[3] = value; + output.phi3[9] = value; + break; + + /* + * xxz: + * (x,x,z), (x,z,x), (z,x,x) + */ + case 14: + output.phi3[2] = value; + output.phi3[6] = value; + output.phi3[18] = value; + break; + + /* + * xyy: + * (x,y,y), (y,x,y), (y,y,x) + */ + case 15: + output.phi3[4] = value; + output.phi3[10] = value; + output.phi3[12] = value; + break; + + /* + * yyz: + * (y,y,z), (y,z,y), (z,y,y) + */ + case 16: + output.phi3[14] = value; + output.phi3[16] = value; + output.phi3[22] = value; + break; + + /* + * xzz: + * (x,z,z), (z,x,z), (z,z,x) + */ + case 17: + output.phi3[8] = value; + output.phi3[20] = value; + output.phi3[24] = value; + break; + + /* + * yzz: + * (y,z,z), (z,y,z), (z,z,y) + */ + case 18: + output.phi3[17] = value; + output.phi3[23] = value; + output.phi3[25] = value; + break; + + /* + * xyz: all six permutations. + */ + case 19: + output.phi3[5] = value; + output.phi3[7] = value; + output.phi3[11] = value; + output.phi3[15] = value; + output.phi3[19] = value; + output.phi3[21] = value; + break; + } +} + __device__ int get_pair_index( int n, int group1, @@ -341,6 +619,138 @@ __device__ void accumulate_lrf_direction( } } +__global__ void count_lrf_atom_degree_kernel(int n_group_ranges, const uint8_t* group_pair_modes, const int* group_sizes, int* degrees) { + const int target_range = blockIdx.x * blockDim.x + threadIdx.x; + + if (target_range >= n_group_ranges) { + return; + } + + int atom_count = 0; + 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) { + atom_count += group_sizes[source_range]; + } + } + degrees[target_range] = atom_count; +} + +__global__ void fill_lrf_atom_csr_kernel(int n_group_ranges, const uint8_t* group_pair_modes, const int* group_start_idx, const int* group_sizes, const int* offsets, int* source_atom_slots) { + const int target_range = blockIdx.x * blockDim.x + threadIdx.x; + + if (target_range >= n_group_ranges) { + return; + } + + int output = offsets[target_range]; + + 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; + } + const int source_start = group_start_idx[source_range]; + const int source_size = group_sizes[source_range]; + for (int local_atom = 0; local_atom < source_size; local_atom++) { + source_atom_slots[output++] = source_start + local_atom; + } + } +} + +__global__ void build_lrf_coefficients_csr_kernel( + int n_group_ranges, + const int* group_indices, + const int* group_start_idx, + const uint8_t* category, + + const int* atom_offsets, + const int* source_atom_slots, + + const real_t* atom_charge, + + const real_t* cx, + const real_t* cy, + const real_t* cz, + LrfCoefficients* coefficients) { + const int target_range = blockIdx.x; + const int phi = blockIdx.y; + const int thread = threadIdx.x; + + if (target_range >= n_group_ranges || phi >= 20) { + return; + } + + const int target_start = group_start_idx[target_range]; + + const uint8_t target_category = category[target_start]; + + constexpr uint8_t P = static_cast(AtomCategory::P); + + constexpr uint8_t W = static_cast(AtomCategory::W); + + /* + * LRF coefficients are only required for non-Q groups. + * This condition is block-uniform. + */ + 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; + + const int begin = atom_offsets[target_range]; + + const int end = atom_offsets[target_range + 1]; + + double local_sum = 0.0; + + for (int entry = begin + thread; entry < end; entry += blockDim.x) { + const int slot = source_atom_slots[entry]; + + const double charge = static_cast(atom_charge[slot]); + + const double x = static_cast(cx[slot]) - target_center.x; + + const double y = static_cast(cy[slot]) - target_center.y; + + const double z = static_cast(cz[slot]) - target_center.z; + + local_sum += calculate_unique_lrf_component(phi, charge, x, y, z); + } + + constexpr unsigned FULL_MASK = 0xffffffffu; + + for (int offset = 16; offset > 0; offset >>= 1) { + local_sum += __shfl_down_sync(FULL_MASK, local_sum, offset); + } + + __shared__ double warp_sums[4]; + + const int lane = thread & 31; + const int warp = thread >> 5; + + if (lane == 0) { + warp_sums[warp] = local_sum; + } + + __syncthreads(); + + if (warp == 0) { + double block_sum = lane < 4 ? warp_sums[lane] : 0.0; + + for (int offset = 16; offset > 0; offset >>= 1) { + block_sum += __shfl_down_sync(FULL_MASK, block_sum, offset); + } + + if (lane == 0) { + write_unique_lrf_component(coefficients[target_group], phi, block_sum); + } + } +} + __global__ void build_lrf_coefficients_kernel( int n_group_ranges, @@ -1029,6 +1439,19 @@ void CudaNonbondedForce::init_backend(Context& ctx) { list_overflow_ = std::make_unique>(1, true, true); lrf_coefficients_ = std::make_unique>(ctx.charge_group_config.charge_groups.size(), false, true); + + lrf_atom_degrees_ = std::make_unique>(n_group_ranges + 1, false, true); + + lrf_atom_offsets_ = std::make_unique>(n_group_ranges + 1, true, true); + + check_cuda(cub::DeviceScan::ExclusiveSum( + nullptr, + lrf_scan_temp_bytes_, + lrf_atom_degrees_->gpu_data_p, + lrf_atom_offsets_->gpu_data_p, + n_group_ranges + 1)); + + lrf_scan_temp_ = std::make_unique>(lrf_scan_temp_bytes_, false, true); } void CudaNonbondedForce::calc_all_direct_pairs(Context& ctx) { @@ -1121,12 +1544,14 @@ void CudaNonbondedForce::init_calculation_groups(Context& ctx) { 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); } + build_lrf_atom_csr(ctx); } void CudaNonbondedForce::calc_exact_tiles(Context& ctx) { @@ -1195,27 +1620,29 @@ void CudaNonbondedForce::init_lrf_coefficients(Context& ctx) { check_cuda(cudaGetLastError()); } - if (n_lrf_pairs_ > 0) { - const int coefficient_grid = (n_group_ranges + warps_per_block - 1) / warps_per_block; + if (n_lrf_source_atom_entries_ > 0) { + constexpr int coefficient_threads = 128; + constexpr int unique_phi_count = 20; - build_lrf_coefficients_kernel<<>>( + dim3 coefficient_grid = (n_group_ranges, unique_phi_count, 1); + + build_lrf_coefficients_csr_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, - data_.category->gpu_data_p, + + lrf_atom_offsets_->gpu_data_p, + lrf_source_atom_slots_->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()); } @@ -1254,6 +1681,61 @@ void CudaNonbondedForce::calc_lrf(Context& ctx) { check_cuda(cudaGetLastError()); } +void CudaNonbondedForce::build_lrf_atom_csr(Context& ctx) { + const int n_group_ranges = data_.group_indices->length; + + n_lrf_source_atom_entries_ = 0; + + lrf_atom_degrees_->zero(); + const int thread_num = 256; + const int grid_size = (n_group_ranges + thread_num - 1) / thread_num; + + count_lrf_atom_degree_kernel<<>>(n_group_ranges, group_pair_modes_->gpu_data_p, data_.group_sizes->gpu_data_p, lrf_atom_degrees_->gpu_data_p); + + check_cuda(cudaGetLastError()); + + /* + * offsets[0] = 0 + * offsets[t+1] = offsets[t] + degrees[t] + * offsets[n] = total CSR entries + */ + + check_cuda(cub::DeviceScan::ExclusiveSum( + lrf_scan_temp_->gpu_data_p, + lrf_scan_temp_bytes_, + lrf_atom_degrees_->gpu_data_p, + lrf_atom_offsets_->gpu_data_p, + n_group_ranges + 1)); + + lrf_atom_offsets_->download(); + n_lrf_source_atom_entries_ = lrf_atom_offsets_->cpu_data_p[n_group_ranges]; + const size_t required_capacity = static_cast(n_lrf_source_atom_entries_); + + if (!lrf_source_atom_slots_ || required_capacity > lrf_source_atom_capacity_) { + size_t new_capacity = required_capacity; + + if (lrf_source_atom_capacity_ != 0) { + const size_t grown_capacity = lrf_source_atom_capacity_ + lrf_source_atom_capacity_ / 2; + + if (grown_capacity > new_capacity) { + new_capacity = grown_capacity; + } + } + + lrf_source_atom_slots_ = std::make_unique>(new_capacity, false, true); + lrf_source_atom_capacity_ = new_capacity; + } + fill_lrf_atom_csr_kernel<<>>( + n_group_ranges, + group_pair_modes_->gpu_data_p, + data_.group_start_idx->gpu_data_p, + data_.group_sizes->gpu_data_p, + lrf_atom_offsets_->gpu_data_p, + lrf_source_atom_slots_->gpu_data_p); + + check_cuda(cudaGetLastError()); +} + void CudaNonbondedForce::calc(Context& ctx) { /* Sync the coords to CudaNonbondedForce::coords first. From 12611276c549b5aa1d9ac9def6eb03c804e6d53e Mon Sep 17 00:00:00 2001 From: shen guo <1257960165@qq.com> Date: Tue, 8 Sep 2026 17:01:52 +0200 Subject: [PATCH 08/12] optimize init lrf part --- .../cuda/include/cuda_nonbonded_force.cuh | 14 +- src/core/cuda/src/cuda_nonbonded_force.cu | 655 ++++++++---------- 2 files changed, 290 insertions(+), 379 deletions(-) diff --git a/src/core/cuda/include/cuda_nonbonded_force.cuh b/src/core/cuda/include/cuda_nonbonded_force.cuh index c73bcea2..11d23940 100644 --- a/src/core/cuda/include/cuda_nonbonded_force.cuh +++ b/src/core/cuda/include/cuda_nonbonded_force.cuh @@ -39,7 +39,7 @@ class CudaNonbondedForce final : public NonbondedForce { void init_lrf_coefficients(Context& ctx); void calc_exact_tiles(Context& ctx); void calc_lrf(Context& ctx); - void build_lrf_atom_csr(Context &ctx); + void build_lrf_atom_csr(Context& ctx); std::unique_ptr> coord_x_, coord_y_, coord_z_; @@ -71,4 +71,16 @@ class CudaNonbondedForce final : public NonbondedForce { size_t lrf_source_atom_capacity_ = 0; size_t lrf_scan_temp_bytes_ = 0; int n_lrf_source_atom_entries_ = 0; + + /* + * Geometry associated with every CSR source-atom entry. + */ + std::unique_ptr> lrf_dx_; + std::unique_ptr> lrf_dy_; + std::unique_ptr> lrf_dz_; + + std::unique_ptr> lrf_q_r1_; + std::unique_ptr> lrf_q_r3_; + std::unique_ptr> lrf_q_r5_; + std::unique_ptr> lrf_q_r7_; }; diff --git a/src/core/cuda/src/cuda_nonbonded_force.cu b/src/core/cuda/src/cuda_nonbonded_force.cu index 02dfa9fa..38b2e099 100644 --- a/src/core/cuda/src/cuda_nonbonded_force.cu +++ b/src/core/cuda/src/cuda_nonbonded_force.cu @@ -5,153 +5,7 @@ #include "geometry.h" namespace { - -/* - * Unique component layout: - * - * 0 phi0 - * - * 1 phi1.x - * 2 phi1.y - * 3 phi1.z - * - * 4 phi2.xxx -> xx - * 5 phi2.yy - * 6 phi2.zz - * 7 phi2.xy - * 8 phi2.xz - * 9 phi2.yz - * - * 10 phi3.xxx - * 11 phi3.yyy - * 12 phi3.zzz - * 13 phi3.xxy - * 14 phi3.xxz - * 15 phi3.xyy - * 16 phi3.yyz - * 17 phi3.xzz - * 18 phi3.yzz - * 19 phi3.xyz - */ -__device__ __forceinline__ double -calculate_unique_lrf_component( - int phi, - double charge, - double x, - double y, - double z) { - const double r2 = - x * x + y * y + z * z; - - if (r2 == 0.0) { - return 0.0; - } - - 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; - - switch (phi) { - case 0: - return charge * inv_r; - - case 1: - return -charge * x * inv_r3; - - case 2: - return -charge * y * inv_r3; - - case 3: - return -charge * z * inv_r3; - - case 4: - return charge * - (3.0 * x * x * inv_r5 - - inv_r3); - - case 5: - return charge * - (3.0 * y * y * inv_r5 - - inv_r3); - - case 6: - return charge * - (3.0 * z * z * inv_r5 - - inv_r3); - - case 7: - return charge * - (3.0 * x * y * inv_r5); - - case 8: - return charge * - (3.0 * x * z * inv_r5); - - case 9: - return charge * - (3.0 * y * z * inv_r5); - - case 10: - return charge * - (9.0 * x * inv_r5 - - 15.0 * x * x * x * inv_r7); - - case 11: - return charge * - (9.0 * y * inv_r5 - - 15.0 * y * y * y * inv_r7); - - case 12: - return charge * - (9.0 * z * inv_r5 - - 15.0 * z * z * z * inv_r7); - - case 13: - return charge * - (3.0 * y * inv_r5 - - 15.0 * x * x * y * inv_r7); - - case 14: - return charge * - (3.0 * z * inv_r5 - - 15.0 * x * x * z * inv_r7); - - case 15: - return charge * - (3.0 * x * inv_r5 - - 15.0 * x * y * y * inv_r7); - - case 16: - return charge * - (3.0 * z * inv_r5 - - 15.0 * y * y * z * inv_r7); - - case 17: - return charge * - (3.0 * x * inv_r5 - - 15.0 * x * z * z * inv_r7); - - case 18: - return charge * - (3.0 * y * inv_r5 - - 15.0 * y * z * z * inv_r7); - - case 19: - return charge * - (-15.0 * x * y * z * inv_r7); - - default: - return 0.0; - } -} - -__device__ __forceinline__ void -write_unique_lrf_component( - LrfCoefficients& output, - int phi, - double value) { +__device__ __forceinline__ void write_unique_lrf_component(LrfCoefficients& output, int phi, double value) { switch (phi) { case 0: output.phi0 = value; @@ -282,6 +136,171 @@ write_unique_lrf_component( } } +template +__global__ void build_lrf_coefficients_order_csr_kernel( + int n_group_ranges, + + const int* __restrict__ group_indices, + const int* __restrict__ group_start_idx, + const uint8_t* __restrict__ category, + const int* __restrict__ atom_offsets, + + const double* __restrict__ dx, + const double* __restrict__ dy, + const double* __restrict__ dz, + + const double* __restrict__ q_r1, + const double* __restrict__ q_r3, + const double* __restrict__ q_r5, + const double* __restrict__ q_r7, + + LrfCoefficients* __restrict__ coefficients) { + constexpr int WARPS_PER_BLOCK = 4; + constexpr unsigned FULL_MASK = 0xffffffffu; + + const int target_range = blockIdx.x; + const int thread = threadIdx.x; + const int lane = thread & 31; + const int warp = thread >> 5; + + if (target_range >= n_group_ranges) { + return; + } + + const int target_start = group_start_idx[target_range]; + const uint8_t target_category = category[target_start]; + + constexpr uint8_t P = static_cast(AtomCategory::P); + constexpr uint8_t W = static_cast(AtomCategory::W); + + if (target_category != P && target_category != W) { + return; + } + + const int begin = atom_offsets[target_range]; + const int end = atom_offsets[target_range + 1]; + + double sums[COMPONENT_COUNT]; + +#pragma unroll + for (int component = 0; component < COMPONENT_COUNT; ++component) { + sums[component] = 0.0; + } + + for (int entry = begin + thread; entry < end; entry += blockDim.x) { + if constexpr (ORDER == 0) { + sums[0] += q_r1[entry]; + } + + if constexpr (ORDER == 1) { + const double x = dx[entry]; + const double y = dy[entry]; + const double z = dz[entry]; + const double qr3 = q_r3[entry]; + + sums[0] -= x * qr3; + sums[1] -= y * qr3; + sums[2] -= z * qr3; + } + + if constexpr (ORDER == 2) { + const double x = dx[entry]; + const double y = dy[entry]; + const double z = dz[entry]; + + const double qr3 = q_r3[entry]; + const double qr5 = q_r5[entry]; + const double three_qr5 = 3.0 * qr5; + + sums[0] += x * x * three_qr5 - qr3; + sums[1] += y * y * three_qr5 - qr3; + sums[2] += z * z * three_qr5 - qr3; + + sums[3] += x * y * three_qr5; + sums[4] += x * z * three_qr5; + sums[5] += y * z * three_qr5; + } + + if constexpr (ORDER == 3) { + const double x = dx[entry]; + const double y = dy[entry]; + const double z = dz[entry]; + + const double qr5 = q_r5[entry]; + const double qr7 = q_r7[entry]; + + const double three_qr5 = 3.0 * qr5; + const double nine_qr5 = 9.0 * qr5; + const double fifteen_qr7 = 15.0 * qr7; + + const double xx = x * x; + const double yy = y * y; + const double zz = z * z; + + sums[0] += x * nine_qr5 - x * xx * fifteen_qr7; + + sums[1] += y * nine_qr5 - y * yy * fifteen_qr7; + + sums[2] += z * nine_qr5 - z * zz * fifteen_qr7; + + sums[3] += y * three_qr5 - xx * y * fifteen_qr7; + + sums[4] += z * three_qr5 - xx * z * fifteen_qr7; + + sums[5] += x * three_qr5 - x * yy * fifteen_qr7; + + sums[6] += z * three_qr5 - yy * z * fifteen_qr7; + + sums[7] += x * three_qr5 - x * zz * fifteen_qr7; + + sums[8] += y * three_qr5 - y * zz * fifteen_qr7; + + sums[9] -= x * y * z * fifteen_qr7; + } + } + +#pragma unroll + for (int component = 0; component < COMPONENT_COUNT; ++component) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + sums[component] += __shfl_down_sync(FULL_MASK, sums[component], offset); + } + } + + __shared__ double warp_sums[COMPONENT_COUNT][WARPS_PER_BLOCK]; + + if (lane == 0) { +#pragma unroll + for (int component = 0; component < COMPONENT_COUNT; ++component) { + warp_sums[component][warp] = sums[component]; + } + } + + __syncthreads(); + + if (warp != 0) { + return; + } + +#pragma unroll + for (int component = 0; component < COMPONENT_COUNT; ++component) { + double block_sum = lane < WARPS_PER_BLOCK ? warp_sums[component][lane] : 0.0; + +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + block_sum += __shfl_down_sync(FULL_MASK, block_sum, offset); + } + + if (lane == 0) { + constexpr int first_component = ORDER == 0 ? 0 : ORDER == 1 ? 1 + : ORDER == 2 ? 4 + : 10; + + write_unique_lrf_component(coefficients[group_indices[target_range]], first_component + component, block_sum); + } + } +} + __device__ int get_pair_index( int n, int group1, @@ -658,11 +677,10 @@ __global__ void fill_lrf_atom_csr_kernel(int n_group_ranges, const uint8_t* grou } } -__global__ void build_lrf_coefficients_csr_kernel( +__global__ void build_lrf_geometry_csr_kernel( int n_group_ranges, + const int* group_indices, - const int* group_start_idx, - const uint8_t* category, const int* atom_offsets, const int* source_atom_slots, @@ -672,28 +690,21 @@ __global__ void build_lrf_coefficients_csr_kernel( const real_t* cx, const real_t* cy, const real_t* cz, - LrfCoefficients* coefficients) { - const int target_range = blockIdx.x; - const int phi = blockIdx.y; - const int thread = threadIdx.x; - - if (target_range >= n_group_ranges || phi >= 20) { - return; - } - - const int target_start = group_start_idx[target_range]; + const LrfCoefficients* coefficients, - const uint8_t target_category = category[target_start]; + double* output_dx, + double* output_dy, + double* output_dz, - constexpr uint8_t P = static_cast(AtomCategory::P); + double* output_q_r1, + double* output_q_r3, + double* output_q_r5, + double* output_q_r7 - constexpr uint8_t W = static_cast(AtomCategory::W); +) { + const int target_range = blockIdx.x; - /* - * LRF coefficients are only required for non-Q groups. - * This condition is block-uniform. - */ - if (target_category != P && target_category != W) { + if (target_range >= n_group_ranges) { return; } @@ -705,9 +716,7 @@ __global__ void build_lrf_coefficients_csr_kernel( const int end = atom_offsets[target_range + 1]; - double local_sum = 0.0; - - for (int entry = begin + thread; entry < end; entry += blockDim.x) { + for (int entry = begin + threadIdx.x; entry < end; entry += blockDim.x) { const int slot = source_atom_slots[entry]; const double charge = static_cast(atom_charge[slot]); @@ -718,222 +727,42 @@ __global__ void build_lrf_coefficients_csr_kernel( const double z = static_cast(cz[slot]) - target_center.z; - local_sum += calculate_unique_lrf_component(phi, charge, x, y, z); - } - - constexpr unsigned FULL_MASK = 0xffffffffu; - - for (int offset = 16; offset > 0; offset >>= 1) { - local_sum += __shfl_down_sync(FULL_MASK, local_sum, offset); - } - - __shared__ double warp_sums[4]; - - const int lane = thread & 31; - const int warp = thread >> 5; - - if (lane == 0) { - warp_sums[warp] = local_sum; - } - - __syncthreads(); - - if (warp == 0) { - double block_sum = lane < 4 ? warp_sums[lane] : 0.0; - - for (int offset = 16; offset > 0; offset >>= 1) { - block_sum += __shfl_down_sync(FULL_MASK, block_sum, offset); - } - - if (lane == 0) { - write_unique_lrf_component(coefficients[target_group], phi, block_sum); - } - } -} - -__global__ void build_lrf_coefficients_kernel( - 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, - const real_t* cy, - const real_t* cz, - - LrfCoefficients* coefficients) { - 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 double r2 = x * x + y * y + z * z; - if (target_range >= n_group_ranges) { - return; - } - const int target_start = group_start_idx[target_range]; - const uint8_t target_category = category[target_start]; + output_dx[entry] = x; + output_dy[entry] = y; + output_dz[entry] = z; - 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] = {}; + if (r2 == 0.0) { + output_q_r1[entry] = 0.0; + output_q_r3[entry] = 0.0; + output_q_r5[entry] = 0.0; + output_q_r7[entry] = 0.0; - 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; } - const int source_start = group_start_idx[source_range]; - const int source_size = group_sizes[source_range]; - - 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; - - /* - * 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; + const double inv_r = rsqrt(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); - } - } + const double inv_r2 = 1.0 / r2; - /* - * 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 inv_r3 = inv_r * inv_r2; - 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 inv_r5 = inv_r3 * inv_r2; - const double v1 = 3.0 * (delta_ab * r[c] + delta_ac * r[b] + delta_bc * r[a]) * inv_r5; + const double inv_r7 = inv_r5 * inv_r2; - const double v2 = -15.0 * r[a] * r[b] * r[c] * inv_r7; + output_q_r1[entry] = charge * inv_r; - local_phi3[index] += charge * (v1 + v2); - } - } - } - } - } + output_q_r3[entry] = charge * inv_r3; - 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); - } + output_q_r5[entry] = charge * inv_r5; -#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; - -#pragma unroll - for (int i = 0; i < 3; ++i) { - output.phi1[i] = local_phi1[i]; - } - -#pragma unroll - for (int i = 0; i < 9; ++i) { - output.phi2[i] = local_phi2[i]; - } - -#pragma unroll - for (int i = 0; i < 27; ++i) { - output.phi3[i] = local_phi3[i]; - } + output_q_r7[entry] = charge * inv_r7; } } + __global__ void calc_lrf_kernel( int n_slots, @@ -1621,17 +1450,12 @@ void CudaNonbondedForce::init_lrf_coefficients(Context& ctx) { check_cuda(cudaGetLastError()); } if (n_lrf_source_atom_entries_ > 0) { - constexpr int coefficient_threads = 128; - constexpr int unique_phi_count = 20; - - dim3 coefficient_grid = (n_group_ranges, unique_phi_count, 1); + constexpr int geometry_threads = 256; - build_lrf_coefficients_csr_kernel<<>>( + build_lrf_geometry_csr_kernel<<>>( n_group_ranges, data_.group_indices->gpu_data_p, - data_.group_start_idx->gpu_data_p, - data_.category->gpu_data_p, lrf_atom_offsets_->gpu_data_p, lrf_source_atom_slots_->gpu_data_p, @@ -1642,9 +1466,60 @@ void CudaNonbondedForce::init_lrf_coefficients(Context& ctx) { coord_y_->gpu_data_p, coord_z_->gpu_data_p, - lrf_coefficients_->gpu_data_p); + lrf_coefficients_->gpu_data_p, + + lrf_dx_->gpu_data_p, + lrf_dy_->gpu_data_p, + lrf_dz_->gpu_data_p, + + lrf_q_r1_->gpu_data_p, + lrf_q_r3_->gpu_data_p, + lrf_q_r5_->gpu_data_p, + lrf_q_r7_->gpu_data_p); + + check_cuda(cudaGetLastError()); + + constexpr int coefficient_threads = 128; + +#define LRF_COEFFICIENT_ARGUMENTS \ + n_group_ranges, \ + data_.group_indices->gpu_data_p, \ + data_.group_start_idx->gpu_data_p, \ + data_.category->gpu_data_p, \ + lrf_atom_offsets_->gpu_data_p, \ + lrf_dx_->gpu_data_p, \ + lrf_dy_->gpu_data_p, \ + lrf_dz_->gpu_data_p, \ + lrf_q_r1_->gpu_data_p, \ + lrf_q_r3_->gpu_data_p, \ + lrf_q_r5_->gpu_data_p, \ + lrf_q_r7_->gpu_data_p, \ + lrf_coefficients_->gpu_data_p + + build_lrf_coefficients_order_csr_kernel<0, 1> + <<>>( + LRF_COEFFICIENT_ARGUMENTS); + + check_cuda(cudaGetLastError()); + + build_lrf_coefficients_order_csr_kernel<1, 3> + <<>>( + LRF_COEFFICIENT_ARGUMENTS); check_cuda(cudaGetLastError()); + + build_lrf_coefficients_order_csr_kernel<2, 6> + <<>>( + LRF_COEFFICIENT_ARGUMENTS); + + check_cuda(cudaGetLastError()); + + build_lrf_coefficients_order_csr_kernel<3, 10> + <<>>( + LRF_COEFFICIENT_ARGUMENTS); + + check_cuda(cudaGetLastError()); +#undef LRF_COEFFICIENT_ARGUMENTS } } @@ -1709,6 +1584,15 @@ void CudaNonbondedForce::build_lrf_atom_csr(Context& ctx) { lrf_atom_offsets_->download(); n_lrf_source_atom_entries_ = lrf_atom_offsets_->cpu_data_p[n_group_ranges]; + + if (n_lrf_source_atom_entries_ < 0) { + throw std::runtime_error("Negative CUDA LRF CSR entry count"); + } + + if (n_lrf_source_atom_entries_ == 0) { + return; + } + const size_t required_capacity = static_cast(n_lrf_source_atom_entries_); if (!lrf_source_atom_slots_ || required_capacity > lrf_source_atom_capacity_) { @@ -1723,6 +1607,21 @@ void CudaNonbondedForce::build_lrf_atom_csr(Context& ctx) { } lrf_source_atom_slots_ = std::make_unique>(new_capacity, false, true); + + lrf_dx_ = std::make_unique>(new_capacity, false, true); + + lrf_dy_ = std::make_unique>(new_capacity, false, true); + + lrf_dz_ = std::make_unique>(new_capacity, false, true); + + lrf_q_r1_ = std::make_unique>(new_capacity, false, true); + + lrf_q_r3_ = std::make_unique>(new_capacity, false, true); + + lrf_q_r5_ = std::make_unique>(new_capacity, false, true); + + lrf_q_r7_ = std::make_unique>(new_capacity, false, true); + lrf_source_atom_capacity_ = new_capacity; } fill_lrf_atom_csr_kernel<<>>( From 7b2b68cab172a21cf755585f70a8d5d5c7415684 Mon Sep 17 00:00:00 2001 From: shen guo <1257960165@qq.com> Date: Wed, 9 Sep 2026 15:48:43 +0200 Subject: [PATCH 09/12] optimize exact entry calculation part --- .../cuda/include/cuda_nonbonded_force.cuh | 13 + src/core/cuda/src/cuda_nonbonded_force.cu | 421 +++++++++++++++--- 2 files changed, 363 insertions(+), 71 deletions(-) diff --git a/src/core/cuda/include/cuda_nonbonded_force.cuh b/src/core/cuda/include/cuda_nonbonded_force.cuh index 11d23940..33fe5c6b 100644 --- a/src/core/cuda/include/cuda_nonbonded_force.cuh +++ b/src/core/cuda/include/cuda_nonbonded_force.cuh @@ -13,7 +13,10 @@ struct ExactEntry { int x_len; // <= 32 int y_start; int y_len; // <= 32 + uint8_t diagonal; + uint8_t y_indirect; + int y_type_slot; }; struct LrfPairEntry { @@ -40,6 +43,7 @@ class CudaNonbondedForce final : public NonbondedForce { void calc_exact_tiles(Context& ctx); void calc_lrf(Context& ctx); void build_lrf_atom_csr(Context& ctx); + void build_exact_atom_tiles(Context& ctx); std::unique_ptr> coord_x_, coord_y_, coord_z_; @@ -83,4 +87,13 @@ class CudaNonbondedForce final : public NonbondedForce { std::unique_ptr> lrf_q_r3_; std::unique_ptr> lrf_q_r5_; std::unique_ptr> lrf_q_r7_; + + std::unique_ptr> exact_atom_degrees_; + std::unique_ptr> exact_atom_offsets_; + std::unique_ptr> exact_entry_degrees_; + std::unique_ptr> exact_entry_offsets_; + + std::unique_ptr> exact_source_atom_slots_; + size_t exact_source_atom_capacity_ = 0; + int n_exact_source_atoms_ = 0; }; diff --git a/src/core/cuda/src/cuda_nonbonded_force.cu b/src/core/cuda/src/cuda_nonbonded_force.cu index 38b2e099..7f701d7f 100644 --- a/src/core/cuda/src/cuda_nonbonded_force.cu +++ b/src/core/cuda/src/cuda_nonbonded_force.cu @@ -5,6 +5,53 @@ #include "geometry.h" namespace { + +__device__ __forceinline__ void emit_exact_packed_segment( + int segment_begin, + int segment_atoms, + int segment_type_slot, + + int target_start, + int target_size, + + int& entry_cursor, + ExactEntry* __restrict__ exact_entries) { + if (segment_atoms == 0) { + return; + } + + /* + * X is the fixed, contiguous target group. + * Y is the packed indirect atom list that rotates. + */ + for (int x_offset = 0; x_offset < target_size; x_offset += 32) { + const int x_len = min(32, target_size - x_offset); + + for (int y_offset = 0; y_offset < segment_atoms; y_offset += 32) { + const int y_len = min(32, segment_atoms - y_offset); + + ExactEntry entry{}; + + entry.x_start = target_start + x_offset; + entry.x_len = x_len; + + entry.y_start = segment_begin + y_offset; + entry.y_len = y_len; + + entry.y_type_slot = segment_type_slot; + + entry.diagonal = 0; + entry.y_indirect = 1; + + exact_entries[entry_cursor++] = entry; + } + } +} + +__device__ bool same_exact_energy_class(int slot1, int slot2, const uint8_t* category, const int* q_state) { + return category[slot1] == category[slot2] && q_state[slot1] == q_state[slot2]; +} + __device__ __forceinline__ void write_unique_lrf_component(LrfCoefficients& output, int phi, double value) { switch (phi) { case 0: @@ -638,6 +685,95 @@ __device__ void accumulate_lrf_direction( } } +__global__ void count_exact_atom_tiles_kernel( + int n_group_ranges, + + const uint8_t* group_pair_modes, + const int* group_start_idx, + const int* group_sizes, + const uint8_t* category, + const int* q_state, + + int* atom_degrees, + int* entry_degrees) { + const int target_range = blockIdx.x * blockDim.x + threadIdx.x; + + if (target_range >= n_group_ranges) { + return; + } + + const int target_size = group_sizes[target_range]; + + /* + * The target group is the fixed, contiguous X axis. + */ + const int target_x_tiles = (target_size + 31) / 32; + + int atom_count = 0; + int entry_count = 0; + + int segment_atoms = 0; + int segment_type_slot = -1; + + /* + * source_range < target_range: + * only cross-group atoms are packed into the indirect Y list. + */ + for (int source_range = 0; source_range < target_range; ++source_range) { + const int pair_index = get_pair_index(n_group_ranges, source_range, target_range); + + if (group_pair_modes[pair_index] != GROUP_PAIR_EXACT) { + continue; + } + + const int source_start = group_start_idx[source_range]; + + const int source_size = group_sizes[source_range]; + + /* + * Finish the current packed Y segment when its + * category or Q state changes. + */ + if (segment_atoms != 0 && !same_exact_energy_class(segment_type_slot, source_start, category, q_state)) { + const int segment_y_tiles = (segment_atoms + 31) / 32; + + entry_count += target_x_tiles * segment_y_tiles; + + segment_atoms = 0; + segment_type_slot = -1; + } + + if (segment_atoms == 0) { + segment_type_slot = source_start; + } + + segment_atoms += source_size; + atom_count += source_size; + } + + /* + * Count the final packed Y segment. + */ + if (segment_atoms != 0) { + const int segment_y_tiles = (segment_atoms + 31) / 32; + + entry_count += target_x_tiles * segment_y_tiles; + } + + /* + * Count the target group's separate diagonal entries. + * These entries do not use the packed Y list. + */ + const int diagonal_pair = get_pair_index(n_group_ranges, target_range, target_range); + + if (group_pair_modes[diagonal_pair] == GROUP_PAIR_EXACT) { + entry_count += target_x_tiles * (target_x_tiles + 1) / 2; + } + + atom_degrees[target_range] = atom_count; + entry_degrees[target_range] = entry_count; +} + __global__ void count_lrf_atom_degree_kernel(int n_group_ranges, const uint8_t* group_pair_modes, const int* group_sizes, int* degrees) { const int target_range = blockIdx.x * blockDim.x + threadIdx.x; @@ -762,7 +898,6 @@ __global__ void build_lrf_geometry_csr_kernel( } } - __global__ void calc_lrf_kernel( int n_slots, @@ -930,12 +1065,6 @@ __global__ void build_pair_lists_kernel( 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, @@ -958,70 +1087,23 @@ __global__ void build_pair_lists_kernel( } 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], - }; + if (mode != GROUP_PAIR_LRF) { 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 dst = atomicAdd(lrf_pair_count, 1); - 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) { + if (dst >= lrf_pair_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; - } - } + lrf_pairs[dst] = { + group1, + group2, + group_indices[group1], + group_indices[group2], + }; } __global__ void classify_group_pairs_by_switch_kernel( @@ -1107,6 +1189,7 @@ __global__ void classify_group_pairs_by_switch_kernel( __global__ void exact_tiles_nonbonded_force_kernel( int n_exact_tiles, const ExactEntry* exact_entries, + const int* exact_source_atom_slots, 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 @@ -1150,9 +1233,9 @@ __global__ void exact_tiles_nonbonded_force_kernel( * 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; + const int y_idx = lane < tile.y_len ? tile.y_indirect == 1 ? exact_source_atom_slots[tile.y_start + lane] : 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, + nonbonded_force_calculation(x_idx, y_idx, tile.diagonal == 1, tile.x_start, tile.y_type_slot, 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); } @@ -1224,6 +1307,98 @@ __global__ void compute_lrf_centers_kernel( } } +__global__ void fill_exact_atom_tiles_kernel( + int n_group_ranges, + + const uint8_t* group_pair_modes, + const int* group_start_idx, + const int* group_sizes, + const uint8_t* category, + const int* q_state, + + const int* atom_offsets, + const int* entry_offsets, + + int* source_atom_slots, + ExactEntry* exact_entries) { + const int target_range = blockIdx.x * blockDim.x + threadIdx.x; + if (target_range >= n_group_ranges) { + return; + } + + const int target_start = group_start_idx[target_range]; + const int target_size = group_sizes[target_range]; + + int atom_cursor = atom_offsets[target_range]; + int entry_cursor = entry_offsets[target_range]; + + int segment_begin = atom_cursor; + int segment_atoms = 0; + int segment_type_slot = -1; + + for (int source_range = 0; source_range < target_range; source_range++) { + const int pair_index = get_pair_index(n_group_ranges, source_range, target_range); + + if (group_pair_modes[pair_index] != GROUP_PAIR_EXACT) { + continue; + } + + const int source_start = group_start_idx[source_range]; + + const int source_size = group_sizes[source_range]; + + if (segment_atoms != 0 && !same_exact_energy_class(segment_type_slot, source_start, category, q_state)) { + emit_exact_packed_segment(segment_begin, segment_atoms, segment_type_slot, target_start, target_size, entry_cursor, exact_entries); + segment_begin = atom_cursor; + segment_atoms = 0; + segment_type_slot = -1; + } + + if (segment_atoms == 0) { + segment_begin = atom_cursor; + segment_type_slot = source_start; + } + + for (int local_atom = 0; local_atom < source_size; local_atom++) { + source_atom_slots[atom_cursor++] = source_start + local_atom; + segment_atoms++; + } + } + + emit_exact_packed_segment(segment_begin, segment_atoms, segment_type_slot, target_start, target_size, entry_cursor, exact_entries); + + const int diagonal_pair = get_pair_index(n_group_ranges, target_range, target_range); + + if (group_pair_modes[diagonal_pair] == GROUP_PAIR_EXACT) { + const int n = (target_size + 31) / 32; + + for (int ix = 0; ix < n; ix++) { + const int x_offset = ix * 32; + const int x_len = min(32, target_size - x_offset); + + for (int iy = ix; iy < n; ++iy) { + const int y_offset = iy * 32; + const int y_len = min(32, target_size - y_offset); + + ExactEntry entry{}; + + entry.x_start = target_start + x_offset; + + entry.x_len = x_len; + entry.y_start = target_start + y_offset; + entry.y_type_slot = entry.y_start; + + entry.y_len = y_len; + + entry.diagonal = (ix == iy); + entry.y_indirect = 0; + + exact_entries[entry_cursor++] = entry; + } + } + } +} + } // namespace void CudaNonbondedForce::init_backend(Context& ctx) { @@ -1281,6 +1456,14 @@ void CudaNonbondedForce::init_backend(Context& ctx) { n_group_ranges + 1)); lrf_scan_temp_ = std::make_unique>(lrf_scan_temp_bytes_, false, true); + + exact_atom_degrees_ = std::make_unique>(n_group_ranges + 1, false, true); + + exact_atom_offsets_ = std::make_unique>(n_group_ranges + 1, true, true); + + exact_entry_degrees_ = std::make_unique>(n_group_ranges + 1, false, true); + + exact_entry_offsets_ = std::make_unique>(n_group_ranges + 1, true, true); } void CudaNonbondedForce::calc_all_direct_pairs(Context& ctx) { @@ -1333,11 +1516,6 @@ void CudaNonbondedForce::init_calculation_groups_by_switch(Context& ctx) { 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, @@ -1352,7 +1530,6 @@ void CudaNonbondedForce::init_calculation_groups_by_switch(Context& ctx) { 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_) { @@ -1362,6 +1539,8 @@ void CudaNonbondedForce::init_calculation_groups_by_switch(Context& ctx) { if (n_lrf_pairs_ < 0 || static_cast(n_lrf_pairs_) > lrf_pair_capacity_) { throw std::runtime_error("Invalid CUDA LRF pair count"); } + + build_exact_atom_tiles(ctx); } void CudaNonbondedForce::init_calculation_groups_by_all_atoms(Context& ctx) { @@ -1397,6 +1576,8 @@ void CudaNonbondedForce::calc_exact_tiles(Context& ctx) { n_exact_tiles_, exact_tiles_->gpu_data_p, + exact_source_atom_slots_->gpu_data_p, + ctx.n_lambdas(), ctx.n_atoms_solute, @@ -1635,6 +1816,104 @@ void CudaNonbondedForce::build_lrf_atom_csr(Context& ctx) { check_cuda(cudaGetLastError()); } +void CudaNonbondedForce::build_exact_atom_tiles(Context& ctx) { + const int n_group_ranges = data_.group_indices->length; + + n_exact_tiles_ = 0; + n_exact_source_atoms_ = 0; + if (n_group_ranges == 0) return; + + exact_atom_degrees_->zero(); + exact_atom_offsets_->zero(); + + exact_entry_degrees_->zero(); + exact_entry_offsets_->zero(); + + const int threads = 256; + + const int blocks = (n_group_ranges + threads - 1) / threads; + + count_exact_atom_tiles_kernel<<>>(n_group_ranges, + group_pair_modes_->gpu_data_p, + data_.group_start_idx->gpu_data_p, + data_.group_sizes->gpu_data_p, + data_.category->gpu_data_p, + data_.q_state->gpu_data_p, + exact_atom_degrees_->gpu_data_p, + exact_entry_degrees_->gpu_data_p); + check_cuda(cudaGetLastError()); + + check_cuda(cub::DeviceScan::ExclusiveSum( + lrf_scan_temp_->gpu_data_p, + lrf_scan_temp_bytes_, + + exact_atom_degrees_->gpu_data_p, + exact_atom_offsets_->gpu_data_p, + + n_group_ranges + 1)); + + check_cuda(cub::DeviceScan::ExclusiveSum( + lrf_scan_temp_->gpu_data_p, + lrf_scan_temp_bytes_, + + exact_entry_degrees_->gpu_data_p, + exact_entry_offsets_->gpu_data_p, + + n_group_ranges + 1)); + + exact_atom_offsets_->download(); + exact_entry_offsets_->download(); + + n_exact_source_atoms_ = exact_atom_offsets_->cpu_data_p[n_group_ranges]; + + n_exact_tiles_ = exact_entry_offsets_->cpu_data_p[n_group_ranges]; + + if (n_exact_source_atoms_ < 0 || n_exact_tiles_ < 0) { + throw std::runtime_error("Negative CUDA exact-list size"); + } + + if (static_cast(n_exact_tiles_) > exact_tile_capacity_) { + throw std::runtime_error("CUDA packed exact-tile capacity exceeded"); + } + + const size_t required_atom_capacity = static_cast(n_exact_source_atoms_); + + if (required_atom_capacity > exact_source_atom_capacity_) { + size_t new_capacity = required_atom_capacity; + + if (exact_source_atom_capacity_ != 0) { + const size_t grown_capacity = exact_source_atom_capacity_ + exact_source_atom_capacity_ / 2; + + new_capacity = std::max(new_capacity, grown_capacity); + } + + exact_source_atom_slots_ = std::make_unique>(new_capacity, false, true); + + exact_source_atom_capacity_ = new_capacity; + } + + if (n_exact_tiles_ == 0) { + return; + } + + fill_exact_atom_tiles_kernel<<>>( + n_group_ranges, + + group_pair_modes_->gpu_data_p, + data_.group_start_idx->gpu_data_p, + data_.group_sizes->gpu_data_p, + data_.category->gpu_data_p, + data_.q_state->gpu_data_p, + + exact_atom_offsets_->gpu_data_p, + exact_entry_offsets_->gpu_data_p, + + exact_source_atom_slots_->gpu_data_p, + exact_tiles_->gpu_data_p); + + check_cuda(cudaGetLastError()); +} + void CudaNonbondedForce::calc(Context& ctx) { /* Sync the coords to CudaNonbondedForce::coords first. From fd69941dd639d32d444af87c50571647b0b68874 Mon Sep 17 00:00:00 2001 From: shen guo <1257960165@qq.com> Date: Fri, 11 Sep 2026 14:14:57 +0200 Subject: [PATCH 10/12] optimize lrf coefficient --- .../cuda/include/cuda_nonbonded_force.cuh | 2 + src/core/cuda/src/cuda_nonbonded_force.cu | 685 ++++++++---------- 2 files changed, 310 insertions(+), 377 deletions(-) diff --git a/src/core/cuda/include/cuda_nonbonded_force.cuh b/src/core/cuda/include/cuda_nonbonded_force.cuh index 33fe5c6b..662c9532 100644 --- a/src/core/cuda/include/cuda_nonbonded_force.cuh +++ b/src/core/cuda/include/cuda_nonbonded_force.cuh @@ -96,4 +96,6 @@ class CudaNonbondedForce final : public NonbondedForce { std::unique_ptr> exact_source_atom_slots_; size_t exact_source_atom_capacity_ = 0; int n_exact_source_atoms_ = 0; + + std::unique_ptr> lrf_slot_to_group_range_; }; diff --git a/src/core/cuda/src/cuda_nonbonded_force.cu b/src/core/cuda/src/cuda_nonbonded_force.cu index 7f701d7f..2697feea 100644 --- a/src/core/cuda/src/cuda_nonbonded_force.cu +++ b/src/core/cuda/src/cuda_nonbonded_force.cu @@ -542,149 +542,6 @@ __device__ void nonbonded_force_calculation( } } -__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 count_exact_atom_tiles_kernel( int n_group_ranges, @@ -813,91 +670,6 @@ __global__ void fill_lrf_atom_csr_kernel(int n_group_ranges, const uint8_t* grou } } -__global__ void build_lrf_geometry_csr_kernel( - int n_group_ranges, - - const int* group_indices, - - const int* atom_offsets, - const int* source_atom_slots, - - const real_t* atom_charge, - - const real_t* cx, - const real_t* cy, - const real_t* cz, - const LrfCoefficients* coefficients, - - double* output_dx, - double* output_dy, - double* output_dz, - - double* output_q_r1, - double* output_q_r3, - double* output_q_r5, - double* output_q_r7 - -) { - const int target_range = blockIdx.x; - - if (target_range >= n_group_ranges) { - return; - } - - const int target_group = group_indices[target_range]; - - const coord_t target_center = coefficients[target_group].center; - - const int begin = atom_offsets[target_range]; - - const int end = atom_offsets[target_range + 1]; - - for (int entry = begin + threadIdx.x; entry < end; entry += blockDim.x) { - const int slot = source_atom_slots[entry]; - - const double charge = static_cast(atom_charge[slot]); - - const double x = static_cast(cx[slot]) - target_center.x; - - const double y = static_cast(cy[slot]) - target_center.y; - - const double z = static_cast(cz[slot]) - target_center.z; - - const double r2 = x * x + y * y + z * z; - - output_dx[entry] = x; - output_dy[entry] = y; - output_dz[entry] = z; - - if (r2 == 0.0) { - output_q_r1[entry] = 0.0; - output_q_r3[entry] = 0.0; - output_q_r5[entry] = 0.0; - output_q_r7[entry] = 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; - - output_q_r1[entry] = charge * inv_r; - - output_q_r3[entry] = charge * inv_r3; - - output_q_r5[entry] = charge * inv_r5; - - output_q_r7[entry] = charge * inv_r7; - } -} - __global__ void calc_lrf_kernel( int n_slots, @@ -1060,52 +832,6 @@ __global__ void nonbonded_kernel( 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 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_LRF) { - return; - } - - 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], - }; -} - __global__ void classify_group_pairs_by_switch_kernel( int n_groups_ranges, @@ -1399,6 +1125,261 @@ __global__ void fill_exact_atom_tiles_kernel( } } +constexpr int LRF_UNIQUE_COMPONENTS = 20; +constexpr int LRF_COEFFICIENT_THREADS = 256; + +__global__ void build_lrf_coefficients_dense_kernel( + int n_group_ranges, + int n_slots, + + const int* __restrict__ group_indices, + const int* __restrict__ group_start_idx, + const uint8_t* __restrict__ category, + const uint8_t* __restrict__ group_pair_modes, + const int* __restrict__ slot_to_group_range, + + const real_t* __restrict__ atom_charge, + const real_t* __restrict__ cx, + const real_t* __restrict__ cy, + const real_t* __restrict__ cz, + + LrfCoefficients* __restrict__ coefficients) { + constexpr int WARPS_PER_BLOCK = LRF_COEFFICIENT_THREADS / 32; + + constexpr unsigned FULL_MASK = 0xffffffffu; + + /* + * Exactly one block processes one target range. + */ + const int target_range = blockIdx.x; + + if (target_range >= n_group_ranges) { + return; + } + + const int target_start = group_start_idx[target_range]; + + const uint8_t target_category = category[target_start]; + + constexpr uint8_t P = static_cast(AtomCategory::P); + + constexpr uint8_t W = static_cast(AtomCategory::W); + + /* + * LRF is only evaluated for P and W atoms. + * This condition is uniform across the whole block. + */ + if (target_category != P && target_category != W) { + return; + } + + const int target_group = group_indices[target_range]; + + const coord_t center = coefficients[target_group].center; + + const int thread = threadIdx.x; + const int lane = thread & 31; + const int warp = thread >> 5; + + /* + * The 20 independent components are: + * + * 0: phi0 + * + * 1: x + * 2: y + * 3: z + * + * 4: xx + * 5: yy + * 6: zz + * 7: xy + * 8: xz + * 9: yz + * + * 10: xxx + * 11: yyy + * 12: zzz + * 13: xxy + * 14: xxz + * 15: xyy + * 16: yyz + * 17: xzz + * 18: yzz + * 19: xyz + */ + double sums[LRF_UNIQUE_COMPONENTS]; + +#pragma unroll + for (int component = 0; component < LRF_UNIQUE_COMPONENTS; ++component) { + sums[component] = 0.0; + } + + /* + * Threads collectively scan the complete contiguous source-slot + * array. Padding slots have source_range == -1. + */ + for (int source_slot = thread; source_slot < n_slots; source_slot += blockDim.x) { + const int source_range = slot_to_group_range[source_slot]; + + if (source_range < 0) { + continue; + } + + const int pair_index = get_pair_index(n_group_ranges, target_range, source_range); + + /* + * This preserves the original CSR semantics exactly: + * + * - GROUP_PAIR_LRF is included. + * - GROUP_PAIR_EXACT is excluded. + * - GROUP_PAIR_IGNORE is excluded. + */ + if (group_pair_modes[pair_index] != GROUP_PAIR_LRF) { + continue; + } + + const double charge = static_cast(atom_charge[source_slot]); + + const double x = static_cast(cx[source_slot]) - center.x; + + const double y = static_cast(cy[source_slot]) - center.y; + + const double z = static_cast(cz[source_slot]) - center.z; + + const double r2 = x * x + y * y + z * z; + + /* + * A self relationship should not be LRF, but guard against + * singular values in case classification changes later. + */ + if (r2 == 0.0) { + continue; + } + + const 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; + + const double q_r1 = charge * inv_r; + + const double q_r3 = charge * inv_r3; + + const double q_r5 = charge * inv_r5; + + const double q_r7 = charge * inv_r7; + + const double three_q_r5 = 3.0 * q_r5; + + const double nine_q_r5 = 9.0 * q_r5; + + const double fifteen_q_r7 = 15.0 * q_r7; + + const double xx = x * x; + const double yy = y * y; + const double zz = z * z; + + /* + * phi0 + */ + sums[0] += q_r1; + + /* + * phi1 + */ + sums[1] -= x * q_r3; + sums[2] -= y * q_r3; + sums[3] -= z * q_r3; + + /* + * Unique phi2 components. + */ + sums[4] += xx * three_q_r5 - q_r3; + + sums[5] += yy * three_q_r5 - q_r3; + + sums[6] += zz * three_q_r5 - q_r3; + + sums[7] += x * y * three_q_r5; + + sums[8] += x * z * three_q_r5; + + sums[9] += y * z * three_q_r5; + + /* + * Unique phi3 components. + */ + sums[10] += x * nine_q_r5 - x * xx * fifteen_q_r7; + + sums[11] += y * nine_q_r5 - y * yy * fifteen_q_r7; + + sums[12] += z * nine_q_r5 - z * zz * fifteen_q_r7; + + sums[13] += y * three_q_r5 - xx * y * fifteen_q_r7; + + sums[14] += z * three_q_r5 - xx * z * fifteen_q_r7; + + sums[15] += x * three_q_r5 - x * yy * fifteen_q_r7; + + sums[16] += z * three_q_r5 - yy * z * fifteen_q_r7; + + sums[17] += x * three_q_r5 - x * zz * fifteen_q_r7; + + sums[18] += y * three_q_r5 - y * zz * fifteen_q_r7; + + sums[19] -= x * y * z * fifteen_q_r7; + } + + /* + * Reduce each component within every warp. + * + * This is performed once after all source atoms have been + * processed, instead of once per source chunk. + */ +#pragma unroll + for (int component = 0; component < LRF_UNIQUE_COMPONENTS; ++component) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + sums[component] += __shfl_down_sync(FULL_MASK, sums[component], offset); + } + } + + __shared__ double warp_sums[LRF_UNIQUE_COMPONENTS][WARPS_PER_BLOCK]; + + if (lane == 0) { +#pragma unroll + for (int component = 0; component < LRF_UNIQUE_COMPONENTS; ++component) { + warp_sums[component][warp] = sums[component]; + } + } + + __syncthreads(); + + /* + * Warp zero reduces the eight warp results. + */ + if (warp == 0) { +#pragma unroll + for (int component = 0; component < LRF_UNIQUE_COMPONENTS; ++component) { + double value = lane < WARPS_PER_BLOCK ? warp_sums[component][lane] : 0.0; + +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value += __shfl_down_sync(FULL_MASK, value, offset); + } + + if (lane == 0) { + write_unique_lrf_component(coefficients[target_group], component, value); + } + } + } +} + } // namespace void CudaNonbondedForce::init_backend(Context& ctx) { @@ -1464,6 +1445,20 @@ void CudaNonbondedForce::init_backend(Context& ctx) { exact_entry_degrees_ = std::make_unique>(n_group_ranges + 1, false, true); exact_entry_offsets_ = std::make_unique>(n_group_ranges + 1, true, true); + + std::vector slot_to_group_range(data_.n_total, -1); + + for (int range = 0; range < n_group_ranges; ++range) { + const int begin = data_.group_start_idx->cpu_data_p[range]; + + const int size = data_.group_sizes->cpu_data_p[range]; + + for (int local_atom = 0; local_atom < size; ++local_atom) { + slot_to_group_range[begin + local_atom] = range; + } + } + + lrf_slot_to_group_range_ = HostDeviceBuffer::from_vector(slot_to_group_range, ctx.command_info.requested_gpu); } void CudaNonbondedForce::calc_all_direct_pairs(Context& ctx) { @@ -1513,32 +1508,6 @@ void CudaNonbondedForce::init_calculation_groups_by_switch(Context& ctx) { 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, - 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_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"); - } build_exact_atom_tiles(ctx); } @@ -1559,7 +1528,7 @@ void CudaNonbondedForce::init_calculation_groups(Context& ctx) { // todo: now alwasys use switch to test init_calculation_groups_by_switch(ctx); } - build_lrf_atom_csr(ctx); + // build_lrf_atom_csr(ctx); } void CudaNonbondedForce::calc_exact_tiles(Context& ctx) { @@ -1608,100 +1577,62 @@ void CudaNonbondedForce::calc_exact_tiles(Context& ctx) { 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); + const int n_group_ranges = static_cast(data_.group_indices->length); - check_cuda(cudaGetLastError()); + if (n_group_ranges <= 0) { + return; } - if (n_lrf_source_atom_entries_ > 0) { - constexpr int geometry_threads = 256; - - build_lrf_geometry_csr_kernel<<>>( - n_group_ranges, - - data_.group_indices->gpu_data_p, - - lrf_atom_offsets_->gpu_data_p, - lrf_source_atom_slots_->gpu_data_p, - - data_.atom_charge->gpu_data_p, - coord_x_->gpu_data_p, - coord_y_->gpu_data_p, - coord_z_->gpu_data_p, + /* + * Calculate target-group centers first. + */ + constexpr int center_threads = 256; + constexpr int center_warps_per_block = center_threads / 32; - lrf_coefficients_->gpu_data_p, + const int center_grid = (n_group_ranges + center_warps_per_block - 1) / center_warps_per_block; - lrf_dx_->gpu_data_p, - lrf_dy_->gpu_data_p, - lrf_dz_->gpu_data_p, + compute_lrf_centers_kernel<<>>( + n_group_ranges, - lrf_q_r1_->gpu_data_p, - lrf_q_r3_->gpu_data_p, - lrf_q_r5_->gpu_data_p, - lrf_q_r7_->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, - check_cuda(cudaGetLastError()); + coord_x_->gpu_data_p, + coord_y_->gpu_data_p, + coord_z_->gpu_data_p, - constexpr int coefficient_threads = 128; + lrf_coefficients_->gpu_data_p); -#define LRF_COEFFICIENT_ARGUMENTS \ - n_group_ranges, \ - data_.group_indices->gpu_data_p, \ - data_.group_start_idx->gpu_data_p, \ - data_.category->gpu_data_p, \ - lrf_atom_offsets_->gpu_data_p, \ - lrf_dx_->gpu_data_p, \ - lrf_dy_->gpu_data_p, \ - lrf_dz_->gpu_data_p, \ - lrf_q_r1_->gpu_data_p, \ - lrf_q_r3_->gpu_data_p, \ - lrf_q_r5_->gpu_data_p, \ - lrf_q_r7_->gpu_data_p, \ - lrf_coefficients_->gpu_data_p + check_cuda(cudaGetLastError()); - build_lrf_coefficients_order_csr_kernel<0, 1> - <<>>( - LRF_COEFFICIENT_ARGUMENTS); + /* + * Exactly one block per target range. + */ + build_lrf_coefficients_dense_kernel<<>>( + n_group_ranges, + data_.n_total, - check_cuda(cudaGetLastError()); + data_.group_indices->gpu_data_p, + data_.group_start_idx->gpu_data_p, + data_.category->gpu_data_p, - build_lrf_coefficients_order_csr_kernel<1, 3> - <<>>( - LRF_COEFFICIENT_ARGUMENTS); + group_pair_modes_->gpu_data_p, - check_cuda(cudaGetLastError()); + lrf_slot_to_group_range_ + ->gpu_data_p, - build_lrf_coefficients_order_csr_kernel<2, 6> - <<>>( - LRF_COEFFICIENT_ARGUMENTS); + data_.atom_charge->gpu_data_p, - check_cuda(cudaGetLastError()); + coord_x_->gpu_data_p, + coord_y_->gpu_data_p, + coord_z_->gpu_data_p, - build_lrf_coefficients_order_csr_kernel<3, 10> - <<>>( - LRF_COEFFICIENT_ARGUMENTS); + lrf_coefficients_->gpu_data_p); - check_cuda(cudaGetLastError()); -#undef LRF_COEFFICIENT_ARGUMENTS - } + check_cuda(cudaGetLastError()); } void CudaNonbondedForce::calc_lrf(Context& ctx) { From 836a488c51173bfbfe226b0a96cac19d98c3447b Mon Sep 17 00:00:00 2001 From: shen guo <1257960165@qq.com> Date: Tue, 15 Sep 2026 10:41:31 +0200 Subject: [PATCH 11/12] add z order when building atom list --- src/core/common/include/geometry.h | 39 +- src/core/common/src/nonbonded_force.cpp | 116 ++ .../cuda/include/cuda_nonbonded_force.cuh | 56 +- src/core/cuda/src/cuda_nonbonded_force.cu | 1018 ++++++----------- 4 files changed, 505 insertions(+), 724 deletions(-) diff --git a/src/core/common/include/geometry.h b/src/core/common/include/geometry.h index 3c4d0a5c..981a3c62 100644 --- a/src/core/common/include/geometry.h +++ b/src/core/common/include/geometry.h @@ -1,9 +1,8 @@ #pragma once #include -#include "md_types.h" #include "cuda_runtime_utility.h" - +#include "md_types.h" template HD inline Real3 operator-(const Real3& a, const Real3& b) { @@ -58,3 +57,39 @@ HD inline Real3 real3_cast(const Real3& a) { static_cast(a.y), static_cast(a.z)}; } + +inline uint32_t expand_morton_bits(uint32_t value) { + value &= 0x000003ffu; + value = (value | (value << 16)) & 0x030000ffu; + + value = (value | (value << 8)) & 0x0300f00fu; + + value = (value | (value << 4)) & 0x030c30c3u; + + value = (value | (value << 2)) & 0x09249249u; + + return value; +} + +inline uint32_t morton_code(uint32_t x, uint32_t y, uint32_t z) { + return expand_morton_bits(x) | (expand_morton_bits(y) << 1) | (expand_morton_bits(z) << 2); +} +inline uint32_t quantize_morton_coordinate(double value, double minimum, double maximum) { + const double extent = maximum - minimum; + + if (!(extent > 0.0)) { + return 0; + } + + double normalized = (value - minimum) / extent; + normalized = std::max(0.0, std::min(1.0, normalized)); + + return static_cast(normalized * 1023.0 + 0.5); +} + +inline uint32_t get_morton_code(const coord_t& p, double x_min, double x_max, double y_min, double y_max, double z_min, double z_max) { + auto x = quantize_morton_coordinate(p.x, x_min, x_max); + auto y = quantize_morton_coordinate(p.y, y_min, y_max); + auto z = quantize_morton_coordinate(p.z, z_min, z_max); + return morton_code(x, y, z); +} \ No newline at end of file diff --git a/src/core/common/src/nonbonded_force.cpp b/src/core/common/src/nonbonded_force.cpp index e9080b05..6e0ab2f2 100644 --- a/src/core/common/src/nonbonded_force.cpp +++ b/src/core/common/src/nonbonded_force.cpp @@ -5,8 +5,114 @@ #include #include "constants.h" +#include "geometry.h" #include "vdw_rules.h" +namespace { + +struct SpatialGroupEntry { + uint32_t morton; + int group; +}; + +coord_t group_spatial_position(const Context& ctx, int group_index, uint8_t group_category) { + const auto& config = ctx.charge_group_config; + + const auto& group = config.charge_groups[group_index]; + + const coord_t* coords = ctx.coords->cpu_data_p; + + constexpr uint8_t W = static_cast(AtomCategory::W); + + /* + * Switch mode always sorts using the switch atom. + * + * Water also uses its switch atom when + * iuse_switch_atom == 0, matching the CPU + * W-W and P-W distance definitions. + */ + if (config.iuse_switch_atom == 1 || group_category == W) { + const int switch_atom = group.iswitch - 1; + + return coords[switch_atom]; + } + + /* + * In all-atom mode, use the centroid for P and Q + * groups. This is only a spatial ordering key; it + * does not change group-pair classification. + */ + coord_t center{0.0, 0.0, 0.0}; + + for (int atom_1based : group.atoms) { + const int atom = atom_1based - 1; + + center.x += coords[atom].x; + center.y += coords[atom].y; + center.z += coords[atom].z; + } + + const double inverse_size = 1.0 / static_cast(group.atoms.size()); + + center.x *= inverse_size; + center.y *= inverse_size; + center.z *= inverse_size; + + return center; +} + +void spatially_sort_group_indices(const Context& ctx, uint8_t group_category, std::vector& group_indices) { + if (group_indices.size() < 2) { + return; + } + + std::vector positions; + positions.reserve(group_indices.size()); + + for (int group : group_indices) { + positions.push_back(group_spatial_position(ctx, group, group_category)); + } + + double minimum_x = positions[0].x; + double minimum_y = positions[0].y; + double minimum_z = positions[0].z; + + double maximum_x = positions[0].x; + double maximum_y = positions[0].y; + double maximum_z = positions[0].z; + + for (const coord_t& position : positions) { + minimum_x = std::min(minimum_x, position.x); + minimum_y = std::min(minimum_y, position.y); + minimum_z = std::min(minimum_z, position.z); + + maximum_x = std::max(maximum_x, position.x); + maximum_y = std::max(maximum_y, position.y); + maximum_z = std::max(maximum_z, position.z); + } + + std::vector entries; + entries.reserve(group_indices.size()); + + for (size_t i = 0; i < group_indices.size(); ++i) { + const coord_t& position = positions[i]; + entries.push_back({get_morton_code(position, minimum_x, maximum_x, minimum_y, maximum_y, minimum_z, maximum_z), group_indices[i]}); + } + + /* + * stable_sort keeps the original topology ordering + * for groups with identical Morton keys. + */ + std::stable_sort(entries.begin(), entries.end(), [](const SpatialGroupEntry& lhs, const SpatialGroupEntry& rhs) { + return lhs.morton < rhs.morton; + }); + + for (size_t i = 0; i < entries.size(); ++i) { + group_indices[i] = entries[i].group; + } +} +} // namespace + void NonbondedForce::init(Context& ctx) { build_combinded_list(ctx); build_charge_table(ctx); @@ -73,6 +179,16 @@ void NonbondedForce::build_combinded_list(Context& ctx) { category_groups[atom_type[atom]].push_back(i); } + if (ctx.command_info.requested_gpu) { + constexpr uint8_t P = static_cast(AtomCategory::P); + constexpr uint8_t Q = static_cast(AtomCategory::Q); + constexpr uint8_t W = static_cast(AtomCategory::W); + + spatially_sort_group_indices(ctx, P, category_groups[P]); + spatially_sort_group_indices(ctx, Q, category_groups[Q]); + spatially_sort_group_indices(ctx, W, category_groups[W]); + } + // P for (int i = 0; i < category_groups[0].size(); i++) { int group_idx = category_groups[0][i]; diff --git a/src/core/cuda/include/cuda_nonbonded_force.cuh b/src/core/cuda/include/cuda_nonbonded_force.cuh index 662c9532..115b9cb8 100644 --- a/src/core/cuda/include/cuda_nonbonded_force.cuh +++ b/src/core/cuda/include/cuda_nonbonded_force.cuh @@ -10,21 +10,7 @@ enum GroupPairMode : uint8_t { struct ExactEntry { int x_start; - int x_len; // <= 32 int y_start; - int y_len; // <= 32 - - uint8_t diagonal; - uint8_t y_indirect; - int y_type_slot; -}; - -struct LrfPairEntry { - int range1; - int range2; - - int group1; - int group2; }; class CudaNonbondedForce final : public NonbondedForce { @@ -37,12 +23,9 @@ class CudaNonbondedForce final : public NonbondedForce { private: 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); - void build_lrf_atom_csr(Context& ctx); void build_exact_atom_tiles(Context& ctx); std::unique_ptr> coord_x_, coord_y_, coord_z_; @@ -50,52 +33,17 @@ class CudaNonbondedForce final : public NonbondedForce { 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; - - std::unique_ptr> lrf_atom_degrees_; - - std::unique_ptr> lrf_atom_offsets_; - - std::unique_ptr> lrf_source_atom_slots_; - - std::unique_ptr> lrf_scan_temp_; - - size_t lrf_source_atom_capacity_ = 0; - size_t lrf_scan_temp_bytes_ = 0; - int n_lrf_source_atom_entries_ = 0; - - /* - * Geometry associated with every CSR source-atom entry. - */ - std::unique_ptr> lrf_dx_; - std::unique_ptr> lrf_dy_; - std::unique_ptr> lrf_dz_; - - std::unique_ptr> lrf_q_r1_; - std::unique_ptr> lrf_q_r3_; - std::unique_ptr> lrf_q_r5_; - std::unique_ptr> lrf_q_r7_; - - std::unique_ptr> exact_atom_degrees_; - std::unique_ptr> exact_atom_offsets_; - std::unique_ptr> exact_entry_degrees_; - std::unique_ptr> exact_entry_offsets_; - - std::unique_ptr> exact_source_atom_slots_; - size_t exact_source_atom_capacity_ = 0; - int n_exact_source_atoms_ = 0; std::unique_ptr> lrf_slot_to_group_range_; + + std::unique_ptr> exact_pair_masks_; }; diff --git a/src/core/cuda/src/cuda_nonbonded_force.cu b/src/core/cuda/src/cuda_nonbonded_force.cu index 2697feea..bc313e73 100644 --- a/src/core/cuda/src/cuda_nonbonded_force.cu +++ b/src/core/cuda/src/cuda_nonbonded_force.cu @@ -6,50 +6,32 @@ namespace { -__device__ __forceinline__ void emit_exact_packed_segment( - int segment_begin, - int segment_atoms, - int segment_type_slot, +__device__ __forceinline__ double combined_slot_distance2( + int slot1, + int slot2, - int target_start, - int target_size, - - int& entry_cursor, - ExactEntry* __restrict__ exact_entries) { - if (segment_atoms == 0) { - return; - } - - /* - * X is the fixed, contiguous target group. - * Y is the packed indirect atom list that rotates. - */ - for (int x_offset = 0; x_offset < target_size; x_offset += 32) { - const int x_len = min(32, target_size - x_offset); - - for (int y_offset = 0; y_offset < segment_atoms; y_offset += 32) { - const int y_len = min(32, segment_atoms - y_offset); - - ExactEntry entry{}; - - entry.x_start = target_start + x_offset; - entry.x_len = x_len; - - entry.y_start = segment_begin + y_offset; - entry.y_len = y_len; + const real_t* __restrict__ cx, + const real_t* __restrict__ cy, + const real_t* __restrict__ cz) { + const double dx = static_cast(cx[slot1]) - static_cast(cx[slot2]); + const double dy = static_cast(cy[slot1]) - static_cast(cy[slot2]); + const double dz = static_cast(cz[slot1]) - static_cast(cz[slot2]); - entry.y_type_slot = segment_type_slot; + return dx * dx + dy * dy + dz * dz; +} - entry.diagonal = 0; - entry.y_indirect = 1; +__device__ __forceinline__ double combined_slot_center_distance2( + int slot, + const coord_t center, - exact_entries[entry_cursor++] = entry; - } - } -} + const real_t* __restrict__ cx, + const real_t* __restrict__ cy, + const real_t* __restrict__ cz) { + const double dx = static_cast(cx[slot]) - center.x; + const double dy = static_cast(cy[slot]) - center.y; + const double dz = static_cast(cz[slot]) - center.z; -__device__ bool same_exact_energy_class(int slot1, int slot2, const uint8_t* category, const int* q_state) { - return category[slot1] == category[slot2] && q_state[slot1] == q_state[slot2]; + return dx * dx + dy * dy + dz * dz; } __device__ __forceinline__ void write_unique_lrf_component(LrfCoefficients& output, int phi, double value) { @@ -183,171 +165,6 @@ __device__ __forceinline__ void write_unique_lrf_component(LrfCoefficients& outp } } -template -__global__ void build_lrf_coefficients_order_csr_kernel( - int n_group_ranges, - - const int* __restrict__ group_indices, - const int* __restrict__ group_start_idx, - const uint8_t* __restrict__ category, - const int* __restrict__ atom_offsets, - - const double* __restrict__ dx, - const double* __restrict__ dy, - const double* __restrict__ dz, - - const double* __restrict__ q_r1, - const double* __restrict__ q_r3, - const double* __restrict__ q_r5, - const double* __restrict__ q_r7, - - LrfCoefficients* __restrict__ coefficients) { - constexpr int WARPS_PER_BLOCK = 4; - constexpr unsigned FULL_MASK = 0xffffffffu; - - const int target_range = blockIdx.x; - const int thread = threadIdx.x; - const int lane = thread & 31; - const int warp = thread >> 5; - - if (target_range >= n_group_ranges) { - return; - } - - const int target_start = group_start_idx[target_range]; - const uint8_t target_category = category[target_start]; - - constexpr uint8_t P = static_cast(AtomCategory::P); - constexpr uint8_t W = static_cast(AtomCategory::W); - - if (target_category != P && target_category != W) { - return; - } - - const int begin = atom_offsets[target_range]; - const int end = atom_offsets[target_range + 1]; - - double sums[COMPONENT_COUNT]; - -#pragma unroll - for (int component = 0; component < COMPONENT_COUNT; ++component) { - sums[component] = 0.0; - } - - for (int entry = begin + thread; entry < end; entry += blockDim.x) { - if constexpr (ORDER == 0) { - sums[0] += q_r1[entry]; - } - - if constexpr (ORDER == 1) { - const double x = dx[entry]; - const double y = dy[entry]; - const double z = dz[entry]; - const double qr3 = q_r3[entry]; - - sums[0] -= x * qr3; - sums[1] -= y * qr3; - sums[2] -= z * qr3; - } - - if constexpr (ORDER == 2) { - const double x = dx[entry]; - const double y = dy[entry]; - const double z = dz[entry]; - - const double qr3 = q_r3[entry]; - const double qr5 = q_r5[entry]; - const double three_qr5 = 3.0 * qr5; - - sums[0] += x * x * three_qr5 - qr3; - sums[1] += y * y * three_qr5 - qr3; - sums[2] += z * z * three_qr5 - qr3; - - sums[3] += x * y * three_qr5; - sums[4] += x * z * three_qr5; - sums[5] += y * z * three_qr5; - } - - if constexpr (ORDER == 3) { - const double x = dx[entry]; - const double y = dy[entry]; - const double z = dz[entry]; - - const double qr5 = q_r5[entry]; - const double qr7 = q_r7[entry]; - - const double three_qr5 = 3.0 * qr5; - const double nine_qr5 = 9.0 * qr5; - const double fifteen_qr7 = 15.0 * qr7; - - const double xx = x * x; - const double yy = y * y; - const double zz = z * z; - - sums[0] += x * nine_qr5 - x * xx * fifteen_qr7; - - sums[1] += y * nine_qr5 - y * yy * fifteen_qr7; - - sums[2] += z * nine_qr5 - z * zz * fifteen_qr7; - - sums[3] += y * three_qr5 - xx * y * fifteen_qr7; - - sums[4] += z * three_qr5 - xx * z * fifteen_qr7; - - sums[5] += x * three_qr5 - x * yy * fifteen_qr7; - - sums[6] += z * three_qr5 - yy * z * fifteen_qr7; - - sums[7] += x * three_qr5 - x * zz * fifteen_qr7; - - sums[8] += y * three_qr5 - y * zz * fifteen_qr7; - - sums[9] -= x * y * z * fifteen_qr7; - } - } - -#pragma unroll - for (int component = 0; component < COMPONENT_COUNT; ++component) { -#pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) { - sums[component] += __shfl_down_sync(FULL_MASK, sums[component], offset); - } - } - - __shared__ double warp_sums[COMPONENT_COUNT][WARPS_PER_BLOCK]; - - if (lane == 0) { -#pragma unroll - for (int component = 0; component < COMPONENT_COUNT; ++component) { - warp_sums[component][warp] = sums[component]; - } - } - - __syncthreads(); - - if (warp != 0) { - return; - } - -#pragma unroll - for (int component = 0; component < COMPONENT_COUNT; ++component) { - double block_sum = lane < WARPS_PER_BLOCK ? warp_sums[component][lane] : 0.0; - -#pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) { - block_sum += __shfl_down_sync(FULL_MASK, block_sum, offset); - } - - if (lane == 0) { - constexpr int first_component = ORDER == 0 ? 0 : ORDER == 1 ? 1 - : ORDER == 2 ? 4 - : 10; - - write_unique_lrf_component(coefficients[group_indices[target_range]], first_component + component, block_sum); - } - } -} - __device__ int get_pair_index( int n, int group1, @@ -450,6 +267,7 @@ __device__ void nonbonded_force_calculation( int x_idx, int y_idx, bool is_diag, + uint32_t pair_mask, int base_x, int base_y, @@ -501,8 +319,10 @@ __device__ void nonbonded_force_calculation( real_t3 atom2_force = {0, 0, 0}; real_t local_e_coul = 0, local_e_vdw = 0; + int y_lane = lane; for (int i = 0; i < 32; i++) { - if (!is_diag || atom1 < atom2) { + const bool mask_enabled = (pair_mask & (uint32_t{1} << y_lane)) != 0; + if (mask_enabled && (!is_diag || atom1 < atom2)) { compute_pair(atom1, atom1_type, atom1_state, atom1_charge, atom1_vdw, atom1_lambda, atom1_coord, atom2, atom2_type, atom2_state, atom2_charge, atom2_vdw, atom2_lambda, atom2_coord, n_atoms_solute, LJ_matrix, @@ -511,6 +331,7 @@ __device__ void nonbonded_force_calculation( local_e_coul, local_e_vdw); } shuffle(atom2, atom2_type, atom2_state, atom2_charge, atom2_vdw, atom2_lambda, atom2_force, atom2_coord); + y_lane = (y_lane + 1) & 31; } if (atom1 >= 0) { @@ -542,131 +363,84 @@ __device__ void nonbonded_force_calculation( } } -__global__ void count_exact_atom_tiles_kernel( +__global__ void build_exact_cluster_tiles_kernel( + int n_clusters, int n_group_ranges, + int max_exact_tiles, const uint8_t* group_pair_modes, - const int* group_start_idx, - const int* group_sizes, - const uint8_t* category, - const int* q_state, - - int* atom_degrees, - int* entry_degrees) { - const int target_range = blockIdx.x * blockDim.x + threadIdx.x; - - if (target_range >= n_group_ranges) { - return; - } + const int* slot_to_group_range, - const int target_size = group_sizes[target_range]; + int* exact_tile_count, + int* list_overflow, - /* - * The target group is the fixed, contiguous X axis. - */ - const int target_x_tiles = (target_size + 31) / 32; + ExactEntry* exact_entries, + uint32_t* exact_pair_masks) { + 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; - int atom_count = 0; - int entry_count = 0; + const int candidate_index = blockIdx.x * warps_per_block + warp_in_block; + const int candidate_count = n_clusters * (n_clusters + 1) / 2; - int segment_atoms = 0; - int segment_type_slot = -1; + if (candidate_index >= candidate_count) { + return; + } - /* - * source_range < target_range: - * only cross-group atoms are packed into the indirect Y list. - */ - for (int source_range = 0; source_range < target_range; ++source_range) { - const int pair_index = get_pair_index(n_group_ranges, source_range, target_range); + const int2 cluster_pair = get_tile_idx(n_clusters, candidate_index); + const int x_cluster = cluster_pair.x; + const int y_cluster = cluster_pair.y; - if (group_pair_modes[pair_index] != GROUP_PAIR_EXACT) { - continue; - } + const int x_start = x_cluster * 32; + const int y_start = y_cluster * 32; - const int source_start = group_start_idx[source_range]; + const int x_range = slot_to_group_range[x_start + lane]; - const int source_size = group_sizes[source_range]; + uint32_t pair_mask = 0; - /* - * Finish the current packed Y segment when its - * category or Q state changes. - */ - if (segment_atoms != 0 && !same_exact_energy_class(segment_type_slot, source_start, category, q_state)) { - const int segment_y_tiles = (segment_atoms + 31) / 32; + if (x_range >= 0) { +#pragma unroll + for (int y_lane = 0; y_lane < 32; y_lane++) { + if (x_cluster == y_cluster && y_lane <= lane) { + continue; + } - entry_count += target_x_tiles * segment_y_tiles; + const int y_range = slot_to_group_range[y_start + y_lane]; + if (y_range < 0) { + continue; + } - segment_atoms = 0; - segment_type_slot = -1; - } + const int group_pair_index = get_pair_index(n_group_ranges, x_range, y_range); - if (segment_atoms == 0) { - segment_type_slot = source_start; + if (group_pair_modes[group_pair_index] == GROUP_PAIR_EXACT) { + pair_mask |= uint32_t{1} << y_lane; + } } - - segment_atoms += source_size; - atom_count += source_size; - } - - /* - * Count the final packed Y segment. - */ - if (segment_atoms != 0) { - const int segment_y_tiles = (segment_atoms + 31) / 32; - - entry_count += target_x_tiles * segment_y_tiles; } - /* - * Count the target group's separate diagonal entries. - * These entries do not use the packed Y list. - */ - const int diagonal_pair = get_pair_index(n_group_ranges, target_range, target_range); - - if (group_pair_modes[diagonal_pair] == GROUP_PAIR_EXACT) { - entry_count += target_x_tiles * (target_x_tiles + 1) / 2; - } - - atom_degrees[target_range] = atom_count; - entry_degrees[target_range] = entry_count; -} + const bool tile_is_active = __any_sync(FULL_MASK, pair_mask != 0); -__global__ void count_lrf_atom_degree_kernel(int n_group_ranges, const uint8_t* group_pair_modes, const int* group_sizes, int* degrees) { - const int target_range = blockIdx.x * blockDim.x + threadIdx.x; - - if (target_range >= n_group_ranges) { + if (!tile_is_active) { return; } - int atom_count = 0; - 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) { - atom_count += group_sizes[source_range]; - } - } - degrees[target_range] = atom_count; -} + int output_index = -1; -__global__ void fill_lrf_atom_csr_kernel(int n_group_ranges, const uint8_t* group_pair_modes, const int* group_start_idx, const int* group_sizes, const int* offsets, int* source_atom_slots) { - const int target_range = blockIdx.x * blockDim.x + threadIdx.x; + if (lane == 0) { + output_index = atomicAdd(exact_tile_count, 1); - if (target_range >= n_group_ranges) { - return; + if (output_index >= max_exact_tiles) { + atomicExch(list_overflow, 1); + } else { + exact_entries[output_index] = {x_start, y_start}; + } } - int output = offsets[target_range]; + output_index = __shfl_sync(FULL_MASK, output_index, 0); - 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; - } - const int source_start = group_start_idx[source_range]; - const int source_size = group_sizes[source_range]; - for (int local_atom = 0; local_atom < source_size; local_atom++) { - source_atom_slots[output++] = source_start + local_atom; - } + if (output_index < max_exact_tiles) { + exact_pair_masks[static_cast(output_index) * 32 + lane] = pair_mask; } } @@ -828,12 +602,13 @@ __global__ void nonbonded_kernel( 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, + nonbonded_force_calculation(x_idx, y_idx, tile_x == tile_y, 0xffffffffu, base_x, base_y, n_states, n_atoms_solute, atom_idx, category, q_state, atom_lambdas, atom_charge, atom_vdw, LJ_matrix, el14_scale, coulomb_constant, vdw_rule, cx, cy, cz, dvelocities, e); } -__global__ void classify_group_pairs_by_switch_kernel( - int n_groups_ranges, +__global__ void classify_group_pairs_kernel( + int n_group_ranges, + int use_switch_atom, double solute_solute_cutoff2, double solute_solvent_cutoff2, @@ -842,61 +617,123 @@ __global__ void classify_group_pairs_by_switch_kernel( 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* __restrict__ group_start_idx, + const int* __restrict__ group_sizes, + + const uint8_t* __restrict__ category, + const int* __restrict__ q_state, + + const real_t* __restrict__ cx, + const real_t* __restrict__ cy, + const real_t* __restrict__ cz, + + uint8_t* __restrict__ group_pair_modes) { const int pair_index = blockIdx.x * blockDim.x + threadIdx.x; - const int total_pairs = n_groups_ranges * (n_groups_ranges + 1) / 2; + + const int total_pairs = n_group_ranges * (n_group_ranges + 1) / 2; if (pair_index >= total_pairs) { return; } - const int2 pair = get_tile_idx(n_groups_ranges, pair_index); - const int group1 = pair.x; - const int group2 = pair.y; + const int2 pair = get_tile_idx(n_group_ranges, pair_index); - 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 int range1 = pair.x; + const int range2 = pair.y; + + const int start1 = group_start_idx[range1]; + const int start2 = group_start_idx[range2]; + + const int size1 = group_sizes[range1]; + const int size2 = group_sizes[range2]; - const uint8_t category1 = category[switch_atom1]; - const uint8_t category2 = category[switch_atom2]; + const uint8_t category1 = category[start1]; + const uint8_t category2 = category[start2]; constexpr uint8_t P = static_cast(AtomCategory::P); + constexpr uint8_t Q = static_cast(AtomCategory::Q); + constexpr uint8_t W = static_cast(AtomCategory::W); - const bool group1_is_q = category1 == Q; - const bool group2_is_q = category2 == Q; + /* + * Default classification. + */ + group_pair_modes[pair_index] = GROUP_PAIR_IGNORE; + + /* + * A group interacts exactly with itself. + */ + if (range1 == range2) { + group_pair_modes[pair_index] = GROUP_PAIR_EXACT; + return; + } + + const bool range1_is_q = category1 == Q; + const bool range2_is_q = category2 == Q; - if (group1_is_q && group2_is_q) { - const int state1 = q_state[switch_atom1]; - const int state2 = q_state[switch_atom2]; - if (state1 == state2) { + /* + * Q-Q is exact only for the same duplicated Q state. + */ + if (range1_is_q && range2_is_q) { + if (q_state[start1] == q_state[start2]) { group_pair_modes[pair_index] = GROUP_PAIR_EXACT; } + return; } - 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); + /* + * Q-P and Q-W use the rcq rule. + * + * With switch atoms: + * test the environment switch atom. + * + * Without switch atoms: + * W still uses its switch atom; + * P is exact when any atom is inside rcq. + */ + if (range1_is_q || range2_is_q) { + const int environment_range = range1_is_q ? range2 : range1; + + const int environment_start = group_start_idx[environment_range]; + + const int environment_size = group_sizes[environment_range]; + + const uint8_t environment_category = category[environment_start]; + + bool inside_rcq = false; - if (environment_distance2 <= rcq2) { + if (use_switch_atom != 0 || environment_category == W) { + inside_rcq = combined_slot_center_distance2(environment_start, solute_center, cx, cy, cz) <= rcq2; + } else { + /* + * The environment is a P group. Check every atom, + * matching CpuNonbondedForce::inside_rcq(). + */ + for (int local_atom = 0; local_atom < environment_size; ++local_atom) { + const int slot = environment_start + local_atom; + + if (combined_slot_center_distance2(slot, solute_center, cx, cy, cz) <= rcq2) { + inside_rcq = true; + break; + } + } + } + + if (inside_rcq) { group_pair_modes[pair_index] = GROUP_PAIR_EXACT; } + return; } - // P-P, P-W, or W-W - const double group_distance2 = norm2(coords[atom_idx[switch_atom1]] - coords[atom_idx[switch_atom2]]); + /* + * At this point the pair is P-P, P-W, or W-W. + */ double normal_cutoff2; + if (category1 == P && category2 == P) { normal_cutoff2 = solute_solute_cutoff2; } else if (category1 == W && category2 == W) { @@ -905,9 +742,100 @@ __global__ void classify_group_pairs_by_switch_kernel( normal_cutoff2 = solute_solvent_cutoff2; } - if (group_distance2 <= normal_cutoff2) { - group_pair_modes[pair_index] = GROUP_PAIR_EXACT; - } else if (group_distance2 <= lrf_cutoff2) { + /* + * Switch-atom mode uses the first slot because the combined + * list always places the switch atom first in each group. + */ + if (use_switch_atom != 0) { + const double distance2 = combined_slot_distance2(start1, start2, cx, cy, cz); + + if (distance2 <= normal_cutoff2) { + group_pair_modes[pair_index] = GROUP_PAIR_EXACT; + } else if (distance2 <= lrf_cutoff2) { + group_pair_modes[pair_index] = GROUP_PAIR_LRF; + } + + return; + } + + /* + * Without switch atoms, W-W still uses switch atoms. + */ + if (category1 == W && category2 == W) { + const double distance2 = combined_slot_distance2(start1, start2, cx, cy, cz); + + if (distance2 <= normal_cutoff2) { + group_pair_modes[pair_index] = GROUP_PAIR_EXACT; + } else if (distance2 <= lrf_cutoff2) { + group_pair_modes[pair_index] = GROUP_PAIR_LRF; + } + + return; + } + + /* + * Without switch atoms, P-W uses the minimum distance from + * every atom in the P group to the W switch atom. + */ + if (category1 == W || category2 == W) { + const int solute_start = category1 == P ? start1 : start2; + + const int solute_size = category1 == P ? size1 : size2; + + const int water_switch = category1 == W ? start1 : start2; + + double minimum_distance2 = lrf_cutoff2 + 1.0; + + for (int local_atom = 0; local_atom < solute_size; ++local_atom) { + const double distance2 = combined_slot_distance2(solute_start + local_atom, water_switch, cx, cy, cz); + + /* + * Exact is the highest-priority classification, so + * stop as soon as any pair enters the exact cutoff. + */ + if (distance2 <= normal_cutoff2) { + group_pair_modes[pair_index] = GROUP_PAIR_EXACT; + return; + } + + if (distance2 < minimum_distance2) { + minimum_distance2 = distance2; + } + } + + if (minimum_distance2 <= lrf_cutoff2) { + group_pair_modes[pair_index] = GROUP_PAIR_LRF; + } + + return; + } + + /* + * P-P without switch atoms: + * minimum distance over every atom pair. + */ + double minimum_distance2 = lrf_cutoff2 + 1.0; + + for (int local1 = 0; local1 < size1; ++local1) { + const int slot1 = start1 + local1; + + for (int local2 = 0; local2 < size2; ++local2) { + const int slot2 = start2 + local2; + + const double distance2 = combined_slot_distance2(slot1, slot2, cx, cy, cz); + + if (distance2 <= normal_cutoff2) { + group_pair_modes[pair_index] = GROUP_PAIR_EXACT; + return; + } + + if (distance2 < minimum_distance2) { + minimum_distance2 = distance2; + } + } + } + + if (minimum_distance2 <= lrf_cutoff2) { group_pair_modes[pair_index] = GROUP_PAIR_LRF; } } @@ -915,7 +843,7 @@ __global__ void classify_group_pairs_by_switch_kernel( __global__ void exact_tiles_nonbonded_force_kernel( int n_exact_tiles, const ExactEntry* exact_entries, - const int* exact_source_atom_slots, + const uint32_t* exact_pair_masks, int n_states, // ctx.n_lambdas, used by nb_coul_slot int n_atoms_solute, // ctx.n_atoms_solute, water grouping + LJ_matrix row stride @@ -953,15 +881,16 @@ __global__ void exact_tiles_nonbonded_force_kernel( } const ExactEntry tile = exact_entries[tile_index]; + const uint32_t pair_mask = exact_pair_masks[static_cast(tile_index) * 32 + lane]; /* * Each lane loads one X atom and one Y atom. * Invalid lanes use atom == -1 and still participate in all shuffles. */ - const int x_idx = lane < tile.x_len ? tile.x_start + lane : -1; - const int y_idx = lane < tile.y_len ? tile.y_indirect == 1 ? exact_source_atom_slots[tile.y_start + lane] : tile.y_start + lane : -1; + const int x_idx = tile.x_start + lane; + const int y_idx = tile.y_start + lane; - nonbonded_force_calculation(x_idx, y_idx, tile.diagonal == 1, tile.x_start, tile.y_type_slot, n_states, n_atoms_solute, atom_idx, + nonbonded_force_calculation(x_idx, y_idx, false, pair_mask, tile.x_start, tile.y_start, n_states, n_atoms_solute, atom_idx, category, q_state, atom_lambdas, atom_charge, atom_vdw, LJ_matrix, el14_scale, coulomb_constant, vdw_rule, cx, cy, cz, dvelocities, e); } @@ -1033,98 +962,6 @@ __global__ void compute_lrf_centers_kernel( } } -__global__ void fill_exact_atom_tiles_kernel( - int n_group_ranges, - - const uint8_t* group_pair_modes, - const int* group_start_idx, - const int* group_sizes, - const uint8_t* category, - const int* q_state, - - const int* atom_offsets, - const int* entry_offsets, - - int* source_atom_slots, - ExactEntry* exact_entries) { - const int target_range = blockIdx.x * blockDim.x + threadIdx.x; - if (target_range >= n_group_ranges) { - return; - } - - const int target_start = group_start_idx[target_range]; - const int target_size = group_sizes[target_range]; - - int atom_cursor = atom_offsets[target_range]; - int entry_cursor = entry_offsets[target_range]; - - int segment_begin = atom_cursor; - int segment_atoms = 0; - int segment_type_slot = -1; - - for (int source_range = 0; source_range < target_range; source_range++) { - const int pair_index = get_pair_index(n_group_ranges, source_range, target_range); - - if (group_pair_modes[pair_index] != GROUP_PAIR_EXACT) { - continue; - } - - const int source_start = group_start_idx[source_range]; - - const int source_size = group_sizes[source_range]; - - if (segment_atoms != 0 && !same_exact_energy_class(segment_type_slot, source_start, category, q_state)) { - emit_exact_packed_segment(segment_begin, segment_atoms, segment_type_slot, target_start, target_size, entry_cursor, exact_entries); - segment_begin = atom_cursor; - segment_atoms = 0; - segment_type_slot = -1; - } - - if (segment_atoms == 0) { - segment_begin = atom_cursor; - segment_type_slot = source_start; - } - - for (int local_atom = 0; local_atom < source_size; local_atom++) { - source_atom_slots[atom_cursor++] = source_start + local_atom; - segment_atoms++; - } - } - - emit_exact_packed_segment(segment_begin, segment_atoms, segment_type_slot, target_start, target_size, entry_cursor, exact_entries); - - const int diagonal_pair = get_pair_index(n_group_ranges, target_range, target_range); - - if (group_pair_modes[diagonal_pair] == GROUP_PAIR_EXACT) { - const int n = (target_size + 31) / 32; - - for (int ix = 0; ix < n; ix++) { - const int x_offset = ix * 32; - const int x_len = min(32, target_size - x_offset); - - for (int iy = ix; iy < n; ++iy) { - const int y_offset = iy * 32; - const int y_len = min(32, target_size - y_offset); - - ExactEntry entry{}; - - entry.x_start = target_start + x_offset; - - entry.x_len = x_len; - entry.y_start = target_start + y_offset; - entry.y_type_slot = entry.y_start; - - entry.y_len = y_len; - - entry.diagonal = (ix == iy); - entry.y_indirect = 0; - - exact_entries[entry_cursor++] = entry; - } - } - } -} - constexpr int LRF_UNIQUE_COMPONENTS = 20; constexpr int LRF_COEFFICIENT_THREADS = 256; @@ -1392,60 +1229,29 @@ void CudaNonbondedForce::init_backend(Context& ctx) { 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; - } - } + + if ((data_.n_total & 31) != 0) { + throw std::runtime_error("CUDA exact cluster list requires n_total to be padded to 32"); } + const int n_exact_clusters = data_.n_total / 32; + + const size_t max_exact_tiles = static_cast(n_exact_clusters) * static_cast(n_exact_clusters + 1) / 2; + exact_tile_capacity_ = max_exact_tiles; - 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_pair_masks_ = std::make_unique>(exact_tile_capacity_ * 32, 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); - lrf_atom_degrees_ = std::make_unique>(n_group_ranges + 1, false, true); - - lrf_atom_offsets_ = std::make_unique>(n_group_ranges + 1, true, true); - - check_cuda(cub::DeviceScan::ExclusiveSum( - nullptr, - lrf_scan_temp_bytes_, - lrf_atom_degrees_->gpu_data_p, - lrf_atom_offsets_->gpu_data_p, - n_group_ranges + 1)); - - lrf_scan_temp_ = std::make_unique>(lrf_scan_temp_bytes_, false, true); - - exact_atom_degrees_ = std::make_unique>(n_group_ranges + 1, false, true); - - exact_atom_offsets_ = std::make_unique>(n_group_ranges + 1, true, true); - - exact_entry_degrees_ = std::make_unique>(n_group_ranges + 1, false, true); - - exact_entry_offsets_ = std::make_unique>(n_group_ranges + 1, true, true); - std::vector slot_to_group_range(data_.n_total, -1); for (int range = 0; range < n_group_ranges; ++range) { @@ -1459,6 +1265,7 @@ void CudaNonbondedForce::init_backend(Context& ctx) { } lrf_slot_to_group_range_ = HostDeviceBuffer::from_vector(slot_to_group_range, ctx.command_info.requested_gpu); + } void CudaNonbondedForce::calc_all_direct_pairs(Context& ctx) { @@ -1477,58 +1284,60 @@ void CudaNonbondedForce::calc_all_direct_pairs(Context& ctx) { 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(); - +void CudaNonbondedForce::init_calculation_groups(Context& ctx) { const double solute_solute_cutoff2 = ctx.md.solute_solute * ctx.md.solute_solute; + const double solute_solvent_cutoff2 = ctx.md.solute_solvent * ctx.md.solute_solvent; + const double solvent_solvent_cutoff2 = ctx.md.solvent_solvent * ctx.md.solvent_solvent; + const double rcq2 = ctx.md.q_atom * ctx.md.q_atom; + const double lrf_cutoff2 = ctx.md.lrf_cutoff * ctx.md.lrf_cutoff; - const int 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; + const int n_group_ranges = static_cast(data_.group_indices->length); - 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()); + if (n_group_ranges == 0) { + n_exact_tiles_ = 0; + return; + } - build_exact_atom_tiles(ctx); -} + constexpr int threads = 256; -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"); -} + const int total_pairs = n_group_ranges * (n_group_ranges + 1) / 2; -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); + const int blocks = (total_pairs + threads - 1) / threads; - } 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); - } - // build_lrf_atom_csr(ctx); + int use_switch_atom = ctx.charge_group_config.iuse_switch_atom == 1 ? 1 : 0; + use_switch_atom = 1; + + classify_group_pairs_kernel<<>>( + n_group_ranges, + use_switch_atom, + + solute_solute_cutoff2, + solute_solvent_cutoff2, + solvent_solvent_cutoff2, + rcq2, + lrf_cutoff2, + + ctx.topo.solute_center, + + data_.group_start_idx->gpu_data_p, + data_.group_sizes->gpu_data_p, + + data_.category->gpu_data_p, + data_.q_state->gpu_data_p, + + coord_x_->gpu_data_p, + coord_y_->gpu_data_p, + coord_z_->gpu_data_p, + + group_pair_modes_->gpu_data_p); + + check_cuda(cudaGetLastError()); + + build_exact_atom_tiles(ctx); } void CudaNonbondedForce::calc_exact_tiles(Context& ctx) { @@ -1544,8 +1353,7 @@ void CudaNonbondedForce::calc_exact_tiles(Context& ctx) { exact_tiles_nonbonded_force_kernel<<>>( n_exact_tiles_, exact_tiles_->gpu_data_p, - - exact_source_atom_slots_->gpu_data_p, + exact_pair_masks_->gpu_data_p, ctx.n_lambdas(), ctx.n_atoms_solute, @@ -1668,181 +1476,55 @@ void CudaNonbondedForce::calc_lrf(Context& ctx) { check_cuda(cudaGetLastError()); } -void CudaNonbondedForce::build_lrf_atom_csr(Context& ctx) { - const int n_group_ranges = data_.group_indices->length; - - n_lrf_source_atom_entries_ = 0; - - lrf_atom_degrees_->zero(); - const int thread_num = 256; - const int grid_size = (n_group_ranges + thread_num - 1) / thread_num; - - count_lrf_atom_degree_kernel<<>>(n_group_ranges, group_pair_modes_->gpu_data_p, data_.group_sizes->gpu_data_p, lrf_atom_degrees_->gpu_data_p); - - check_cuda(cudaGetLastError()); - - /* - * offsets[0] = 0 - * offsets[t+1] = offsets[t] + degrees[t] - * offsets[n] = total CSR entries - */ - - check_cuda(cub::DeviceScan::ExclusiveSum( - lrf_scan_temp_->gpu_data_p, - lrf_scan_temp_bytes_, - lrf_atom_degrees_->gpu_data_p, - lrf_atom_offsets_->gpu_data_p, - n_group_ranges + 1)); +void CudaNonbondedForce::build_exact_atom_tiles(Context& ctx) { + n_exact_tiles_ = 0; - lrf_atom_offsets_->download(); - n_lrf_source_atom_entries_ = lrf_atom_offsets_->cpu_data_p[n_group_ranges]; + const int n_group_ranges = static_cast(data_.group_indices->length); - if (n_lrf_source_atom_entries_ < 0) { - throw std::runtime_error("Negative CUDA LRF CSR entry count"); - } + const int n_clusters = data_.n_total / 32; - if (n_lrf_source_atom_entries_ == 0) { + if (n_group_ranges == 0 || n_clusters == 0) { return; } - const size_t required_capacity = static_cast(n_lrf_source_atom_entries_); - - if (!lrf_source_atom_slots_ || required_capacity > lrf_source_atom_capacity_) { - size_t new_capacity = required_capacity; - - if (lrf_source_atom_capacity_ != 0) { - const size_t grown_capacity = lrf_source_atom_capacity_ + lrf_source_atom_capacity_ / 2; - - if (grown_capacity > new_capacity) { - new_capacity = grown_capacity; - } - } - - lrf_source_atom_slots_ = std::make_unique>(new_capacity, false, true); - - lrf_dx_ = std::make_unique>(new_capacity, false, true); - - lrf_dy_ = std::make_unique>(new_capacity, false, true); - - lrf_dz_ = std::make_unique>(new_capacity, false, true); - - lrf_q_r1_ = std::make_unique>(new_capacity, false, true); + exact_tile_count_->zero(); + list_overflow_->zero(); - lrf_q_r3_ = std::make_unique>(new_capacity, false, true); + constexpr int threads = 256; + constexpr int warps_per_block = threads / 32; - lrf_q_r5_ = std::make_unique>(new_capacity, false, true); + const int candidate_count = n_clusters * (n_clusters + 1) / 2; - lrf_q_r7_ = std::make_unique>(new_capacity, false, true); + const int blocks = (candidate_count + warps_per_block - 1) / warps_per_block; - lrf_source_atom_capacity_ = new_capacity; - } - fill_lrf_atom_csr_kernel<<>>( + build_exact_cluster_tiles_kernel<<>>( + n_clusters, n_group_ranges, - group_pair_modes_->gpu_data_p, - data_.group_start_idx->gpu_data_p, - data_.group_sizes->gpu_data_p, - lrf_atom_offsets_->gpu_data_p, - lrf_source_atom_slots_->gpu_data_p); + static_cast(exact_tile_capacity_), - check_cuda(cudaGetLastError()); -} - -void CudaNonbondedForce::build_exact_atom_tiles(Context& ctx) { - const int n_group_ranges = data_.group_indices->length; - - n_exact_tiles_ = 0; - n_exact_source_atoms_ = 0; - if (n_group_ranges == 0) return; - - exact_atom_degrees_->zero(); - exact_atom_offsets_->zero(); - - exact_entry_degrees_->zero(); - exact_entry_offsets_->zero(); + group_pair_modes_->gpu_data_p, + lrf_slot_to_group_range_->gpu_data_p, - const int threads = 256; + exact_tile_count_->gpu_data_p, + list_overflow_->gpu_data_p, - const int blocks = (n_group_ranges + threads - 1) / threads; + exact_tiles_->gpu_data_p, + exact_pair_masks_->gpu_data_p); - count_exact_atom_tiles_kernel<<>>(n_group_ranges, - group_pair_modes_->gpu_data_p, - data_.group_start_idx->gpu_data_p, - data_.group_sizes->gpu_data_p, - data_.category->gpu_data_p, - data_.q_state->gpu_data_p, - exact_atom_degrees_->gpu_data_p, - exact_entry_degrees_->gpu_data_p); check_cuda(cudaGetLastError()); - check_cuda(cub::DeviceScan::ExclusiveSum( - lrf_scan_temp_->gpu_data_p, - lrf_scan_temp_bytes_, - - exact_atom_degrees_->gpu_data_p, - exact_atom_offsets_->gpu_data_p, - - n_group_ranges + 1)); - - check_cuda(cub::DeviceScan::ExclusiveSum( - lrf_scan_temp_->gpu_data_p, - lrf_scan_temp_bytes_, - - exact_entry_degrees_->gpu_data_p, - exact_entry_offsets_->gpu_data_p, + exact_tile_count_->download(); + list_overflow_->download(); - n_group_ranges + 1)); - - exact_atom_offsets_->download(); - exact_entry_offsets_->download(); - - n_exact_source_atoms_ = exact_atom_offsets_->cpu_data_p[n_group_ranges]; - - n_exact_tiles_ = exact_entry_offsets_->cpu_data_p[n_group_ranges]; - - if (n_exact_source_atoms_ < 0 || n_exact_tiles_ < 0) { - throw std::runtime_error("Negative CUDA exact-list size"); + if (list_overflow_->cpu_data_p[0] != 0) { + throw std::runtime_error("CUDA exact cluster tile capacity exceeded"); } - if (static_cast(n_exact_tiles_) > exact_tile_capacity_) { - throw std::runtime_error("CUDA packed exact-tile capacity exceeded"); - } - - const size_t required_atom_capacity = static_cast(n_exact_source_atoms_); - - if (required_atom_capacity > exact_source_atom_capacity_) { - size_t new_capacity = required_atom_capacity; - - if (exact_source_atom_capacity_ != 0) { - const size_t grown_capacity = exact_source_atom_capacity_ + exact_source_atom_capacity_ / 2; - - new_capacity = std::max(new_capacity, grown_capacity); - } - - exact_source_atom_slots_ = std::make_unique>(new_capacity, false, true); - - exact_source_atom_capacity_ = new_capacity; - } + n_exact_tiles_ = exact_tile_count_->cpu_data_p[0]; - if (n_exact_tiles_ == 0) { - return; + if (n_exact_tiles_ < 0 || static_cast(n_exact_tiles_) > exact_tile_capacity_) { + throw std::runtime_error("Invalid CUDA exact cluster tile count"); } - - fill_exact_atom_tiles_kernel<<>>( - n_group_ranges, - - group_pair_modes_->gpu_data_p, - data_.group_start_idx->gpu_data_p, - data_.group_sizes->gpu_data_p, - data_.category->gpu_data_p, - data_.q_state->gpu_data_p, - - exact_atom_offsets_->gpu_data_p, - exact_entry_offsets_->gpu_data_p, - - exact_source_atom_slots_->gpu_data_p, - exact_tiles_->gpu_data_p); - - check_cuda(cudaGetLastError()); } void CudaNonbondedForce::calc(Context& ctx) { From a2532c8897c81cdc38a2af0d2210c6d0058621d7 Mon Sep 17 00:00:00 2001 From: shen guo <1257960165@qq.com> Date: Tue, 15 Sep 2026 14:36:59 +0200 Subject: [PATCH 12/12] add 14 mask --- .../cuda/include/cuda_nonbonded_force.cuh | 1 + src/core/cuda/src/cuda_nonbonded_force.cu | 218 ++++++++++++------ 2 files changed, 143 insertions(+), 76 deletions(-) diff --git a/src/core/cuda/include/cuda_nonbonded_force.cuh b/src/core/cuda/include/cuda_nonbonded_force.cuh index 115b9cb8..a528de64 100644 --- a/src/core/cuda/include/cuda_nonbonded_force.cuh +++ b/src/core/cuda/include/cuda_nonbonded_force.cuh @@ -46,4 +46,5 @@ class CudaNonbondedForce final : public NonbondedForce { std::unique_ptr> lrf_slot_to_group_range_; std::unique_ptr> exact_pair_masks_; + std::unique_ptr> exact_pair_14_masks_; }; diff --git a/src/core/cuda/src/cuda_nonbonded_force.cu b/src/core/cuda/src/cuda_nonbonded_force.cu index bc313e73..b0dd2c0c 100644 --- a/src/core/cuda/src/cuda_nonbonded_force.cu +++ b/src/core/cuda/src/cuda_nonbonded_force.cu @@ -6,30 +6,30 @@ namespace { -__device__ __forceinline__ double combined_slot_distance2( +__device__ __forceinline__ real_t combined_slot_distance2( int slot1, int slot2, const real_t* __restrict__ cx, const real_t* __restrict__ cy, const real_t* __restrict__ cz) { - const double dx = static_cast(cx[slot1]) - static_cast(cx[slot2]); - const double dy = static_cast(cy[slot1]) - static_cast(cy[slot2]); - const double dz = static_cast(cz[slot1]) - static_cast(cz[slot2]); + const real_t dx = cx[slot1] - cx[slot2]; + const real_t dy = cy[slot1] - cy[slot2]; + const real_t dz = cz[slot1] - cz[slot2]; return dx * dx + dy * dy + dz * dz; } -__device__ __forceinline__ double combined_slot_center_distance2( +__device__ __forceinline__ real_t combined_slot_center_distance2( int slot, const coord_t center, const real_t* __restrict__ cx, const real_t* __restrict__ cy, const real_t* __restrict__ cz) { - const double dx = static_cast(cx[slot]) - center.x; - const double dy = static_cast(cy[slot]) - center.y; - const double dz = static_cast(cz[slot]) - center.z; + const real_t dx = cx[slot] - center.x; + const real_t dy = cy[slot] - center.y; + const real_t dz = cz[slot] - center.z; return dx * dx + dy * dy + dz * dz; } @@ -188,6 +188,7 @@ __device__ int2 get_tile_idx(int n, int t) { return {x, y}; } +template __device__ void compute_pair( // atom1 int atom1, uint8_t atom1_type, int atom1_state, @@ -201,14 +202,16 @@ __device__ void compute_pair( real_t el14_scale, real_t coulomb_constant, int vdw_rule, // output real_t3& atom1_force, real_t3& atom2_force, - real_t& e_coul, real_t& e_vdw) { - if (atom1 == -1 || atom2 == -1 || atom1 == atom2) return; - auto bond_type = get_bond_type(n_atoms_solute, LJ_matrix, atom1, atom1_type, atom2, atom2_type); - if (bond_type == BondType::Bond23) return; - - constexpr uint8_t Q = static_cast(AtomCategory::Q); - if (atom1_type == Q && atom2_type == Q && atom1_state != atom2_state) return; - + real_t& e_coul, real_t& e_vdw, bool is_14 = false) { + if constexpr (!BONDS_PRECLASSIFIED) { + if (atom1 == -1 || atom2 == -1 || atom1 == atom2) return; + auto bond_type = get_bond_type(n_atoms_solute, LJ_matrix, atom1, atom1_type, atom2, atom2_type); + if (bond_type == BondType::Bond23) return; + + constexpr uint8_t Q = static_cast(AtomCategory::Q); + if (atom1_type == Q && atom2_type == Q && atom1_state != atom2_state) return; + is_14 = bond_type == BondType::Bond14; + } real_t dx = atom2_coord.x - atom1_coord.x; real_t dy = atom2_coord.y - atom1_coord.y; real_t dz = atom2_coord.z - atom1_coord.z; @@ -216,7 +219,6 @@ __device__ void compute_pair( real_t inv_dis2 = static_cast(1.0) / dis2; real_t inv_dis = sqrt(inv_dis2); - bool is_14 = (bond_type == BondType::Bond14); real_t qij = atom1_charge * atom2_charge; real_t scaling = is_14 ? el14_scale : 1; real_t2 pair = (is_14) ? combine_vdw(vdw_rule, atom1_vdw.aii_14, atom1_vdw.bii_14, atom2_vdw.aii_14, atom2_vdw.bii_14) : combine_vdw(vdw_rule, atom1_vdw.aii_normal, atom1_vdw.bii_normal, atom2_vdw.aii_normal, atom2_vdw.bii_normal); @@ -263,11 +265,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); } +template __device__ void nonbonded_force_calculation( int x_idx, int y_idx, bool is_diag, uint32_t pair_mask, + uint32_t pair_14_mask, int base_x, int base_y, @@ -321,14 +325,16 @@ __device__ void nonbonded_force_calculation( real_t local_e_coul = 0, local_e_vdw = 0; int y_lane = lane; for (int i = 0; i < 32; i++) { - const bool mask_enabled = (pair_mask & (uint32_t{1} << y_lane)) != 0; + const uint32_t bit = uint32_t{1} << y_lane; + const bool mask_enabled = (pair_mask & bit) != 0; if (mask_enabled && (!is_diag || atom1 < atom2)) { - compute_pair(atom1, atom1_type, atom1_state, atom1_charge, atom1_vdw, atom1_lambda, atom1_coord, - atom2, atom2_type, atom2_state, atom2_charge, atom2_vdw, atom2_lambda, atom2_coord, - n_atoms_solute, LJ_matrix, - el14_scale, coulomb_constant, vdw_rule, - atom1_force, atom2_force, - local_e_coul, local_e_vdw); + const bool is_14 = (pair_14_mask & bit) != 0; + compute_pair(atom1, atom1_type, atom1_state, atom1_charge, atom1_vdw, atom1_lambda, atom1_coord, + atom2, atom2_type, atom2_state, atom2_charge, atom2_vdw, atom2_lambda, atom2_coord, + n_atoms_solute, LJ_matrix, + el14_scale, coulomb_constant, vdw_rule, + atom1_force, atom2_force, + local_e_coul, local_e_vdw, is_14); } shuffle(atom2, atom2_type, atom2_state, atom2_charge, atom2_vdw, atom2_lambda, atom2_force, atom2_coord); y_lane = (y_lane + 1) & 31; @@ -371,57 +377,103 @@ __global__ void build_exact_cluster_tiles_kernel( const uint8_t* group_pair_modes, const int* slot_to_group_range, + int n_atoms_solute, + const int* atom_idx, + const uint8_t* category, + const int* q_state, + const int* LJ_matrix, + int* exact_tile_count, int* list_overflow, ExactEntry* exact_entries, - uint32_t* exact_pair_masks) { + uint32_t* exact_pair_masks, + uint32_t* exact_pair_14_masks) { constexpr unsigned FULL_MASK = 0xffffffffu; + constexpr uint8_t Q = static_cast(AtomCategory::Q); + const int lane = threadIdx.x & 31; const int warp_in_block = threadIdx.x >> 5; const int warps_per_block = blockDim.x >> 5; const int candidate_index = blockIdx.x * warps_per_block + warp_in_block; + const int candidate_count = n_clusters * (n_clusters + 1) / 2; + // Uniform return for the entire warp. if (candidate_index >= candidate_count) { return; } const int2 cluster_pair = get_tile_idx(n_clusters, candidate_index); + const int x_cluster = cluster_pair.x; const int y_cluster = cluster_pair.y; const int x_start = x_cluster * 32; const int y_start = y_cluster * 32; - const int x_range = slot_to_group_range[x_start + lane]; + const int x_slot = x_start + lane; + const int x_range = slot_to_group_range[x_slot]; + const int x_atom = atom_idx[x_slot]; uint32_t pair_mask = 0; + uint32_t pair_14_mask = 0; + + if (x_range >= 0 && x_atom >= 0) { + const uint8_t x_category = category[x_slot]; + const int x_state = q_state[x_slot]; - if (x_range >= 0) { #pragma unroll - for (int y_lane = 0; y_lane < 32; y_lane++) { + for (int y_lane = 0; y_lane < 32; ++y_lane) { + // Unique ownership within a diagonal tile. if (x_cluster == y_cluster && y_lane <= lane) { continue; } - const int y_range = slot_to_group_range[y_start + y_lane]; + const int y_slot = y_start + y_lane; + const int y_range = slot_to_group_range[y_slot]; + if (y_range < 0) { continue; } const int group_pair_index = get_pair_index(n_group_ranges, x_range, y_range); - if (group_pair_modes[group_pair_index] == GROUP_PAIR_EXACT) { - pair_mask |= uint32_t{1} << y_lane; + if (group_pair_modes[group_pair_index] != GROUP_PAIR_EXACT) { + continue; + } + + const int y_atom = atom_idx[y_slot]; + + if (y_atom < 0 || x_atom == y_atom) { + continue; + } + + const uint8_t y_category = category[y_slot]; + + if (x_category == Q && y_category == Q && x_state != q_state[y_slot]) { + continue; + } + + const BondType bond_type = get_bond_type(n_atoms_solute, LJ_matrix, x_atom, x_category, y_atom, y_category); + + if (bond_type == BondType::Bond23) { + continue; + } + + const uint32_t bit = uint32_t{1} << y_lane; + + pair_mask |= bit; + + if (bond_type == BondType::Bond14) { + pair_14_mask |= bit; } } } - const bool tile_is_active = __any_sync(FULL_MASK, pair_mask != 0); - - if (!tile_is_active) { + // Test activity after removing exclusions. + if (!__any_sync(FULL_MASK, pair_mask != 0)) { return; } @@ -433,14 +485,18 @@ __global__ void build_exact_cluster_tiles_kernel( if (output_index >= max_exact_tiles) { atomicExch(list_overflow, 1); } else { - exact_entries[output_index] = {x_start, y_start}; + exact_entries[output_index] = { + x_start, y_start}; } } output_index = __shfl_sync(FULL_MASK, output_index, 0); if (output_index < max_exact_tiles) { - exact_pair_masks[static_cast(output_index) * 32 + lane] = pair_mask; + const size_t mask_index = static_cast(output_index) * 32 + lane; + + exact_pair_masks[mask_index] = pair_mask; + exact_pair_14_masks[mask_index] = pair_14_mask; } } @@ -602,19 +658,19 @@ __global__ void nonbonded_kernel( x_idx = x_idx < sz ? x_idx : -1; y_idx = y_idx < sz ? y_idx : -1; - nonbonded_force_calculation(x_idx, y_idx, tile_x == tile_y, 0xffffffffu, 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); + nonbonded_force_calculation(x_idx, y_idx, tile_x == tile_y, 0xffffffffu, 0, base_x, base_y, n_states, n_atoms_solute, + atom_idx, category, q_state, atom_lambdas, atom_charge, atom_vdw, LJ_matrix, el14_scale, coulomb_constant, vdw_rule, cx, cy, cz, dvelocities, e); } __global__ void classify_group_pairs_kernel( int n_group_ranges, int use_switch_atom, - double solute_solute_cutoff2, - double solute_solvent_cutoff2, - double solvent_solvent_cutoff2, - double rcq2, - double lrf_cutoff2, + real_t solute_solute_cutoff2, + real_t solute_solvent_cutoff2, + real_t solvent_solvent_cutoff2, + real_t rcq2, + real_t lrf_cutoff2, const coord_t solute_center, @@ -732,7 +788,7 @@ __global__ void classify_group_pairs_kernel( /* * At this point the pair is P-P, P-W, or W-W. */ - double normal_cutoff2; + real_t normal_cutoff2; if (category1 == P && category2 == P) { normal_cutoff2 = solute_solute_cutoff2; @@ -747,7 +803,7 @@ __global__ void classify_group_pairs_kernel( * list always places the switch atom first in each group. */ if (use_switch_atom != 0) { - const double distance2 = combined_slot_distance2(start1, start2, cx, cy, cz); + const real_t distance2 = combined_slot_distance2(start1, start2, cx, cy, cz); if (distance2 <= normal_cutoff2) { group_pair_modes[pair_index] = GROUP_PAIR_EXACT; @@ -762,7 +818,7 @@ __global__ void classify_group_pairs_kernel( * Without switch atoms, W-W still uses switch atoms. */ if (category1 == W && category2 == W) { - const double distance2 = combined_slot_distance2(start1, start2, cx, cy, cz); + const real_t distance2 = combined_slot_distance2(start1, start2, cx, cy, cz); if (distance2 <= normal_cutoff2) { group_pair_modes[pair_index] = GROUP_PAIR_EXACT; @@ -784,10 +840,10 @@ __global__ void classify_group_pairs_kernel( const int water_switch = category1 == W ? start1 : start2; - double minimum_distance2 = lrf_cutoff2 + 1.0; + real_t minimum_distance2 = lrf_cutoff2 + 1.0; for (int local_atom = 0; local_atom < solute_size; ++local_atom) { - const double distance2 = combined_slot_distance2(solute_start + local_atom, water_switch, cx, cy, cz); + const real_t distance2 = combined_slot_distance2(solute_start + local_atom, water_switch, cx, cy, cz); /* * Exact is the highest-priority classification, so @@ -814,7 +870,7 @@ __global__ void classify_group_pairs_kernel( * P-P without switch atoms: * minimum distance over every atom pair. */ - double minimum_distance2 = lrf_cutoff2 + 1.0; + real_t minimum_distance2 = lrf_cutoff2 + 1.0; for (int local1 = 0; local1 < size1; ++local1) { const int slot1 = start1 + local1; @@ -822,7 +878,7 @@ __global__ void classify_group_pairs_kernel( for (int local2 = 0; local2 < size2; ++local2) { const int slot2 = start2 + local2; - const double distance2 = combined_slot_distance2(slot1, slot2, cx, cy, cz); + const real_t distance2 = combined_slot_distance2(slot1, slot2, cx, cy, cz); if (distance2 <= normal_cutoff2) { group_pair_modes[pair_index] = GROUP_PAIR_EXACT; @@ -844,6 +900,7 @@ __global__ void exact_tiles_nonbonded_force_kernel( int n_exact_tiles, const ExactEntry* exact_entries, const uint32_t* exact_pair_masks, + const uint32_t* exact_pair_14_masks, int n_states, // ctx.n_lambdas, used by nb_coul_slot int n_atoms_solute, // ctx.n_atoms_solute, water grouping + LJ_matrix row stride @@ -882,6 +939,7 @@ __global__ void exact_tiles_nonbonded_force_kernel( const ExactEntry tile = exact_entries[tile_index]; const uint32_t pair_mask = exact_pair_masks[static_cast(tile_index) * 32 + lane]; + const uint32_t pair_14_mask = exact_pair_14_masks[static_cast(tile_index) * 32 + lane]; /* * Each lane loads one X atom and one Y atom. @@ -890,8 +948,8 @@ __global__ void exact_tiles_nonbonded_force_kernel( const int x_idx = tile.x_start + lane; const int y_idx = tile.y_start + lane; - nonbonded_force_calculation(x_idx, y_idx, false, pair_mask, 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); + nonbonded_force_calculation(x_idx, y_idx, false, pair_mask, pair_14_mask, tile.x_start, tile.y_start, n_states, n_atoms_solute, atom_idx, + category, q_state, atom_lambdas, atom_charge, atom_vdw, LJ_matrix, el14_scale, coulomb_constant, vdw_rule, cx, cy, cz, dvelocities, e); } __global__ void compute_lrf_centers_kernel( @@ -1045,7 +1103,7 @@ __global__ void build_lrf_coefficients_dense_kernel( * 18: yzz * 19: xyz */ - double sums[LRF_UNIQUE_COMPONENTS]; + real_t sums[LRF_UNIQUE_COMPONENTS]; #pragma unroll for (int component = 0; component < LRF_UNIQUE_COMPONENTS; ++component) { @@ -1076,15 +1134,15 @@ __global__ void build_lrf_coefficients_dense_kernel( continue; } - const double charge = static_cast(atom_charge[source_slot]); + const real_t charge = atom_charge[source_slot]; - const double x = static_cast(cx[source_slot]) - center.x; + const real_t x = cx[source_slot] - center.x; - const double y = static_cast(cy[source_slot]) - center.y; + const real_t y = cy[source_slot] - center.y; - const double z = static_cast(cz[source_slot]) - center.z; + const real_t z = cz[source_slot] - center.z; - const double r2 = x * x + y * y + z * z; + const real_t r2 = x * x + y * y + z * z; /* * A self relationship should not be LRF, but guard against @@ -1094,32 +1152,32 @@ __global__ void build_lrf_coefficients_dense_kernel( continue; } - const double inv_r = rsqrt(r2); - const double inv_r2 = 1.0 / r2; + const real_t inv_r = rsqrt(r2); + const real_t inv_r2 = 1.0 / r2; - const double inv_r3 = inv_r * inv_r2; + const real_t inv_r3 = inv_r * inv_r2; - const double inv_r5 = inv_r3 * inv_r2; + const real_t inv_r5 = inv_r3 * inv_r2; - const double inv_r7 = inv_r5 * inv_r2; + const real_t inv_r7 = inv_r5 * inv_r2; - const double q_r1 = charge * inv_r; + const real_t q_r1 = charge * inv_r; - const double q_r3 = charge * inv_r3; + const real_t q_r3 = charge * inv_r3; - const double q_r5 = charge * inv_r5; + const real_t q_r5 = charge * inv_r5; - const double q_r7 = charge * inv_r7; + const real_t q_r7 = charge * inv_r7; - const double three_q_r5 = 3.0 * q_r5; + const real_t three_q_r5 = 3.0 * q_r5; - const double nine_q_r5 = 9.0 * q_r5; + const real_t nine_q_r5 = 9.0 * q_r5; - const double fifteen_q_r7 = 15.0 * q_r7; + const real_t fifteen_q_r7 = 15.0 * q_r7; - const double xx = x * x; - const double yy = y * y; - const double zz = z * z; + const real_t xx = x * x; + const real_t yy = y * y; + const real_t zz = z * z; /* * phi0 @@ -1246,6 +1304,8 @@ void CudaNonbondedForce::init_backend(Context& ctx) { exact_pair_masks_ = std::make_unique>(exact_tile_capacity_ * 32, false, true); + exact_pair_14_masks_ = std::make_unique>(exact_tile_capacity_ * 32, false, true); + exact_tile_count_ = std::make_unique>(1, true, true); list_overflow_ = std::make_unique>(1, true, true); @@ -1265,7 +1325,6 @@ void CudaNonbondedForce::init_backend(Context& ctx) { } lrf_slot_to_group_range_ = HostDeviceBuffer::from_vector(slot_to_group_range, ctx.command_info.requested_gpu); - } void CudaNonbondedForce::calc_all_direct_pairs(Context& ctx) { @@ -1309,7 +1368,6 @@ void CudaNonbondedForce::init_calculation_groups(Context& ctx) { const int blocks = (total_pairs + threads - 1) / threads; int use_switch_atom = ctx.charge_group_config.iuse_switch_atom == 1 ? 1 : 0; - use_switch_atom = 1; classify_group_pairs_kernel<<>>( n_group_ranges, @@ -1354,6 +1412,7 @@ void CudaNonbondedForce::calc_exact_tiles(Context& ctx) { n_exact_tiles_, exact_tiles_->gpu_data_p, exact_pair_masks_->gpu_data_p, + exact_pair_14_masks_->gpu_data_p, ctx.n_lambdas(), ctx.n_atoms_solute, @@ -1505,11 +1564,18 @@ void CudaNonbondedForce::build_exact_atom_tiles(Context& ctx) { group_pair_modes_->gpu_data_p, lrf_slot_to_group_range_->gpu_data_p, + ctx.n_atoms_solute, + data_.atom_idx->gpu_data_p, + data_.category->gpu_data_p, + data_.q_state->gpu_data_p, + ctx.LJ_matrix->gpu_data_p, + exact_tile_count_->gpu_data_p, list_overflow_->gpu_data_p, exact_tiles_->gpu_data_p, - exact_pair_masks_->gpu_data_p); + exact_pair_masks_->gpu_data_p, + exact_pair_14_masks_->gpu_data_p); check_cuda(cudaGetLastError());