Optimize KOKKOS package for small system sizes (71a62272) · Commits · 郑智淋 / lammps

src/KOKKOS/atom_vec_atomic_kokkos.cpp

+3 −2

Changes for src/KOKKOS/atom_vec_atomic_kokkos.cpp: 3 added lines, 2 removed lines.
Original line number	Diff line number	Diff line
		@@ -24,7 +24,7 @@

		using namespace LAMMPS_NS;

		#define DELTA 10000
		#define DELTA 10

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

		@@ -55,7 +55,8 @@ AtomVecAtomicKokkos::AtomVecAtomicKokkos(LAMMPS *lmp) : AtomVecKokkos(lmp)

		void AtomVecAtomicKokkos::grow(int n)
		{
		if (n == 0) nmax += DELTA;
		int step = MAX(DELTA,nmax*0.01);
		if (n == 0) nmax += step;
		else nmax = n;
		atomKK->nmax = nmax;
		if (nmax < 0 \|\| nmax > MAXSMALLINT)

src/KOKKOS/domain_kokkos.cpp

+12 −0

Changes for src/KOKKOS/domain_kokkos.cpp: 12 added lines, 0 removed lines.
Original line number	Diff line number	Diff line
		@@ -17,6 +17,7 @@
		#include "error.h"
		#include "force.h"
		#include "kspace.h"
		#include "kokkos.h"

		using namespace LAMMPS_NS;

		@@ -339,6 +340,17 @@ struct DomainPBCFunctor {

		void DomainKokkos::pbc()
		{

		if (lmp->kokkos->exchange_comm_classic) {

		// reduce GPU data movement

		atomKK->sync(Host,X_MASK\|V_MASK\|MASK_MASK\|IMAGE_MASK);
		Domain::pbc();
		atomKK->modified(Host,X_MASK\|V_MASK\|MASK_MASK\|IMAGE_MASK);
		return;
		}

		double lo,hi,*period;
		int nlocal = atomKK->nlocal;

src/KOKKOS/npair_kokkos.cpp

+18 −6

Changes for src/KOKKOS/npair_kokkos.cpp: 18 added lines, 6 removed lines.
Original line number	Diff line number	Diff line
		@@ -15,6 +15,7 @@
		#include "atom_kokkos.h"
		#include "atom_masks.h"
		#include "domain_kokkos.h"
		#include "update.h"
		#include "neighbor_kokkos.h"
		#include "nbin_kokkos.h"
		#include "nstencil.h"
		@@ -27,6 +28,16 @@ namespace LAMMPS_NS {
		template<class DeviceType, int HALF_NEIGH, int GHOST, int TRI, int SIZE>
		NPairKokkos<DeviceType,HALF_NEIGH,GHOST,TRI,SIZE>::NPairKokkos(LAMMPS *lmp) : NPair(lmp) {

		// use 1D view for scalars to reduce GPU memory operations

		d_scalars = typename AT::t_int_1d("neighbor:scalars",2);
		h_scalars = HAT::t_int_1d("neighbor:scalars_mirror",2);

		d_resize = Kokkos::subview(d_scalars,0);
		d_new_maxneighs = Kokkos::subview(d_scalars,1);

		h_resize = Kokkos::subview(h_scalars,0);
		h_new_maxneighs = Kokkos::subview(h_scalars,1);
		}

		/* ----------------------------------------------------------------------
		@@ -84,9 +95,11 @@ template<class DeviceType, int HALF_NEIGH, int GHOST, int TRI, int SIZE>
		void NPairKokkos<DeviceType,HALF_NEIGH,GHOST,TRI,SIZE>::copy_stencil_info()
		{
		NPair::copy_stencil_info();

		nstencil = ns->nstencil;

		if (neighbor->last_setup_bins == update->ntimestep) {
		// copy stencil to device as it may have changed

		int maxstencil = ns->get_maxstencil();

		if (maxstencil > k_stencil.extent(0))
		@@ -107,6 +120,7 @@ void NPairKokkos<DeviceType,HALF_NEIGH,GHOST,TRI,SIZE>::copy_stencil_info()
		k_stencilxyz.sync<DeviceType>();
		}
		}
		}

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

		@@ -157,7 +171,7 @@ void NPairKokkos<DeviceType,HALF_NEIGH,GHOST,TRI,SIZE>::build(NeighList *list_)
		bboxhi,bboxlo,
		domain->xperiodic,domain->yperiodic,domain->zperiodic,
		domain->xprd_half,domain->yprd_half,domain->zprd_half,
		skin);
		skin,d_resize,h_resize,d_new_maxneighs,h_new_maxneighs);

		k_cutneighsq.sync<DeviceType>();
		k_ex1_type.sync<DeviceType>();
		@@ -185,8 +199,7 @@ void NPairKokkos<DeviceType,HALF_NEIGH,GHOST,TRI,SIZE>::build(NeighList *list_)
		data.h_new_maxneighs() = list->maxneighs;
		data.h_resize() = 0;

		Kokkos::deep_copy(data.resize, data.h_resize);
		Kokkos::deep_copy(data.new_maxneighs, data.h_new_maxneighs);
		Kokkos::deep_copy(d_scalars, h_scalars);
		#ifdef KOKKOS_ENABLE_CUDA
		#define BINS_PER_BLOCK 2
		const int factor = atoms_per_bin<64?2:1;
		@@ -245,10 +258,9 @@ void NPairKokkos<DeviceType,HALF_NEIGH,GHOST,TRI,SIZE>::build(NeighList *list_)
		}
		}
		}
		deep_copy(data.h_resize, data.resize);
		Kokkos::deep_copy(h_scalars, d_scalars);

		if(data.h_resize()) {
		deep_copy(data.h_new_maxneighs, data.new_maxneighs);
		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);
		data.neigh_list.d_neighbors = list->d_neighbors;

src/KOKKOS/npair_kokkos.h

+15 −15

Changes for src/KOKKOS/npair_kokkos.h: 15 added lines, 15 removed lines.
Original line number	Diff line number	Diff line
		@@ -95,6 +95,8 @@ namespace LAMMPS_NS {

		template<class DeviceType, int HALF_NEIGH, int GHOST, int TRI, int SIZE>
		class NPairKokkos : public NPair {
		typedef ArrayTypes<DeviceType> AT;

		public:
		NPairKokkos(class LAMMPS *);
		~NPairKokkos() {}
		@@ -105,6 +107,12 @@ class NPairKokkos : public NPair {

		private:
		int newton_pair;
		typename AT::t_int_1d d_scalars;
		HAT::t_int_1d h_scalars;
		typename AT::t_int_scalar d_resize;
		typename AT::t_int_scalar d_new_maxneighs;
		HAT::t_int_scalar h_resize;
		HAT::t_int_scalar h_new_maxneighs;

		// data from Neighbor class

		@@ -251,7 +259,11 @@ class NeighborKokkosExecute
		const X_FLOAT _bboxhi, const X_FLOAT _bboxlo,
		const int & _xperiodic, const int & _yperiodic, const int & _zperiodic,
		const int & _xprd_half, const int & _yprd_half, const int & _zprd_half,
		const X_FLOAT _skin):
		const X_FLOAT _skin,
		const typename AT::t_int_scalar _resize,
		const typename ArrayTypes<LMPHostType>::t_int_scalar _h_resize,
		const typename AT::t_int_scalar _new_maxneighs,
		const typename ArrayTypes<LMPHostType>::t_int_scalar _h_new_maxneighs):
		neigh_list(_neigh_list), cutneighsq(_cutneighsq),
		bincount(_bincount),c_bincount(_bincount),bins(_bins),c_bins(_bins),
		atom2bin(_atom2bin),c_atom2bin(_atom2bin),
		@@ -272,7 +284,8 @@ class NeighborKokkosExecute
		ex_mol_intra(_ex_mol_intra),
		xperiodic(_xperiodic),yperiodic(_yperiodic),zperiodic(_zperiodic),
		xprd_half(_xprd_half),yprd_half(_yprd_half),zprd_half(_zprd_half),
		skin(_skin) {
		skin(_skin),resize(_resize),h_resize(_h_resize),
		new_maxneighs(_new_maxneighs),h_new_maxneighs(_h_new_maxneighs) {

		if (molecular == 2) moltemplate = 1;
		else moltemplate = 0;
		@@ -280,20 +293,7 @@ class NeighborKokkosExecute
		bboxlo[0] = _bboxlo[0]; bboxlo[1] = _bboxlo[1]; bboxlo[2] = _bboxlo[2];
		bboxhi[0] = _bboxhi[0]; bboxhi[1] = _bboxhi[1]; bboxhi[2] = _bboxhi[2];

		resize = typename AT::t_int_scalar(Kokkos::view_alloc("NeighborKokkosFunctor::resize",Kokkos::WithoutInitializing));
		#ifndef KOKKOS_USE_CUDA_UVM
		h_resize = Kokkos::create_mirror_view(resize);
		#else
		h_resize = resize;
		#endif
		h_resize() = 1;
		new_maxneighs = typename AT::
		t_int_scalar(Kokkos::view_alloc("NeighborKokkosFunctor::new_maxneighs",Kokkos::WithoutInitializing));
		#ifndef KOKKOS_USE_CUDA_UVM
		h_new_maxneighs = Kokkos::create_mirror_view(new_maxneighs);
		#else
		h_new_maxneighs = new_maxneighs;
		#endif
		h_new_maxneighs() = neigh_list.maxneighs;
		};

src/KOKKOS/pair_kokkos.h

+4 −5

Changes for src/KOKKOS/pair_kokkos.h: 4 added lines, 5 removed lines.
Original line number	Diff line number	Diff line
		@@ -388,7 +388,6 @@ struct PairComputeFunctor {
		const int lastatom = firstatom + atoms_per_team < inum ? firstatom + atoms_per_team : inum;
		Kokkos::parallel_for(Kokkos::TeamThreadRange(team, firstatom, lastatom), [&] (const int &ii) {


		const int i = list.d_ilist[ii];
		const X_FLOAT xtmp = c.x(i,0);
		const X_FLOAT ytmp = c.x(i,1);
		@@ -875,14 +874,14 @@ EV_FLOAT pair_compute_neighlist (PairStyle* fpair, typename Kokkos::Impl::enable
		PairComputeFunctor<PairStyle,NEIGHFLAG,false,Specialisation > ff(fpair,list);
		atoms_per_team = GetTeamSize(ff, atoms_per_team, vector_length);
		Kokkos::TeamPolicy<Kokkos::IndexType<int> > policy(list->inum,atoms_per_team,vector_length);
		if (fpair->eflag \|\| fpair->vflag) Kokkos::parallel_reduce(policy,ff,ev);
		else Kokkos::parallel_for(policy,ff);
		if (fpair->eflag \|\| fpair->vflag) Kokkos::parallel_reduce(Kokkos::Experimental::require(policy,Kokkos::Experimental::WorkItemProperty::HintLightWeight),ff,ev);
		else Kokkos::parallel_for(Kokkos::Experimental::require(policy,Kokkos::Experimental::WorkItemProperty::HintLightWeight),ff);
		} else {
		PairComputeFunctor<PairStyle,NEIGHFLAG,true,Specialisation > ff(fpair,list);
		atoms_per_team = GetTeamSize(ff, atoms_per_team, vector_length);
		Kokkos::TeamPolicy<Kokkos::IndexType<int> > policy(list->inum,atoms_per_team,vector_length);
		if (fpair->eflag \|\| fpair->vflag) Kokkos::parallel_reduce(policy,ff,ev);
		else Kokkos::parallel_for(policy,ff);
		if (fpair->eflag \|\| fpair->vflag) Kokkos::parallel_reduce(Kokkos::Experimental::require(policy,Kokkos::Experimental::WorkItemProperty::HintLightWeight),ff,ev);
		else Kokkos::parallel_for(Kokkos::Experimental::require(policy,Kokkos::Experimental::WorkItemProperty::HintLightWeight),ff);
		}
		} else {
		if(fpair->atom->ntypes > MAX_TYPES_STACKPARAMS) {

Admin message