Porting fix nve/sphere to Kokkos.

src/KOKKOS/Install.sh

@@ -109,6 +109,8 @@ action fix_npt_kokkos.cpp
 action fix_npt_kokkos.h
 action fix_nve_kokkos.cpp
 action fix_nve_kokkos.h
+action fix_nve_sphere_kokkos.cpp
+action fix_nve_sphere_kokkos.h
 action fix_nvt_kokkos.cpp
 action fix_nvt_kokkos.h
 action fix_property_atom_kokkos.cpp

src/KOKKOS/fix_nve_sphere_kokkos.cpp

+/* ----------------------------------------------------------------------
+   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.
+------------------------------------------------------------------------- */
+
+#include "fix_nve_sphere_kokkos.h"
+#include "atom_masks.h"
+#include "atom_kokkos.h"
+#include "error.h"
+
+using namespace LAMMPS_NS;
+
+enum{NONE,DIPOLE};
+
+/* ---------------------------------------------------------------------- */
+
+template<class DeviceType>
+FixNVESphereKokkos<DeviceType>::FixNVESphereKokkos(LAMMPS *lmp, int narg, char **arg) :
+  FixNVESphere(lmp, narg, arg)
+{
+  kokkosable = 1;
+  atomKK = (AtomKokkos *)atom;
+  execution_space = ExecutionSpaceFromDevice<DeviceType>::space;
+
+  datamask_read = F_MASK | TORQUE_MASK | RMASS_MASK | RADIUS_MASK | MASK_MASK;
+  datamask_modify = X_MASK | V_MASK | OMEGA_MASK;
+}
+
+/* ---------------------------------------------------------------------- */
+
+template<class DeviceType>
+void FixNVESphereKokkos<DeviceType>::cleanup_copy()
+{
+  id = style = NULL;
+  vatom = NULL;
+}
+
+/* ---------------------------------------------------------------------- */
+
+template<class DeviceType>
+void FixNVESphereKokkos<DeviceType>::init()
+{
+  FixNVESphere::init();
+
+  if (extra == DIPOLE) {
+    error->all(FLERR,"Fix nve/sphere/kk doesn't yet support dipole");
+  }
+}
+
+/* ---------------------------------------------------------------------- */
+
+template<class DeviceType>
+void FixNVESphereKokkos<DeviceType>::initial_integrate(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>();
+  omega = atomKK->k_omega.view<DeviceType>();
+  f = atomKK->k_f.view<DeviceType>();
+  torque = atomKK->k_torque.view<DeviceType>();
+  mask = atomKK->k_mask.view<DeviceType>();
+  rmass = atomKK->k_rmass.view<DeviceType>();
+  radius = atomKK->k_radius.view<DeviceType>();
+
+  int nlocal = atom->nlocal;
+  if (igroup == atom->firstgroup) nlocal = atom->nfirst;
+
+  FixNVESphereKokkosInitialIntegrateFunctor<DeviceType> f(this); 
+  Kokkos::parallel_for(nlocal,f);
+}
+
+/* ---------------------------------------------------------------------- */
+
+template <class DeviceType>
+KOKKOS_INLINE_FUNCTION
+void FixNVESphereKokkos<DeviceType>::initial_integrate_item(const int i) const
+{
+  const double dtfrotate = dtf / inertia;
+
+  if (mask(i) & groupbit) {
+    const double dtfm = dtf / rmass(i);
+    v(i,0) += dtfm * f(i,0);
+    v(i,1) += dtfm * f(i,1);
+    v(i,2) += dtfm * f(i,2);
+    x(i,0) += dtv * v(i,0);
+    x(i,1) += dtv * v(i,1);
+    x(i,2) += dtv * v(i,2);
+      
+    const double dtirotate = dtfrotate / (radius(i)*radius(i)*rmass(i));
+    omega(i,0) += dtirotate * torque(i,0);
+    omega(i,1) += dtirotate * torque(i,1);
+    omega(i,2) += dtirotate * torque(i,2);
+  }
+}
+
+/* ---------------------------------------------------------------------- */
+
+template<class DeviceType>
+void FixNVESphereKokkos<DeviceType>::final_integrate()
+{
+  atomKK->sync(execution_space,datamask_read);
+  atomKK->modified(execution_space,datamask_modify);
+
+  v = atomKK->k_v.view<DeviceType>();
+  omega = atomKK->k_omega.view<DeviceType>();
+  f = atomKK->k_f.view<DeviceType>();
+  torque = atomKK->k_torque.view<DeviceType>();
+  mask = atomKK->k_mask.view<DeviceType>();
+  rmass = atomKK->k_rmass.view<DeviceType>();
+  radius = atomKK->k_radius.view<DeviceType>();
+
+  int nlocal = atom->nlocal;
+  if (igroup == atom->firstgroup) nlocal = atom->nfirst;
+
+  FixNVESphereKokkosFinalIntegrateFunctor<DeviceType> f(this);
+  Kokkos::parallel_for(nlocal,f);
+}
+
+/* ---------------------------------------------------------------------- */
+
+template <class DeviceType>
+KOKKOS_INLINE_FUNCTION
+void FixNVESphereKokkos<DeviceType>::final_integrate_item(const int i) const
+{
+  const double dtfrotate = dtf / inertia;
+
+  if (mask(i) & groupbit) {
+    const double dtfm = dtf / rmass(i);
+    v(i,0) += dtfm * f(i,0);
+    v(i,1) += dtfm * f(i,1);
+    v(i,2) += dtfm * f(i,2);
+      
+    const double dtirotate = dtfrotate / (radius(i)*radius(i)*rmass(i));
+    omega(i,0) += dtirotate * torque(i,0);
+    omega(i,1) += dtirotate * torque(i,1);
+    omega(i,2) += dtirotate * torque(i,2);
+  }
+}
+
+namespace LAMMPS_NS {
+template class FixNVESphereKokkos<LMPDeviceType>;
+#ifdef KOKKOS_HAVE_CUDA
+template class FixNVESphereKokkos<LMPHostType>;
+#endif
+}

