Add method to copy Kokkos neighbor list to CPU list

This commit is contained in:
Stan Moore
2020-02-26 14:26:46 -07:00
parent a828facbd8
commit eb80fd3031
10 changed files with 133 additions and 33 deletions

View File

@ -19,20 +19,20 @@ using namespace LAMMPS_NS;
/* ---------------------------------------------------------------------- */
template<class Device>
NeighListKokkos<Device>::NeighListKokkos(class LAMMPS *lmp):NeighList(lmp)
template<class DeviceType>
NeighListKokkos<DeviceType>::NeighListKokkos(class LAMMPS *lmp):NeighList(lmp)
{
_stride = 1;
maxneighs = 16;
kokkos = 1;
maxatoms = 0;
execution_space = ExecutionSpaceFromDevice<Device>::space;
execution_space = ExecutionSpaceFromDevice<DeviceType>::space;
};
/* ---------------------------------------------------------------------- */
template<class Device>
void NeighListKokkos<Device>::grow(int nmax)
template<class DeviceType>
void NeighListKokkos<DeviceType>::grow(int nmax)
{
// skip if this list is already long enough to store nmax atoms
// and maxneighs neighbors
@ -40,14 +40,12 @@ void NeighListKokkos<Device>::grow(int nmax)
if (nmax <= maxatoms && d_neighbors.extent(1) >= maxneighs) return;
maxatoms = nmax;
k_ilist =
DAT::tdual_int_1d("neighlist:ilist",maxatoms);
d_ilist = k_ilist.view<Device>();
d_numneigh =
typename ArrayTypes<Device>::t_int_1d("neighlist:numneigh",maxatoms);
d_neighbors =
typename ArrayTypes<Device>::t_neighbors_2d("neighlist:neighbors",
maxatoms,maxneighs);
k_ilist = DAT::tdual_int_1d("neighlist:ilist",maxatoms);
d_ilist = k_ilist.view<DeviceType>();
k_numneigh = DAT::tdual_int_1d("neighlist:numneigh",maxatoms);
d_numneigh = k_numneigh.view<DeviceType>();
k_neighbors = DAT::tdual_neighbors_2d("neighlist:neighbors",maxatoms,maxneighs);
d_neighbors = k_neighbors.view<DeviceType>();
}
/* ---------------------------------------------------------------------- */

View File

