Kokkos version of compute reaxff/bonds
This commit is contained in:
Richard Berger
192
src/KOKKOS/compute_reaxff_bonds_kokkos.cpp
Normal file
192
src/KOKKOS/compute_reaxff_bonds_kokkos.cpp
Normal file
@ -0,0 +1,192 @@
|
||||
// clang-format off
|
||||
/* ----------------------------------------------------------------------
|
||||
LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
|
||||
https://www.lammps.org/, Sandia National Laboratories
|
||||
LAMMPS development team: developers@lammps.org
|
||||
|
||||
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.
|
||||
------------------------------------------------------------------------- */
|
||||
|
||||
/* ----------------------------------------------------------------------
|
||||
Contributing author: Richard Berger (LANL)
|
||||
------------------------------------------------------------------------- */
|
||||
|
||||
#include "compute_reaxff_bonds_kokkos.h"
|
||||
#include "atom.h"
|
||||
#include "molecule.h"
|
||||
#include "update.h"
|
||||
#include "force.h"
|
||||
#include "memory.h"
|
||||
#include "error.h"
|
||||
#include "neigh_list.h"
|
||||
|
||||
#include "memory_kokkos.h"
|
||||
#include "pair_reaxff_kokkos.h"
|
||||
#include "reaxff_api.h"
|
||||
|
||||
using namespace LAMMPS_NS;
|
||||
using namespace ReaxFF;
|
||||
|
||||
/* ---------------------------------------------------------------------- */
|
||||
|
||||
template<class DeviceType>
|
||||
ComputeReaxFFBondsKokkos<DeviceType>::ComputeReaxFFBondsKokkos(LAMMPS *lmp, int narg, char **arg) :
|
||||
ComputeReaxFFBonds(lmp, narg, arg),
|
||||
nbuf(-1), buf(nullptr)
|
||||
{
|
||||
}
|
||||
|
||||
|
||||
/* ---------------------------------------------------------------------- */
|
||||
|
||||
template<class DeviceType>
|
||||
ComputeReaxFFBondsKokkos<DeviceType>::~ComputeReaxFFBondsKokkos()
|
||||
{
|
||||
memoryKK->destroy_kokkos(k_buf, buf);
|
||||
}
|
||||
|
||||
/* ---------------------------------------------------------------------- */
|
||||
|
||||
template<class DeviceType>
|
||||
void ComputeReaxFFBondsKokkos<DeviceType>::init()
|
||||
{
|
||||
reaxff = dynamic_cast<PairReaxFF*>(force->pair_match("^reax../kk",0));
|
||||
if (reaxff == nullptr) error->all(FLERR,"Cannot use compute reaxff/bonds without "
|
||||
"pair_style reaxff/kk");
|
||||
}
|
||||
|
||||
/* ---------------------------------------------------------------------- */
|
||||
template<class DeviceType>
|
||||
void ComputeReaxFFBondsKokkos<DeviceType>::compute_bonds()
|
||||
{
|
||||
if (atom->nlocal > nlocal) {
|
||||
memory->destroy(array_atom);
|
||||
nlocal = atom->nlocal;
|
||||
memory->create(array_atom, nlocal, 7, "reaxff/bonds:array_atom");
|
||||
}
|
||||
|
||||
// retrieve bond information from kokkos pair style. the data potentially
|
||||
// lives on device. it is copied into buf on the host in a condensed format
|
||||
// compute_local and compute_atom then expand the data from this buffer into
|
||||
// appropiate arrays for consumption by others (e.g. dump local, dump custom
|
||||
// or library interface)
|
||||
int maxnumbonds = 0;
|
||||
if (reaxff->execution_space == Device)
|
||||
device_pair()->FindBond(maxnumbonds);
|
||||
else
|
||||
host_pair()->FindBond(maxnumbonds);
|
||||
|
||||
nbuf = 1+(maxnumbonds*2 + 7)*nlocal;
|
||||
|
||||
if(!buf || k_buf.extent(0) < nbuf) {
|
||||
memoryKK->destroy_kokkos(k_buf, buf);
|
||||
memoryKK->create_kokkos(k_buf, buf, nbuf, "reaxff/bonds:buf");
|
||||
}
|
||||
|
||||
// Pass information to buffer, will sync to host
|
||||
int nbuf_local;
|
||||
if (reaxff->execution_space == Device)
|
||||
device_pair()->PackBondBuffer(k_buf, nbuf_local);
|
||||
else
|
||||
host_pair()->PackBondBuffer(k_buf, nbuf_local);
|
||||
buf[0] = nlocal;
|
||||
|
||||
// Extract number of bonds from buffer
|
||||
nbonds = 0;
|
||||
int j = 1;
|
||||
for (int i = 0; i < nlocal; i++) {
|
||||
int numbonds = static_cast<int>(buf[j+5]);
|
||||
nbonds += numbonds;
|
||||
j += 2*numbonds + 7;
|
||||
}
|
||||
}
|
||||
|
||||
/* ---------------------------------------------------------------------- */
|
||||
|
||||
template<class DeviceType>
|
||||
void ComputeReaxFFBondsKokkos<DeviceType>::compute_local()
|
||||
{
|
||||
invoked_local = update->ntimestep;
|
||||
|
||||
if(invoked_bonds < update->ntimestep) {
|
||||
compute_bonds();
|
||||
}
|
||||
|
||||
if(nbonds > prev_nbonds) {
|
||||
// grow array_local
|
||||
memory->destroy(array_local);
|
||||
memory->create(array_local, nbonds, 3, "reaxff/bonds:array_local");
|
||||
prev_nbonds = nbonds;
|
||||
}
|
||||
|
||||
size_local_rows = nbonds;
|
||||
|
||||
// extract local bond information from buffer
|
||||
int b = 0;
|
||||
int j = 1;
|
||||
|
||||
for (int i = 0; i < nlocal; ++i) {
|
||||
const int numbonds = static_cast<int>(buf[j+5]);
|
||||
const int neigh_offset = j + 6;
|
||||
const int bo_offset = neigh_offset + numbonds + 1;
|
||||
for (int k = 0; k < numbonds; k++) {
|
||||
auto bond = array_local[b++];
|
||||
bond[0] = i;
|
||||
bond[1] = static_cast<tagint> (buf[neigh_offset+k]);
|
||||
bond[2] = buf[bo_offset+k];
|
||||
}
|
||||
j += 2*numbonds + 7;
|
||||
}
|
||||
}
|
||||
|
||||
/* ---------------------------------------------------------------------- */
|
||||
|
||||
template<class DeviceType>
|
||||
void ComputeReaxFFBondsKokkos<DeviceType>::compute_peratom()
|
||||
{
|
||||
invoked_peratom = update->ntimestep;
|
||||
|
||||
if(invoked_bonds < update->ntimestep) {
|
||||
compute_bonds();
|
||||
}
|
||||
|
||||
// extract peratom bond information from buffer
|
||||
int j = 1;
|
||||
for (int i = 0; i < nlocal; ++i) {
|
||||
auto ptr = array_atom[i];
|
||||
int numbonds = static_cast<int>(buf[j+5]);
|
||||
const int mol_offset = j + 6 + numbonds;
|
||||
ptr[0] = buf[j]; // jtag
|
||||
ptr[1] = buf[j+1]; // itype
|
||||
ptr[2] = numbonds;
|
||||
ptr[3] = buf[mol_offset]; // mol
|
||||
ptr[4] = buf[j+2]; // sbo
|
||||
ptr[5] = buf[j+3]; // nlp
|
||||
ptr[6] = buf[j+4]; // q
|
||||
j += 2*numbonds + 7;
|
||||
}
|
||||
}
|
||||
|
||||
/* ----------------------------------------------------------------------
|
||||
memory usage of local data
|
||||
------------------------------------------------------------------------- */
|
||||
|
||||
template<class DeviceType>
|
||||
double ComputeReaxFFBondsKokkos<DeviceType>::memory_usage()
|
||||
{
|
||||
double bytes = (double)(nbonds*3) * sizeof(double);
|
||||
bytes += (double)(nlocal*7) * sizeof(double);
|
||||
return bytes;
|
||||
}
|
||||
|
||||
namespace LAMMPS_NS {
|
||||
template class ComputeReaxFFBondsKokkos<LMPDeviceType>;
|
||||
#ifdef LMP_KOKKOS_GPU
|
||||
template class ComputeReaxFFBondsKokkos<LMPHostType>;
|
||||
#endif
|
||||
}
|
||||
@ -17,41 +17,39 @@
|
||||
|
||||
#ifdef COMPUTE_CLASS
|
||||
// clang-format off
|
||||
ComputeStyle(reaxff/bonds/local/kk,ComputeReaxFFBondsLocalKokkos<LMPDeviceType>);
|
||||
ComputeStyle(reaxff/bonds/local/kk/device,ComputeReaxFFBondsLocalKokkos<LMPDeviceType>);
|
||||
ComputeStyle(reaxff/bonds/local/kk/host,ComputeReaxFFBondsLocalKokkos<LMPHostType>);
|
||||
ComputeStyle(reaxff/bonds/kk,ComputeReaxFFBondsKokkos<LMPDeviceType>);
|
||||
ComputeStyle(reaxff/bonds/kk/device,ComputeReaxFFBondsKokkos<LMPDeviceType>);
|
||||
ComputeStyle(reaxff/bonds/kk/host,ComputeReaxFFBondsKokkos<LMPHostType>);
|
||||
// clang-format on
|
||||
#else
|
||||
|
||||
#ifndef LMP_COMPUTE_REAXFF_BONDS_LOCAL_KOKKOS_H
|
||||
#define LMP_COMPUTE_REAXFF_BONDS_LOCAL_KOKKOS_H
|
||||
#ifndef LMP_COMPUTE_REAXFF_BONDS_KOKKOS_H
|
||||
#define LMP_COMPUTE_REAXFF_BONDS_KOKKOS_H
|
||||
|
||||
#include "compute_reaxff_bonds_local.h"
|
||||
#include "compute_reaxff_bonds.h"
|
||||
#include "pair_reaxff_kokkos.h"
|
||||
#include "kokkos_type.h"
|
||||
|
||||
namespace LAMMPS_NS {
|
||||
|
||||
template<class DeviceType>
|
||||
class ComputeReaxFFBondsLocalKokkos : public Compute {
|
||||
class ComputeReaxFFBondsKokkos : public ComputeReaxFFBonds {
|
||||
public:
|
||||
using device_type = DeviceType;
|
||||
using AT = ArrayTypes<DeviceType>;
|
||||
|
||||
ComputeReaxFFBondsLocalKokkos(class LAMMPS *, int, char **);
|
||||
~ComputeReaxFFBondsLocalKokkos() override;
|
||||
ComputeReaxFFBondsKokkos(class LAMMPS *, int, char **);
|
||||
~ComputeReaxFFBondsKokkos() override;
|
||||
void init() override;
|
||||
void compute_local() override;
|
||||
void compute_peratom() override;
|
||||
void compute_bonds() override;
|
||||
double memory_usage() override;
|
||||
|
||||
private:
|
||||
int nlocal;
|
||||
int nvalues;
|
||||
int prev_nvalues;
|
||||
|
||||
double **alocal;
|
||||
typename AT::tdual_float_2d k_alocal;
|
||||
PairReaxFF *reaxff;
|
||||
int nbuf;
|
||||
double *buf;
|
||||
typename AT::tdual_float_1d k_buf;
|
||||
|
||||
auto device_pair() {
|
||||
return dynamic_cast<PairReaxFFKokkos<LMPDeviceType>*>(reaxff);
|
||||
@ -1,125 +0,0 @@
|
||||
// clang-format off
|
||||
/* ----------------------------------------------------------------------
|
||||
LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
|
||||
https://www.lammps.org/, Sandia National Laboratories
|
||||
LAMMPS development team: developers@lammps.org
|
||||
|
||||
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.
|
||||
------------------------------------------------------------------------- */
|
||||
|
||||
/* ----------------------------------------------------------------------
|
||||
Contributing author: Richard Berger (LANL)
|
||||
------------------------------------------------------------------------- */
|
||||
|
||||
#include "compute_reaxff_bonds_local_kokkos.h"
|
||||
#include "atom.h"
|
||||
#include "molecule.h"
|
||||
#include "update.h"
|
||||
#include "force.h"
|
||||
#include "memory.h"
|
||||
#include "error.h"
|
||||
#include "neigh_list.h"
|
||||
|
||||
#include "memory_kokkos.h"
|
||||
#include "pair_reaxff_kokkos.h"
|
||||
#include "reaxff_api.h"
|
||||
|
||||
using namespace LAMMPS_NS;
|
||||
using namespace ReaxFF;
|
||||
|
||||
/* ---------------------------------------------------------------------- */
|
||||
|
||||
template<class DeviceType>
|
||||
ComputeReaxFFBondsLocalKokkos<DeviceType>::ComputeReaxFFBondsLocalKokkos(LAMMPS *lmp, int narg, char **arg) :
|
||||
Compute(lmp, narg, arg),
|
||||
alocal(nullptr), reaxff(nullptr)
|
||||
{
|
||||
if (atom->tag_consecutive() == 0)
|
||||
error->all(FLERR,"Atom IDs must be consecutive for compute reaxff/bonds/local");
|
||||
|
||||
local_flag = 1;
|
||||
|
||||
nvalues = 7 + 2*MAXREAXBOND;
|
||||
prev_nvalues = 0;
|
||||
|
||||
// initialize output
|
||||
|
||||
nlocal = atom->nlocal;
|
||||
|
||||
size_local_rows = atom->nlocal;
|
||||
size_local_cols = 7 + 2*MAXREAXBOND;
|
||||
printf("RUNNING KOKKOS VERSION\n");
|
||||
}
|
||||
|
||||
|
||||
/* ---------------------------------------------------------------------- */
|
||||
|
||||
template<class DeviceType>
|
||||
ComputeReaxFFBondsLocalKokkos<DeviceType>::~ComputeReaxFFBondsLocalKokkos()
|
||||
{
|
||||
memoryKK->destroy_kokkos(k_alocal, alocal);
|
||||
}
|
||||
|
||||
/* ---------------------------------------------------------------------- */
|
||||
|
||||
template<class DeviceType>
|
||||
void ComputeReaxFFBondsLocalKokkos<DeviceType>::init()
|
||||
{
|
||||
reaxff = dynamic_cast<PairReaxFF*>(force->pair_match("^reax../kk",0));
|
||||
if (reaxff == nullptr) error->all(FLERR,"Cannot use fix reaxff/bonds without "
|
||||
"pair_style reaxff/kk");
|
||||
}
|
||||
|
||||
/* ---------------------------------------------------------------------- */
|
||||
|
||||
template<class DeviceType>
|
||||
void ComputeReaxFFBondsLocalKokkos<DeviceType>::compute_local()
|
||||
{
|
||||
invoked_local = update->ntimestep;
|
||||
|
||||
int maxnumbonds = 0;
|
||||
if (reaxff->execution_space == Device)
|
||||
device_pair()->FindBond(maxnumbonds);
|
||||
else
|
||||
host_pair()->FindBond(maxnumbonds);
|
||||
nvalues = 7+2*maxnumbonds;
|
||||
|
||||
if(atom->nlocal > nlocal || nvalues > prev_nvalues) {
|
||||
nlocal = atom->nlocal;
|
||||
memoryKK->destroy_kokkos(k_alocal, alocal);
|
||||
memoryKK->create_kokkos(k_alocal, alocal, atom->nlocal, nvalues,"reaxff/bonds/local:alocal");
|
||||
prev_nvalues = nvalues;
|
||||
array_local = alocal;
|
||||
}
|
||||
|
||||
size_local_rows = nlocal;
|
||||
size_local_cols = nvalues;
|
||||
|
||||
if (reaxff->execution_space == Device)
|
||||
device_pair()->PackBondInfo(k_alocal);
|
||||
else
|
||||
host_pair()->PackBondInfo(k_alocal);
|
||||
}
|
||||
|
||||
/* ----------------------------------------------------------------------
|
||||
memory usage of local data
|
||||
------------------------------------------------------------------------- */
|
||||
|
||||
template<class DeviceType>
|
||||
double ComputeReaxFFBondsLocalKokkos<DeviceType>::memory_usage()
|
||||
{
|
||||
double bytes = (double)nlocal*nvalues * sizeof(double);
|
||||
return bytes;
|
||||
}
|
||||
|
||||
namespace LAMMPS_NS {
|
||||
template class ComputeReaxFFBondsLocalKokkos<LMPDeviceType>;
|
||||
#ifdef LMP_KOKKOS_GPU
|
||||
template class ComputeReaxFFBondsLocalKokkos<LMPHostType>;
|
||||
#endif
|
||||
}
|
||||
@ -4290,61 +4290,6 @@ void PairReaxFFKokkos<DeviceType>::pack_bond_buffer_item(int i, int &j, const bo
|
||||
|
||||
/* ---------------------------------------------------------------------- */
|
||||
|
||||
template<class DeviceType>
|
||||
void PairReaxFFKokkos<DeviceType>::PackBondInfo(DAT::tdual_float_2d k_alocal)
|
||||
{
|
||||
d_alocal = k_alocal.view<DeviceType>();
|
||||
k_params_sing.template sync<DeviceType>();
|
||||
atomKK->sync(execution_space,TAG_MASK|TYPE_MASK|Q_MASK|MOLECULE_MASK);
|
||||
|
||||
tag = atomKK->k_tag.view<DeviceType>();
|
||||
type = atomKK->k_type.view<DeviceType>();
|
||||
q = atomKK->k_q.view<DeviceType>();
|
||||
if (atom->molecule)
|
||||
molecule = atomKK->k_molecule.view<DeviceType>();
|
||||
|
||||
copymode = 1;
|
||||
nlocal = atomKK->nlocal;
|
||||
PairReaxKokkosPackBondInfoFunctor<DeviceType> pack_bond_info_functor(this);
|
||||
Kokkos::parallel_for(nlocal, pack_bond_info_functor);
|
||||
copymode = 0;
|
||||
|
||||
k_alocal.modify<DeviceType>();
|
||||
k_alocal.sync<LMPHostType>();
|
||||
}
|
||||
|
||||
template<class DeviceType>
|
||||
KOKKOS_INLINE_FUNCTION
|
||||
void PairReaxFFKokkos<DeviceType>::pack_bond_info_item(const int i) const
|
||||
{
|
||||
const int numbonds = d_numneigh_bonds[i];
|
||||
d_alocal(i,0) = tag[i];
|
||||
d_alocal(i,1) = type[i];
|
||||
d_alocal(i,2) = numbonds;
|
||||
|
||||
int j = 3;
|
||||
|
||||
for (int k = 0; k < numbonds; k++) {
|
||||
d_alocal(i,j++) = d_neighid(i,k);
|
||||
}
|
||||
|
||||
d_alocal(i,j++) = molecule.data() ? molecule[i] : 0.0;
|
||||
|
||||
for (int k = 0; k < numbonds; k++) {
|
||||
d_alocal(i,j++) = d_abo(i,k);
|
||||
}
|
||||
|
||||
d_alocal(i,j++) = d_total_bo[i];
|
||||
d_alocal(i,j++) = paramssing(type[i]).nlp_opt - d_Delta_lp[i];
|
||||
d_alocal(i,j++) = q[i];
|
||||
|
||||
for(; j < d_alocal.extent(1); ++j) {
|
||||
d_alocal(i,j) = 0.0;
|
||||
}
|
||||
}
|
||||
|
||||
/* ---------------------------------------------------------------------- */
|
||||
|
||||
template<class DeviceType>
|
||||
void PairReaxFFKokkos<DeviceType>::FindBondSpecies()
|
||||
{
|
||||
|
||||
@ -135,7 +135,6 @@ class PairReaxFFKokkos : public PairReaxFF {
|
||||
double memory_usage();
|
||||
void FindBond(int &);
|
||||
void PackBondBuffer(DAT::tdual_ffloat_1d, int &);
|
||||
void PackBondInfo(DAT::tdual_float_2d);
|
||||
void FindBondSpecies();
|
||||
|
||||
template<int NEIGHFLAG>
|
||||
@ -293,9 +292,6 @@ class PairReaxFFKokkos : public PairReaxFF {
|
||||
KOKKOS_INLINE_FUNCTION
|
||||
void pack_bond_buffer_item(int, int&, const bool&) const;
|
||||
|
||||
KOKKOS_INLINE_FUNCTION
|
||||
void pack_bond_info_item(const int) const;
|
||||
|
||||
KOKKOS_INLINE_FUNCTION
|
||||
void operator()(TagPairReaxFindBondSpeciesZero, const int&) const;
|
||||
|
||||
@ -510,8 +506,6 @@ class PairReaxFFKokkos : public PairReaxFF {
|
||||
typename AT::t_ffloat_1d d_buf;
|
||||
DAT::tdual_int_scalar k_nbuf_local;
|
||||
|
||||
typename AT::t_float_2d d_alocal;
|
||||
|
||||
typedef Kokkos::View<reax_int4**, LMPDeviceType::array_layout, DeviceType> t_reax_int4_2d;
|
||||
|
||||
t_reax_int4_2d d_angular_pack, d_torsion_pack;
|
||||
@ -555,19 +549,6 @@ struct PairReaxKokkosPackBondBufferFunctor {
|
||||
}
|
||||
};
|
||||
|
||||
template <class DeviceType>
|
||||
struct PairReaxKokkosPackBondInfoFunctor {
|
||||
using device_type = DeviceType;
|
||||
using value_type = int;
|
||||
PairReaxFFKokkos<DeviceType> c;
|
||||
PairReaxKokkosPackBondInfoFunctor(PairReaxFFKokkos<DeviceType>* c_ptr):c(*c_ptr) {};
|
||||
|
||||
KOKKOS_INLINE_FUNCTION
|
||||
void operator()(const int i) const {
|
||||
c.pack_bond_info_item(i);
|
||||
}
|
||||
};
|
||||
|
||||
}
|
||||
|
||||
#endif
|
||||
|
||||
Reference in New Issue
Block a user