src/KOKKOS/fix_nve_sphere_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(nve/sphere/kk,FixNVESphereKokkos<LMPDeviceType>)
+FixStyle(nve/sphere/kk/device,FixNVESphereKokkos<LMPDeviceType>)
+FixStyle(nve/sphere/kk/host,FixNVESphereKokkos<LMPHostType>)
+
+#else
+
+#ifndef LMP_FIX_NVE_SPHERE_KOKKOS_H
+#define LMP_FIX_NVE_SPHERE_KOKKOS_H
+
+#include "fix_nve_sphere.h"
+#include "kokkos_type.h"
+
+namespace LAMMPS_NS {
+  
+template<class DeviceType>
+class FixNVESphereKokkos : public FixNVESphere {
+  public:
+    FixNVESphereKokkos(class LAMMPS *, int, char **);
+    virtual ~FixNVESphereKokkos() {}
+    void cleanup_copy();
+    void init();
+    void initial_integrate(int);
+    void final_integrate();
+  
+    KOKKOS_INLINE_FUNCTION
+    void initial_integrate_item(const int i) const;
+    KOKKOS_INLINE_FUNCTION
+    void final_integrate_item(const int i) const;
+
+  private:
+    typename ArrayTypes<DeviceType>::t_x_array x;
+    typename ArrayTypes<DeviceType>::t_v_array v;
+    typename ArrayTypes<DeviceType>::t_v_array omega;
+    typename ArrayTypes<DeviceType>::t_f_array f;
+    typename ArrayTypes<DeviceType>::t_f_array torque;
+    typename ArrayTypes<DeviceType>::t_float_1d rmass;
+    typename ArrayTypes<DeviceType>::t_float_1d radius;
+    typename ArrayTypes<DeviceType>::t_int_1d mask;
+};
+
+template <class DeviceType>
+struct FixNVESphereKokkosInitialIntegrateFunctor {
+  FixNVESphereKokkos<DeviceType> c;
+  FixNVESphereKokkosInitialIntegrateFunctor(FixNVESphereKokkos<DeviceType> *c_ptr): c(*c_ptr) { c.cleanup_copy(); }
+  KOKKOS_INLINE_FUNCTION
+  void operator()(const int i) const {
+    c.initial_integrate_item(i);
+  }
+};
+
+template <class DeviceType>
+struct FixNVESphereKokkosFinalIntegrateFunctor {
+  FixNVESphereKokkos<DeviceType> c;
+  FixNVESphereKokkosFinalIntegrateFunctor(FixNVESphereKokkos<DeviceType> *c_ptr): c(*c_ptr) { c.cleanup_copy(); }
+  KOKKOS_INLINE_FUNCTION
+  void operator()(const int i) const {
+    c.final_integrate_item(i);
+  }
+};
+
+} // namespace LAMMPS_NS
+
+#endif // LMP_FIX_NVE_SPHERE_KOKKOS_H
+#endif // FIX_CLASS