From 62555d0fcfcc0dfcf0b6f17d4ef270e8b38acd56 Mon Sep 17 00:00:00 2001 From: Stan Moore Date: Fri, 14 Aug 2026 20:59:51 +0000 Subject: [PATCH 1/2] KOKKOS: sync per-atom data on demand during variable evaluation The KOKKOS versions of fix setforce, fix addforce, and fix efield had to call atomKK->sync(Host,ALL_MASK) before evaluating their force variables, since variable evaluation happens on the host. Copying all per-atom data back from the device every step is far more than these formulas need, and adding the same boilerplate to every caller of compute_atom() would not scale: there are more than fifty of them, and a caller that forgets an entry gets a silently incomplete sync. Let the Variable class request the data instead, at the points where it reads per-atom arrays on the host. The functions that do so become virtual, and VariableKokkos overrides them to make the data current and then delegate to the base class: compute_atom() for the group test, atom_vector() for an atom vector in a formula, group_function() for xcm() and friends, special_function() for gmask, rmask, and grmask, and peratom2global() and custom2global() for access to a single atom by ID. Per-atom data reached through a compute, fix, or custom property has no such seam, since evaluate() handles it inline, so those three branches call a small virtual sync_peratom() instead. This keeps all knowledge of the bitmasks in VariableKokkos. atom_masks.h is not needed in variable.cpp at all, and variable.cpp gains only three one-line calls plus the virtual keywords. Input instantiates VariableKokkos when KOKKOS is active. The three fixes drop their blanket sync with no replacement, and every other caller of an atom-style variable gets the same treatment for free. The mapping is fail-safe: atom_vector() starts at ALL_MASK and only narrows, so an atom vector added to the base class without a matching entry in VariableKokkos syncs more than it needs to rather than reading stale data. The compute, fix, and group function paths request ALL_MASK for the same reason. Three cases need care: "mass" without per-atom rmass is a per-type array that eval_tree() indexes with atom->type[i], so it needs TYPE_MASK; peratom2global() reads the named vector as well as the atom map; and group_function() is tried for every function word in a formula and returns 0 for anything that is not a group function, so the guard is factored out as is_group_function() and shared rather than syncing on every function call. The modify_host() and sync() calls on the result arrays in the three fixes stay. Those are a real host to device copy of the variable result, not a stale flag, and can only go away if variables are evaluated on the device. Correct the comment that said otherwise. Add a unit test that swaps in a Variable subclass recording which seams a formula routes through, and asserts them for the core atom vectors, gmask/rmask/grmask, indexed per-atom access, custom properties, and the compute and group function fallbacks. --- cmake/Modules/Packages/KOKKOS.cmake | 3 +- src/KOKKOS/Install.sh | 2 + src/KOKKOS/fix_addforce_kokkos.cpp | 8 +- src/KOKKOS/fix_efield_kokkos.cpp | 8 +- src/KOKKOS/fix_setforce_kokkos.cpp | 8 +- src/KOKKOS/variable_kokkos.cpp | 157 +++++++++++++++++++++++++++ src/KOKKOS/variable_kokkos.h | 44 ++++++++ src/accelerator_kokkos.h | 7 ++ src/input.cpp | 5 +- src/variable.cpp | 34 ++++-- src/variable.h | 46 +++++--- unittest/commands/test_variables.cpp | 153 ++++++++++++++++++++++++++ 12 files changed, 440 insertions(+), 35 deletions(-) create mode 100644 src/KOKKOS/variable_kokkos.cpp create mode 100644 src/KOKKOS/variable_kokkos.h diff --git a/cmake/Modules/Packages/KOKKOS.cmake b/cmake/Modules/Packages/KOKKOS.cmake index 9043c1ae70a..1a1f640fb4c 100644 --- a/cmake/Modules/Packages/KOKKOS.cmake +++ b/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 diff --git a/src/KOKKOS/Install.sh b/src/KOKKOS/Install.sh index 465c1480ccf..c2cd28af2ee 100755 --- a/src/KOKKOS/Install.sh +++ b/src/KOKKOS/Install.sh @@ -563,6 +563,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 diff --git a/src/KOKKOS/fix_addforce_kokkos.cpp b/src/KOKKOS/fix_addforce_kokkos.cpp index d5f5dbb5c1e..b830dad0206 100644 --- a/src/KOKKOS/fix_addforce_kokkos.cpp +++ b/src/KOKKOS/fix_addforce_kokkos.cpp @@ -132,8 +132,6 @@ void FixAddForceKokkos::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::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(); } diff --git a/src/KOKKOS/fix_efield_kokkos.cpp b/src/KOKKOS/fix_efield_kokkos.cpp index 90b5ea72b0a..848e81ee889 100644 --- a/src/KOKKOS/fix_efield_kokkos.cpp +++ b/src/KOKKOS/fix_efield_kokkos.cpp @@ -150,11 +150,13 @@ void FixEfieldKokkos::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(); } diff --git a/src/KOKKOS/fix_setforce_kokkos.cpp b/src/KOKKOS/fix_setforce_kokkos.cpp index 568831fccce..2d0054ab60e 100644 --- a/src/KOKKOS/fix_setforce_kokkos.cpp +++ b/src/KOKKOS/fix_setforce_kokkos.cpp @@ -114,8 +114,6 @@ void FixSetForceKokkos::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::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(); } diff --git a/src/KOKKOS/variable_kokkos.cpp b/src/KOKKOS/variable_kokkos.cpp new file mode 100644 index 00000000000..89d83051b5c --- /dev/null +++ b/src/KOKKOS/variable_kokkos.cpp @@ -0,0 +1,157 @@ +/* ---------------------------------------------------------------------- + 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 + +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(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. the other special functions read 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) +{ + 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); + + 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); +} diff --git a/src/KOKKOS/variable_kokkos.h b/src/KOKKOS/variable_kokkos.h new file mode 100644 index 00000000000..566e31a9a1b --- /dev/null +++ b/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 diff --git a/src/accelerator_kokkos.h b/src/accelerator_kokkos.h index c05e958b85c..94d399e47d4 100644 --- a/src/accelerator_kokkos.h +++ b/src/accelerator_kokkos.h @@ -27,6 +27,7 @@ #include "kokkos.h" // IWYU pragma: export #include "memory_kokkos.h" // IWYU pragma: export #include "modify_kokkos.h" // IWYU pragma: export +#include "variable_kokkos.h" // IWYU pragma: export #include "neighbor_kokkos.h" // IWYU pragma: export #define LAMMPS_INLINE KOKKOS_INLINE_FUNCTION @@ -43,6 +44,7 @@ #include "group.h" #include "memory.h" #include "modify.h" +#include "variable.h" #include "neighbor.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: diff --git a/src/input.cpp b/src/input.cpp index a4ef7835a20..2081a8c9951 100644 --- a/src/input.cpp +++ b/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 diff --git a/src/variable.cpp b/src/variable.cpp index 8597dbc43ec..b91bc8bd12c 100644 --- a/src/variable.cpp +++ b/src/variable.cpp @@ -1651,6 +1651,10 @@ double Variable::evaluate(char *str, Tree **tree, int ivar) print_var_error(FLERR,"Variable evaluation before simulation box is defined" + utils::errorurl(30),ivar); + // the compute is invoked below and reads per-atom data on the host + + sync_peratom(nullptr); + // uppercase used to access of peratom data by equal-style var int lowercase = 1; @@ -1938,6 +1942,10 @@ double Variable::evaluate(char *str, Tree **tree, int ivar) print_var_error(FLERR,"Variable evaluation before simulation box is defined" + utils::errorurl(30),ivar); + // the fix supplies per-atom data it stores on the host + + sync_peratom(nullptr); + // uppercase used to force access of // global vector vs global scalar, and global array vs global vector @@ -2360,6 +2368,10 @@ double Variable::evaluate(char *str, Tree **tree, int ivar) print_var_error(FLERR,"Variable evaluation before simulation box is defined" + utils::errorurl(30),ivar); + // the i_ / d_ / i2_ / d2_ prefix says which custom array is read + + sync_peratom(word); + int index_custom,type_custom,cols_custom; if (word[1] == '2') index_custom = atom->find_custom(word+3,type_custom,cols_custom); else index_custom = atom->find_custom(word+2,type_custom,cols_custom); @@ -4335,19 +4347,25 @@ int Variable::math_function(char *word, char *contents, Tree **tree, Tree **tree torque(group,dim),inertia(group,dim),omega(group,dim) ------------------------------------------------------------------------- */ +int Variable::is_group_function(const char *word) +{ + return (strcmp(word,"count") == 0) || (strcmp(word,"mass") == 0) || + (strcmp(word,"charge") == 0) || (strcmp(word,"xcm") == 0) || + (strcmp(word,"vcm") == 0) || (strcmp(word,"fcm") == 0) || + (strcmp(word,"bound") == 0) || (strcmp(word,"gyration") == 0) || + (strcmp(word,"ke") == 0) || (strcmp(word,"angmom") == 0) || + (strcmp(word,"torque") == 0) || (strcmp(word,"inertia") == 0) || + (strcmp(word,"omega") == 0); +} + +/* ---------------------------------------------------------------------- */ + int Variable::group_function(char *word, char *contents, Tree **tree, Tree **treestack, int &ntreestack, double *argstack, int &nargstack, int ivar) { // word not a match to any group function - if (strcmp(word,"count") != 0 && strcmp(word,"mass") && - strcmp(word,"charge") != 0 && strcmp(word,"xcm") != 0 && - strcmp(word,"vcm") != 0 && strcmp(word,"fcm") != 0 && - strcmp(word,"bound") != 0 && strcmp(word,"gyration") != 0 && - strcmp(word,"ke") != 0 && strcmp(word,"angmom") != 0 && - strcmp(word,"torque") != 0 && strcmp(word,"inertia") != 0 && - strcmp(word,"omega") != 0) - return 0; + if (!is_group_function(word)) return 0; // parse contents for comma-separated args // narg = number of args, args = strings between commas diff --git a/src/variable.h b/src/variable.h index 211e646b07d..ce76182e287 100644 --- a/src/variable.h +++ b/src/variable.h @@ -51,7 +51,7 @@ class Variable : protected Pointers { char *retrieve(const char *); double compute_equal(int); double compute_equal(const std::string &); - void compute_atom(int, int, double *, int, int); + virtual void compute_atom(int, int, double *, int, int); int compute_vector(int, double **); void internal_set(int, double); int internal_create(char *, double); @@ -92,15 +92,6 @@ class Variable : protected Pointers { }; std::vector variables; - private: - int treetype; // ATOM or VECTOR flag for formula evaluation - - class RanMars *randomequal; // random number generator for equal-style vars - class RanMars *randomatom; // random number generator for atom-style vars - - int precedence[18]; // precedence level of math operators - // set length to include up to XOR in enum - struct Tree { // parse tree for atom-style or vector-style vars double value; // single scalar double *array; // per-atom or per-type list of doubles @@ -127,6 +118,35 @@ class Variable : protected Pointers { } }; + // seams for accelerator packages that keep a second copy of the per-atom + // data. each of these reads per-atom arrays on the host, so VariableKokkos + // overrides them to make the data it needs current and then delegates here. + + virtual void atom_vector(char *, Tree **, Tree **, int &); + static int is_group_function(const char *); + virtual int group_function(char *, char *, Tree **, Tree **, int &, double *, int &, int); + virtual int special_function(const std::string &, char *, Tree **, Tree **, int &, double *, + int &, int, char *, int &, char *&); + virtual void peratom2global(int, char *, double *, int, tagint, Tree **, Tree **, int &, + double *, int &); + virtual void custom2global(int *, double *, int, tagint, Tree **, Tree **, int &, double *, + int &); + + // per-atom data reached through a compute or fix cannot be identified from + // the formula, so evaluate() passes a null pointer and asks for all of it. + // custom properties pass the i_ / d_ / i2_ / d2_ name they were given. + + virtual void sync_peratom(const char *) {} + + private: + int treetype; // ATOM or VECTOR flag for formula evaluation + + class RanMars *randomequal; // random number generator for equal-style vars + class RanMars *randomatom; // random number generator for atom-style vars + + int precedence[18]; // precedence level of math operators + // set length to include up to XOR in enum + int compute_python(int); void remove(int); int recycle(); @@ -140,15 +160,9 @@ class Variable : protected Pointers { void free_tree(Tree *); int find_matching_paren(char *, int, char *&, int); int math_function(char *, char *, Tree **, Tree **, int &, double *, int &, int); - int group_function(char *, char *, Tree **, Tree **, int &, double *, int &, int); Region *region_function(char *, int); - int special_function(const std::string &, char *, Tree **, Tree **, int &, double *, int &, int, - char *, int &, char *&); int feature_function(char *, char *, Tree **, Tree **, int &, double *, int &, int); - void peratom2global(int, char *, double *, int, tagint, Tree **, Tree **, int &, double *, int &); - void custom2global(int *, double *, int, tagint, Tree **, Tree **, int &, double *, int &); int is_atom_vector(char *); - void atom_vector(char *, Tree **, Tree **, int &); int parse_args(char *, char **); void parse_vector(int, char *); char *find_next_comma(char *); diff --git a/unittest/commands/test_variables.cpp b/unittest/commands/test_variables.cpp index 95ddb0d89da..12b36635310 100644 --- a/unittest/commands/test_variables.cpp +++ b/unittest/commands/test_variables.cpp @@ -14,6 +14,7 @@ #include "lammps.h" #include "atom.h" +#include "atom_masks.h" #include "domain.h" #include "group.h" #include "info.h" @@ -27,6 +28,8 @@ #include "gtest/gtest.h" #include +#include +#include #include // whether to print verbose output (i.e. not capturing LAMMPS screen output). @@ -867,6 +870,156 @@ TEST_F(VariableTest, Set) variable->internal_set(variable->find("ten"), -2.5); ASSERT_THAT(variable->retrieve("ten"), StrEq("-2.5")); } +// Records which of the accelerator seams in Variable a formula routes +// through, and with what name. VariableKokkos overrides exactly these to +// decide what to copy back from the device, so this pins down the behavior +// that matters without needing a device. + +class RecordingVariable : public Variable { +public: + RecordingVariable(LAMMPS *lmp) : Variable(lmp) {} + std::set seen; + void reset() { seen.clear(); } + + void compute_atom(int ivar, int igroup, double *result, int stride, int sumflag) override + { + seen.insert("mask"); + Variable::compute_atom(ivar, igroup, result, stride, sumflag); + } + +protected: + void atom_vector(char *word, Tree **tree, Tree **treestack, int &ntreestack) override + { + seen.insert(word); + Variable::atom_vector(word, tree, treestack, ntreestack); + } + + int group_function(char *word, char *contents, Tree **tree, Tree **treestack, int &ntreestack, + double *argstack, int &nargstack, int ivar) override + { + if (is_group_function(word)) seen.insert(""); + return Variable::group_function(word, contents, tree, treestack, ntreestack, argstack, + nargstack, ivar); + } + + int 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) override + { + if ((word == "gmask") || (word == "rmask") || (word == "grmask")) seen.insert(word); + return Variable::special_function(word, contents, tree, treestack, ntreestack, argstack, + nargstack, ivar, str, i, ptr); + } + + void peratom2global(int flag, char *word, double *vector, int nstride, tagint id, Tree **tree, + Tree **treestack, int &ntreestack, double *argstack, int &nargstack) override + { + seen.insert(""); + Variable::peratom2global(flag, word, vector, nstride, id, tree, treestack, ntreestack, + argstack, nargstack); + } + + void custom2global(int *ivector, double *dvector, int nstride, tagint id, Tree **tree, + Tree **treestack, int &ntreestack, double *argstack, int &nargstack) override + { + seen.insert(""); + Variable::custom2global(ivector, dvector, nstride, id, tree, treestack, ntreestack, + argstack, nargstack); + } + + void sync_peratom(const char *word) override { seen.insert(word ? word : ""); } +}; + +class VariableSyncTest : public LAMMPSTest { +protected: + RecordingVariable *rec; + Group *group; + + void SetUp() override + { + testbinary = "VariableSyncTest"; + args = {"-log", "none", "-echo", "screen", "-nocite"}; + LAMMPSTest::SetUp(); + group = lmp->group; + + // swap in the recording subclass before any variable is defined + delete lmp->input->variable; + rec = new RecordingVariable(lmp); + lmp->input->variable = rec; + + BEGIN_HIDE_OUTPUT(); + command("fix props all property/atom mol rmass q d_dm_val"); + command("atom_modify map array"); // needed for x[N] style access + command("units real"); + command("lattice sc 1.0 origin 0.125 0.125 0.125"); + command("region box block -2 2 -2 2 -2 2"); + command("create_box 8 box"); + command("create_atoms 1 box"); + command("mass * 1.0"); + command("region left block -2.0 -1.0 INF INF INF INF"); + command("compute dm_ke all ke/atom"); + command("run 0 post no"); + END_HIDE_OUTPUT(); + } + + std::set seams_for(const std::string &formula) + { + static int n = 0; + std::string name = fmt::format("vsync{}", n++); + BEGIN_HIDE_OUTPUT(); + command(fmt::format("variable {} atom \"{}\"", name, formula)); + END_HIDE_OUTPUT(); + const int ivar = rec->find(name.c_str()); + const int nlocal = lmp->atom->nlocal; + std::vector buf(nlocal > 0 ? nlocal : 1); + rec->reset(); + rec->compute_atom(ivar, group->find("all"), buf.data(), 1, 0); + return rec->seen; + } +}; + +using StrSet = std::set; + +TEST_F(VariableSyncTest, AcceleratorSeams) +{ + // compute_atom() always reads atom->mask for the group test + + EXPECT_EQ(seams_for("x"), (StrSet{"x", "mask"})); + EXPECT_EQ(seams_for("x*y+z"), (StrSet{"x", "y", "z", "mask"})); + EXPECT_EQ(seams_for("vx*vy+vz"), (StrSet{"vx", "vy", "vz", "mask"})); + EXPECT_EQ(seams_for("fx+fy+fz"), (StrSet{"fx", "fy", "fz", "mask"})); + EXPECT_EQ(seams_for("q"), (StrSet{"q", "mask"})); + EXPECT_EQ(seams_for("type"), (StrSet{"type", "mask"})); + EXPECT_EQ(seams_for("id"), (StrSet{"id", "mask"})); + EXPECT_EQ(seams_for("mol"), (StrSet{"mol", "mask"})); + EXPECT_EQ(seams_for("mass"), (StrSet{"mass", "mask"})); + + // only the arrays the formula actually names + + EXPECT_EQ(seams_for("x*vy"), (StrSet{"x", "vy", "mask"})); + EXPECT_EQ(seams_for("sqrt(x*x)+q*type"), (StrSet{"x", "q", "type", "mask"})); + + // group and region tests route through special_function() + + EXPECT_EQ(seams_for("gmask(all)"), (StrSet{"gmask", "mask"})); + EXPECT_EQ(seams_for("rmask(left)"), (StrSet{"rmask", "mask"})); + EXPECT_EQ(seams_for("grmask(all,left)"), (StrSet{"grmask", "mask"})); + + // custom per-atom properties are named by their prefix + + EXPECT_EQ(seams_for("d_dm_val"), (StrSet{"d_dm_val", "mask"})); + + // data that cannot be accounted for falls back to everything + + EXPECT_THAT(seams_for("c_dm_ke"), ::testing::Contains("")); + EXPECT_THAT(seams_for("x*count(all)"), ::testing::Contains("")); + EXPECT_THAT(seams_for("x[1]+y[2]"), ::testing::Contains("")); + + // constant-folded formulas touch no per-atom data beyond the group test + + EXPECT_EQ(seams_for("1.0+2.0"), (StrSet{"mask"})); +} + } // namespace LAMMPS_NS int main(int argc, char **argv) From 3bda429aa6a9139f061ef077a543966f8babe752 Mon Sep 17 00:00:00 2001 From: Stan Moore Date: Tue, 25 Aug 2026 18:56:00 +0000 Subject: [PATCH 2/2] KOKKOS: add sync seams for reducing special functions and thermo keywords Two paths in evaluate() read per-atom data on the host without going through any of the seams added in the previous commit, so removing the blanket atomKK->sync(Host,ALL_MASK) from the three fixes left them reading stale data. special_function() handles sum(), min(), max(), ave(), trap(), slope(), sort(), and rsort(), which take a compute or a fix as their argument and invoke it. The override only synced for gmask, rmask, and grmask, so a formula such as sum(c_msd) driving fix setforce/kk read a stale host copy of the atom coordinates. Factor the map lookup that guards special_function() out as is_special_function() so the override can use the same test rather than a second copy of the name list, and have the override fall back to ALL_MASK for any special function that is not on an explicit list of ones known to touch no per-atom data. A special function added later therefore syncs too much rather than too little. The thermo keyword branch of evaluate() calls Thermo::evaluate_keyword(), which invokes the thermo computes, with no seam at all. Add one. Extend the unit test to cover both paths, and record the name of any special function the formula reaches rather than just gmask, rmask, and grmask, so the test measures which seams are used instead of repeating the mask policy that belongs in VariableKokkos. Also drop an unused include from the test, keep the new includes in accelerator_kokkos.h in alphabetical order, and move is_group_function() above the comment block that describes group_function(). --- src/KOKKOS/variable_kokkos.cpp | 23 +++++++++++++--- src/accelerator_kokkos.h | 4 +-- src/variable.cpp | 40 ++++++++++++++++++---------- src/variable.h | 1 + unittest/commands/test_variables.cpp | 17 ++++++++++-- 5 files changed, 63 insertions(+), 22 deletions(-) diff --git a/src/KOKKOS/variable_kokkos.cpp b/src/KOKKOS/variable_kokkos.cpp index 89d83051b5c..1dea2610bf2 100644 --- a/src/KOKKOS/variable_kokkos.cpp +++ b/src/KOKKOS/variable_kokkos.cpp @@ -97,16 +97,31 @@ int VariableKokkos::group_function(char *word, char *contents, Tree **tree, Tree /* ---------------------------------------------------------------------- gmask() tests atom->mask, rmask() and grmask() also match a region - against atom->x. the other special functions read no per-atom data. + 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) { - 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); + // 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); diff --git a/src/accelerator_kokkos.h b/src/accelerator_kokkos.h index 94d399e47d4..aa1f3072db8 100644 --- a/src/accelerator_kokkos.h +++ b/src/accelerator_kokkos.h @@ -27,8 +27,8 @@ #include "kokkos.h" // IWYU pragma: export #include "memory_kokkos.h" // IWYU pragma: export #include "modify_kokkos.h" // IWYU pragma: export -#include "variable_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,8 +44,8 @@ #include "group.h" #include "memory.h" #include "modify.h" -#include "variable.h" #include "neighbor.h" +#include "variable.h" #define LAMMPS_INLINE inline diff --git a/src/variable.cpp b/src/variable.cpp index 8b128dc758b..62bc3904c4e 100644 --- a/src/variable.cpp +++ b/src/variable.cpp @@ -2546,6 +2546,11 @@ double Variable::evaluate(char *str, Tree **tree, int ivar) print_var_error(FLERR,"Variable evaluation before simulation box is defined" + utils::errorurl(30),ivar); + // thermo keywords invoke the thermo computes, which read per-atom + // data on the host + + sync_peratom(nullptr); + int flag = output->thermo->evaluate_keyword(word,&value1); if (flag) print_var_error(FLERR,fmt::format("Invalid thermo keyword '{}' in variable formula", @@ -4337,19 +4342,6 @@ int Variable::math_function(char *word, char *contents, Tree **tree, Tree **tree return 1; } -/* ---------------------------------------------------------------------- - process a group function in formula with optional region arg - push result onto tree or arg stack - word = group function - contents = str between parentheses with one,two,three args - return 0 if not a match, 1 if successfully processed - customize by adding a group function with optional region arg: - count(group),mass(group),charge(group), - xcm(group,dim),vcm(group,dim),fcm(group,dim), - bound(group,xmin),gyration(group),ke(group),angmom(group,dim), - torque(group,dim),inertia(group,dim),omega(group,dim) -------------------------------------------------------------------------- */ - int Variable::is_group_function(const char *word) { return (strcmp(word,"count") == 0) || (strcmp(word,"mass") == 0) || @@ -4363,6 +4355,19 @@ int Variable::is_group_function(const char *word) /* ---------------------------------------------------------------------- */ +/* ---------------------------------------------------------------------- + process a group function in formula with optional region arg + push result onto tree or arg stack + word = group function + contents = str between parentheses with one,two,three args + return 0 if not a match, 1 if successfully processed + customize by adding a group function with optional region arg: + count(group),mass(group),charge(group), + xcm(group,dim),vcm(group,dim),fcm(group,dim), + bound(group,xmin),gyration(group),ke(group),angmom(group,dim), + torque(group,dim),inertia(group,dim),omega(group,dim) +------------------------------------------------------------------------- */ + int Variable::group_function(char *word, char *contents, Tree **tree, Tree **treestack, int &ntreestack, double *argstack, int &nargstack, int ivar) { @@ -4616,6 +4621,13 @@ const std::unordered_map special_function_map = { // NOLINTEND } +int Variable::is_special_function(const std::string &word) +{ + return special_function_map.find(word) != special_function_map.end(); +} + +/* ---------------------------------------------------------------------- */ + int Variable::special_function(const std::string &word, char *contents, Tree **tree, Tree **treestack, int &ntreestack, double *argstack, int &nargstack, int ivar, char *str, int &istr, char *&ptr) @@ -4624,7 +4636,7 @@ int Variable::special_function(const std::string &word, char *contents, Tree **t double value,sy,sxy; // return if "word" is not a match to any special function - if (special_function_map.find(word) == special_function_map.end()) return 0; + if (!is_special_function(word)) return 0; // process label2type() separately b/c its label arg can have commas in it diff --git a/src/variable.h b/src/variable.h index ce76182e287..d3aea6ba899 100644 --- a/src/variable.h +++ b/src/variable.h @@ -125,6 +125,7 @@ class Variable : protected Pointers { virtual void atom_vector(char *, Tree **, Tree **, int &); static int is_group_function(const char *); virtual int group_function(char *, char *, Tree **, Tree **, int &, double *, int &, int); + static int is_special_function(const std::string &); virtual int special_function(const std::string &, char *, Tree **, Tree **, int &, double *, int &, int, char *, int &, char *&); virtual void peratom2global(int, char *, double *, int, tagint, Tree **, Tree **, int &, diff --git a/unittest/commands/test_variables.cpp b/unittest/commands/test_variables.cpp index a04bc21eb66..a8c706eed03 100644 --- a/unittest/commands/test_variables.cpp +++ b/unittest/commands/test_variables.cpp @@ -14,7 +14,6 @@ #include "lammps.h" #include "atom.h" -#include "atom_masks.h" #include "domain.h" #include "group.h" #include "info.h" @@ -943,7 +942,10 @@ class RecordingVariable : public Variable { int &ntreestack, double *argstack, int &nargstack, int ivar, char *str, int &i, char *&ptr) override { - if ((word == "gmask") || (word == "rmask") || (word == "grmask")) seen.insert(word); + // record the name of any special function reached, so this test + // measures which seams a formula uses rather than duplicating the + // mask policy that lives in VariableKokkos + if (is_special_function(word)) seen.insert(word); return Variable::special_function(word, contents, tree, treestack, ntreestack, argstack, nargstack, ivar, str, i, ptr); } @@ -995,6 +997,7 @@ class VariableSyncTest : public LAMMPSTest { command("mass * 1.0"); command("region left block -2.0 -1.0 INF INF INF INF"); command("compute dm_ke all ke/atom"); + command("compute dm_msd all msd"); command("run 0 post no"); END_HIDE_OUTPUT(); } @@ -1052,6 +1055,16 @@ TEST_F(VariableSyncTest, AcceleratorSeams) EXPECT_THAT(seams_for("x*count(all)"), ::testing::Contains("")); EXPECT_THAT(seams_for("x[1]+y[2]"), ::testing::Contains("")); + // special functions that reduce a compute or fix invoke it, and thermo + // keywords invoke the thermo computes; both read per-atom data on the host + + EXPECT_THAT(seams_for("sum(c_dm_msd)"), ::testing::Contains("sum")); + EXPECT_THAT(seams_for("x*temp"), ::testing::Contains("")); + + // special functions that touch no per-atom data must not force a sync + + EXPECT_EQ(seams_for("x*is_os(^Linux)"), (StrSet{"x", "is_os", "mask"})); + // constant-folded formulas touch no per-atom data beyond the group test EXPECT_EQ(seams_for("1.0+2.0"), (StrSet{"mask"}));