Some tweaks (9b955a6c) · Commits · 郑智淋 / lammps

src/KOKKOS/pair_snap_kokkos_impl.h

+3 −0

Original line number	Diff line number	Diff line
		@@ -234,6 +234,9 @@ void PairSNAPKokkos<DeviceType>::compute(int eflag_in, int vflag_in)
		typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPComputeUi> policy_ui(((inum+team_size-1)/team_size)*max_neighs,team_size,ui_vector_length);
		Kokkos::parallel_for("ComputeUi",policy_ui,*this);

		//Ulisttot transpose
		snaKK.transpose_ulisttot();

		//Compute bispectrum
		if (quadraticflag \|\| eflag) {
		//ComputeZi

src/KOKKOS/sna_kokkos.h

+26 −3

Original line number	Diff line number	Diff line
		@@ -53,12 +53,32 @@ public:
		typedef Kokkos::View<SNAcomplex*, DeviceType> t_sna_1c;
		typedef Kokkos::View<SNAcomplex*, DeviceType, Kokkos::MemoryTraits<Kokkos::Atomic> > t_sna_1c_atomic;
		typedef Kokkos::View<SNAcomplex**, DeviceType> t_sna_2c;
		typedef Kokkos::View<SNAcomplex**, Kokkos::LayoutRight, DeviceType> t_sna_2c_cpu;
		typedef Kokkos::View<SNAcomplex**, Kokkos::LayoutRight, DeviceType> t_sna_2c_lr;
		typedef Kokkos::View<SNAcomplex***, DeviceType> t_sna_3c;
		typedef Kokkos::View<SNAcomplex***[3], DeviceType> t_sna_4c;
		typedef Kokkos::View<SNAcomplex**[3], DeviceType> t_sna_3c3;
		typedef Kokkos::View<SNAcomplex*****, DeviceType> t_sna_5c;

		// Helper class to get ulisttot_r

		template<typename DeviceLayout, typename T1, typename T2>
		class UlisttotHelper {
		public:
		inline
		static void transpose(T1 &ulisttot_lr, const T2 &ulisttot) {
		Kokkos::deep_copy(ulisttot_lr,ulisttot);
		}
		};

		template<typename T1, typename T2>
		class UlisttotHelper<Kokkos::LayoutRight,T1,T2> {
		public:
		inline
		static void transpose(T1 &ulisttot_lr, const T2 &ulisttot) {
		ulisttot_lr = ulisttot;
		}
		};

		inline
		SNAKokkos() {};
		KOKKOS_INLINE_FUNCTION
		@@ -80,6 +100,9 @@ inline

		int ncoeff;

		inline
		void transpose_ulisttot();

		// functions for bispectrum coefficients
		KOKKOS_INLINE_FUNCTION
		void pre_ui(const typename Kokkos::TeamPolicy<DeviceType>::member_type& team, int); // ForceSNAP
		@@ -130,11 +153,11 @@ inline

		t_sna_2d blist;
		t_sna_2c ulisttot;
		t_sna_2c_cpu ulisttot_lr;
		t_sna_2c_lr ulisttot_lr;
		t_sna_2c zlist;

		t_sna_3c ulist;
		t_sna_2c ylist;
		t_sna_2c_lr ylist;

		// derivatives of data
		t_sna_4c dulist;

src/KOKKOS/sna_kokkos_impl.h

+11 −8

Original line number	Diff line number	Diff line
		@@ -226,11 +226,12 @@ void SNAKokkos<DeviceType>::grow_rij(int newnatom, int newnmax)

		blist = t_sna_2d("sna:blist",natom,idxb_max);
		ulisttot = t_sna_2c("sna:ulisttot",natom,idxu_max);
		ulisttot_lr = t_sna_2c_cpu("sna:ulisttot_lr",natom,idxu_max);
		if (!Kokkos::Impl::is_same<typename DeviceType::array_layout,Kokkos::LayoutRight>::value)
		ulisttot_lr = t_sna_2c_lr("sna:ulisttot_lr",natom,idxu_max);
		zlist = t_sna_2c("sna:zlist",natom,idxz_max);

		ulist = t_sna_3c("sna:ulist",natom,nmax,idxu_max);
		ylist = t_sna_2c("sna:ylist",natom,idxu_max);
		ylist = t_sna_2c_lr("sna:ylist",natom,idxu_max);

		dulist = t_sna_4c("sna:dulist",natom,nmax,idxu_max);
		}
		@@ -282,11 +283,6 @@ void SNAKokkos<DeviceType>::compute_ui(const typename Kokkos::TeamPolicy<DeviceT
		//Kokkos::single(Kokkos::PerThread(team), [&] (){
		add_uarraytot(team, iatom, jnbor, r, wj(iatom,jnbor), rcutij(iatom,jnbor));
		//});

		if (Kokkos::Impl::is_same<typename DeviceType::array_layout,Kokkos::LayoutRight>::value)
		ulisttot_lr = ulisttot;
		else
		Kokkos::deep_copy(ulisttot_lr,ulisttot);
		}

		template<class DeviceType>
		@@ -788,6 +784,12 @@ void SNAKokkos<DeviceType>::compute_uarray(const typename Kokkos::TeamPolicy<Dev
		}
		}

		template<class DeviceType>
		void SNAKokkos<DeviceType>::transpose_ulisttot()
		{
		UlisttotHelper<typename DeviceType::array_layout,decltype(ulisttot_lr),decltype(ulisttot)>::transpose(ulisttot_lr,ulisttot);
		}

		/* ----------------------------------------------------------------------
		compute derivatives of Wigner U-functions for one neighbor
		see comments in compute_uarray()
		@@ -1320,6 +1322,7 @@ double SNAKokkos<DeviceType>::memory_usage()

		bytes += natom * idxu_max * sizeof(double) * 2; // ulist
		bytes += natom * idxu_max * sizeof(double) * 2; // ulisttot
		if (!Kokkos::Impl::is_same<typename DeviceType::array_layout,Kokkos::LayoutRight>::value)
		bytes += natom * idxu_max * sizeof(double) * 2; // ulisttot_lr
		bytes += natom * idxu_max * 3 * sizeof(double) * 2; // dulist

Admin message