Merge pull request #3522 from mkanski/fix_viscous_kokkos

KOKKOS version of fix viscous and fix dt/reset

doc/src/Commands_fix.rst

@@ -65,7 +65,7 @@ OPT.
    * :doc:`drude <fix_drude>`
    * :doc:`drude/transform/direct <fix_drude_transform>`
    * :doc:`drude/transform/inverse <fix_drude_transform>`
-   * :doc:`dt/reset <fix_dt_reset>`
+   * :doc:`dt/reset (k) <fix_dt_reset>`
    * :doc:`edpd/source <fix_dpd_source>`
    * :doc:`efield <fix_efield>`
    * :doc:`ehex <fix_ehex>`
@@ -250,7 +250,7 @@ OPT.
    * :doc:`tune/kspace <fix_tune_kspace>`
    * :doc:`vector <fix_vector>`
    * :doc:`viscosity <fix_viscosity>`
-   * :doc:`viscous <fix_viscous>`
+   * :doc:`viscous (k) <fix_viscous>`
    * :doc:`viscous/sphere <fix_viscous_sphere>`
    * :doc:`wall/body/polygon <fix_wall_body_polygon>`
    * :doc:`wall/body/polyhedron <fix_wall_body_polyhedron>`

doc/src/fix_dt_reset.rst

@@ -3,6 +3,8 @@
 fix dt/reset command
 ====================
 
+Accelerator Variants: *dt/reset/kk*
+
 Syntax
 """"""
 
@@ -86,6 +88,10 @@ allows dump files to be written at intervals specified by simulation
 time, rather than by timesteps.  Simulation time is in time units;
 see the :doc:`units <units>` doc page for details.
 
+----------
+
+.. include:: accel_styles.rst
+
 Restart, fix_modify, output, run start/stop, minimize info
 """""""""""""""""""""""""""""""""""""""""""""""""""""""""""
 

doc/src/fix_viscous.rst

@@ -3,6 +3,8 @@
 fix viscous command
 ===================
 
+Accelerator Variants: *viscous/kk*
+
 Syntax
 """"""
 
@@ -84,6 +86,10 @@ more easily be used as a thermostat.
 
 ----------
 
+.. include:: accel_styles.rst
+
+----------
+
 Restart, fix_modify, output, run start/stop, minimize info
 """""""""""""""""""""""""""""""""""""""""""""""""""""""""""
 

src/KOKKOS/Install.sh

@@ -117,6 +117,8 @@ action fix_acks2_reaxff_kokkos.cpp fix_acks2_reaxff.cpp
 action fix_acks2_reaxff_kokkos.h fix_acks2_reaxff.h
 action fix_deform_kokkos.cpp
 action fix_deform_kokkos.h
+action fix_dt_reset_kokkos.cpp
+action fix_dt_reset_kokkos.h
 action fix_enforce2d_kokkos.cpp
 action fix_enforce2d_kokkos.h
 action fix_eos_table_rx_kokkos.cpp fix_eos_table_rx.cpp
@@ -165,6 +167,8 @@ action fix_wall_lj93_kokkos.cpp
 action fix_wall_lj93_kokkos.h
 action fix_wall_reflect_kokkos.cpp
 action fix_wall_reflect_kokkos.h
+action fix_viscous_kokkos.cpp
+action fix_viscous_kokkos.h
 action fix_dpd_energy_kokkos.cpp fix_dpd_energy.cpp
 action fix_dpd_energy_kokkos.h fix_dpd_energy.h
 action fix_rx_kokkos.cpp fix_rx.cpp

src/KOKKOS/fix_dt_reset_kokkos.cpp

