first pass at implementing atom2bin for KOKKOS neighbor lists (32e0de7a) · Commits · 郑智淋 / lammps

src/KOKKOS/nbin_kokkos.cpp

+5 −0

Original line number	Diff line number	Diff line
		@@ -75,6 +75,10 @@ void NBinKokkos<DeviceType>::bin_atoms_setup(int nall)
		k_bincount = DAT::tdual_int_1d("Neighbor::d_bincount",mbins);
		bincount = k_bincount.view<DeviceType>();
		}
		if (nall > k_atom2bin.d_view.dimension_0()) {
		k_atom2bin = DAT::tdual_int_1d("Neighbor::d_atom2bin",nall);
		atom2bin = k_atom2bin.view<DeviceType>();
		}
		}

		/* ----------------------------------------------------------------------
		@@ -125,6 +129,7 @@ void NBinKokkos<DeviceType>::binatomsItem(const int &i) const
		{
		const int ibin = coord2bin(x(i, 0), x(i, 1), x(i, 2));

		atom2bin(i) = ibin;
		const int ac = Kokkos::atomic_fetch_add(&bincount[ibin], (int)1);
		if(ac < bins.dimension_1()) {
		bins(ibin, ac) = i;

src/KOKKOS/nbin_kokkos.h

+2 −0

Original line number	Diff line number	Diff line
		@@ -44,11 +44,13 @@ class NBinKokkos : public NBinStandard {
		int atoms_per_bin;
		DAT::tdual_int_1d k_bincount;
		DAT::tdual_int_2d k_bins;
		DAT::tdual_int_1d k_atom2bin;

		typename AT::t_int_1d bincount;
		const typename AT::t_int_1d_const c_bincount;
		typename AT::t_int_2d bins;
		typename AT::t_int_2d_const c_bins;
		typename AT::t_int_1d atom2bin;
		typename AT::t_int_scalar d_resize;
		typename ArrayTypes<LMPHostType>::t_int_scalar h_resize;
		typename AT::t_x_array_randomread x;

src/KOKKOS/npair_kokkos.cpp

+5 −2

Original line number	Diff line number	Diff line
		@@ -73,6 +73,7 @@ void NPairKokkos<DeviceType,HALF_NEIGH,GHOST,TRI>::copy_bin_info()
		atoms_per_bin = nbKK->atoms_per_bin;
		k_bincount = nbKK->k_bincount;
		k_bins = nbKK->k_bins;
		k_atom2bin = nbKK->k_atom2bin;
		}

		/* ----------------------------------------------------------------------
		@@ -122,6 +123,7 @@ void NPairKokkos<DeviceType,HALF_NEIGH,GHOST,TRI>::build(NeighList *list_)
		k_cutneighsq.view<DeviceType>(),
		k_bincount.view<DeviceType>(),
		k_bins.view<DeviceType>(),
		k_atom2bin.view<DeviceType>(),
		nstencil,
		k_stencil.view<DeviceType>(),
		k_stencilxyz.view<DeviceType>(),
		@@ -166,6 +168,7 @@ void NPairKokkos<DeviceType,HALF_NEIGH,GHOST,TRI>::build(NeighList *list_)
		k_ex_mol_intra.sync<DeviceType>();
		k_bincount.sync<DeviceType>();
		k_bins.sync<DeviceType>();
		k_atom2bin.sync<DeviceType>();
		atomKK->sync(Device,X_MASK\|TYPE_MASK\|MASK_MASK\|MOLECULE_MASK\|TAG_MASK\|SPECIAL_MASK);

		data.special_flag[0] = special_flag[0];
		@@ -317,7 +320,7 @@ void NeighborKokkosExecute<DeviceType>::
		const X_FLOAT ztmp = x(i, 2);
		const int itype = type(i);

		const int ibin = coord2bin(xtmp, ytmp, ztmp);
		const int ibin = c_atom2bin(i);

		const typename ArrayTypes<DeviceType>::t_int_1d_const_um stencil
		= d_stencil;
		@@ -678,7 +681,7 @@ void NeighborKokkosExecute<DeviceType>::
		// no molecular test when i = ghost atom

		if (i < nlocal) {
		const int ibin = coord2bin(xtmp, ytmp, ztmp);
		const int ibin = c_atom2bin(i);
		for (int k = 0; k < nstencil; k++) {
		const int jbin = ibin + stencil[k];
		for(int m = 0; m < c_bincount(jbin); m++) {

src/KOKKOS/npair_kokkos.h

+5 −32

Original line number	Diff line number	Diff line
		@@ -105,6 +105,7 @@ class NPairKokkos : public NPair {
		int atoms_per_bin;
		DAT::tdual_int_1d k_bincount;
		DAT::tdual_int_2d k_bins;
		DAT::tdual_int_1d k_atom2bin;

		// data from NStencil class

		@@ -148,6 +149,8 @@ class NeighborKokkosExecute
		const typename AT::t_int_1d_const c_bincount;
		typename AT::t_int_2d bins;
		typename AT::t_int_2d_const c_bins;
		const typename AT::t_int_1d atom2bin;
		const typename AT::t_int_1d_const c_atom2bin;


		// data from NStencil class
		@@ -190,6 +193,7 @@ class NeighborKokkosExecute
		const typename AT::t_xfloat_2d_randomread &_cutneighsq,
		const typename AT::t_int_1d &_bincount,
		const typename AT::t_int_2d &_bins,
		const typename AT::t_int_1d &_atom2bin,
		const int _nstencil,
		const typename AT::t_int_1d &_d_stencil,
		const typename AT::t_int_1d_3 &_d_stencilxyz,
		@@ -224,6 +228,7 @@ class NeighborKokkosExecute
		const int & _xprd_half, const int & _yprd_half, const int & _zprd_half):
		neigh_list(_neigh_list), cutneighsq(_cutneighsq),
		bincount(_bincount),c_bincount(_bincount),bins(_bins),c_bins(_bins),
		atom2bin(_atom2bin),c_atom2bin(_atom2bin),
		nstencil(_nstencil),d_stencil(_d_stencil),d_stencilxyz(_d_stencilxyz),
		nlocal(_nlocal),
		x(_x),type(_type),mask(_mask),molecule(_molecule),
		@@ -281,38 +286,6 @@ class NeighborKokkosExecute
		void build_ItemCuda(typename Kokkos::TeamPolicy<DeviceType>::member_type dev) const;
		#endif

		KOKKOS_INLINE_FUNCTION
		int coord2bin(const X_FLOAT & x,const X_FLOAT & y,const X_FLOAT & z) const
		{
		int ix,iy,iz;

		if (x >= bboxhi[0])
		ix = static_cast<int> ((x-bboxhi[0])*bininvx) + nbinx;
		else if (x >= bboxlo[0]) {
		ix = static_cast<int> ((x-bboxlo[0])*bininvx);
		ix = MIN(ix,nbinx-1);
		} else
		ix = static_cast<int> ((x-bboxlo[0])*bininvx) - 1;

		if (y >= bboxhi[1])
		iy = static_cast<int> ((y-bboxhi[1])*bininvy) + nbiny;
		else if (y >= bboxlo[1]) {
		iy = static_cast<int> ((y-bboxlo[1])*bininvy);
		iy = MIN(iy,nbiny-1);
		} else
		iy = static_cast<int> ((y-bboxlo[1])*bininvy) - 1;

		if (z >= bboxhi[2])
		iz = static_cast<int> ((z-bboxhi[2])*bininvz) + nbinz;
		else if (z >= bboxlo[2]) {
		iz = static_cast<int> ((z-bboxlo[2])*bininvz);
		iz = MIN(iz,nbinz-1);
		} else
		iz = static_cast<int> ((z-bboxlo[2])*bininvz) - 1;

		return (iz-mbinzlo)mbinymbinx + (iy-mbinylo)*mbinx + (ix-mbinxlo);
		}

		KOKKOS_INLINE_FUNCTION
		int coord2bin(const X_FLOAT & x,const X_FLOAT & y,const X_FLOAT & z, int* i) const
		{

Admin message