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; }