From a72a3ed50d12365ebb41a02001454fe667036b23 Mon Sep 17 00:00:00 2001 From: Richard Berger Date: Tue, 17 Oct 2023 16:24:14 -0600 Subject: [PATCH] Kokkos version of compute reaxff/bonds --- src/KOKKOS/compute_reaxff_bonds_kokkos.cpp | 192 ++++++++++++++++++ ...kokkos.h => compute_reaxff_bonds_kokkos.h} | 30 ++- .../compute_reaxff_bonds_local_kokkos.cpp | 125 ------------ src/KOKKOS/pair_reaxff_kokkos.cpp | 55 ----- src/KOKKOS/pair_reaxff_kokkos.h | 19 -- 5 files changed, 206 insertions(+), 215 deletions(-) create mode 100644 src/KOKKOS/compute_reaxff_bonds_kokkos.cpp rename src/KOKKOS/{compute_reaxff_bonds_local_kokkos.h => compute_reaxff_bonds_kokkos.h} (67%) delete mode 100644 src/KOKKOS/compute_reaxff_bonds_local_kokkos.cpp diff --git a/src/KOKKOS/compute_reaxff_bonds_kokkos.cpp b/src/KOKKOS/compute_reaxff_bonds_kokkos.cpp new file mode 100644 index 0000000000..921acb9193 --- /dev/null +++ b/src/KOKKOS/compute_reaxff_bonds_kokkos.cpp @@ -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 +ComputeReaxFFBondsKokkos::ComputeReaxFFBondsKokkos(LAMMPS *lmp, int narg, char **arg) : + ComputeReaxFFBonds(lmp, narg, arg), + nbuf(-1), buf(nullptr) +{ +} + + +/* ---------------------------------------------------------------------- */ + +template +ComputeReaxFFBondsKokkos::~ComputeReaxFFBondsKokkos() +{ + memoryKK->destroy_kokkos(k_buf, buf); +} + +/* ---------------------------------------------------------------------- */ + +template +void ComputeReaxFFBondsKokkos::init() +{ + reaxff = dynamic_cast(force->pair_match("^reax../kk",0)); + if (reaxff == nullptr) error->all(FLERR,"Cannot use compute reaxff/bonds without " + "pair_style reaxff/kk"); +} + +/* ---------------------------------------------------------------------- */ +template +void ComputeReaxFFBondsKokkos::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(buf[j+5]); + nbonds += numbonds; + j += 2*numbonds + 7; + } +} + +/* ---------------------------------------------------------------------- */ + +template +void ComputeReaxFFBondsKokkos::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(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 (buf[neigh_offset+k]); + bond[2] = buf[bo_offset+k]; + } + j += 2*numbonds + 7; + } +} + +/* ---------------------------------------------------------------------- */ + +template +void ComputeReaxFFBondsKokkos::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(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 +double ComputeReaxFFBondsKokkos::memory_usage() +{ + double bytes = (double)(nbonds*3) * sizeof(double); + bytes += (double)(nlocal*7) * sizeof(double); + return bytes; +} + +namespace LAMMPS_NS { +template class ComputeReaxFFBondsKokkos; +#ifdef LMP_KOKKOS_GPU +template class ComputeReaxFFBondsKokkos; +#endif +} diff --git a/src/KOKKOS/compute_reaxff_bonds_local_kokkos.h b/src/KOKKOS/compute_reaxff_bonds_kokkos.h similarity index 67% rename from src/KOKKOS/compute_reaxff_bonds_local_kokkos.h rename to src/KOKKOS/compute_reaxff_bonds_kokkos.h index bfb5521199..45020ffa81 100644 --- a/src/KOKKOS/compute_reaxff_bonds_local_kokkos.h +++ b/src/KOKKOS/compute_reaxff_bonds_kokkos.h @@ -17,41 +17,39 @@ #ifdef COMPUTE_CLASS // clang-format off -ComputeStyle(reaxff/bonds/local/kk,ComputeReaxFFBondsLocalKokkos); -ComputeStyle(reaxff/bonds/local/kk/device,ComputeReaxFFBondsLocalKokkos); -ComputeStyle(reaxff/bonds/local/kk/host,ComputeReaxFFBondsLocalKokkos); +ComputeStyle(reaxff/bonds/kk,ComputeReaxFFBondsKokkos); +ComputeStyle(reaxff/bonds/kk/device,ComputeReaxFFBondsKokkos); +ComputeStyle(reaxff/bonds/kk/host,ComputeReaxFFBondsKokkos); // 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 ComputeReaxFFBondsLocalKokkos : public Compute { +class ComputeReaxFFBondsKokkos : public ComputeReaxFFBonds { public: using device_type = DeviceType; using AT = ArrayTypes; - 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*>(reaxff); diff --git a/src/KOKKOS/compute_reaxff_bonds_local_kokkos.cpp b/src/KOKKOS/compute_reaxff_bonds_local_kokkos.cpp deleted file mode 100644 index a56d243ed6..0000000000 --- a/src/KOKKOS/compute_reaxff_bonds_local_kokkos.cpp +++ /dev/null @@ -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 -ComputeReaxFFBondsLocalKokkos::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 -ComputeReaxFFBondsLocalKokkos::~ComputeReaxFFBondsLocalKokkos() -{ - memoryKK->destroy_kokkos(k_alocal, alocal); -} - -/* ---------------------------------------------------------------------- */ - -template -void ComputeReaxFFBondsLocalKokkos::init() -{ - reaxff = dynamic_cast(force->pair_match("^reax../kk",0)); - if (reaxff == nullptr) error->all(FLERR,"Cannot use fix reaxff/bonds without " - "pair_style reaxff/kk"); -} - -/* ---------------------------------------------------------------------- */ - -template -void ComputeReaxFFBondsLocalKokkos::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 -double ComputeReaxFFBondsLocalKokkos::memory_usage() -{ - double bytes = (double)nlocal*nvalues * sizeof(double); - return bytes; -} - -namespace LAMMPS_NS { -template class ComputeReaxFFBondsLocalKokkos; -#ifdef LMP_KOKKOS_GPU -template class ComputeReaxFFBondsLocalKokkos; -#endif -} diff --git a/src/KOKKOS/pair_reaxff_kokkos.cpp b/src/KOKKOS/pair_reaxff_kokkos.cpp index fbf90f140d..c7d54b80cd 100644 --- a/src/KOKKOS/pair_reaxff_kokkos.cpp +++ b/src/KOKKOS/pair_reaxff_kokkos.cpp @@ -4290,61 +4290,6 @@ void PairReaxFFKokkos::pack_bond_buffer_item(int i, int &j, const bo /* ---------------------------------------------------------------------- */ -template -void PairReaxFFKokkos::PackBondInfo(DAT::tdual_float_2d k_alocal) -{ - d_alocal = k_alocal.view(); - k_params_sing.template sync(); - atomKK->sync(execution_space,TAG_MASK|TYPE_MASK|Q_MASK|MOLECULE_MASK); - - tag = atomKK->k_tag.view(); - type = atomKK->k_type.view(); - q = atomKK->k_q.view(); - if (atom->molecule) - molecule = atomKK->k_molecule.view(); - - copymode = 1; - nlocal = atomKK->nlocal; - PairReaxKokkosPackBondInfoFunctor pack_bond_info_functor(this); - Kokkos::parallel_for(nlocal, pack_bond_info_functor); - copymode = 0; - - k_alocal.modify(); - k_alocal.sync(); -} - -template -KOKKOS_INLINE_FUNCTION -void PairReaxFFKokkos::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 void PairReaxFFKokkos::FindBondSpecies() { diff --git a/src/KOKKOS/pair_reaxff_kokkos.h b/src/KOKKOS/pair_reaxff_kokkos.h index 32007e9970..1ad0955a1e 100644 --- a/src/KOKKOS/pair_reaxff_kokkos.h +++ b/src/KOKKOS/pair_reaxff_kokkos.h @@ -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 @@ -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 t_reax_int4_2d; t_reax_int4_2d d_angular_pack, d_torsion_pack; @@ -555,19 +549,6 @@ struct PairReaxKokkosPackBondBufferFunctor { } }; -template -struct PairReaxKokkosPackBondInfoFunctor { - using device_type = DeviceType; - using value_type = int; - PairReaxFFKokkos c; - PairReaxKokkosPackBondInfoFunctor(PairReaxFFKokkos* c_ptr):c(*c_ptr) {}; - - KOKKOS_INLINE_FUNCTION - void operator()(const int i) const { - c.pack_bond_info_item(i); - } -}; - } #endif