@ -59,7 +59,7 @@ class AtomNeighborsConst
const int _stride;
};
template<class Device>
template<class DeviceType>
class NeighListKokkos: public NeighList {
int _stride;
@ -67,10 +67,12 @@ public:
int maxneighs;
void grow(int nmax);
typename ArrayTypes<Device>::t_neighbors_2d d_neighbors;
typename DAT::tdual_int_1d k_ilist; // local indices of I atoms
typename ArrayTypes<Device>::t_int_1d d_ilist;
typename ArrayTypes<Device>::t_int_1d d_numneigh; // # of J neighs for each I
DAT::tdual_neighbors_2d k_neighbors;
typename ArrayTypes<DeviceType>::t_neighbors_2d d_neighbors;
DAT::tdual_int_1d k_ilist; // local indices of I atoms
typename ArrayTypes<DeviceType>::t_int_1d d_ilist;
DAT::tdual_int_1d k_numneigh; // # of J neighs for each I
typename ArrayTypes<DeviceType>::t_int_1d d_numneigh;
NeighListKokkos(class LAMMPS *lmp);
@ -82,8 +84,8 @@ public:
KOKKOS_INLINE_FUNCTION
static AtomNeighborsConst static_neighbors_const(int i,
typename ArrayTypes<Device>::t_neighbors_2d_const const& d_neighbors,
typename ArrayTypes<Device>::t_int_1d_const const& d_numneigh) {
typename ArrayTypes<DeviceType>::t_neighbors_2d_const const& d_neighbors,
typename ArrayTypes<DeviceType>::t_int_1d_const const& d_numneigh) {
return AtomNeighborsConst(&d_neighbors(i,0),d_numneigh(i),
&d_neighbors(i,1)-&d_neighbors(i,0));
}

View File

@ -13,6 +13,8 @@
#include "npair_copy_kokkos.h"
#include "neigh_list_kokkos.h"
#include "my_page.h"
#include "error.h"
using namespace LAMMPS_NS;
@ -30,6 +32,24 @@ void NPairCopyKokkos<DeviceType>::build(NeighList *list)
{
NeighList *listcopy = list->listcopy;
if (list->kokkos) {
if (!listcopy->kokkos)
error->all(FLERR,"Cannot copy non-Kokkos neighbor list to Kokkos neigh list");
copy_to_kokkos(list);
} else {
if (!listcopy->kokkos)
error->all(FLERR,"Missing Kokkos neighbor list for copy");
copy_to_cpu(list);
}
}
/* ---------------------------------------------------------------------- */
template<class DeviceType>
void NPairCopyKokkos<DeviceType>::copy_to_kokkos(NeighList *list)
{
NeighList *listcopy = list->listcopy;
list->inum = listcopy->inum;
list->gnum = listcopy->gnum;
list->ilist = listcopy->ilist;
@ -44,6 +64,62 @@ void NPairCopyKokkos<DeviceType>::build(NeighList *list)
list_kk->d_neighbors = listcopy_kk->d_neighbors;
}
/* ---------------------------------------------------------------------- */
template<class DeviceType>
void NPairCopyKokkos<DeviceType>::copy_to_cpu(NeighList *list)
{
NeighList *listcopy = list->listcopy;
NeighListKokkos<DeviceType>* listcopy_kk = (NeighListKokkos<DeviceType>*) listcopy;
listcopy_kk->k_ilist.template sync<LMPHostType>();
listcopy_kk->k_numneigh.template sync<LMPHostType>();
listcopy_kk->k_neighbors.template sync<LMPHostType>();
int inum = listcopy->inum;
int gnum = listcopy->gnum;
int inum_all = inum;
if (list->ghost) inum_all += gnum;
auto h_ilist = listcopy_kk->k_ilist.h_view;
auto h_numneigh = listcopy_kk->k_numneigh.h_view;
auto h_neighbors = listcopy_kk->k_neighbors.h_view;
list->inum = inum;
list->gnum = gnum;
auto ilist = list->ilist;
auto numneigh = list->numneigh;
// Kokkos neighbor data is stored differently than regular CPU,
// must copy element by element
int *neighptr;
int **firstneigh = list->firstneigh;
MyPage<int> *ipage = list->ipage;
ipage->reset();
for (int ii = 0; ii < inum_all; ii++) {
neighptr = ipage->vget();
const int i = h_ilist[ii];
ilist[ii] = i;
// loop over Kokkos neighbor list
const int jnum = h_numneigh[i];
numneigh[i] = jnum;
for (int jj = 0; jj < jnum; jj++) {
const int joriginal = h_neighbors(i,jj);
neighptr[jj] = joriginal;
}
firstneigh[i] = neighptr;
ipage->vgot(jnum);
if (ipage->status())
error->one(FLERR,"Neighbor list overflow, boost neigh_modify one");
}
}
namespace LAMMPS_NS {
template class NPairCopyKokkos<LMPDeviceType>;
#ifdef KOKKOS_ENABLE_CUDA

View File

@ -36,6 +36,9 @@ class NPairCopyKokkos : public NPair {
NPairCopyKokkos(class LAMMPS *);
~NPairCopyKokkos() {}
void build(class NeighList *);
private:
void copy_to_kokkos(class NeighList *);
void copy_to_cpu(class NeighList *);
};
}

View File

@ -68,12 +68,14 @@ void NPairHalffullKokkos<DeviceType,NEWTON>::build(NeighList *list)
copymode = 1;
Kokkos::parallel_for(Kokkos::RangePolicy<DeviceType, TagNPairHalffullCompute>(0,inum_full),*this);
copymode = 0;
list->inum = k_list_full->inum;
list->gnum = k_list_full->gnum;
k_list->k_ilist.template modify<DeviceType>();
copymode = 0;
k_list->k_ilist.template modify<DeviceType>();
k_list->k_numneigh.template modify<DeviceType>();
k_list->k_neighbors.template modify<DeviceType>();
}
template<class DeviceType, int NEWTON>

View File

@ -273,7 +273,8 @@ void NPairKokkos<DeviceType,HALF_NEIGH,GHOST,TRI,SIZE>::build(NeighList *list_)
if(data.h_resize()) {
list->maxneighs = data.h_new_maxneighs() * 1.2;
list->d_neighbors = typename ArrayTypes<DeviceType>::t_neighbors_2d("neighbors", list->d_neighbors.extent(0), list->maxneighs);
list->k_neighbors = DAT::tdual_neighbors_2d("neighbors", list->d_neighbors.extent(0), list->maxneighs);
list->d_neighbors = list->k_neighbors.template view<DeviceType>();
data.neigh_list.d_neighbors = list->d_neighbors;
data.neigh_list.maxneighs = list->maxneighs;
}
@ -288,6 +289,8 @@ void NPairKokkos<DeviceType,HALF_NEIGH,GHOST,TRI,SIZE>::build(NeighList *list_)
}
list->k_ilist.template modify<DeviceType>();
list->k_numneigh.template modify<DeviceType>();
list->k_neighbors.template modify<DeviceType>();
}
/* ---------------------------------------------------------------------- */

