Add Kokkos version of compute coord/atom (a44e49e2) · Commits · 郑智淋 / lammps

src/KOKKOS/Install.sh

+2 −0

Original line number	Diff line number	Diff line
		@@ -85,6 +85,8 @@ action comm_kokkos.cpp
		action comm_kokkos.h
		action comm_tiled_kokkos.cpp
		action comm_tiled_kokkos.h
		action compute_coord_atom_kokkos.cpp
		action compute_coord_atom_kokkos.h
		action compute_orientorder_atom_kokkos.cpp
		action compute_orientorder_atom_kokkos.h
		action compute_temp_kokkos.cpp

src/KOKKOS/compute_coord_atom_kokkos.cpp

0 → 100644

+249 −0

Original line number	Diff line number	Diff line
		/* ----------------------------------------------------------------------
		LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
		http://lammps.sandia.gov, Sandia National Laboratories
		Steve Plimpton, sjplimp@sandia.gov

		Copyright (2003) Sandia Corporation. Under the terms of Contract
		DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government retains
		certain rights in this software. This software is distributed under
		the GNU General Public License.

		See the README file in the top-level LAMMPS directory.
		------------------------------------------------------------------------- */

		#include "compute_coord_atom_kokkos.h"
		#include <cmath>
		#include <cstring>
		#include "compute_orientorder_atom_kokkos.h"
		#include "atom_kokkos.h"
		#include "update.h"
		#include "modify.h"
		#include "neighbor_kokkos.h"
		#include "neigh_list.h"
		#include "neigh_request.h"
		#include "force.h"
		#include "pair.h"
		#include "comm.h"
		#include "group.h"
		#include "memory_kokkos.h"
		#include "error.h"
		#include "atom_masks.h"

		using namespace LAMMPS_NS;

		#define INVOKED_PERATOM 8

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

		template<class DeviceType>
		ComputeCoordAtomKokkos<DeviceType>::ComputeCoordAtomKokkos(LAMMPS lmp, int narg, char *arg) :
		ComputeCoordAtom(lmp, narg, arg)
		{
		kokkosable = 1;
		atomKK = (AtomKokkos *) atom;
		execution_space = ExecutionSpaceFromDevice<DeviceType>::space;
		datamask_read = EMPTY_MASK;
		datamask_modify = EMPTY_MASK;

		d_typelo = typename AT::t_int_1d("coord/atom:typelo",ncol);
		d_typehi = typename AT::t_int_1d("coord/atom:typehi",ncol);

		auto h_typelo = Kokkos::create_mirror_view(d_typelo);
		auto h_typehi = Kokkos::create_mirror_view(d_typehi);

		for (int i = 0; i < ncol; i++) {
		h_typelo(i) = typelo[i];
		h_typehi(i) = typehi[i];
		}

		Kokkos::deep_copy(d_typelo,h_typelo);
		Kokkos::deep_copy(d_typehi,h_typehi);
		}

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

		template<class DeviceType>
		ComputeCoordAtomKokkos<DeviceType>::~ComputeCoordAtomKokkos<DeviceType>()
		{
		if (copymode) return;

		memoryKK->destroy_kokkos(k_cvec,cvec);
		memoryKK->destroy_kokkos(k_carray,carray);
		}

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

		template<class DeviceType>
		void ComputeCoordAtomKokkos<DeviceType>::init()
		{
		ComputeCoordAtom::init();

		// need an occasional full neighbor list

		// irequest = neigh request made by parent class

		int irequest = neighbor->nrequest - 1;

		neighbor->requests[irequest]->
		kokkos_host = Kokkos::Impl::is_same<DeviceType,LMPHostType>::value &&
		!Kokkos::Impl::is_same<DeviceType,LMPDeviceType>::value;
		neighbor->requests[irequest]->
		kokkos_device = Kokkos::Impl::is_same<DeviceType,LMPDeviceType>::value;
		}

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

		template<class DeviceType>
		void ComputeCoordAtomKokkos<DeviceType>::compute_peratom()
		{
		invoked_peratom = update->ntimestep;

		// grow coordination array if necessary

		if (atom->nmax > nmax) {
		if (ncol == 1) {
		memoryKK->destroy_kokkos(k_cvec,cvec);
		nmax = atom->nmax;
		memoryKK->create_kokkos(k_cvec,cvec,nmax,"coord/atom:cvec");
		vector_atom = cvec;
		d_cvec = k_cvec.template view<DeviceType>();
		} else {
		memoryKK->destroy_kokkos(k_carray,carray);
		nmax = atom->nmax;
		memoryKK->create_kokkos(k_carray,carray,nmax,ncol,"coord/atom:carray");
		array_atom = carray;
		d_carray = k_carray.template view<DeviceType>();
		}
		}

		if (cstyle == ORIENT) {
		if (!(c_orientorder->invoked_flag & INVOKED_PERATOM)) {
		c_orientorder->compute_peratom();
		c_orientorder->invoked_flag \|= INVOKED_PERATOM;
		}
		nqlist = c_orientorder->nqlist;
		normv = c_orientorder->array_atom;
		comm->forward_comm_compute(this);

		if (!c_orientorder->kokkosable)
		error->all(FLERR,"Must use compute orientorder/atom/kk with compute coord/atom/kk");

		if (c_orientorder->execution_space == Host) {
		ComputeOrientOrderAtomKokkos<LMPHostType>* c_orientorder_kk;
		c_orientorder_kk = (ComputeOrientOrderAtomKokkos<LMPHostType>*) c_orientorder;
		c_orientorder_kk->k_qnarray.modify<LMPHostType>();
		c_orientorder_kk->k_qnarray.sync<DeviceType>();
		d_normv = c_orientorder_kk->k_qnarray.view<DeviceType>();
		} else {
		ComputeOrientOrderAtomKokkos<LMPDeviceType>* c_orientorder_kk;
		c_orientorder_kk = (ComputeOrientOrderAtomKokkos<LMPDeviceType>*) c_orientorder;
		c_orientorder_kk->k_qnarray.modify<LMPHostType>();
		c_orientorder_kk->k_qnarray.sync<DeviceType>();
		d_normv = c_orientorder_kk->k_qnarray.view<DeviceType>();
		}
		}

		// invoke full neighbor list (will copy or build if necessary)

		neighbor->build_one(list);

		inum = list->inum;
		NeighListKokkos<DeviceType>* k_list = static_cast<NeighListKokkos<DeviceType>*>(list);
		d_numneigh = k_list->d_numneigh;
		d_neighbors = k_list->d_neighbors;
		d_ilist = k_list->d_ilist;

		// compute coordination number(s) for each atom in group
		// use full neighbor list to count atoms less than cutoff

		atomKK->sync(execution_space,X_MASK\|TYPE_MASK\|MASK_MASK);
		x = atomKK->k_x.view<DeviceType>();
		type = atomKK->k_type.view<DeviceType>();
		mask = atomKK->k_mask.view<DeviceType>();

		copymode = 1;
		if (cstyle == CUTOFF) {
		if (ncol == 1) {
		typename Kokkos::RangePolicy<DeviceType, TagComputeCoordAtom<CUTOFF,1> > policy(0,inum);
		Kokkos::parallel_for("ComputeCoordAtom",policy,*this);
		} else {
		typename Kokkos::RangePolicy<DeviceType, TagComputeCoordAtom<CUTOFF,0> > policy(0,inum);
		Kokkos::parallel_for("ComputeCoordAtom",policy,*this);
		}
		} else if (cstyle == ORIENT) {
		typename Kokkos::RangePolicy<DeviceType, TagComputeCoordAtom<ORIENT,1> > policy(0,inum);
		Kokkos::parallel_for("ComputeCoordAtom",policy,*this);
		}
		copymode = 0;

		if (ncol == 1 \|\| cstyle == ORIENT) {
		k_cvec.modify<DeviceType>();
		k_cvec.sync<LMPHostType>();
		} else {
		k_carray.modify<DeviceType>();
		k_carray.sync<LMPHostType>();
		}

		}

		template<class DeviceType>
		template<int CSTYLE, int NCOL>
		KOKKOS_INLINE_FUNCTION
		void ComputeCoordAtomKokkos<DeviceType>::operator()(TagComputeCoordAtom<CSTYLE,NCOL>, const int &ii) const
		{
		const int i = d_ilist[ii];
		if (NCOL == 1 \|\| CSTYLE == ORIENT)
		d_cvec(i) = 0.0;
		else
		for (int m = 0; m < ncol; m++) d_carray(i,m) = 0.0;
		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 jnum = d_numneigh[i];

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

		if (NCOL == 1)
		if (!(mask[j] & jgroupbit)) continue;

		const int jtype = type[j];
		const F_FLOAT delx = x(j,0) - xtmp;
		const F_FLOAT dely = x(j,1) - ytmp;
		const F_FLOAT delz = x(j,2) - ztmp;
		const F_FLOAT rsq = delxdelx + delydely + delz*delz;
		if (rsq < cutsq) {
		if (CSTYLE == CUTOFF) {
		if (NCOL == 1) {
		if (jtype >= d_typelo[0] && jtype <= d_typehi[0])
		n++;
		} else {
		for (int m = 0; m < ncol; m++)
		if (jtype >= d_typelo[m] && jtype <= d_typehi[m])
		d_carray(i,m) += 1.0;
		}
		} else if (CSTYLE == ORIENT) {
		double dot_product = 0.0;
		for (int m=0; m < 2(2l+1); m++) {
		dot_product += d_normv(i,nqlist+m)*d_normv(j,nqlist+m);
		}
		if (dot_product > threshold) n++;
		}
		}
		}

		if (NCOL == 1 \|\| CSTYLE == ORIENT)
		d_cvec[i] = n;
		}

		}

		namespace LAMMPS_NS {
		template class ComputeCoordAtomKokkos<LMPDeviceType>;
		#ifdef KOKKOS_ENABLE_CUDA
		template class ComputeCoordAtomKokkos<LMPHostType>;
		#endif
		}

