Kokkos version of compute reaxff/bonds

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<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
+}

src/KOKKOS/compute_reaxff_bonds_kokkos.h

@@ -17,41 +17,39 @@
 
 #ifdef COMPUTE_CLASS
 // clang-format off
-ComputeStyle(reaxff/bonds/local/kk,ComputeReaxFFBondsLocalKokkos<LMPDeviceType>);
+ComputeStyle(reaxff/bonds/kk,ComputeReaxFFBondsKokkos<LMPDeviceType>);
-ComputeStyle(reaxff/bonds/local/kk/device,ComputeReaxFFBondsLocalKokkos<LMPDeviceType>);
+ComputeStyle(reaxff/bonds/kk/device,ComputeReaxFFBondsKokkos<LMPDeviceType>);
-ComputeStyle(reaxff/bonds/local/kk/host,ComputeReaxFFBondsLocalKokkos<LMPHostType>);
+ComputeStyle(reaxff/bonds/kk/host,ComputeReaxFFBondsKokkos<LMPHostType>);
 // clang-format on
 #else
 
-#ifndef LMP_COMPUTE_REAXFF_BONDS_LOCAL_KOKKOS_H
+#ifndef LMP_COMPUTE_REAXFF_BONDS_KOKKOS_H
-#define LMP_COMPUTE_REAXFF_BONDS_LOCAL_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 **);
+  ComputeReaxFFBondsKokkos(class LAMMPS *, int, char **);
-  ~ComputeReaxFFBondsLocalKokkos() override;
+  ~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 nbuf;
-  int nvalues;
+  double *buf;
-  int prev_nvalues;
+  typename AT::tdual_float_1d k_buf;
-
-  double **alocal;
-  typename AT::tdual_float_2d k_alocal;
-  PairReaxFF *reaxff;
 
   auto device_pair() {
     return dynamic_cast<PairReaxFFKokkos<LMPDeviceType>*>(reaxff);

src/KOKKOS/compute_reaxff_bonds_local_kokkos.cpp

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

src/KOKKOS/pair_reaxff_kokkos.cpp

@@ -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()
 {

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