KOKKOS: sync per-atom data on demand during variable evaluation

cmake/Modules/Packages/KOKKOS.cmake

@@ -176,7 +176,8 @@ set(KOKKOS_PKG_SOURCES ${KOKKOS_PKG_SOURCES_DIR}/kokkos.cpp
                        ${KOKKOS_PKG_SOURCES_DIR}/domain_kokkos.cpp
                        ${KOKKOS_PKG_SOURCES_DIR}/modify_kokkos.cpp
                        ${KOKKOS_PKG_SOURCES_DIR}/rand_pool_wrap_kokkos.cpp
-                       ${KOKKOS_PKG_SOURCES_DIR}/tune_kokkos.cpp)
+                       ${KOKKOS_PKG_SOURCES_DIR}/tune_kokkos.cpp
+                       ${KOKKOS_PKG_SOURCES_DIR}/variable_kokkos.cpp)
 
 
 # fix wall/gran has been refactored in an incompatible way. Use old version of base class for now

src/KOKKOS/Install.sh

@@ -574,6 +574,8 @@ action sna_kokkos.h sna.h
 action third_order_kokkos.cpp dynamical_matrix.cpp
 action third_order_kokkos.h dynamical_matrix.h
 action transpose_helper_kokkos.h
+action variable_kokkos.cpp
+action variable_kokkos.h
 action verlet_kokkos.cpp
 action verlet_kokkos.h
 

src/KOKKOS/fix_addforce_kokkos.cpp