src/KOKKOS/compute_coord_atom_kokkos.h

0 → 100644

+78 −0

Original line number	Diff line number	Diff line
		/* -- c++ -- ----------------------------------------------------------
		LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
		http://lammps.sandia.gov, Sandia National Laboratories
		Steve Plimpton, sjplimp@sandia.gov

		Copyright (2003) Sandia Corporation. Under the terms of Contract
		DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government retains
		certain rights in this software. This software is distributed under
		the GNU General Public License.

		See the README file in the top-level LAMMPS directory.
		------------------------------------------------------------------------- */

		#ifdef COMPUTE_CLASS

		ComputeStyle(orientorder/atom/kk,ComputeCoordAtomKokkos<LMPDeviceType>)
		ComputeStyle(orientorder/atom/kk/device,ComputeCoordAtomKokkos<LMPDeviceType>)
		ComputeStyle(orientorder/atom/kk/host,ComputeCoordAtomKokkos<LMPHostType>)

		#else

		#ifndef LMP_COMPUTE_COORD_ATOM_KOKKOS_H
		#define LMP_COMPUTE_COORD_ATOM_KOKKOS_H

		#include "compute_coord_atom.h"
		#include "kokkos_type.h"

		namespace LAMMPS_NS {

		template<int CSTYLE, int NCOL>
		struct TagComputeCoordAtom{};

		template<class DeviceType>
		class ComputeCoordAtomKokkos : public ComputeCoordAtom {
		public:
		typedef DeviceType device_type;
		typedef ArrayTypes<DeviceType> AT;

		ComputeCoordAtomKokkos(class LAMMPS , int, char *);
		virtual ~ComputeCoordAtomKokkos();
		void init();
		void compute_peratom();
		enum {NONE,CUTOFF,ORIENT};

		template<int CSTYLE, int NCOL>
		KOKKOS_INLINE_FUNCTION
		void operator()(TagComputeCoordAtom<CSTYLE,NCOL>, const int&) const;

		private:
		int inum;

		typename AT::t_x_array_randomread x;
		typename ArrayTypes<DeviceType>::t_int_1d_randomread type;
		typename ArrayTypes<DeviceType>::t_int_1d mask;

		typename AT::t_neighbors_2d d_neighbors;
		typename AT::t_int_1d_randomread d_ilist;
		typename AT::t_int_1d_randomread d_numneigh;

		typename AT::t_int_1d d_typelo;
		typename AT::t_int_1d d_typehi;

		DAT::tdual_float_1d k_cvec;
		typename AT::t_float_1d d_cvec;
		DAT::tdual_float_2d k_carray;
		typename AT::t_float_2d d_carray;

		typename AT::t_float_2d d_normv;
		};

		}

		#endif
		#endif

		/* ERROR/WARNING messages:

		*/

src/compute_coord_atom.cpp

+2 −0

Original line number	Diff line number	Diff line
		@@ -119,6 +119,8 @@ ComputeCoordAtom::ComputeCoordAtom(LAMMPS lmp, int narg, char *arg) :

		ComputeCoordAtom::~ComputeCoordAtom()
		{
		if (copymode) return;

		delete [] group2;
		delete [] typelo;
		delete [] typehi;

src/compute_coord_atom.h

+4 −4

Original line number	Diff line number	Diff line
		@@ -27,16 +27,16 @@ namespace LAMMPS_NS {
		class ComputeCoordAtom : public Compute {
		public:
		ComputeCoordAtom(class LAMMPS , int, char *);
		~ComputeCoordAtom();
		void init();
		virtual ~ComputeCoordAtom();
		virtual void init();
		void init_list(int, class NeighList *);
		void compute_peratom();
		virtual void compute_peratom();
		int pack_forward_comm(int, int , double , int, int *);
		void unpack_forward_comm(int, int, double *);
		double memory_usage();
		enum {NONE,CUTOFF,ORIENT};

		private:
		protected:
		int nmax,ncol;
		double cutsq;
		class NeighList *list;

Admin message