Collapse loop in compute_zi (67a284c2) · Commits · 郑智淋 / lammps

src/KOKKOS/pair_snap_kokkos.h

+1 −1

Original line number	Diff line number	Diff line
		@@ -80,7 +80,7 @@ public:
		void operator() (TagPairSNAPComputeUi,const typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPComputeUi>::member_type& team) const;

		KOKKOS_INLINE_FUNCTION
		void operator() (TagPairSNAPComputeZi,const typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPComputeZi>::member_type& team) const;
		void operator() (TagPairSNAPComputeZi,const int& ii) const;

		KOKKOS_INLINE_FUNCTION
		void operator() (TagPairSNAPComputeBi,const typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPComputeBi>::member_type& team) const;

+4 −4

Original line number	Diff line number	Diff line
		@@ -230,7 +230,8 @@ void PairSNAPKokkos<DeviceType>::compute(int eflag_in, int vflag_in)
		//Compute bispectrum
		if (quadraticflag \|\| eflag) {
		//ComputeZi
		typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPComputeZi> policy_zi(chunk_size,team_size,vector_length);
		int idxz_max = snaKK.idxz_max;
		typename Kokkos::RangePolicy<DeviceType, TagPairSNAPComputeZi> policy_zi(0,chunk_size*idxz_max);
		Kokkos::parallel_for("ComputeZi",policy_zi,*this);

		//ComputeBi
		@@ -540,10 +541,9 @@ void PairSNAPKokkos<DeviceType>::operator() (TagPairSNAPComputeYi,const int& ii)

		template<class DeviceType>
		KOKKOS_INLINE_FUNCTION
		void PairSNAPKokkos<DeviceType>::operator() (TagPairSNAPComputeZi,const typename Kokkos::TeamPolicy<DeviceType, TagPairSNAPComputeZi>::member_type& team) const {
		int ii = team.league_rank();
		void PairSNAPKokkos<DeviceType>::operator() (TagPairSNAPComputeZi,const int& ii) const {
		SNAKokkos<DeviceType> my_sna = snaKK;
		my_sna.compute_zi(team,ii);
		my_sna.compute_zi(ii);
		}

		template<class DeviceType>

+10 −22

Original line number	Diff line number	Diff line
		@@ -33,22 +33,12 @@ struct alignas(2*sizeof(SNAreal)) SNAcomplex{

		KOKKOS_INLINE_FUNCTION
		SNAcomplex() : re(0),im(0)
		{;}
		{}

		KOKKOS_INLINE_FUNCTION
		SNAcomplex(SNAreal real_in, SNAreal imag_in)
		:re(real_in),im(imag_in)
		{;}

		#ifdef __CUDACC__
		KOKKOS_INLINE_FUNCTION
		SNAcomplex(double2 ref)
		:re(ref.x),im(ref.y)
		{;}

		operator double2() { return make_double2(re,im); }
		#endif

		{}
		};

		//struct SNAKK_ZINDICES {
		@@ -125,8 +115,6 @@ inline

		inline
		void transpose_ulisttot();
		inline
		void zero_yi();

		// functions for bispectrum coefficients
		KOKKOS_INLINE_FUNCTION
		@@ -136,7 +124,7 @@ inline
		KOKKOS_INLINE_FUNCTION
		void compute_ui_orig(const typename Kokkos::TeamPolicy<DeviceType>::member_type& team, int, int); // ForceSNAP
		KOKKOS_INLINE_FUNCTION
		void compute_zi(const typename Kokkos::TeamPolicy<DeviceType>::member_type& team, int); // ForceSNAP
		void compute_zi(const int&); // ForceSNAP
		KOKKOS_INLINE_FUNCTION
		void zero_yi(const int&);
		KOKKOS_INLINE_FUNCTION

+45 −47

Original line number	Diff line number	Diff line
		@@ -289,11 +289,11 @@ void SNAKokkos<DeviceType>::compute_ui(const typename Kokkos::TeamPolicy<DeviceT

		template<class DeviceType>
		KOKKOS_INLINE_FUNCTION
		void SNAKokkos<DeviceType>::compute_zi(const typename Kokkos::TeamPolicy<DeviceType>::member_type& team, int iatom)
		void SNAKokkos<DeviceType>::compute_zi(const int& iter)
		{
		Kokkos::parallel_for(Kokkos::TeamThreadRange(team,idxz_max),
		[&] (const int& jjz) {
		//for(int jjz = 0; jjz < idxz_max; jjz++) {
		const int iatom = iter / idxz_max;
		const int jjz = iter % idxz_max;

		const int j1 = idxz(jjz,0);
		const int j2 = idxz(jjz,1);
		const int j = idxz(jjz,2);
		@@ -335,8 +335,6 @@ void SNAKokkos<DeviceType>::compute_zi(const typename Kokkos::TeamPolicy<DeviceT
		jju2 -= j2+1;
		icgb += j2;
		} // end loop over ib

		}); // end loop over jjz
		}

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