int gmx_pmeonly(struct gmx_pme_t *pme,
t_commrec *cr, t_nrnb *mynrnb,
gmx_wallcycle_t wcycle,
gmx_walltime_accounting_t walltime_accounting,
real ewaldcoeff_q, real ewaldcoeff_lj,
t_inputrec *ir)
{
int npmedata;
struct gmx_pme_t **pmedata;
gmx_pme_pp_t pme_pp;
int ret;
int natoms;
matrix box;
rvec *x_pp = NULL, *f_pp = NULL;
real *chargeA = NULL, *chargeB = NULL;
real *c6A = NULL, *c6B = NULL;
real *sigmaA = NULL, *sigmaB = NULL;
real lambda_q = 0;
real lambda_lj = 0;
int maxshift_x = 0, maxshift_y = 0;
real energy_q, energy_lj, dvdlambda_q, dvdlambda_lj;
matrix vir_q, vir_lj;
float cycles;
int count;
gmx_bool bEnerVir;
int pme_flags;
gmx_int64_t step, step_rel;
ivec grid_switch;
/* This data will only use with PME tuning, i.e. switching PME grids */
npmedata = 1;
snew(pmedata, npmedata);
pmedata[0] = pme;
pme_pp = gmx_pme_pp_init(cr);
init_nrnb(mynrnb);
count = 0;
do /****** this is a quasi-loop over time steps! */
{
/* The reason for having a loop here is PME grid tuning/switching */
do
{
/* Domain decomposition */
ret = gmx_pme_recv_coeffs_coords(pme_pp,
&natoms,
&chargeA, &chargeB,
&c6A, &c6B,
&sigmaA, &sigmaB,
box, &x_pp, &f_pp,
&maxshift_x, &maxshift_y,
&pme->bFEP_q, &pme->bFEP_lj,
&lambda_q, &lambda_lj,
&bEnerVir,
&pme_flags,
&step,
grid_switch, &ewaldcoeff_q, &ewaldcoeff_lj);
if (ret == pmerecvqxSWITCHGRID)
{
/* Switch the PME grid to grid_switch */
gmx_pmeonly_switch(&npmedata, &pmedata, grid_switch, cr, ir, &pme);
}
if (ret == pmerecvqxRESETCOUNTERS)
{
/* Reset the cycle and flop counters */
reset_pmeonly_counters(wcycle, walltime_accounting, mynrnb, ir, step);
}
}
while (ret == pmerecvqxSWITCHGRID || ret == pmerecvqxRESETCOUNTERS);
if (ret == pmerecvqxFINISH)
{
/* We should stop: break out of the loop */
break;
}
step_rel = step - ir->init_step;
if (count == 0)
{
wallcycle_start(wcycle, ewcRUN);
walltime_accounting_start(walltime_accounting);
}
wallcycle_start(wcycle, ewcPMEMESH);
dvdlambda_q = 0;
dvdlambda_lj = 0;
clear_mat(vir_q);
clear_mat(vir_lj);
gmx_pme_do(pme, 0, natoms, x_pp, f_pp,
chargeA, chargeB, c6A, c6B, sigmaA, sigmaB, box,
cr, maxshift_x, maxshift_y, mynrnb, wcycle,
vir_q, ewaldcoeff_q, vir_lj, ewaldcoeff_lj,
&energy_q, &energy_lj, lambda_q, lambda_lj, &dvdlambda_q, &dvdlambda_lj,
pme_flags | GMX_PME_DO_ALL_F | (bEnerVir ? GMX_PME_CALC_ENER_VIR : 0));
cycles = wallcycle_stop(wcycle, ewcPMEMESH);
gmx_pme_send_force_vir_ener(pme_pp,
f_pp, vir_q, energy_q, vir_lj, energy_lj,
dvdlambda_q, dvdlambda_lj, cycles);
count++;
} /***** end of quasi-loop, we stop with the break above */
while (TRUE);
walltime_accounting_end(walltime_accounting);
return 0;
}
int gmx_pmeonly(struct gmx_pme_t *pme,
const t_commrec *cr, t_nrnb *mynrnb,
gmx_wallcycle *wcycle,
gmx_walltime_accounting_t walltime_accounting,
t_inputrec *ir, PmeRunMode runMode)
{
int ret;
int natoms = 0;
matrix box;
real lambda_q = 0;
real lambda_lj = 0;
int maxshift_x = 0, maxshift_y = 0;
real energy_q, energy_lj, dvdlambda_q, dvdlambda_lj;
matrix vir_q, vir_lj;
float cycles;
int count;
gmx_bool bEnerVir = FALSE;
int64_t step;
/* This data will only use with PME tuning, i.e. switching PME grids */
std::vector<gmx_pme_t *> pmedata;
pmedata.push_back(pme);
auto pme_pp = gmx_pme_pp_init(cr);
//TODO the variable below should be queried from the task assignment info
const bool useGpuForPme = (runMode == PmeRunMode::GPU) || (runMode == PmeRunMode::Mixed);
if (useGpuForPme)
{
changePinningPolicy(&pme_pp->chargeA, pme_get_pinning_policy());
changePinningPolicy(&pme_pp->x, pme_get_pinning_policy());
}
init_nrnb(mynrnb);
count = 0;
do /****** this is a quasi-loop over time steps! */
{
/* The reason for having a loop here is PME grid tuning/switching */
do
{
/* Domain decomposition */
ivec newGridSize;
bool atomSetChanged = false;
real ewaldcoeff_q = 0, ewaldcoeff_lj = 0;
ret = gmx_pme_recv_coeffs_coords(pme_pp.get(),
&natoms,
box,
&maxshift_x, &maxshift_y,
&lambda_q, &lambda_lj,
&bEnerVir,
&step,
&newGridSize,
&ewaldcoeff_q,
&ewaldcoeff_lj,
&atomSetChanged);
if (ret == pmerecvqxSWITCHGRID)
{
/* Switch the PME grid to newGridSize */
pme = gmx_pmeonly_switch(&pmedata, newGridSize, ewaldcoeff_q, ewaldcoeff_lj, cr, ir);
}
if (atomSetChanged)
{
gmx_pme_reinit_atoms(pme, natoms, pme_pp->chargeA.data());
}
if (ret == pmerecvqxRESETCOUNTERS)
{
/* Reset the cycle and flop counters */
reset_pmeonly_counters(wcycle, walltime_accounting, mynrnb, step, useGpuForPme);
}
}
while (ret == pmerecvqxSWITCHGRID || ret == pmerecvqxRESETCOUNTERS);
if (ret == pmerecvqxFINISH)
{
/* We should stop: break out of the loop */
break;
}
if (count == 0)
{
wallcycle_start(wcycle, ewcRUN);
walltime_accounting_start_time(walltime_accounting);
}
wallcycle_start(wcycle, ewcPMEMESH);
dvdlambda_q = 0;
dvdlambda_lj = 0;
clear_mat(vir_q);
clear_mat(vir_lj);
energy_q = 0;
energy_lj = 0;
// TODO Make a struct of array refs onto these per-atom fields
// of pme_pp (maybe box, energy and virial, too; and likewise
// from mdatoms for the other call to gmx_pme_do), so we have
// fewer lines of code and less parameter passing.
const int pmeFlags = GMX_PME_DO_ALL_F | (bEnerVir ? GMX_PME_CALC_ENER_VIR : 0);
gmx::ArrayRef<const gmx::RVec> forces;
if (useGpuForPme)
{
const bool boxChanged = false;
//TODO this should be set properly by gmx_pme_recv_coeffs_coords,
// or maybe use inputrecDynamicBox(ir), at the very least - change this when this codepath is tested!
pme_gpu_prepare_computation(pme, boxChanged, box, wcycle, pmeFlags);
pme_gpu_launch_spread(pme, pme_pp->x.rvec_array(), wcycle);
pme_gpu_launch_complex_transforms(pme, wcycle);
pme_gpu_launch_gather(pme, wcycle, PmeForceOutputHandling::Set);
pme_gpu_wait_finish_task(pme, wcycle, &forces, vir_q, &energy_q);
pme_gpu_reinit_computation(pme, wcycle);
}
else
{
gmx_pme_do(pme, 0, natoms, pme_pp->x.rvec_array(), as_rvec_array(pme_pp->f.data()),
pme_pp->chargeA.data(), pme_pp->chargeB.data(),
pme_pp->sqrt_c6A.data(), pme_pp->sqrt_c6B.data(),
pme_pp->sigmaA.data(), pme_pp->sigmaB.data(), box,
cr, maxshift_x, maxshift_y, mynrnb, wcycle,
vir_q, vir_lj,
&energy_q, &energy_lj, lambda_q, lambda_lj, &dvdlambda_q, &dvdlambda_lj,
pmeFlags);
forces = pme_pp->f;
}
cycles = wallcycle_stop(wcycle, ewcPMEMESH);
gmx_pme_send_force_vir_ener(pme_pp.get(), as_rvec_array(forces.data()),
vir_q, energy_q, vir_lj, energy_lj,
dvdlambda_q, dvdlambda_lj, cycles);
count++;
} /***** end of quasi-loop, we stop with the break above */
while (TRUE);
walltime_accounting_end_time(walltime_accounting);
return 0;
}
