Merge pull request #1906 from stanmoore1/kk_copycpu (f5a8d40a) · Commits · 郑智淋 / lammps

src/KOKKOS/neigh_list_kokkos.cpp

+11 −13

Original line number	Diff line number	Diff line
		@@ -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>();
		}

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

src/KOKKOS/neigh_list_kokkos.h

+9 −7

Original line number	Diff line number	Diff line
		@@ -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));
		}

src/KOKKOS/npair_copy_kokkos.cpp

+76 −0

Original line number	Diff line number	Diff line
		@@ -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 neighbor 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

src/KOKKOS/npair_copy_kokkos.h

+3 −0

Original line number	Diff line number	Diff line
		@@ -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 *);
		};

		}

src/KOKKOS/npair_halffull_kokkos.cpp

+4 −2

Original line number	Diff line number	Diff line
		@@ -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>

Admin message