Base for KOKKOS dt/reset

src/KOKKOS/fix_dt_reset_kokkos.cpp

@@ -0,0 +1,136 @@
+// 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 "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;
+
+#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_params = Kokkos::DualView<double*, Kokkos::LayoutRight, DeviceType>("FixDtResetKokkos:gamma", 1);
+
+  k_params.h_view(0) = emax;
+
+
+  k_params.template modify<LMPHostType>();
+  k_params.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 dtmin = BIG;
+
+//  Kokkos::DualView<double*, Kokkos::LayoutRight, DeviceType> dt =
+//      Kokkos::DualView<double*, Kokkos::LayoutRight, DeviceType>("FixViscousKokkos:gamma", 1);
+
+  double dt;
+
+  copymode = 1;
+  Kokkos::parallel_reduce(Kokkos::RangePolicy<DeviceType, TagFixDtResetMass>(0,nlocal), *this, Kokkos::Min<double>(dt));
+//  Kokkos::parallel_reduce(Kokkos::RangePolicy<DeviceType, TagFixDtResetMass>(0,nlocal), *this, dt);
+  copymode = 0;
+
+  printf ("Dt: %f \n", dt);
+  atomKK->modified(execution_space, F_MASK);
+}
+
+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 k_emax = k_params.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 ((k_emax > 0.0) && (fsq * vsq > 0.0)) {
+      dte = k_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,63 @@
+/* -*- 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{};
+
+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;
+
+ 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_params;
+};
+
+}
+
+#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;