Add Kokkos version of compute coord/atom

This commit is contained in:
Stan Moore 2020-02-25 10:22:30 -07:00
parent a3c5c49a6b
commit a44e49e276
5 changed files with 335 additions and 4 deletions

View File

@ -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

View File

@ -0,0 +1,249 @@
/* ----------------------------------------------------------------------
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 = delx*delx + dely*dely + 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*(2*l+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
}

View File

@ -0,0 +1,78 @@
/* -*- 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:
*/

View File

@ -119,6 +119,8 @@ ComputeCoordAtom::ComputeCoordAtom(LAMMPS *lmp, int narg, char **arg) :
ComputeCoordAtom::~ComputeCoordAtom()
{
if (copymode) return;
delete [] group2;
delete [] typelo;
delete [] typehi;

View File

@ -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;