Optimize Kokkos comm for small systems (b50ef59a) · Commits · 郑智淋 / lammps

src/KOKKOS/comm_kokkos.cpp

+20 −14

Changes for src/KOKKOS/comm_kokkos.cpp: 20 added lines, 14 removed lines.
Original line number	Diff line number	Diff line
		@@ -57,10 +57,9 @@ CommKokkos::CommKokkos(LAMMPS *lmp) : CommBrick(lmp)
		memory->destroy(buf_recv);
		buf_recv = NULL;

		k_exchange_lists = DAT::
		tdual_int_1d("comm:k_exchange_lists",2,100);
		k_exchange_sendlist = Kokkos::subview(k_exchange_lists,0,KOKKOS::ALL);
		k_exchange_copylist = Kokkos::subview(k_exchange_lists,1,KOKKOS::ALL);
		k_exchange_lists = DAT::tdual_int_2d("comm:k_exchange_lists",2,100);
		k_exchange_sendlist = Kokkos::subview(k_exchange_lists,0,Kokkos::ALL);
		k_exchange_copylist = Kokkos::subview(k_exchange_lists,1,Kokkos::ALL);
		k_count = DAT::tdual_int_scalar("comm:k_count");
		k_sendflag = DAT::tdual_int_1d("comm:k_sendflag",100);

		@@ -188,6 +187,8 @@ void CommKokkos::forward_comm_device(int dummy)
		atomKK->sync(ExecutionSpaceFromDevice<DeviceType>::space,X_MASK);

		if (comm->nprocs == 1) {
		k_swap.sync<DeviceType>();
		k_pbc.sync<DeviceType>();
		n = avec->pack_comm_self_fused(totalsend,k_sendlist,k_sendnum_scan,
		k_firstrecv,k_pbc_flag,k_pbc);
		} else {
		@@ -620,8 +621,7 @@ void CommKokkos::exchange_device()
		}
		}

		auto k_exchange_lists_short = Kokkos::subview(k_exchange_lists,KOKKOS::ALL,k_count.h_view());
		k_exchange_lists_short.template sync<LMPHostType>();
		k_exchange_lists.sync<LMPHostType>();
		k_sendflag.sync<LMPHostType>();

		int sendpos = nlocal-1;
		@@ -635,6 +635,7 @@ void CommKokkos::exchange_device()
		k_exchange_copylist.h_view(i) = -1;
		}

		auto k_exchange_copylist_short = Kokkos::subview(k_exchange_copylist,k_count.h_view());
		k_exchange_copylist_short.modify<LMPHostType>();
		k_exchange_copylist_short.sync<DeviceType>();
		nsend = k_count.h_view();
		@@ -749,9 +750,7 @@ void CommKokkos::borders()
		if (!exchange_comm_classic) {
		if (exchange_comm_on_host) borders_device<LMPHostType>();
		else borders_device<LMPDeviceType>();
		return;
		}

		} else {
		atomKK->sync(Host,ALL_MASK);
		k_sendlist.sync<LMPHostType>();
		CommBrick::borders();
		@@ -759,6 +758,10 @@ void CommKokkos::borders()
		atomKK->modified(Host,ALL_MASK);
		}

		if (comm->nprocs == 1 && !forward_comm_classic)
		copy_pbc_info();
		}

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

		template<class DeviceType>
		@@ -1043,8 +1046,14 @@ void CommKokkos::borders_device() {
		atomKK->sync(Host,TAG_MASK);
		atom->map_set();
		}
		}

		if (comm->nprocs == 1) {
		/* ----------------------------------------------------------------------
		copy pbc info
		------------------------------------------------------------------------- */

		void CommKokkos::copy_pbc_info()
		{
		if (nswap > k_pbc.extent(0)) {
		k_pbc = DAT::tdual_int_2d("comm:pbc",nswap,6);
		k_swap = DAT::tdual_int_2d("comm:swap",3,nswap);
		@@ -1072,11 +1081,8 @@ void CommKokkos::borders_device() {

		k_swap.modify<LMPHostType>();
		k_pbc.modify<LMPHostType>();

		k_swap.sync<LMPDeviceType>();
		k_pbc.sync<LMPDeviceType>();
		}
		}

		/* ----------------------------------------------------------------------
		realloc the size of the send buffer as needed with BUFFACTOR and bufextra
		if flag = 1, realloc

src/KOKKOS/comm_kokkos.h

+2 −0

Changes for src/KOKKOS/comm_kokkos.h: 2 added lines, 0 removed lines.
Original line number	Diff line number	Diff line
		@@ -58,6 +58,7 @@ class CommKokkos : public CommBrick {
		DAT::tdual_int_2d k_sendlist;
		DAT::tdual_int_scalar k_total_send;
		DAT::tdual_xfloat_2d k_buf_send,k_buf_recv;
		DAT::tdual_int_2d k_exchange_lists;
		DAT::tdual_int_1d k_exchange_sendlist,k_exchange_copylist,k_sendflag;
		DAT::tdual_int_scalar k_count;
		//double *buf_send; // send buffer for all comm
		@@ -81,6 +82,7 @@ class CommKokkos : public CommBrick {
		void grow_recv_kokkos(int, ExecutionSpace space = Host);
		void grow_list(int, int);
		void grow_swap(int);
		void copy_pbc_info();
		};

		}

Admin message