diff --git a/src/KOKKOS/neigh_list_kokkos.cpp b/src/KOKKOS/neigh_list_kokkos.cpp index 2b9c5ef645..afaf0dd1b8 100644 --- a/src/KOKKOS/neigh_list_kokkos.cpp +++ b/src/KOKKOS/neigh_list_kokkos.cpp @@ -19,20 +19,20 @@ using namespace LAMMPS_NS; /* ---------------------------------------------------------------------- */ -template -NeighListKokkos::NeighListKokkos(class LAMMPS *lmp):NeighList(lmp) +template +NeighListKokkos::NeighListKokkos(class LAMMPS *lmp):NeighList(lmp) { _stride = 1; maxneighs = 16; kokkos = 1; maxatoms = 0; - execution_space = ExecutionSpaceFromDevice::space; + execution_space = ExecutionSpaceFromDevice::space; }; /* ---------------------------------------------------------------------- */ -template -void NeighListKokkos::grow(int nmax) +template +void NeighListKokkos::grow(int nmax) { // skip if this list is already long enough to store nmax atoms // and maxneighs neighbors @@ -40,14 +40,12 @@ void NeighListKokkos::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(); - d_numneigh = - typename ArrayTypes::t_int_1d("neighlist:numneigh",maxatoms); - d_neighbors = - typename ArrayTypes::t_neighbors_2d("neighlist:neighbors", - maxatoms,maxneighs); + k_ilist = DAT::tdual_int_1d("neighlist:ilist",maxatoms); + d_ilist = k_ilist.view(); + k_numneigh = DAT::tdual_int_1d("neighlist:numneigh",maxatoms); + d_numneigh = k_numneigh.view(); + k_neighbors = DAT::tdual_neighbors_2d("neighlist:neighbors",maxatoms,maxneighs); + d_neighbors = k_neighbors.view(); } /* ---------------------------------------------------------------------- */ diff --git a/src/KOKKOS/neigh_list_kokkos.h b/src/KOKKOS/neigh_list_kokkos.h index 585422c54f..f8195a01f4 100644 --- a/src/KOKKOS/neigh_list_kokkos.h +++ b/src/KOKKOS/neigh_list_kokkos.h @@ -59,7 +59,7 @@ class AtomNeighborsConst const int _stride; }; -template +template class NeighListKokkos: public NeighList { int _stride; @@ -67,10 +67,12 @@ public: int maxneighs; void grow(int nmax); - typename ArrayTypes::t_neighbors_2d d_neighbors; - typename DAT::tdual_int_1d k_ilist; // local indices of I atoms - typename ArrayTypes::t_int_1d d_ilist; - typename ArrayTypes::t_int_1d d_numneigh; // # of J neighs for each I + DAT::tdual_neighbors_2d k_neighbors; + typename ArrayTypes::t_neighbors_2d d_neighbors; + DAT::tdual_int_1d k_ilist; // local indices of I atoms + typename ArrayTypes::t_int_1d d_ilist; + DAT::tdual_int_1d k_numneigh; // # of J neighs for each I + typename ArrayTypes::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::t_neighbors_2d_const const& d_neighbors, - typename ArrayTypes::t_int_1d_const const& d_numneigh) { + typename ArrayTypes::t_neighbors_2d_const const& d_neighbors, + typename ArrayTypes::t_int_1d_const const& d_numneigh) { return AtomNeighborsConst(&d_neighbors(i,0),d_numneigh(i), &d_neighbors(i,1)-&d_neighbors(i,0)); } diff --git a/src/KOKKOS/npair_copy_kokkos.cpp b/src/KOKKOS/npair_copy_kokkos.cpp index 0ce0f4d3ff..b0554db9ea 100644 --- a/src/KOKKOS/npair_copy_kokkos.cpp +++ b/src/KOKKOS/npair_copy_kokkos.cpp @@ -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::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 +void NPairCopyKokkos::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::build(NeighList *list) list_kk->d_neighbors = listcopy_kk->d_neighbors; } +/* ---------------------------------------------------------------------- */ + +template +void NPairCopyKokkos::copy_to_cpu(NeighList *list) +{ + NeighList *listcopy = list->listcopy; + NeighListKokkos* listcopy_kk = (NeighListKokkos*) listcopy; + + listcopy_kk->k_ilist.template sync(); + listcopy_kk->k_numneigh.template sync(); + listcopy_kk->k_neighbors.template sync(); + + 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 *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; #ifdef KOKKOS_ENABLE_CUDA diff --git a/src/KOKKOS/npair_copy_kokkos.h b/src/KOKKOS/npair_copy_kokkos.h index 84eb10b204..4bbb2749e5 100644 --- a/src/KOKKOS/npair_copy_kokkos.h +++ b/src/KOKKOS/npair_copy_kokkos.h @@ -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 *); }; } diff --git a/src/KOKKOS/npair_halffull_kokkos.cpp b/src/KOKKOS/npair_halffull_kokkos.cpp index cc8f663ef2..754a5ca010 100644 --- a/src/KOKKOS/npair_halffull_kokkos.cpp +++ b/src/KOKKOS/npair_halffull_kokkos.cpp @@ -68,12 +68,14 @@ void NPairHalffullKokkos::build(NeighList *list) copymode = 1; Kokkos::parallel_for(Kokkos::RangePolicy(0,inum_full),*this); + copymode = 0; list->inum = k_list_full->inum; list->gnum = k_list_full->gnum; - k_list->k_ilist.template modify(); - copymode = 0; + k_list->k_ilist.template modify(); + k_list->k_numneigh.template modify(); + k_list->k_neighbors.template modify(); } template diff --git a/src/KOKKOS/npair_kokkos.cpp b/src/KOKKOS/npair_kokkos.cpp index aa2a1fc5ff..8dce226ef0 100644 --- a/src/KOKKOS/npair_kokkos.cpp +++ b/src/KOKKOS/npair_kokkos.cpp @@ -273,7 +273,8 @@ void NPairKokkos::build(NeighList *list_) if(data.h_resize()) { list->maxneighs = data.h_new_maxneighs() * 1.2; - list->d_neighbors = typename ArrayTypes::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(); data.neigh_list.d_neighbors = list->d_neighbors; data.neigh_list.maxneighs = list->maxneighs; } @@ -288,6 +289,8 @@ void NPairKokkos::build(NeighList *list_) } list->k_ilist.template modify(); + list->k_numneigh.template modify(); + list->k_neighbors.template modify(); } /* ---------------------------------------------------------------------- */ diff --git a/src/KOKKOS/npair_ssa_kokkos.cpp b/src/KOKKOS/npair_ssa_kokkos.cpp index 5234c4e6d2..b5f892958b 100644 --- a/src/KOKKOS/npair_ssa_kokkos.cpp +++ b/src/KOKKOS/npair_ssa_kokkos.cpp @@ -519,6 +519,8 @@ fprintf(stdout, "Fina%03d %6d inum %6d gnum, total used %6d, allocated %6d\n" #endif list->k_ilist.template modify(); + list->k_numneigh.template modify(); + list->k_neighbors.template modify(); } diff --git a/src/neigh_list.cpp b/src/neigh_list.cpp index 1c53b2f7a4..c01f8ca595 100644 --- a/src/neigh_list.cpp +++ b/src/neigh_list.cpp @@ -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]; diff --git a/src/neigh_list.h b/src/neigh_list.h index 146a01e7b5..501188a724 100644 --- a/src/neigh_list.h +++ b/src/neigh_list.h @@ -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 diff --git a/src/neighbor.cpp b/src/neighbor.cpp index 3d1f37a03b..5783c24d2d 100644 --- a/src/neighbor.cpp +++ b/src/neighbor.cpp @@ -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); }