Merge pull request #940 from Pakketeretet2/port-enforce2d-kokkos

Port enforce2d kokkos

doc/src/Section_commands.txt

@@ -582,7 +582,7 @@ USER-INTEL, k = KOKKOS, o = USER-OMP, t = OPT.
 "dt/reset"_fix_dt_reset.html,
 "efield"_fix_efield.html,
 "ehex"_fix_ehex.html,
-"enforce2d"_fix_enforce2d.html,
+"enforce2d (k)"_fix_enforce2d.html,
 "evaporate"_fix_evaporate.html,
 "external"_fix_external.html,
 "freeze"_fix_freeze.html,

doc/src/fix_enforce2d.txt

@@ -7,6 +7,7 @@
 :line
 
 fix enforce2d command :h3
+fix enforce2d/kk command :h3
 
 [Syntax:]
 
@@ -27,12 +28,13 @@ not move from their initial z coordinate.
 
 :line
 
-Styles with a suffix are functionally the same as the corresponding
-style without the suffix.  They have been optimized to run faster,
-depending on your available hardware, as discussed in
-"Section 5"_Section_accelerate.html of the manual.  The
-accelerated styles take the same arguments and should produce the same
-results, except for round-off and precision issues.
+Styles with a {gpu}, {intel}, {kk}, {omp}, or {opt} suffix are
+functionally the same as the corresponding style without the suffix.
+They have been optimized to run faster, depending on your available
+hardware, as discussed in "Section 5"_Section_accelerate.html
+of the manual.  The accelerated styles take the same arguments and
+should produce the same results, except for round-off and precision
+issues.
 
 These accelerated styles are part of the GPU, USER-INTEL, KOKKOS,
 USER-OMP and OPT packages, respectively.  They are only enabled if

src/KOKKOS/Install.sh

@@ -93,6 +93,8 @@ action domain_kokkos.cpp
 action domain_kokkos.h
 action fix_deform_kokkos.cpp
 action fix_deform_kokkos.h
+action fix_enforce2d_kokkos.cpp
+action fix_enforce2d_kokkos.h
 action fix_eos_table_rx_kokkos.cpp fix_eos_table_rx.cpp
 action fix_eos_table_rx_kokkos.h fix_eos_table_rx.h  
 action fix_langevin_kokkos.cpp

src/KOKKOS/fix_enforce2d_kokkos.cpp

