Add squashed comm forward for Kokkos (0a02097e) · Commits · 郑智淋 / lammps

src/KOKKOS/atom_vec_kokkos.cpp

+100 −0

Changes for src/KOKKOS/atom_vec_kokkos.cpp: 100 added lines, 0 removed lines.
Original line number	Diff line number	Diff line
		@@ -267,6 +267,106 @@ int AtomVecKokkos::pack_comm_self(const int &n, const DAT::tdual_int_2d &list, c
		return n*3;
		}


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

		template<class DeviceType,int TRICLINIC>
		struct AtomVecKokkos_PackCommSelfSquash {
		typedef DeviceType device_type;

		typename ArrayTypes<DeviceType>::t_x_array_randomread _x;
		typename ArrayTypes<DeviceType>::t_x_array _xw;
		typename ArrayTypes<DeviceType>::t_int_2d_const _list;
		typename ArrayTypes<DeviceType>::t_int_2d_const _pbc;
		typename ArrayTypes<DeviceType>::t_int_1d_const _pbc_flag;
		typename ArrayTypes<DeviceType>::t_int_1d_const _firstrecv;
		typename ArrayTypes<DeviceType>::t_int_1d_const _sendnum_scan;
		X_FLOAT _xprd,_yprd,_zprd,_xy,_xz,_yz;

		AtomVecKokkos_PackCommSelfSquash(
		const typename DAT::tdual_x_array &x,
		const typename DAT::tdual_int_2d &list,
		const typename DAT::tdual_int_2d &pbc,
		const typename DAT::tdual_int_1d &pbc_flag,
		const typename DAT::tdual_int_1d &firstrecv,
		const typename DAT::tdual_int_1d &sendnum_scan,
		const X_FLOAT &xprd, const X_FLOAT &yprd, const X_FLOAT &zprd,
		const X_FLOAT &xy, const X_FLOAT &xz, const X_FLOAT &yz):
		_x(x.view<DeviceType>()),_xw(x.view<DeviceType>()),
		_list(list.view<DeviceType>()),
		_pbc(pbc.view<DeviceType>()),
		_pbc_flag(pbc_flag.view<DeviceType>()),
		_firstrecv(firstrecv.view<DeviceType>()),
		_sendnum_scan(sendnum_scan.view<DeviceType>()),
		_xprd(xprd),_yprd(yprd),_zprd(zprd),
		_xy(xy),_xz(xz),_yz(yz) {};

		KOKKOS_INLINE_FUNCTION
		void operator() (const int& ii) const {

		int iswap = 0;
		while (ii >= _sendnum_scan[iswap]) iswap++;
		int i = ii;
		if (iswap > 1)
		i = ii - _sendnum_scan[iswap-1];
		const int _nfirst = _firstrecv[iswap];

		const int j = _list(iswap,i);
		if (_pbc_flag(iswap) == 0) {
		_xw(i+_nfirst,0) = _x(j,0);
		_xw(i+_nfirst,1) = _x(j,1);
		_xw(i+_nfirst,2) = _x(j,2);
		} else {
		if (TRICLINIC == 0) {
		_xw(i+_nfirst,0) = _x(j,0) + _pbc(iswap,0)*_xprd;
		_xw(i+_nfirst,1) = _x(j,1) + _pbc(iswap,1)*_yprd;
		_xw(i+_nfirst,2) = _x(j,2) + _pbc(iswap,2)*_zprd;
		} else {
		_xw(i+_nfirst,0) = _x(j,0) + _pbc(iswap,0)_xprd + _pbc(iswap,5)_xy + _pbc(iswap,4)*_xz;
		_xw(i+_nfirst,1) = _x(j,1) + _pbc(iswap,1)_yprd + _pbc(iswap,3)_yz;
		_xw(i+_nfirst,2) = _x(j,2) + _pbc(iswap,2)*_zprd;
		}
		}

		}
		};

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

		int AtomVecKokkos::pack_comm_self_squash(const int &n, const DAT::tdual_int_2d &list, const DAT::tdual_int_1d &sendnum_scan,
		const DAT::tdual_int_1d &firstrecv, const DAT::tdual_int_1d &pbc_flag, const DAT::tdual_int_2d &pbc) {
		if(commKK->forward_comm_on_host) {
		sync(Host,X_MASK);
		modified(Host,X_MASK);
		if(domain->triclinic) {
		struct AtomVecKokkos_PackCommSelfSquash<LMPHostType,1> f(atomKK->k_x,list,pbc,pbc_flag,firstrecv,sendnum_scan,
		domain->xprd,domain->yprd,domain->zprd,
		domain->xy,domain->xz,domain->yz);
		Kokkos::parallel_for(n,f);
		} else {
		struct AtomVecKokkos_PackCommSelfSquash<LMPHostType,0> f(atomKK->k_x,list,pbc,pbc_flag,firstrecv,sendnum_scan,
		domain->xprd,domain->yprd,domain->zprd,
		domain->xy,domain->xz,domain->yz);
		Kokkos::parallel_for(n,f);
		}
		} else {
		sync(Device,X_MASK);
		modified(Device,X_MASK);
		if(domain->triclinic) {
		struct AtomVecKokkos_PackCommSelfSquash<LMPDeviceType,1> f(atomKK->k_x,list,pbc,pbc_flag,firstrecv,sendnum_scan,
		domain->xprd,domain->yprd,domain->zprd,
		domain->xy,domain->xz,domain->yz);
		Kokkos::parallel_for(n,f);
		} else {
		struct AtomVecKokkos_PackCommSelfSquash<LMPDeviceType,0> f(atomKK->k_x,list,pbc,pbc_flag,firstrecv,sendnum_scan,
		domain->xprd,domain->yprd,domain->zprd,
		domain->xy,domain->xz,domain->yz);
		Kokkos::parallel_for(n,f);
		}
		}
		return n*3;
		}

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

		template<class DeviceType>

