Merge pull request #4012 from stanmoore1/kk_fix_temp

Port Fix temp berendsen and rescale to use Kokkos

doc/src/Commands_fix.rst

@@ -239,10 +239,10 @@ OPT.
    * :doc:`store/force <fix_store_force>`
    * :doc:`store/state <fix_store_state>`
    * :doc:`tdpd/source <fix_dpd_source>`
-   * :doc:`temp/berendsen <fix_temp_berendsen>`
+   * :doc:`temp/berendsen (k) <fix_temp_berendsen>`
    * :doc:`temp/csld <fix_temp_csvr>`
    * :doc:`temp/csvr <fix_temp_csvr>`
-   * :doc:`temp/rescale <fix_temp_rescale>`
+   * :doc:`temp/rescale (k) <fix_temp_rescale>`
    * :doc:`temp/rescale/eff <fix_temp_rescale_eff>`
    * :doc:`tfmc <fix_tfmc>`
    * :doc:`tgnpt/drude <fix_tgnh_drude>`

doc/src/fix_temp_berendsen.rst

@@ -1,8 +1,11 @@
 .. index:: fix temp/berendsen
+.. index:: fix temp/berendsen/kk
 
 fix temp/berendsen command
 ==========================
 
+Accelerator Variants: *temp/berendsen/kk*
+
 Syntax
 """"""
 
@@ -118,6 +121,10 @@ remaining thermal degrees of freedom, and the bias is added back in.
 
 ----------
 
+.. include:: accel_styles.rst
+
+----------
+
 Restart, fix_modify, output, run start/stop, minimize info
 """""""""""""""""""""""""""""""""""""""""""""""""""""""""""
 

doc/src/fix_temp_rescale.rst

@@ -1,8 +1,11 @@
 .. index:: fix temp/rescale
+.. index:: fix temp/rescale/kk
 
 fix temp/rescale command
 ========================
 
+Accelerator Variants: *temp/rescale/kk*
+
 Syntax
 """"""
 
@@ -125,6 +128,10 @@ remaining thermal degrees of freedom, and the bias is added back in.
 
 ----------
 
+.. include:: accel_styles.rst
+
+----------
+
 Restart, fix_modify, output, run start/stop, minimize info
 """""""""""""""""""""""""""""""""""""""""""""""""""""""""""
 

src/KOKKOS/Install.sh

@@ -177,6 +177,10 @@ action fix_shardlow_kokkos.cpp fix_shardlow.cpp
 action fix_shardlow_kokkos.h fix_shardlow.h
 action fix_spring_self_kokkos.cpp
 action fix_spring_self_kokkos.h
+action fix_temp_berendsen_kokkos.cpp
+action fix_temp_berendsen_kokkos.h
+action fix_temp_rescale_kokkos.cpp
+action fix_temp_rescale_kokkos.h
 action fix_viscous_kokkos.cpp
 action fix_viscous_kokkos.h
 action fix_wall_gran_kokkos.cpp fix_wall_gran.cpp

src/KOKKOS/fix_temp_berendsen_kokkos.cpp