@@ -0,0 +1,168 @@
+/* -*- 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 & Matthew Peterson (Brandeis University)
+------------------------------------------------------------------------- */
+
+#include "atom_masks.h"
+#include "atom_kokkos.h"
+#include "comm.h"
+#include "error.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   = V_MASK | F_MASK | OMEGA_MASK | MASK_MASK
+	  | TORQUE_MASK | ANGMOM_MASK;
+
+  datamask_modify = V_MASK | F_MASK | OMEGA_MASK
+	  | TORQUE_MASK | ANGMOM_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);
+
+  v = atomKK->k_v.view<DeviceType>();
+  f = atomKK->k_f.view<DeviceType>();
+
+  if( atomKK->omega_flag )
+    omega  = atomKK->k_omega.view<DeviceType>();
+
+  if( atomKK->angmom_flag )
+    angmom = atomKK->k_angmom.view<DeviceType>();
+
+  if( atomKK->torque_flag )
+    torque = atomKK->k_torque.view<DeviceType>();
+
+
+  mask = atomKK->k_mask.view<DeviceType>();
+
+  int nlocal = atomKK->nlocal;
+  if (igroup == atomKK->firstgroup) nlocal = atomKK->nfirst;
+
+  int flag_mask = 0;
+  if( atomKK->omega_flag ) flag_mask  |= 1;
+  if( atomKK->angmom_flag ) flag_mask |= 2;
+  if( atomKK->torque_flag ) flag_mask |= 4;
+
+  copymode = 1;
+  switch( flag_mask ){
+    case 0:{
+      FixEnforce2DKokkosPostForceFunctor<DeviceType,0,0,0> functor(this);
+      Kokkos::parallel_for(nlocal,functor);
+      break;
+    }
+    case 1:{
+      FixEnforce2DKokkosPostForceFunctor<DeviceType,1,0,0> functor(this);
+      Kokkos::parallel_for(nlocal,functor);
+      break;
+    }
+    case 2:{
+      FixEnforce2DKokkosPostForceFunctor<DeviceType,0,1,0> functor(this);
+      Kokkos::parallel_for(nlocal,functor);
+      break;
+    }
+    case 3:{
+      FixEnforce2DKokkosPostForceFunctor<DeviceType,1,1,0> functor(this);
+      Kokkos::parallel_for(nlocal,functor);
+      break;
+    }
+    case 4:{
+      FixEnforce2DKokkosPostForceFunctor<DeviceType,0,0,1> functor(this);
+      Kokkos::parallel_for(nlocal,functor);
+      break;
+    }
+    case 5:{
+      FixEnforce2DKokkosPostForceFunctor<DeviceType,1,0,1> functor(this);
+      Kokkos::parallel_for(nlocal,functor);
+      break;
+    }
+    case 6:{
+      FixEnforce2DKokkosPostForceFunctor<DeviceType,0,1,1> functor(this);
+      Kokkos::parallel_for(nlocal,functor);
+      break;
+    }
+    case 7:{
+      FixEnforce2DKokkosPostForceFunctor<DeviceType,1,1,1> functor(this);
+      Kokkos::parallel_for(nlocal,functor);
+      break;
+    }
+    default:
+      error->all(FLERR, "Flag in fix_enforce2d_kokkos outside of what it should be");
+  }
+  copymode = 0;
+
+  atomKK->modified(execution_space,datamask_modify);
+
+  for (int m = 0; m < nfixlist; m++) {
+    atomKK->sync(flist[m]->execution_space,flist[m]->datamask_read);
+    flist[m]->enforce2d();
+    atomKK->modified(flist[m]->execution_space,flist[m]->datamask_modify);
+  }
+
+}
+
+
+template <class DeviceType>
+template <int omega_flag, int angmom_flag, int torque_flag>
+void FixEnforce2DKokkos<DeviceType>::post_force_item( int i ) const
+{
+  if (mask[i] & groupbit){
+    v(i,2) = 0.0;
+    f(i,2) = 0.0;
+
+    if(omega_flag){
+      omega(i,0) = 0.0;
+      omega(i,1) = 0.0;
+    }
+
+    if(angmom_flag){
+      angmom(i,0) = 0.0;
+      angmom(i,1) = 0.0;
+    }
+
+    if(torque_flag){
+      torque(i,0) = 0.0;
+      torque(i,1) = 0.0;
+    }
+  }
+}
+
+
+namespace LAMMPS_NS {
+template class FixEnforce2DKokkos<LMPDeviceType>;
+#ifdef KOKKOS_HAVE_CUDA
+template class FixEnforce2DKokkos<LMPHostType>;
+#endif
+}

src/KOKKOS/fix_enforce2d_kokkos.h

@@ -0,0 +1,84 @@
+/* -*- 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 setup(int);
+  void post_force(int);
+
+  template <int omega_flag, int angmom_flag, int torque_flag>
+  KOKKOS_INLINE_FUNCTION
+  void post_force_item(const int i) 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_v_array v;
+  typename ArrayTypes<DeviceType>::t_f_array f;
+
+  typename ArrayTypes<DeviceType>::t_v_array omega;
+  typename ArrayTypes<DeviceType>::t_v_array angmom;
+  typename ArrayTypes<DeviceType>::t_f_array torque;
+
+  typename ArrayTypes<DeviceType>::t_int_1d mask;
+};
+
+
+template <class DeviceType, int omega_flag, int angmom_flag, int torque_flag>
+struct FixEnforce2DKokkosPostForceFunctor {
+  typedef DeviceType device_type;
+  FixEnforce2DKokkos<DeviceType> c;
+
+  FixEnforce2DKokkosPostForceFunctor(FixEnforce2DKokkos<DeviceType>* c_ptr):
+    c(*c_ptr) {};
+
+  KOKKOS_INLINE_FUNCTION
+  void operator()(const int i) const {
+    c.template post_force_item <omega_flag, angmom_flag, torque_flag>(i);
+  }
+};
+
+
+}
+
+#endif
+#endif
+
+/* ERROR/WARNING messages:
+
+E: Flag in fix_enforce2d_kokkos outside of what it should be
+
+LAMMPS developer-only error.
+
+*/

src/fix_enforce2d.cpp

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

src/fix_enforce2d.h

@@ -36,7 +36,7 @@ class FixEnforce2D : public Fix {
   void post_force_respa(int, int, int);
   void min_post_force(int);
 
- private:
+ protected:
   int nfixlist;
   class Fix **flist;
 };