src/KOKKOS/atom_vec_kokkos.h

+7 −0

Changes for src/KOKKOS/atom_vec_kokkos.h: 7 added lines, 0 removed lines.
Original line number	Diff line number	Diff line
		@@ -51,6 +51,13 @@ class AtomVecKokkos : public AtomVec {
		const int & iswap, const int nfirst,
		const int &pbc_flag, const int pbc[]);

		virtual int
		pack_comm_self_squash(const int &n, const DAT::tdual_int_2d &list,
		const DAT::tdual_int_1d &sendnum_scan,
		const DAT::tdual_int_1d &firstrecv,
		const DAT::tdual_int_1d &pbc_flag,
		const DAT::tdual_int_2d &pbc);

		virtual int
		pack_comm_kokkos(const int &n, const DAT::tdual_int_2d &list,
		const int & iswap, const DAT::tdual_xfloat_2d &buf,

src/KOKKOS/comm_kokkos.cpp

+35 −0

Changes for src/KOKKOS/comm_kokkos.cpp: 35 added lines, 0 removed lines.
Original line number	Diff line number	Diff line
		@@ -242,10 +242,15 @@ void CommKokkos::forward_comm_device(int dummy)
		}
		} else {
		if (!ghost_velocity) {
		if (1) {
		n = avec->pack_comm_self_squash(totalsend,k_sendlist,k_sendnum_scan,
		k_firstrecv,k_pbc_flag,k_pbc);
		} else {
		if (sendnum[iswap])
		n = avec->pack_comm_self(sendnum[iswap],k_sendlist,iswap,
		firstrecv[iswap],pbc_flag[iswap],pbc[iswap]);
		DeviceType::fence();
		}
		} else {
		n = avec->pack_comm_vel_kokkos(sendnum[iswap],k_sendlist,iswap,
		k_buf_send,pbc_flag[iswap],pbc[iswap]);
		@@ -1036,6 +1041,36 @@ void CommKokkos::borders_device() {
		atomKK->sync(Host,TAG_MASK);
		atom->map_set();
		}

		if (1) {
		k_pbc = DAT::tdual_int_2d("comm:pbc",nswap,6);
		k_pbc_flag = DAT::tdual_int_1d("comm:pbc_flag",nswap);
		k_firstrecv = DAT::tdual_int_1d("comm:firstrecv",nswap);
		k_sendnum_scan = DAT::tdual_int_1d("comm:sendnum_scan",nswap);
		int scan = 0;
		for (int iswap = 0; iswap < nswap; iswap++) {
		scan += sendnum[iswap];
		k_sendnum_scan.h_view[iswap] = scan;
		k_firstrecv.h_view[iswap] = firstrecv[iswap];
		k_pbc_flag.h_view[iswap] = pbc_flag[iswap];
		k_pbc.h_view(iswap,0) = pbc[iswap][0];
		k_pbc.h_view(iswap,1) = pbc[iswap][1];
		k_pbc.h_view(iswap,2) = pbc[iswap][2];
		k_pbc.h_view(iswap,3) = pbc[iswap][3];
		k_pbc.h_view(iswap,4) = pbc[iswap][4];
		k_pbc.h_view(iswap,5) = pbc[iswap][5];
		}
		totalsend = scan;
		k_pbc .modify<LMPHostType>();
		k_pbc_flag .modify<LMPHostType>();
		k_firstrecv .modify<LMPHostType>();
		k_sendnum_scan.modify<LMPHostType>();

		k_pbc .sync<LMPDeviceType>();
		k_pbc_flag .sync<LMPDeviceType>();
		k_firstrecv .sync<LMPDeviceType>();
		k_sendnum_scan.sync<LMPDeviceType>();
		}
		}
		/* ----------------------------------------------------------------------
		realloc the size of the send buffer as needed with BUFFACTOR and bufextra

src/KOKKOS/comm_kokkos.h

+6 −0

Changes for src/KOKKOS/comm_kokkos.h: 6 added lines, 0 removed lines.
Original line number	Diff line number	Diff line
		@@ -63,6 +63,12 @@ class CommKokkos : public CommBrick {
		//double *buf_send; // send buffer for all comm
		//double *buf_recv; // recv buffer for all comm

		DAT::tdual_int_2d k_pbc;
		DAT::tdual_int_1d k_pbc_flag;
		DAT::tdual_int_1d k_firstrecv;
		DAT::tdual_int_1d k_sendnum_scan;
		int totalsend;

		int max_buf_pair;
		DAT::tdual_xfloat_1d k_buf_send_pair;
		DAT::tdual_xfloat_1d k_buf_recv_pair;

Admin message