Make Kokkos pair_snap chunksize user-settable (1b7a8d88) · Commits · 郑智淋 / lammps

src/KOKKOS/pair_snap_kokkos_impl.h

+16 −13

Original line number	Diff line number	Diff line
		@@ -198,7 +198,7 @@ void PairSNAPKokkos<DeviceType>::compute(int eflag_in, int vflag_in)
		d_ninside = Kokkos::View<int*, DeviceType>("PairSNAPKokkos:ninside",inum);
		}

		int chunk_size = MIN(2000,inum);
		chunk_size = MIN(chunk_size,inum);
		chunk_offset = 0;

		snaKK.grow_rij(chunk_size,max_neighs);
		@@ -221,7 +221,7 @@ void PairSNAPKokkos<DeviceType>::compute(int eflag_in, int vflag_in)
		Kokkos::parallel_for("PreUi",policy_preui,*this);

		//ComputeUi
		typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPComputeUi> policy_ui(((inum+team_size-1)/team_size)*max_neighs,team_size,vector_length);
		typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPComputeUi> policy_ui(((chunk_size+team_size-1)/team_size)*max_neighs,team_size,vector_length);
		Kokkos::parallel_for("ComputeUi",policy_ui,*this);

		//Ulisttot transpose
		@@ -253,11 +253,11 @@ void PairSNAPKokkos<DeviceType>::compute(int eflag_in, int vflag_in)
		Kokkos::parallel_for("ComputeYi",policy_yi,*this);

		//ComputeDuidrj
		typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPComputeDuidrj> policy_duidrj(((inum+team_size-1)/team_size)*max_neighs,team_size,vector_length);
		typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPComputeDuidrj> policy_duidrj(((chunk_size+team_size-1)/team_size)*max_neighs,team_size,vector_length);
		Kokkos::parallel_for("ComputeDuidrj",policy_duidrj,*this);

		//ComputeDeidrj
		typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPComputeDeidrj> policy_deidrj(((inum+team_size-1)/team_size)*max_neighs,team_size,vector_length);
		typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPComputeDeidrj> policy_deidrj(((chunk_size+team_size-1)/team_size)*max_neighs,team_size,vector_length);
		Kokkos::parallel_for("ComputeDeidrj",policy_deidrj,*this);

		//ComputeForce
		@@ -514,11 +514,12 @@ void PairSNAPKokkos<DeviceType>::operator() (TagPairSNAPComputeUi,const typename
		SNAKokkos<DeviceType> my_sna = snaKK;

		// Extract the atom number
		int ii = team.team_rank() + team.team_size() * (team.league_rank() % ((inum+team.team_size()-1)/team.team_size()));
		if (ii >= inum) return;
		int ii = team.team_rank() + team.team_size() * (team.league_rank() %
		((chunk_size+team.team_size()-1)/team.team_size()));
		if (ii >= chunk_size) return;

		// Extract the neighbor number
		const int jj = team.league_rank() / ((inum+team.team_size()-1)/team.team_size());
		const int jj = team.league_rank() / ((chunk_size+team.team_size()-1)/team.team_size());
		const int ninside = d_ninside(ii);
		if (jj >= ninside) return;

		@@ -560,11 +561,12 @@ void PairSNAPKokkos<DeviceType>::operator() (TagPairSNAPComputeDuidrj,const type
		SNAKokkos<DeviceType> my_sna = snaKK;

		// Extract the atom number
		int ii = team.team_rank() + team.team_size() * (team.league_rank() % ((inum+team.team_size()-1)/team.team_size()));
		if (ii >= inum) return;
		int ii = team.team_rank() + team.team_size() * (team.league_rank() %
		((chunk_size+team.team_size()-1)/team.team_size()));
		if (ii >= chunk_size) return;

		// Extract the neighbor number
		const int jj = team.league_rank() / ((inum+team.team_size()-1)/team.team_size());
		const int jj = team.league_rank() / ((chunk_size+team.team_size()-1)/team.team_size());
		const int ninside = d_ninside(ii);
		if (jj >= ninside) return;

		@@ -577,11 +579,12 @@ void PairSNAPKokkos<DeviceType>::operator() (TagPairSNAPComputeDeidrj,const type
		SNAKokkos<DeviceType> my_sna = snaKK;

		// Extract the atom number
		int ii = team.team_rank() + team.team_size() * (team.league_rank() % ((inum+team.team_size()-1)/team.team_size()));
		if (ii >= inum) return;
		int ii = team.team_rank() + team.team_size() * (team.league_rank() %
		((chunk_size+team.team_size()-1)/team.team_size()));
		if (ii >= chunk_size) return;

		// Extract the neighbor number
		const int jj = team.league_rank() / ((inum+team.team_size()-1)/team.team_size());
		const int jj = team.league_rank() / ((chunk_size+team.team_size()-1)/team.team_size());
		const int ninside = d_ninside(ii);
		if (jj >= ninside) return;

src/SNAP/pair_snap.cpp

+3 −0

Original line number	Diff line number	Diff line
		@@ -635,6 +635,7 @@ void PairSNAP::read_files(char coefffilename, char paramfilename)
		switchflag = 1;
		bzeroflag = 1;
		quadraticflag = 0;
		chunk_size = 2000;

		// open SNAP parameter file on proc 0

		@@ -698,6 +699,8 @@ void PairSNAP::read_files(char coefffilename, char paramfilename)
		bzeroflag = atoi(keyval);
		else if (strcmp(keywd,"quadraticflag") == 0)
		quadraticflag = atoi(keyval);
		else if (strcmp(keywd,"chunksize") == 0)
		chunk_size = atoi(keyval);
		else
		error->all(FLERR,"Incorrect SNAP parameter file");
		}

src/SNAP/pair_snap.h

+1 −0

Original line number	Diff line number	Diff line
		@@ -59,6 +59,7 @@ protected:
		double** bispectrum; // bispectrum components for all atoms in list
		int *map; // mapping from atom types to elements
		int twojmax, switchflag, bzeroflag;
		int chunk_size;
		double rfac0, rmin0, wj1, wj2;
		int rcutfacflag, twojmaxflag; // flags for required parameters
		int beta_max; // length of beta

Admin message