@@ -0,0 +1,135 @@
+// 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.
+------------------------------------------------------------------------- */
+
+#include "fix_temp_berendsen_kokkos.h"
+
+#include "atom_kokkos.h"
+#include "comm.h"
+#include "compute.h"
+#include "error.h"
+#include "force.h"
+#include "group.h"
+#include "input.h"
+#include "modify.h"
+#include "update.h"
+#include "variable.h"
+#include "atom_masks.h"
+
+#include <cmath>
+#include <cstring>
+
+using namespace LAMMPS_NS;
+using namespace FixConst;
+
+enum{NOBIAS,BIAS};
+enum{CONSTANT,EQUAL};
+
+/* ---------------------------------------------------------------------- */
+
+template<class DeviceType>
+FixTempBerendsenKokkos<DeviceType>::FixTempBerendsenKokkos(LAMMPS *lmp, int narg, char **arg) :
+  FixTempBerendsen(lmp, narg, arg)
+{
+  kokkosable = 1;
+  atomKK = (AtomKokkos *)atom;
+  execution_space = ExecutionSpaceFromDevice<DeviceType>::space;
+
+  datamask_read = EMPTY_MASK;
+  datamask_modify = EMPTY_MASK;
+}
+
+/* ---------------------------------------------------------------------- */
+
+template<class DeviceType>
+void FixTempBerendsenKokkos<DeviceType>::end_of_step()
+{
+  atomKK->sync(temperature->execution_space,temperature->datamask_read);
+  double t_current = temperature->compute_scalar();
+  atomKK->modified(temperature->execution_space,temperature->datamask_modify);
+  atomKK->sync(execution_space,temperature->datamask_modify);
+
+  double tdof = temperature->dof;
+
+  // there is nothing to do, if there are no degrees of freedom
+
+  if (tdof < 1) return;
+
+  if (t_current == 0.0)
+    error->all(FLERR, "Computed current temperature for fix temp/berendsen must not be 0.0");
+
+  double delta = update->ntimestep - update->beginstep;
+  if (delta != 0.0) delta /= update->endstep - update->beginstep;
+
+  // set current t_target
+  // if variable temp, evaluate variable, wrap with clear/add
+
+  if (tstyle == CONSTANT)
+    t_target = t_start + delta * (t_stop-t_start);
+  else {
+    modify->clearstep_compute();
+    t_target = input->variable->compute_equal(tvar);
+    if (t_target < 0.0)
+      error->one(FLERR, "Fix temp/berendsen variable {} returned negative temperature",
+                 input->variable->names[tvar]);
+    modify->addstep_compute(update->ntimestep + nevery);
+  }
+
+  // rescale velocities by lamda
+  // for BIAS:
+  //   temperature is current, so do not need to re-compute
+  //   OK to not test returned v = 0, since lamda is multiplied by v
+
+  double lamda = sqrt(1.0 + update->dt/t_period*(t_target/t_current - 1.0));
+  double efactor = 0.5 * force->boltz * tdof;
+  energy += t_current * (1.0-lamda*lamda) * efactor;
+
+  auto v = atomKK->k_v.view<DeviceType>();
+  auto mask = atomKK->k_mask.view<DeviceType>();
+  int nlocal = atom->nlocal;
+  auto groupbit = this->groupbit;
+
+  if (which == NOBIAS) {
+    atomKK->sync(temperature->execution_space,temperature->datamask_read);
+    temperature->remove_bias_all();
+    atomKK->modified(temperature->execution_space,temperature->datamask_modify);
+    atomKK->sync(execution_space,temperature->datamask_modify);
+  }
+
+  atomKK->sync(execution_space,V_MASK|MASK_MASK);
+
+  Kokkos::parallel_for(Kokkos::RangePolicy<DeviceType>(0,nlocal), LAMMPS_LAMBDA(int i) {
+    if (mask[i] & groupbit) {
+      v(i,0) *= lamda;
+      v(i,1) *= lamda;
+      v(i,2) *= lamda;
+    }
+  });
+
+  atomKK->modified(execution_space,V_MASK);
+
+  if (which == NOBIAS) {
+    atomKK->sync(temperature->execution_space,temperature->datamask_read);
+    temperature->restore_bias_all();
+    atomKK->modified(temperature->execution_space,temperature->datamask_modify);
+    atomKK->sync(execution_space,temperature->datamask_modify);
+  }
+}
+/* ---------------------------------------------------------------------- */
+
+namespace LAMMPS_NS {
+template class FixTempBerendsenKokkos<LMPDeviceType>;
+#ifdef LMP_KOKKOS_GPU
+template class FixTempBerendsenKokkos<LMPHostType>;
+#endif
+}

src/KOKKOS/fix_temp_berendsen_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.
+------------------------------------------------------------------------- */
+
+#ifdef FIX_CLASS
+// clang-format off
+FixStyle(temp/berendsen/kk,FixTempBerendsenKokkos<LMPDeviceType>);
+FixStyle(temp/berendsen/kk/device,FixTempBerendsenKokkos<LMPDeviceType>);
+FixStyle(temp/berendsen/kk/host,FixTempBerendsenKokkos<LMPHostType>);
+// clang-format on
+#else
+
+// clang-format off
+#ifndef LMP_FIX_TEMP_BERENDSEN_KOKKOS_H
+#define LMP_FIX_TEMP_BERENDSEN_KOKKOS_H
+
+#include "fix_temp_berendsen.h"
+#include "kokkos_type.h"
+
+namespace LAMMPS_NS {
+
+template<class DeviceType>
+class FixTempBerendsenKokkos : public FixTempBerendsen {
+ public:
+  typedef DeviceType device_type;
+
+  FixTempBerendsenKokkos(class LAMMPS *, int, char **);
+  ~FixTempBerendsenKokkos() override {}
+  void end_of_step() override;
+};
+
+}    // namespace LAMMPS_NS
+
+#endif
+#endif

src/KOKKOS/fix_temp_rescale_kokkos.cpp

