Skip to content

Commit

Permalink
Merge pull request lammps#3973 from stanmoore1/kk_min
Browse files Browse the repository at this point in the history
Bugfixes for Kokkos minimize
  • Loading branch information
akohlmey authored Nov 16, 2023
2 parents d31c3ab + e648500 commit 3664eaf
Show file tree
Hide file tree
Showing 4 changed files with 46 additions and 7 deletions.
32 changes: 32 additions & 0 deletions src/KOKKOS/atom_vec_spin_kokkos.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -586,6 +586,38 @@ int AtomVecSpinKokkos::unpack_exchange_kokkos(DAT::tdual_xfloat_2d &k_buf, int n
}
}

/* ----------------------------------------------------------------------
clear extra forces starting at atom N
nbytes = # of bytes to clear for a per-atom vector
include f b/c this is invoked from within SPIN pair styles
------------------------------------------------------------------------- */

void AtomVecSpinKokkos::force_clear(int n, size_t nbytes)
{
int nzero = (double)nbytes/sizeof(double);

if (nzero) {
atomKK->k_fm.clear_sync_state(); // will be cleared below
atomKK->k_fm_long.clear_sync_state(); // will be cleared below

// local variables for lambda capture

auto l_fm = atomKK->k_fm.d_view;
auto l_fm_long = atomKK->k_fm_long.d_view;

Kokkos::parallel_for(nzero, LAMMPS_LAMBDA(int i) {
l_fm(i,0) = 0.0;
l_fm(i,1) = 0.0;
l_fm(i,2) = 0.0;
l_fm_long(i,0) = 0.0;
l_fm_long(i,1) = 0.0;
l_fm_long(i,2) = 0.0;
});

atomKK->modified(Device,FM_MASK|FML_MASK);
}
}

/* ---------------------------------------------------------------------- */

void AtomVecSpinKokkos::sync(ExecutionSpace space, unsigned int mask)
Expand Down
1 change: 1 addition & 0 deletions src/KOKKOS/atom_vec_spin_kokkos.h
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ class AtomVecSpinKokkos : public AtomVecKokkos, public AtomVecSpin {
AtomVecSpinKokkos(class LAMMPS *);
void grow(int) override;
void grow_pointers() override;
void force_clear(int, size_t) override;
void sort_kokkos(Kokkos::BinSort<KeyViewType, BinOp> &Sorter) override;
int pack_border_kokkos(int n, DAT::tdual_int_2d k_sendlist,
DAT::tdual_xfloat_2d buf,int iswap,
Expand Down
18 changes: 12 additions & 6 deletions src/KOKKOS/min_kokkos.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -72,8 +72,7 @@ void MinKokkos::init()
void MinKokkos::setup(int flag)
{
if (comm->me == 0 && screen) {
fmt::print(screen,"Setting up {} style minimization ...\n",
update->minimize_style);
fmt::print(screen,"Setting up {} style minimization ...\n", update->minimize_style);
if (flag) {
fmt::print(screen," Unit style : {}\n", update->unit_style);
fmt::print(screen," Current step : {}\n", update->ntimestep);
Expand All @@ -92,14 +91,13 @@ void MinKokkos::setup(int flag)
fextra = new double[nextra_global];
if (comm->me == 0)
error->warning(FLERR, "Energy due to {} extra global DOFs will"
" be included in minimizer energies\n", nextra_global);
" be included in minimizer energies\n",nextra_global);
}

// compute for potential energy

int id = modify->find_compute("thermo_pe");
if (id < 0) error->all(FLERR,"Minimization could not find thermo_pe compute");
pe_compute = modify->compute[id];
pe_compute = modify->get_compute_by_id("thermo_pe");
if (!pe_compute) error->all(FLERR,"Minimization could not find thermo_pe compute");

// style-specific setup does two tasks
// setup extra global dof vectors
Expand Down Expand Up @@ -537,6 +535,7 @@ double MinKokkos::energy_force(int resetflag)
if (resetflag) fix_minimize_kk->reset_coords();
reset_vectors();
}

return energy;
}

Expand Down Expand Up @@ -575,7 +574,14 @@ void MinKokkos::force_clear()
l_torque(i,2) = 0.0;
}
});

if (extraflag) {
size_t nbytes = sizeof(double) * atom->nlocal;
if (force->newton) nbytes += sizeof(double) * atom->nghost;
atom->avec->force_clear(0,nbytes);
}
}

atomKK->modified(Device,F_MASK|TORQUE_MASK);
}

Expand Down
2 changes: 1 addition & 1 deletion src/KOKKOS/npair_kokkos.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -113,7 +113,7 @@ void NPairKokkos<DeviceType,HALF,NEWTON,GHOST,TRI,SIZE>::copy_stencil_info()
NPair::copy_stencil_info();
nstencil = ns->nstencil;

if (ns->last_stencil != last_stencil_old) {
if (ns->last_stencil != last_stencil_old || ns->last_stencil == update->ntimestep) {
// copy stencil to device as it may have changed

last_stencil_old = ns->last_stencil;
Expand Down

0 comments on commit 3664eaf

Please sign in to comment.