@@ -132,8 +132,6 @@ void FixAddForceKokkos<DeviceType>::post_force(int vflag)
 
   } else {
 
-    atomKK->sync(Host,ALL_MASK); // this can be removed when variable class is ported to Kokkos
-
     modify->clearstep_compute();
 
     if (xstyle == EQUAL) xvalue = input->variable->compute_equal(xvar);
@@ -149,7 +147,11 @@ void FixAddForceKokkos<DeviceType>::post_force(int vflag)
 
     modify->addstep_compute(update->ntimestep + 1);
 
-    if (varflag == ATOM) {  // this can be removed when variable class is ported to Kokkos
+    // atom-style variables are evaluated on the host, so the result has to be
+    // copied to the device for the kernel below.  this is a real copy, not a
+    // stale flag: it can only go away if variables are evaluated on the device.
+
+    if (varflag == ATOM) {
       k_sforce.modify_host();
       k_sforce.sync<DeviceType>();
     }

src/KOKKOS/fix_efield_kokkos.cpp

@@ -150,11 +150,13 @@ void FixEfieldKokkos<DeviceType>::post_force(int vflag)
 
   } else {
 
-    atomKK->sync(Host,ALL_MASK); // this can be removed when variable class is ported to Kokkos
-
     FixEfield::update_efield_variables();
 
-    if (varflag == ATOM) {  // this can be removed when variable class is ported to Kokkos
+    // atom-style variables are evaluated on the host, so the result has to be
+    // copied to the device for the kernel below.  this is a real copy, not a
+    // stale flag: it can only go away if variables are evaluated on the device.
+
+    if (varflag == ATOM) {
       k_efield.modify_host();
       k_efield.sync<DeviceType>();
     }

src/KOKKOS/fix_setforce_kokkos.cpp

@@ -114,8 +114,6 @@ void FixSetForceKokkos<DeviceType>::post_force(int /*vflag*/)
 
   } else {
 
-    atomKK->sync(Host,ALL_MASK); // this can be removed when variable class is ported to Kokkos
-
     modify->clearstep_compute();
 
     if (xstyle == EQUAL) xvalue = input->variable->compute_equal(xvar);
@@ -130,7 +128,11 @@ void FixSetForceKokkos<DeviceType>::post_force(int /*vflag*/)
 
     modify->addstep_compute(update->ntimestep + 1);
 
-    if (varflag == ATOM) {  // this can be removed when variable class is ported to Kokkos
+    // atom-style variables are evaluated on the host, so the result has to be
+    // copied to the device for the kernel below.  this is a real copy, not a
+    // stale flag: it can only go away if variables are evaluated on the device.
+
+    if (varflag == ATOM) {
       k_sforce.modify_host();
       k_sforce.sync<DeviceType>();
     }

src/KOKKOS/variable_kokkos.cpp

@@ -0,0 +1,172 @@
+/* ----------------------------------------------------------------------
+   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.
+------------------------------------------------------------------------- */
+
+#include "variable_kokkos.h"
+
+#include "atom_kokkos.h"
+#include "atom_masks.h"
+
+#include <cstring>
+
+using namespace LAMMPS_NS;
+
+/* ----------------------------------------------------------------------
+   make the per-atom arrays named by mask current on the host
+------------------------------------------------------------------------- */
+
+void VariableKokkos::sync_host(uint64_t mask)
+{
+  // Input is created before Atom is replaced by AtomKokkos, so the cast
+  // cannot be done once in the constructor
+
+  auto *atomKK = dynamic_cast<AtomKokkos *>(atom);
+  if (atomKK) atomKK->sync(Host, mask);
+}
+
+/* ----------------------------------------------------------------------
+   compute_atom() reads atom->mask on the host for the group test
+------------------------------------------------------------------------- */
+
+void VariableKokkos::compute_atom(int ivar, int igroup, double *result, int stride, int sumflag)
+{
+  sync_host(MASK_MASK);
+  Variable::compute_atom(ivar, igroup, result, stride, sumflag);
+}
+
+/* ----------------------------------------------------------------------
+   an atom vector in a formula reads exactly one per-atom array.
+   keep this list in step with Variable::atom_vector()
+------------------------------------------------------------------------- */
+
+void VariableKokkos::atom_vector(char *word, Tree **tree, Tree **treestack, int &ntreestack)
+{
+  uint64_t mask = ALL_MASK;
+
+  if (strcmp(word,"id") == 0) mask = TAG_MASK;
+  else if (strcmp(word,"type") == 0) mask = TYPE_MASK;
+  else if (strcmp(word,"mol") == 0) mask = MOLECULE_MASK;
+  else if (strcmp(word,"radius") == 0) mask = RADIUS_MASK;
+  else if (strcmp(word,"q") == 0) mask = Q_MASK;
+  else if (strcmp(word,"x") == 0) mask = X_MASK;
+  else if (strcmp(word,"y") == 0) mask = X_MASK;
+  else if (strcmp(word,"z") == 0) mask = X_MASK;
+  else if (strcmp(word,"vx") == 0) mask = V_MASK;
+  else if (strcmp(word,"vy") == 0) mask = V_MASK;
+  else if (strcmp(word,"vz") == 0) mask = V_MASK;
+  else if (strcmp(word,"fx") == 0) mask = F_MASK;
+  else if (strcmp(word,"fy") == 0) mask = F_MASK;
+  else if (strcmp(word,"fz") == 0) mask = F_MASK;
+
+  // without per-atom rmass, "mass" is a per-type array that eval_tree()
+  // indexes with atom->type[i], so it is atom->type that must be current
+
+  else if (strcmp(word,"mass") == 0) mask = atom->rmass ? RMASS_MASK : TYPE_MASK;
+
+  // an atom vector added to the base class but not listed here still works,
+  // it just syncs more than it needs to
+
+  sync_host(mask);
+  Variable::atom_vector(word, tree, treestack, ntreestack);
+}
+
+/* ----------------------------------------------------------------------
+   group functions such as xcm(), fcm(), and gyration() read a range of
+   per-atom arrays that cannot be narrowed down from the formula
+------------------------------------------------------------------------- */
+
+int VariableKokkos::group_function(char *word, char *contents, Tree **tree, Tree **treestack,
+                                   int &ntreestack, double *argstack, int &nargstack, int ivar)
+{
+  // group_function() is tried for every function word in a formula and returns
+  // 0 for anything that is not one, so only sync when it really is one
+
+  if (is_group_function(word)) sync_host(ALL_MASK);
+  return Variable::group_function(word, contents, tree, treestack, ntreestack, argstack, nargstack,
+                                  ivar);
+}
+
+/* ----------------------------------------------------------------------
+   gmask() tests atom->mask, rmask() and grmask() also match a region
+   against atom->x.  sum(), min(), ave(), sort() and the like take a
+   compute or fix as their argument and invoke it, which reads per-atom
+   data that cannot be narrowed down.  the rest touch no per-atom data.
+------------------------------------------------------------------------- */
+
+int VariableKokkos::special_function(const std::string &word, char *contents, Tree **tree,
+                                     Tree **treestack, int &ntreestack, double *argstack,
+                                     int &nargstack, int ivar, char *str, int &i, char *&ptr)
+{
+  // like group_function(), this is tried for every function word in a formula
+  // and returns 0 for anything that is not a special function
+
+  if (is_special_function(word)) {
+    if (word == "gmask") sync_host(MASK_MASK);
+    else if (word == "rmask") sync_host(X_MASK);
+    else if (word == "grmask") sync_host(X_MASK | MASK_MASK);
+
+    // these read nothing per-atom; anything else, including a special
+    // function added later, falls through to ALL_MASK
+
+    else if ((word != "next") && (word != "is_file") && (word != "is_os") &&
+             (word != "is_timeout") && (word != "extract_setting") &&
+             (word != "label2type") && (word != "is_typelabel"))
+      sync_host(ALL_MASK);
+  }
+
+  return Variable::special_function(word, contents, tree, treestack, ntreestack, argstack,
+                                    nargstack, ivar, str, i, ptr);
+}
+
+/* ----------------------------------------------------------------------
+   access to a single atom by ID, such as x[100] or c_ID[7], reads the atom
+   map as well as whichever array was named
+------------------------------------------------------------------------- */
+
+void VariableKokkos::peratom2global(int flag, char *word, double *vector, int nstride, tagint id,
+                                    Tree **tree, Tree **treestack, int &ntreestack,
+                                    double *argstack, int &nargstack)
+{
+  sync_host(ALL_MASK);
+  Variable::peratom2global(flag, word, vector, nstride, id, tree, treestack, ntreestack, argstack,
+                           nargstack);
+}
+
+/* ---------------------------------------------------------------------- */
+
+void VariableKokkos::custom2global(int *ivector, double *dvector, int nstride, tagint id,
+                                   Tree **tree, Tree **treestack, int &ntreestack,
+                                   double *argstack, int &nargstack)
+{
+  sync_host(ALL_MASK);
+  Variable::custom2global(ivector, dvector, nstride, id, tree, treestack, ntreestack, argstack,
+                          nargstack);
+}
+
+/* ----------------------------------------------------------------------
+   called from evaluate() for per-atom data reached through a compute or a
+   fix, which passes a null pointer, or through a custom property, which
+   passes the i_ / d_ / i2_ / d2_ name the formula used
+------------------------------------------------------------------------- */
+
+void VariableKokkos::sync_peratom(const char *word)
+{
+  uint64_t mask = ALL_MASK;
+
+  if (!word) mask = ALL_MASK;
+  else if (strncmp(word,"i2_",3) == 0) mask = IARRAY_MASK;
+  else if (strncmp(word,"d2_",3) == 0) mask = DARRAY_MASK;
+  else if (strncmp(word,"i_",2) == 0) mask = IVECTOR_MASK;
+  else if (strncmp(word,"d_",2) == 0) mask = DVECTOR_MASK;
+
+  sync_host(mask);
+}

src/KOKKOS/variable_kokkos.h

@@ -0,0 +1,44 @@
+/* -*- c++ -*- ----------------------------------------------------------
+   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.
+------------------------------------------------------------------------- */
+
+#ifndef LMP_VARIABLE_KOKKOS_H
+#define LMP_VARIABLE_KOKKOS_H
+
+#include "variable.h"
+
+namespace LAMMPS_NS {
+
+class VariableKokkos : public Variable {
+ public:
+  VariableKokkos(class LAMMPS *lmp) : Variable(lmp) {}
+
+  void compute_atom(int, int, double *, int, int) override;
+
+ protected:
+  void atom_vector(char *, Tree **, Tree **, int &) override;
+  int group_function(char *, char *, Tree **, Tree **, int &, double *, int &, int) override;
+  int special_function(const std::string &, char *, Tree **, Tree **, int &, double *, int &, int,
+                       char *, int &, char *&) override;
+  void peratom2global(int, char *, double *, int, tagint, Tree **, Tree **, int &, double *,
+                      int &) override;
+  void custom2global(int *, double *, int, tagint, Tree **, Tree **, int &, double *,
+                     int &) override;
+  void sync_peratom(const char *) override;
+
+ private:
+  void sync_host(uint64_t);
+};
+
+}    // namespace LAMMPS_NS
+
+#endif

src/accelerator_kokkos.h

@@ -28,6 +28,7 @@
 #include "memory_kokkos.h"        // IWYU pragma: export
 #include "modify_kokkos.h"        // IWYU pragma: export
 #include "neighbor_kokkos.h"      // IWYU pragma: export
+#include "variable_kokkos.h"      // IWYU pragma: export
 
 #define LAMMPS_INLINE KOKKOS_INLINE_FUNCTION
 
@@ -44,6 +45,7 @@
 #include "memory.h"
 #include "modify.h"
 #include "neighbor.h"
+#include "variable.h"
 
 #define LAMMPS_INLINE inline
 
@@ -110,6 +112,11 @@ class ModifyKokkos : public Modify {
   ModifyKokkos(class LAMMPS *lmp) : Modify(lmp) {}
 };
 
+class VariableKokkos : public Variable {
+ public:
+  VariableKokkos(class LAMMPS *lmp) : Variable(lmp) {}
+};
+
 // NOLINTBEGIN
 class DAT {
  public:

src/input.cpp

@@ -128,7 +128,10 @@ Input::Input(LAMMPS *lmp, int argc, char **argv) :
     inlines = new int[LMP_MAXFILE];
   }
 
-  variable = new Variable(lmp);
+  if (lmp->kokkos && lmp->kokkos->kokkos_exists)
+    variable = new VariableKokkos(lmp);
+  else
+    variable = new Variable(lmp);
 
   // fill map with commands listed in style_command.h