@@ -0,0 +1,140 @@
+// 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.
+------------------------------------------------------------------------- */
+
+#include "fix_temp_rescale_kokkos.h"
+
+#include "atom_kokkos.h"
+#include "comm.h"
+#include "compute.h"
+#include "error.h"
+#include "force.h"
+#include "group.h"
+#include "input.h"
+#include "modify.h"
+#include "update.h"
+#include "variable.h"
+#include "atom_masks.h"
+
+#include <cmath>
+#include <cstring>
+
+using namespace LAMMPS_NS;
+using namespace FixConst;
+
+enum{NOBIAS,BIAS};
+enum{CONSTANT,EQUAL};
+
+/* ---------------------------------------------------------------------- */
+
+template<class DeviceType>
+FixTempRescaleKokkos<DeviceType>::FixTempRescaleKokkos(LAMMPS *lmp, int narg, char **arg) :
+  FixTempRescale(lmp, narg, arg)
+{
+  kokkosable = 1;
+  atomKK = (AtomKokkos *)atom;
+  execution_space = ExecutionSpaceFromDevice<DeviceType>::space;
+
+  datamask_read = EMPTY_MASK;
+  datamask_modify = EMPTY_MASK;
+}
+
+/* ---------------------------------------------------------------------- */
+
+template<class DeviceType>
+void FixTempRescaleKokkos<DeviceType>::end_of_step()
+{
+  atomKK->sync(temperature->execution_space,temperature->datamask_read);
+  double t_current = temperature->compute_scalar();
+  atomKK->modified(temperature->execution_space,temperature->datamask_modify);
+  atomKK->sync(execution_space,temperature->datamask_modify);
+
+  // there is nothing to do, if there are no degrees of freedom
+
+  if (temperature->dof < 1) return;
+
+  // protect against division by zero
+
+  if (t_current == 0.0)
+    error->all(FLERR,"Computed temperature for fix temp/rescale cannot be 0.0");
+
+  double delta = update->ntimestep - update->beginstep;
+  if (delta != 0.0) delta /= update->endstep - update->beginstep;
+
+  // set current t_target
+  // if variable temp, evaluate variable, wrap with clear/add
+
+  if (tstyle == CONSTANT)
+    t_target = t_start + delta * (t_stop-t_start);
+  else {
+    modify->clearstep_compute();
+    t_target = input->variable->compute_equal(tvar);
+    if (t_target < 0.0)
+      error->one(FLERR, "Fix temp/rescale variable returned negative temperature");
+    modify->addstep_compute(update->ntimestep + nevery);
+  }
+
+  // rescale velocity of appropriate atoms if outside window
+  // for BIAS:
+  //   temperature is current, so do not need to re-compute
+  //   OK to not test returned v = 0, since factor is multiplied by v
+
+  if (fabs(t_current-t_target) > t_window) {
+    t_target = t_current - fraction*(t_current-t_target);
+    double factor = sqrt(t_target/t_current);
+    double efactor = 0.5 * force->boltz * temperature->dof;
+
+    energy += (t_current-t_target) * efactor;
+
+    auto v = atomKK->k_v.view<DeviceType>();
+    auto mask = atomKK->k_mask.view<DeviceType>();
+    int nlocal = atom->nlocal;
+    auto groupbit = this->groupbit;
+
+    if (which == NOBIAS) {
+      atomKK->sync(temperature->execution_space,temperature->datamask_read);
+      temperature->remove_bias_all();
+      atomKK->modified(temperature->execution_space,temperature->datamask_modify);
+      atomKK->sync(execution_space,temperature->datamask_modify);
+    }
+
+    atomKK->sync(execution_space,V_MASK|MASK_MASK);
+
+    Kokkos::parallel_for(Kokkos::RangePolicy<DeviceType>(0,nlocal), LAMMPS_LAMBDA(int i) {
+      if (mask[i] & groupbit) {
+        v(i,0) *= factor;
+        v(i,1) *= factor;
+        v(i,2) *= factor;
+      }
+    });
+
+    atomKK->modified(execution_space,V_MASK);
+
+    if (which == NOBIAS) {
+      atomKK->sync(temperature->execution_space,temperature->datamask_read);
+      temperature->restore_bias_all();
+      atomKK->modified(temperature->execution_space,temperature->datamask_modify);
+      atomKK->sync(execution_space,temperature->datamask_modify);
+
+    }
+  }
+}
+
+/* ---------------------------------------------------------------------- */
+
+namespace LAMMPS_NS {
+template class FixTempRescaleKokkos<LMPDeviceType>;
+#ifdef LMP_KOKKOS_GPU
+template class FixTempRescaleKokkos<LMPHostType>;
+#endif
+}

src/KOKKOS/fix_temp_rescale_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.
+------------------------------------------------------------------------- */
+
+#ifdef FIX_CLASS
+// clang-format off
+FixStyle(temp/rescale/kk,FixTempRescaleKokkos<LMPDeviceType>);
+FixStyle(temp/rescale/kk/device,FixTempRescaleKokkos<LMPDeviceType>);
+FixStyle(temp/rescale/kk/host,FixTempRescaleKokkos<LMPHostType>);
+// clang-format on
+#else
+
+// clang-format off
+#ifndef LMP_FIX_TEMP_RESCALE_KOKKOS_H
+#define LMP_FIX_TEMP_RESCALE_KOKKOS_H
+
+#include "fix_temp_rescale.h"
+#include "kokkos_type.h"
+
+namespace LAMMPS_NS {
+
+template<class DeviceType>
+class FixTempRescaleKokkos : public FixTempRescale {
+ public:
+  typedef DeviceType device_type;
+
+  FixTempRescaleKokkos(class LAMMPS *, int, char **);
+  ~FixTempRescaleKokkos() override {}
+  void end_of_step() override;
+};
+
+}    // namespace LAMMPS_NS
+
+#endif
+#endif

src/fix_temp_berendsen.h

@@ -38,7 +38,7 @@ class FixTempBerendsen : public Fix {
   void restart(char *buf) override;
   void *extract(const char *, int &) override;
 
- private:
+ protected:
   int which;
   double t_start, t_stop, t_period, t_target;
   double energy;