@@ -0,0 +1,195 @@
+// 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_dt_reset_kokkos.h"
+
+#include "atom_kokkos.h"
+#include "atom_masks.h"
+#include "error.h"
+#include "force.h"
+#include "input.h"
+#include "integrate.h"
+#include "kokkos_base.h"
+#include "memory_kokkos.h"
+#include "modify.h"
+#include "output.h"
+#include "pair.h"
+#include "update.h"
+
+using namespace LAMMPS_NS;
+using namespace FixConst;
+
+#define BIG 1.0e20
+
+/* ---------------------------------------------------------------------- */
+
+template<class DeviceType>
+FixDtResetKokkos<DeviceType>::FixDtResetKokkos(LAMMPS *lmp, int narg, char **arg) :
+  FixDtReset(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 FixDtResetKokkos<DeviceType>::init()
+{
+  FixDtReset::init();
+
+  k_emax = Kokkos::DualView<double*, Kokkos::LayoutRight, DeviceType>("FixDtResetKokkos:gamma", 1);
+
+  k_emax.h_view(0) = emax;
+
+
+  k_emax.template modify<LMPHostType>();
+  k_emax.template sync<DeviceType>();
+
+  if (utils::strmatch(update->integrate_style,"^respa"))
+    error->all(FLERR,"Cannot (yet) use respa with Kokkos");
+}
+
+/* ---------------------------------------------------------------------- */
+
+template<class DeviceType>
+void FixDtResetKokkos<DeviceType>::end_of_step()
+{
+  atomKK->sync(execution_space, V_MASK | F_MASK | MASK_MASK | TYPE_MASK | RMASS_MASK);
+
+  v = atomKK->k_v.view<DeviceType>();
+  f = atomKK->k_f.view<DeviceType>();
+  mask = atomKK->k_mask.view<DeviceType>();
+  type = atomKK->k_type.view<DeviceType>();
+  if (atomKK->rmass)
+    rmass = atomKK->k_rmass.view<DeviceType>();
+  else
+    mass = atomKK->k_mass.view<DeviceType>();
+
+  int nlocal = atom->nlocal;
+
+  double dt;
+
+  copymode = 1;
+  if (atomKK->rmass)
+    Kokkos::parallel_reduce(Kokkos::RangePolicy<DeviceType, TagFixDtResetRMass>(0,nlocal), *this, Kokkos::Min<double>(dt));
+  else
+    Kokkos::parallel_reduce(Kokkos::RangePolicy<DeviceType, TagFixDtResetMass>(0,nlocal), *this, Kokkos::Min<double>(dt));
+  copymode = 0;
+
+  MPI_Allreduce(MPI_IN_PLACE, &dt, 1, MPI_DOUBLE, MPI_MIN, world);
+
+  atomKK->modified(execution_space, F_MASK);
+
+   if (minbound) dt = MAX(dt, tmin);
+   if (maxbound) dt = MIN(dt, tmax);
+
+   // if timestep didn't change, just return
+   // else reset update->dt and other classes that depend on it
+   // rRESPA, pair style, fixes
+
+   if (dt == update->dt) return;
+
+   laststep = update->ntimestep;
+
+   // calls to other classes that need to know timestep size changed
+   // similar logic is in Input::timestep()
+
+   update->update_time();
+   update->dt = dt;
+   update->dt_default = 0;
+   if (force->pair) force->pair->reset_dt();
+   for (int i = 0; i < modify->nfix; i++) modify->fix[i]->reset_dt();
+   output->reset_dt();
+
+}
+
+/* ---------------------------------------------------------------------- */
+
+template<class DeviceType>
+KOKKOS_INLINE_FUNCTION
+void FixDtResetKokkos<DeviceType>::operator()(TagFixDtResetMass, const int &i, double &k_dt) const {
+
+  double dtv, dtf, dte, dtsq;
+  double vsq, fsq, massinv;
+  double delx, dely, delz, delr;
+
+  double emax = k_emax.d_view(0);
+
+  if (mask[i] & groupbit) {
+
+    massinv = 1.0 / mass[type[i]];
+    vsq = v(i,0) * v(i,0) + v(i,1) * v(i,1) + v(i,2) * v(i,2);
+    fsq = f(i,0) * f(i,0) + f(i,1) * f(i,1) + f(i,2) * f(i,2);
+    dtv = dtf = dte = BIG;
+    if (vsq > 0.0) dtv = xmax / sqrt(vsq);
+    if (fsq > 0.0) dtf = sqrt(2.0 * xmax / (ftm2v * sqrt(fsq) * massinv));
+    k_dt = MIN(dtv, dtf);
+    if ((emax > 0.0) && (fsq * vsq > 0.0)) {
+      dte = emax / sqrt(fsq * vsq) / sqrt(ftm2v * mvv2e);
+      k_dt = MIN(dt, dte);
+    }
+    dtsq = k_dt * k_dt;
+    delx = k_dt * v(i,0) + 0.5 * dtsq * massinv * f(i,0) * ftm2v;
+    dely = k_dt * v(i,1) + 0.5 * dtsq * massinv * f(i,1) * ftm2v;
+    delz = k_dt * v(i,2) + 0.5 * dtsq * massinv * f(i,2) * ftm2v;
+    delr = sqrt(delx * delx + dely * dely + delz * delz);
+    if (delr > xmax) k_dt *= xmax / delr;
+  }
+ }
+
+/* ---------------------------------------------------------------------- */
+
+template<class DeviceType>
+KOKKOS_INLINE_FUNCTION
+void FixDtResetKokkos<DeviceType>::operator()(TagFixDtResetRMass, const int &i, double &k_dt) const {
+
+  double dtv, dtf, dte, dtsq;
+  double vsq, fsq, massinv;
+  double delx, dely, delz, delr;
+
+  double emax = k_emax.d_view(0);
+
+  if (mask[i] & groupbit) {
+
+    massinv = 1.0 / rmass[i];
+    vsq = v(i,0) * v(i,0) + v(i,1) * v(i,1) + v(i,2) * v(i,2);
+    fsq = f(i,0) * f(i,0) + f(i,1) * f(i,1) + f(i,2) * f(i,2);
+    dtv = dtf = dte = BIG;
+    if (vsq > 0.0) dtv = xmax / sqrt(vsq);
+    if (fsq > 0.0) dtf = sqrt(2.0 * xmax / (ftm2v * sqrt(fsq) * massinv));
+    k_dt = MIN(dtv, dtf);
+    if ((emax > 0.0) && (fsq * vsq > 0.0)) {
+      dte = emax / sqrt(fsq * vsq) / sqrt(ftm2v * mvv2e);
+      k_dt = MIN(dt, dte);
+    }
+    dtsq = k_dt * k_dt;
+    delx = k_dt * v(i,0) + 0.5 * dtsq * massinv * f(i,0) * ftm2v;
+    dely = k_dt * v(i,1) + 0.5 * dtsq * massinv * f(i,1) * ftm2v;
+    delz = k_dt * v(i,2) + 0.5 * dtsq * massinv * f(i,2) * ftm2v;
+    delr = sqrt(delx * delx + dely * dely + delz * delz);
+    if (delr > xmax) k_dt *= xmax / delr;
+  }
+}
+
+namespace LAMMPS_NS {
+template class FixDtResetKokkos<LMPDeviceType>;
+#ifdef LMP_KOKKOS_GPU
+template class FixDtResetKokkos<LMPHostType>;
+#endif
+}
+

src/KOKKOS/fix_dt_reset_kokkos.h

@@ -0,0 +1,66 @@
+/* -*- 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(dt/reset/kk,FixDtResetKokkos<LMPDeviceType>);
+FixStyle(dt/reset/kk/device,FixDtResetKokkos<LMPDeviceType>);
+FixStyle(dt/reset/kk/host,FixDtResetKokkos<LMPHostType>);
+// clang-format on
+#else
+
+// clang-format off
+#ifndef LMP_FIX_DT_RESET_KOKKOS_H
+#define LMP_FIX_DT_RESET_KOKKOS_H
+
+#include "fix_dt_reset.h"
+#include "kokkos_type.h"
+
+namespace LAMMPS_NS {
+
+struct TagFixDtResetMass{};
+struct TagFixDtResetRMass{};
+
+template<class DeviceType>
+class FixDtResetKokkos : public FixDtReset {
+ public:
+  typedef DeviceType device_type;
+  typedef ArrayTypes<DeviceType> AT;
+
+  FixDtResetKokkos(class LAMMPS *, int, char **);
+//  ~FixDtResetKokkos() override;
+  void init() override;
+  void end_of_step() override;
+
+  KOKKOS_INLINE_FUNCTION
+  void operator()(TagFixDtResetMass, const int&, double&) const;
+  KOKKOS_INLINE_FUNCTION
+  void operator()(TagFixDtResetRMass, const int&, double&) const;
+
+ private:
+  typename AT::t_v_array v;
+  typename AT::t_f_array f;
+  typename AT::t_int_1d_randomread mask;
+  typename AT::t_int_1d_randomread type;
+  typename ArrayTypes<DeviceType>::t_float_1d_randomread rmass;
+  typename ArrayTypes<DeviceType>::t_float_1d_randomread mass;
+
+
+  Kokkos::DualView<double*, Kokkos::LayoutRight, DeviceType> k_emax;
+};
+
+}
+
+#endif
+#endif
+

src/KOKKOS/fix_viscous_kokkos.cpp

@@ -0,0 +1,106 @@
+// 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_viscous_kokkos.h"
+
+#include "atom_kokkos.h"
+#include "update.h"
+#include "modify.h"
+#include "input.h"
+#include "memory_kokkos.h"
+#include "error.h"
+#include "atom_masks.h"
+#include "kokkos_base.h"
+
+using namespace LAMMPS_NS;
+using namespace FixConst;
+
+/* ---------------------------------------------------------------------- */
+
+template<class DeviceType>
+FixViscousKokkos<DeviceType>::FixViscousKokkos(LAMMPS *lmp, int narg, char **arg) :
+  FixViscous(lmp, narg, arg)
+{
+  kokkosable = 1;
+  atomKK = (AtomKokkos *) atom;
+  execution_space = ExecutionSpaceFromDevice<DeviceType>::space;
+  datamask_read = EMPTY_MASK;
+  datamask_modify = EMPTY_MASK;
+}
+
+/* ---------------------------------------------------------------------- */
+
+template<class DeviceType>
+FixViscousKokkos<DeviceType>::~FixViscousKokkos()
+{
+  if (copymode) return;
+}
+
+/* ---------------------------------------------------------------------- */
+
+template<class DeviceType>
+void FixViscousKokkos<DeviceType>::init()
+{
+  FixViscous::init();
+
+  k_gamma = Kokkos::DualView<double*, Kokkos::LayoutRight, DeviceType>("FixViscousKokkos:gamma",atom->ntypes+1);
+
+  for (int i = 1; i <= atom->ntypes; i++) k_gamma.h_view(i) = gamma[i];
+
+  k_gamma.template modify<LMPHostType>();
+  k_gamma.template sync<DeviceType>();
+
+  if (utils::strmatch(update->integrate_style,"^respa"))
+    error->all(FLERR,"Cannot (yet) use respa with Kokkos");
+}
+
+/* ---------------------------------------------------------------------- */
+
+template<class DeviceType>
+void FixViscousKokkos<DeviceType>::post_force(int /*vflag*/)
+{
+  atomKK->sync(execution_space, V_MASK | F_MASK | MASK_MASK | TYPE_MASK);
+
+  v = atomKK->k_v.view<DeviceType>();
+  f = atomKK->k_f.view<DeviceType>();
+  mask = atomKK->k_mask.view<DeviceType>();
+  type = atomKK->k_type.view<DeviceType>();
+
+  int nlocal = atom->nlocal;
+
+  copymode = 1;
+  Kokkos::parallel_for(Kokkos::RangePolicy<DeviceType, TagFixViscous>(0,nlocal),*this);
+  copymode = 0;
+
+  atomKK->modified(execution_space, F_MASK);
+}
+
+template<class DeviceType>
+KOKKOS_INLINE_FUNCTION
+void FixViscousKokkos<DeviceType>::operator()(TagFixViscous, const int &i) const {
+  if (mask[i] & groupbit) {
+    double drag = k_gamma.d_view(type[i]);
+    f(i,0) -= drag*v(i,0);
+    f(i,1) -= drag*v(i,1);
+    f(i,2) -= drag*v(i,2);
+  }
+}
+
+namespace LAMMPS_NS {
+template class FixViscousKokkos<LMPDeviceType>;
+#ifdef LMP_KOKKOS_GPU
+template class FixViscousKokkos<LMPHostType>;
+#endif
+}
+

src/KOKKOS/fix_viscous_kokkos.h

@@ -0,0 +1,60 @@
+/* -*- 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(viscous/kk,FixViscousKokkos<LMPDeviceType>);
+FixStyle(viscous/kk/device,FixViscousKokkos<LMPDeviceType>);
+FixStyle(viscous/kk/host,FixViscousKokkos<LMPHostType>);
+// clang-format on
+#else
+
+// clang-format off
+#ifndef LMP_FIX_VISCOUS_KOKKOS_H
+#define LMP_FIX_VISCOUS_KOKKOS_H
+
+#include "fix_viscous.h"
+#include "kokkos_type.h"
+
+namespace LAMMPS_NS {
+
+struct TagFixViscous{};
+
+template<class DeviceType>
+class FixViscousKokkos : public FixViscous {
+ public:
+  typedef DeviceType device_type;
+  typedef ArrayTypes<DeviceType> AT;
+
+  FixViscousKokkos(class LAMMPS *, int, char **);
+  ~FixViscousKokkos() override;
+  void init() override;
+  void post_force(int) override;
+
+  KOKKOS_INLINE_FUNCTION
+  void operator()(TagFixViscous, const int&) const;
+
+ private:
+  typename AT::t_v_array v;
+  typename AT::t_f_array f;
+  typename AT::t_int_1d_randomread mask;
+  typename AT::t_int_1d_randomread type;
+
+  Kokkos::DualView<double*, Kokkos::LayoutRight, DeviceType> k_gamma;
+};
+
+}
+
+#endif
+#endif
+

src/fix_dt_reset.h

@@ -34,7 +34,7 @@ class FixDtReset : public Fix {
   void end_of_step() override;
   double compute_scalar() override;
 
- private:
+ protected:
   bigint laststep;
   int minbound, maxbound;
   double tmin, tmax, xmax, emax;

src/fix_viscous.cpp

@@ -61,6 +61,8 @@ FixViscous::FixViscous(LAMMPS *lmp, int narg, char **arg) :
 
 FixViscous::~FixViscous()
 {
+  if (copymode) return;
+
   delete [] gamma;
 }