View File

@ -519,6 +519,8 @@ fprintf(stdout, "Fina%03d %6d inum %6d gnum, total used %6d, allocated %6d\n"
#endif
list->k_ilist.template modify<DeviceType>();
list->k_numneigh.template modify<DeviceType>();
list->k_neighbors.template modify<DeviceType>();
}

View File

@ -80,6 +80,7 @@ NeighList::NeighList(LAMMPS *lmp) : Pointers(lmp)
// Kokkos package
kokkos = 0;
kk2cpu = 0;
execution_space = Host;
// USER-DPD package
@ -143,8 +144,11 @@ void NeighList::post_constructor(NeighRequest *nq)
respainner = nq->respainner;
copy = nq->copy;
if (nq->copy)
if (nq->copy) {
listcopy = neighbor->lists[nq->copylist];
if (listcopy->kokkos && !this->kokkos)
kk2cpu = 1;
}
if (nq->skip) {
listskip = neighbor->lists[nq->skiplist];

View File

@ -42,6 +42,7 @@ class NeighList : protected Pointers {
int respamiddle; // 1 if there is also a rRespa middle list
int respainner; // 1 if there is also a rRespa inner list
int copy; // 1 if this list is copied from another list
int kk2cpu; // 1 if this list is copied from Kokkos to CPU
int copymode; // 1 if this is a Kokkos on-device copy
// data structs to store neighbor pairs I,J and associated values

View File

@ -849,7 +849,8 @@ int Neighbor::init_pair()
// allocate initial pages for each list, except if copy flag set
for (i = 0; i < nlist; i++) {
if (lists[i]->copy) continue;
if (lists[i]->copy && !lists[i]->kk2cpu)
continue;
lists[i]->setup_pages(pgsize,oneatom);
}
@ -860,8 +861,10 @@ int Neighbor::init_pair()
// also Kokkos list initialization
int maxatom = atom->nmax;
for (i = 0; i < nlist; i++)
if (neigh_pair[i] && !lists[i]->copy) lists[i]->grow(maxatom,maxatom);
for (i = 0; i < nlist; i++) {
if (neigh_pair[i] && (!lists[i]->copy || lists[i]->kk2cpu))
lists[i]->grow(maxatom,maxatom);
}
// plist = indices of perpetual NPair classes
// perpetual = non-occasional, re-built at every reneighboring
@ -1257,8 +1260,8 @@ void Neighbor::morph_copy()
if (irq->history != jrq->history) continue;
if (irq->bond != jrq->bond) continue;
if (irq->intel != jrq->intel) continue;
if (irq->kokkos_host != jrq->kokkos_host) continue;
if (irq->kokkos_device != jrq->kokkos_device) continue;
if (irq->kokkos_host && !jrq->kokkos_host) continue;
if (irq->kokkos_device && !jrq->kokkos_device) continue;
if (irq->ssa != jrq->ssa) continue;
if (irq->cut != jrq->cut) continue;
if (irq->cutoff != jrq->cutoff) continue;
@ -1789,8 +1792,12 @@ int Neighbor::choose_pair(NeighRequest *rq)
if (rq->copy) {
if (!(mask & NP_COPY)) continue;
if (!rq->kokkos_device != !(mask & NP_KOKKOS_DEVICE)) continue;
if (!rq->kokkos_host != !(mask & NP_KOKKOS_HOST)) continue;
if (rq->kokkos_device || rq->kokkos_host) {
if (!rq->kokkos_device != !(mask & NP_KOKKOS_DEVICE)) continue;
if (!rq->kokkos_host != !(mask & NP_KOKKOS_HOST)) continue;
}
if (!requests[rq->copylist]->kokkos_device != !(mask & NP_KOKKOS_DEVICE)) continue;
if (!requests[rq->copylist]->kokkos_host != !(mask & NP_KOKKOS_HOST)) continue;
return i+1;
}
@ -2102,7 +2109,8 @@ void Neighbor::build(int topoflag)
for (i = 0; i < npair_perpetual; i++) {
m = plist[i];
if (!lists[m]->copy) lists[m]->grow(nlocal,nall);
if (!lists[i]->copy || lists[i]->kk2cpu)
lists[m]->grow(nlocal,nall);
neigh_pair[m]->build_setup();
neigh_pair[m]->build(lists[m]);
}
@ -2191,7 +2199,8 @@ void Neighbor::build_one(class NeighList *mylist, int preflag)
// build the list
if (!mylist->copy) mylist->grow(atom->nlocal,atom->nlocal+atom->nghost);
if (!mylist->copy || mylist->kk2cpu)
mylist->grow(atom->nlocal,atom->nlocal+atom->nghost);
np->build_setup();
np->build(mylist);
}