got fix_momentum_kokkos to compile (ca30c1ec) · Commits · 郑智淋 / lammps

src/KOKKOS/fix_momentum_kokkos.cpp

+18 −9

Original line number	Diff line number	Diff line
		@@ -15,6 +15,7 @@
		#include <string.h>
		#include "fix_momentum_kokkos.h"
		#include "atom_kokkos.h"
		#include "atom_masks.h"
		#include "domain.h"
		#include "group.h"
		#include "error.h"
		@@ -54,6 +55,7 @@ void FixMomentumKokkos<DeviceType>::init()

		template<class DeviceType>
		double FixMomentumKokkos<DeviceType>::get_kinetic_energy(
		int nlocal,
		typename AT::t_v_array_randomread v,
		typename AT::t_int_1d_randomread mask)
		{
		@@ -61,7 +63,7 @@ double FixMomentumKokkos<DeviceType>::get_kinetic_energy(
		// D.I. : does this rmass check make sense in Kokkos mode ?
		if (atom->rmass) {
		atomKK->sync(execution_space, RMASS_MASK);
		typename AT::t_float_1d_randomread rmass = atomKK->k_rmass;
		typename AT::t_float_1d_randomread rmass = atomKK->k_rmass.view<DeviceType>();
		Kokkos::parallel_reduce(nlocal, LAMMPS_LAMBDA(int i, double& update) {
		if (mask(i) & groupbit)
		update += rmass(i) *
		@@ -70,8 +72,8 @@ double FixMomentumKokkos<DeviceType>::get_kinetic_energy(
		} else {
		// D.I. : why is there no MASS_MASK ?
		atomKK->sync(execution_space, TYPE_MASK);
		typename AT::t_int_1d_randomread type = atomKK->k_type;
		typename AT::t_float_1d_randomread mass = atomKK->k_mass;
		typename AT::t_int_1d_randomread type = atomKK->k_type.view<DeviceType>();
		typename AT::t_float_1d_randomread mass = atomKK->k_mass.view<DeviceType>();
		Kokkos::parallel_reduce(nlocal, LAMMPS_LAMBDA(int i, double& update) {
		if (mask(i) & groupbit)
		update += mass(type(i)) *
		@@ -104,7 +106,7 @@ void FixMomentumKokkos<DeviceType>::end_of_step()

		// compute kinetic energy before momentum removal, if needed

		if (rescale) ekin_old = get_kinetic_energy(v, mask);
		if (rescale) ekin_old = get_kinetic_energy(nlocal, v, mask);

		if (linear) {
		double vcm[3];
		@@ -115,9 +117,9 @@ void FixMomentumKokkos<DeviceType>::end_of_step()

		Kokkos::parallel_for(nlocal, LAMMPS_LAMBDA(int i) {
		if (mask(i) & groupbit) {
		if (xflag) v(i,0) -= vcm(0);
		if (yflag) v(i,1) -= vcm(1);
		if (zflag) v(i,2) -= vcm(2);
		if (xflag) v(i,0) -= vcm[0];
		if (yflag) v(i,1) -= vcm[1];
		if (zflag) v(i,2) -= vcm[2];
		}
		});
		}
		@@ -142,7 +144,7 @@ void FixMomentumKokkos<DeviceType>::end_of_step()
		if (mask[i] & groupbit) {
		double dx,dy,dz;
		double unwrap[3];
		domain->unmap(x[i],image[i],unwrap);
		domain->unmap(&x(i,0),image(i),unwrap); // this will not work in CUDA !!!
		dx = unwrap[0] - xcm[0];
		dy = unwrap[1] - xcm[1];
		dz = unwrap[2] - xcm[2];
		@@ -157,7 +159,7 @@ void FixMomentumKokkos<DeviceType>::end_of_step()

		if (rescale) {

		ekin_new = get_kinetic_energy(v, mask);
		ekin_new = get_kinetic_energy(nlocal, v, mask);

		double factor = 1.0;
		if (ekin_new != 0.0) factor = sqrt(ekin_old/ekin_new);
		@@ -171,3 +173,10 @@ void FixMomentumKokkos<DeviceType>::end_of_step()
		}
		}

		namespace LAMMPS_NS {
		template class FixMomentumKokkos<LMPDeviceType>;
		#ifdef KOKKOS_HAVE_CUDA
		template class FixMomentumKokkos<LMPHostType>;
		#endif
		}

src/KOKKOS/fix_momentum_kokkos.h

+6 −1

Original line number	Diff line number	Diff line
		@@ -23,6 +23,7 @@ FixStyle(momentum/kk/host,FixMomentumKokkos<LMPHostType>)
		#define LMP_FIX_MOMENTUM_KOKKOS_H

		#include "fix_momentum.h"
		#include "kokkos_type.h"

		namespace LAMMPS_NS {

		@@ -31,11 +32,15 @@ class FixMomentumKokkos : public FixMomentum {
		public:
		typedef ArrayTypes<DeviceType> AT;

		FixMomentum(class LAMMPS , int, char *);
		FixMomentumKokkos(class LAMMPS , int, char *);
		void init();
		void end_of_step();

		private:
		double get_kinetic_energy(
		int nlocal,
		typename AT::t_v_array_randomread v,
		typename AT::t_int_1d_randomread mask);
		};

		}

src/domain.cpp

+1 −1

Original line number	Diff line number	Diff line
		@@ -1441,7 +1441,7 @@ void Domain::unmap(double *x, imageint image)
		for triclinic, use h[] to add in tilt factors in other dims as needed
		------------------------------------------------------------------------- */

		void Domain::unmap(double x, imageint image, double y)
		void Domain::unmap(const double x, imageint image, double y)
		{
		int xbox = (image & IMGMASK) - IMGMAX;
		int ybox = (image >> IMGBITS & IMGMASK) - IMGMAX;

src/domain.h

+1 −1

Original line number	Diff line number	Diff line
		@@ -119,7 +119,7 @@ class Domain : protected Pointers {
		void remap(double *);
		void remap_near(double , double );
		void unmap(double *, imageint);
		void unmap(double , imageint, double );
		void unmap(const double , imageint, double );
		void image_flip(int, int, int);
		int ownatom(int, double , imageint , int);

src/fix_momentum.h

+1 −1

Original line number	Diff line number	Diff line
		@@ -31,7 +31,7 @@ class FixMomentum : public Fix {
		void init();
		void end_of_step();

		private:
		protected:
		int linear,angular,rescale;
		int xflag,yflag,zflag;
		int dynamic;

Admin message