Ported fix enforce2d to Kokkos.

doc/src/fix_enforce2d.txt

@@ -7,6 +7,7 @@
 :line
 
 fix enforce2d command :h3
+fix enforce2d/kk command :h3
 
 [Syntax:]
 

src/KOKKOS/fix_enforce2d_kokkos.cpp

+/* -*- c++ -*- ----------------------------------------------------------
+   LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
+   http://lammps.sandia.gov, Sandia National Laboratories
+   Steve Plimpton, sjplimp@sandia.gov
+
+   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.
+------------------------------------------------------------------------- */
+
+/* ----------------------------------------------------------------------
+   Contributing authors: Stefan Paquay (Brandeis University)
+------------------------------------------------------------------------- */
+
+#include "atom_masks.h"
+#include "atom_kokkos.h"
+#include "fix_enforce2d_kokkos.h"
+
+using namespace LAMMPS_NS;
+
+
+template <class DeviceType>
+FixEnforce2DKokkos<DeviceType>::FixEnforce2DKokkos(LAMMPS *lmp, int narg, char **arg) :
+  FixEnforce2D(lmp, narg, arg)
+{
+  kokkosable = 1;
+  atomKK = (AtomKokkos *) atom;
+  execution_space = ExecutionSpaceFromDevice<DeviceType>::space;
+
+  datamask_read   = X_MASK | V_MASK | F_MASK | MASK_MASK;
+  datamask_modify = X_MASK | V_MASK | F_MASK;
+}
+
+
+template <class DeviceType>
+void FixEnforce2DKokkos<DeviceType>::setup(int vflag)
+{
+  post_force(vflag);
+}
+
+
+template <class DeviceType>
+void FixEnforce2DKokkos<DeviceType>::post_force(int vflag)
+{
+  atomKK->sync(execution_space,datamask_read);
+  atomKK->modified(execution_space,datamask_modify);
+
+  x = atomKK->k_x.view<DeviceType>();
+  v = atomKK->k_v.view<DeviceType>();
+  f = atomKK->k_f.view<DeviceType>();
+
+  mask = atomKK->k_mask.view<DeviceType>();
+
+  int nlocal = atomKK->nlocal;
+  if (igroup == atomKK->firstgroup) nlocal = atomKK->nfirst;
+
+  FixEnforce2DKokkosPostForceFunctor<DeviceType> functor(this);
+  Kokkos::parallel_for(nlocal,functor);
+
+  // Probably sync here again?
+  atomKK->sync(execution_space,datamask_read);
+  atomKK->modified(execution_space,datamask_modify);
+
+  for (int m = 0; m < nfixlist; m++)
+    flist[m]->enforce2d();
+
+
+}
+
+
+template <class DeviceType>
+void FixEnforce2DKokkos<DeviceType>::post_force_item( int i ) const
+{
+
+  if (mask[i] & groupbit){
+    v(i,2) = 0;
+    x(i,2) = 0;
+    f(i,2) = 0;
+
+    // Add for omega, angmom, torque...
+  }
+
+}
+
+
+template<class DeviceType>
+void FixEnforce2DKokkos<DeviceType>::cleanup_copy()
+{
+  id = style = NULL;
+  vatom = NULL;
+}
+
+
+namespace LAMMPS_NS {
+template class FixEnforce2DKokkos<LMPDeviceType>;
+#ifdef KOKKOS_HAVE_CUDA
+template class FixEnforce2DKokkos<LMPHostType>;
+#endif
+}

src/KOKKOS/fix_enforce2d_kokkos.h

+/* -*- c++ -*- ----------------------------------------------------------
+   LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
+   http://lammps.sandia.gov, Sandia National Laboratories
+   Steve Plimpton, sjplimp@sandia.gov
+
+   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
+
+FixStyle(enforce2d/kk,FixEnforce2DKokkos<LMPDeviceType>)
+FixStyle(enforce2d/kk/device,FixEnforce2DKokkos<LMPDeviceType>)
+FixStyle(enforce2d/kk/host,FixEnforce2DKokkos<LMPHostType>)
+
+#else
+
+#ifndef LMP_FIX_ENFORCE2D_KOKKOS_H
+#define LMP_FIX_ENFORCE2D_KOKKOS_H
+
+#include "fix_enforce2d.h"
+#include "kokkos_type.h"
+
+namespace LAMMPS_NS {
+
+template<class DeviceType>
+class FixEnforce2DKokkos : public FixEnforce2D {
+ public:
+  FixEnforce2DKokkos(class LAMMPS *, int, char **);
+  // ~FixEnforce2DKokkos() {}
+  // void init();
+  void cleanup_copy();
+  void setup(int);
+  void post_force(int);
+
+  KOKKOS_INLINE_FUNCTION
+  void post_force_item(int) const;
+
+  // void min_setup(int);       Kokkos does not support minimization (yet)
+  // void min_post_force(int);  Kokkos does not support minimization (yet)
+  // void post_force_respa(int, int, int);  No RRESPA support yet.
+
+ private:
+
+  typename ArrayTypes<DeviceType>::t_x_array x;
+  typename ArrayTypes<DeviceType>::t_v_array v;
+  typename ArrayTypes<DeviceType>::t_f_array f;
+
+  typename ArrayTypes<DeviceType>::t_int_1d mask;
+};
+
+
+template <class DeviceType>
+struct FixEnforce2DKokkosPostForceFunctor  {
+  typedef DeviceType device_type;
+  FixEnforce2DKokkos<DeviceType> c;
+
+  FixEnforce2DKokkosPostForceFunctor(FixEnforce2DKokkos<DeviceType>* c_ptr):
+    c(*c_ptr) {c.cleanup_copy();};
+  KOKKOS_INLINE_FUNCTION
+  void operator()(const int i) const {
+    c.post_force_item(i);
+  }
+};
+
+
+}
+
+#endif
+#endif
+
+/* ERROR/WARNING messages:
+
+E: Illegal ... command
+
+Self-explanatory.  Check the input script syntax and compare to the
+documentation for the command.  You can use -echo screen as a
+command-line option when running LAMMPS to see the offending line.
+
+E: Cannot use fix enforce2d with 3d simulation
+
+Self-explanatory.
+
+E: Fix enforce2d must be defined after fix %s
+
+UNDOCUMENTED
+
+*/