Merge pull request #1969 from stanmoore1/team_size (e1924c60) · Commits · 郑智淋 / lammps

src/KOKKOS/npair_kokkos.cpp

+32 −13

Original line number	Diff line number	Diff line
		@@ -214,7 +214,6 @@ void NPairKokkos<DeviceType,HALF_NEIGH,GHOST,TRI,SIZE>::build(NeighList *list_)
		#ifdef KOKKOS_ENABLE_CUDA
		#define BINS_PER_BLOCK 2
		const int factor = atoms_per_bin<64?2:1;
		Kokkos::TeamPolicy<DeviceType> config((mbins+factor-1)/factor,atoms_per_bin*factor);
		#else
		const int factor = 1;
		#endif
		@@ -227,9 +226,14 @@ void NPairKokkos<DeviceType,HALF_NEIGH,GHOST,TRI,SIZE>::build(NeighList *list_)
		if (SIZE) {
		NPairKokkosBuildFunctorSize<DeviceType,TRI?0:HALF_NEIGH,1,TRI> f(data,atoms_per_bin * 5 * sizeof(X_FLOAT) * factor);
		#ifdef KOKKOS_ENABLE_CUDA
		if (ExecutionSpaceFromDevice<DeviceType>::space == Device)
		if (ExecutionSpaceFromDevice<DeviceType>::space == Device) {
		int team_size = atoms_per_bin*factor;
		int team_size_max = Kokkos::TeamPolicy<DeviceType>(team_size,Kokkos::AUTO).team_size_max(f,Kokkos::ParallelForTag());
		if (team_size <= team_size_max) {
		Kokkos::TeamPolicy<DeviceType> config((mbins+factor-1)/factor,team_size);
		Kokkos::parallel_for(config, f);
		else
		} else Kokkos::parallel_for(nall, f); // fall back to flat method
		} else
		Kokkos::parallel_for(nall, f);
		#else
		Kokkos::parallel_for(nall, f);
		@@ -237,9 +241,14 @@ void NPairKokkos<DeviceType,HALF_NEIGH,GHOST,TRI,SIZE>::build(NeighList *list_)
		} else {
		NPairKokkosBuildFunctor<DeviceType,TRI?0:HALF_NEIGH,1,TRI> f(data,atoms_per_bin * 5 * sizeof(X_FLOAT) * factor);
		#ifdef KOKKOS_ENABLE_CUDA
		if (ExecutionSpaceFromDevice<DeviceType>::space == Device)
		if (ExecutionSpaceFromDevice<DeviceType>::space == Device) {
		int team_size = atoms_per_bin*factor;
		int team_size_max = Kokkos::TeamPolicy<DeviceType>(team_size,Kokkos::AUTO).team_size_max(f,Kokkos::ParallelForTag());
		if (team_size <= team_size_max) {
		Kokkos::TeamPolicy<DeviceType> config((mbins+factor-1)/factor,team_size);
		Kokkos::parallel_for(config, f);
		else
		} else Kokkos::parallel_for(nall, f); // fall back to flat method
		} else
		Kokkos::parallel_for(nall, f);
		#else
		Kokkos::parallel_for(nall, f);
		@@ -249,9 +258,14 @@ void NPairKokkos<DeviceType,HALF_NEIGH,GHOST,TRI,SIZE>::build(NeighList *list_)
		if (SIZE) {
		NPairKokkosBuildFunctorSize<DeviceType,HALF_NEIGH,0,0> f(data,atoms_per_bin * 5 * sizeof(X_FLOAT) * factor);
		#ifdef KOKKOS_ENABLE_CUDA
		if (ExecutionSpaceFromDevice<DeviceType>::space == Device)
		if (ExecutionSpaceFromDevice<DeviceType>::space == Device) {
		int team_size = atoms_per_bin*factor;
		int team_size_max = Kokkos::TeamPolicy<DeviceType>(team_size,Kokkos::AUTO).team_size_max(f,Kokkos::ParallelForTag());
		if (team_size <= team_size_max) {
		Kokkos::TeamPolicy<DeviceType> config((mbins+factor-1)/factor,team_size);
		Kokkos::parallel_for(config, f);
		else
		} else Kokkos::parallel_for(nall, f); // fall back to flat method
		} else
		Kokkos::parallel_for(nall, f);
		#else
		Kokkos::parallel_for(nall, f);
		@@ -259,9 +273,14 @@ void NPairKokkos<DeviceType,HALF_NEIGH,GHOST,TRI,SIZE>::build(NeighList *list_)
		} else {
		NPairKokkosBuildFunctor<DeviceType,HALF_NEIGH,0,0> f(data,atoms_per_bin * 5 * sizeof(X_FLOAT) * factor);
		#ifdef KOKKOS_ENABLE_CUDA
		if (ExecutionSpaceFromDevice<DeviceType>::space == Device)
		if (ExecutionSpaceFromDevice<DeviceType>::space == Device) {
		int team_size = atoms_per_bin*factor;
		int team_size_max = Kokkos::TeamPolicy<DeviceType>(team_size,Kokkos::AUTO).team_size_max(f,Kokkos::ParallelForTag());
		if (team_size <= team_size_max) {
		Kokkos::TeamPolicy<DeviceType> config((mbins+factor-1)/factor,team_size);
		Kokkos::parallel_for(config, f);
		else
		} else Kokkos::parallel_for(nall, f); // fall back to flat method
		} else
		Kokkos::parallel_for(nall, f);
		#else
		Kokkos::parallel_for(nall, f);

Admin message