Fix team_size issue in comm_kokkos (5c0f0274) · Commits · 郑智淋 / lammps

src/KOKKOS/comm_kokkos.cpp

+6 −1

Original line number	Diff line number	Diff line
		@@ -830,6 +830,11 @@ void CommKokkos::borders_device() {
		ExecutionSpace exec_space = ExecutionSpaceFromDevice<DeviceType>::space;
		atomKK->sync(exec_space,ALL_MASK);

		int team_size = 1;
		#ifdef KOKKOS_ENABLE_CUDA
		team_size = 128;
		#endif

		// do swaps over all 3 dimensions

		iswap = 0;
		@@ -883,7 +888,7 @@ void CommKokkos::borders_device() {

		BuildBorderListFunctor<DeviceType> f(atomKK->k_x,k_sendlist,
		k_total_send,nfirst,nlast,dim,lo,hi,iswap,maxsendlist[iswap]);
		Kokkos::TeamPolicy<DeviceType> config((nlast-nfirst+127)/128,128);
		Kokkos::TeamPolicy<DeviceType> config((nlast-nfirst+team_size-1)/team_size,team_size);
		Kokkos::parallel_for(config,f);

		k_total_send.template modify<DeviceType>();