Restore original compute_h in fix_qeq_reax_kokkos (85999fc4) · Commits · 郑智淋 / lammps

src/KOKKOS/fix_qeq_reax_kokkos.cpp

+88 −20

Original line number	Diff line number	Diff line
		@@ -227,6 +227,15 @@ void FixQEqReaxKokkos<DeviceType>::pre_force(int vflag)

		// compute_H

		if (lmp->kokkos->ngpus == 0) { // CPU
		if (neighflag == FULL) {
		FixQEqReaxKokkosComputeHFunctor<DeviceType, FULL> computeH_functor(this);
		Kokkos::parallel_scan(inum,computeH_functor);
		} else { // HALF and HALFTHREAD are the same
		FixQEqReaxKokkosComputeHFunctor<DeviceType, HALF> computeH_functor(this);
		Kokkos::parallel_scan(inum,computeH_functor);
		}
		} else { // GPU, use teams
		Kokkos::deep_copy(d_mfill_offset,0);

		int vector_length = 32;
		@@ -239,14 +248,11 @@ void FixQEqReaxKokkos<DeviceType>::pre_force(int vflag)
		FixQEqReaxKokkosComputeHFunctor<DeviceType, FULL> computeH_functor(
		this, atoms_per_team, vector_length);
		Kokkos::parallel_for(policy, computeH_functor);
		} else if (neighflag == HALF) {
		} else { // HALF and HALFTHREAD are the same
		FixQEqReaxKokkosComputeHFunctor<DeviceType, HALF> computeH_functor(
		this, atoms_per_team, vector_length);
		Kokkos::parallel_for(policy, computeH_functor);
		} else {
		FixQEqReaxKokkosComputeHFunctor<DeviceType, HALFTHREAD> computeH_functor(
		this, atoms_per_team, vector_length);
		Kokkos::parallel_for(policy, computeH_functor);
		}
		}

		// init_matvec
		@@ -401,6 +407,68 @@ void FixQEqReaxKokkos<DeviceType>::zero_item(int ii) const

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

		template<class DeviceType>
		template <int NEIGHFLAG>
		KOKKOS_INLINE_FUNCTION
		void FixQEqReaxKokkos<DeviceType>::compute_h_item(int ii, int &m_fill, const bool &final) const
		{
		const int i = d_ilist[ii];
		int j,jj,jtype;

		if (mask[i] & groupbit) {

		const X_FLOAT xtmp = x(i,0);
		const X_FLOAT ytmp = x(i,1);
		const X_FLOAT ztmp = x(i,2);
		const int itype = type(i);
		const tagint itag = tag(i);
		const int jnum = d_numneigh[i];
		if (final)
		d_firstnbr[i] = m_fill;

		for (jj = 0; jj < jnum; jj++) {
		j = d_neighbors(i,jj);
		j &= NEIGHMASK;
		jtype = type(j);

		const X_FLOAT delx = x(j,0) - xtmp;
		const X_FLOAT dely = x(j,1) - ytmp;
		const X_FLOAT delz = x(j,2) - ztmp;

		if (NEIGHFLAG != FULL) {
		// skip half of the interactions
		const tagint jtag = tag(j);
		if (j >= nlocal) {
		if (itag > jtag) {
		if ((itag+jtag) % 2 == 0) continue;
		} else if (itag < jtag) {
		if ((itag+jtag) % 2 == 1) continue;
		} else {
		if (x(j,2) < ztmp) continue;
		if (x(j,2) == ztmp && x(j,1) < ytmp) continue;
		if (x(j,2) == ztmp && x(j,1) == ytmp && x(j,0) < xtmp) continue;
		}
		}
		}

		const F_FLOAT rsq = delxdelx + delydely + delz*delz;
		if (rsq > cutsq) continue;

		if (final) {
		const F_FLOAT r = sqrt(rsq);
		d_jlist(m_fill) = j;
		const F_FLOAT shldij = d_shield(itype,jtype);
		d_val(m_fill) = calculate_H_k(r,shldij);
		}
		m_fill++;
		}
		if (final)
		d_numnbrs[i] = m_fill - d_firstnbr[i];
		}
		}

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

		// Calculate Qeq matrix H where H is a sparse matrix and H[i][j] represents the electrostatic interaction coefficients on atom-i with atom-j
		// d_val - contains the non-zero entries of sparse matrix H
		// d_numnbrs - d_numnbrs[i] contains the # of non-zero entries in the i-th row of H (which also represents the # of neighbor atoms with electrostatic interaction coefficients with atom-i)

src/KOKKOS/fix_qeq_reax_kokkos.h

+14 −1

Original line number	Diff line number	Diff line
		@@ -53,6 +53,10 @@ class FixQEqReaxKokkos : public FixQEqReax {
		KOKKOS_INLINE_FUNCTION
		void zero_item(int) const;

		template<int NEIGHFLAG>
		KOKKOS_INLINE_FUNCTION
		void compute_h_item(int, int &, const bool &) const;

		template<int NEIGHFLAG>
		KOKKOS_INLINE_FUNCTION
		void compute_h_team(const typename Kokkos::TeamPolicy <DeviceType> ::member_type &team, int, int) const;
		@@ -151,7 +155,6 @@ class FixQEqReaxKokkos : public FixQEqReax {
		int allocated_flag;
		int need_dup;

		DAT::tdual_int_scalar k_mfill_offset;
		typename AT::t_int_scalar d_mfill_offset;

		typedef Kokkos::DualView<int***,DeviceType> tdual_int_1d;
		@@ -254,9 +257,14 @@ struct FixQEqReaxKokkosMatVecFunctor {
		template <class DeviceType, int NEIGHFLAG>
		struct FixQEqReaxKokkosComputeHFunctor {
		int atoms_per_team, vector_length;
		typedef int value_type;
		typedef Kokkos::ScratchMemorySpace<DeviceType> scratch_space;
		FixQEqReaxKokkos<DeviceType> c;

		FixQEqReaxKokkosComputeHFunctor(FixQEqReaxKokkos<DeviceType>* c_ptr):c(*c_ptr) {
		c.cleanup_copy();
		};

		FixQEqReaxKokkosComputeHFunctor(FixQEqReaxKokkos<DeviceType> *c_ptr,
		int _atoms_per_team, int _vector_length)
		: c(*c_ptr), atoms_per_team(_atoms_per_team),
		@@ -264,6 +272,11 @@ struct FixQEqReaxKokkosComputeHFunctor {
		c.cleanup_copy();
		};

		KOKKOS_INLINE_FUNCTION
		void operator()(const int ii, int &m_fill, const bool &final) const {
		c.template compute_h_item<NEIGHFLAG>(ii,m_fill,final);
		}

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

Admin message