Add Kokkos support for fix shake and forward comm on Device

src/Depend.sh

@@ -107,6 +107,7 @@ if (test $1 = "PERI") then
 fi
 
 if (test $1 = "RIGID") then
+  depend KOKKOS
   depend USER-OMP
   depend USER-SDPD
 fi

src/KOKKOS/Install.sh

@@ -142,6 +142,8 @@ action fix_reaxc_species_kokkos.cpp fix_reaxc_species.cpp
 action fix_reaxc_species_kokkos.h fix_reaxc_species.h
 action fix_setforce_kokkos.cpp
 action fix_setforce_kokkos.h
+action fix_shake_kokkos.cpp fix_shake.cpp
+action fix_shake_kokkos.h fix_shake.h
 action fix_shardlow_kokkos.cpp fix_shardlow.cpp
 action fix_shardlow_kokkos.h fix_shardlow.h
 action fix_momentum_kokkos.cpp

src/KOKKOS/comm_kokkos.cpp

@@ -75,6 +75,11 @@ CommKokkos::CommKokkos(LAMMPS *lmp) : CommBrick(lmp)
   max_buf_pair = 0;
   k_buf_send_pair = DAT::tdual_xfloat_1d("comm:k_buf_send_pair",1);
   k_buf_recv_pair = DAT::tdual_xfloat_1d("comm:k_recv_send_pair",1);
+
+  max_buf_fix = 0;
+  k_buf_send_fix = DAT::tdual_xfloat_1d("comm:k_buf_send_fix",1);
+  k_buf_recv_fix = DAT::tdual_xfloat_1d("comm:k_recv_send_fix",1);
+
 }
 
 /* ---------------------------------------------------------------------- */
@@ -356,10 +361,79 @@ void CommKokkos::reverse_comm_device()
 
 void CommKokkos::forward_comm_fix(Fix *fix, int size)
 {
-  k_sendlist.sync<LMPHostType>();
+  if (fix->execution_space == Device && fix->forward_comm_device) {
-  CommBrick::forward_comm_fix(fix,size);
+    k_sendlist.sync<LMPDeviceType>();
+    forward_comm_fix_device<LMPDeviceType>(fix,size);
+  } else {
+    k_sendlist.sync<LMPHostType>();
+    CommBrick::forward_comm_fix(fix,size);
+  }
 }
 
+template<class DeviceType>
+void CommKokkos::forward_comm_fix_device(Fix *fix, int size)
+{
+  int iswap,n,nsize;
+  MPI_Request request;
+
+  if (size) nsize = size;
+  else nsize = fix->comm_forward;
+  KokkosBase* fixKKBase = dynamic_cast<KokkosBase*>(fix);
+
+  for (iswap = 0; iswap < nswap; iswap++) {
+    int n = MAX(max_buf_fix,nsize*sendnum[iswap]);
+    n = MAX(n,nsize*recvnum[iswap]);
+    if (n > max_buf_fix)
+      grow_buf_fix(n);
+  }
+
+  for (iswap = 0; iswap < nswap; iswap++) {
+
+    // pack buffer
+
+    n = fixKKBase->pack_forward_comm_fix_kokkos(sendnum[iswap],k_sendlist,
+                                      iswap,k_buf_send_fix,pbc_flag[iswap],pbc[iswap]);
+    DeviceType().fence();
+
+    // exchange with another proc
+    // if self, set recv buffer to send buffer
+
+    if (sendproc[iswap] != me) {
+      double* buf_send_fix;
+      double* buf_recv_fix;
+      if (lmp->kokkos->gpu_aware_flag) {
+        buf_send_fix = k_buf_send_fix.view<DeviceType>().data();
+        buf_recv_fix = k_buf_recv_fix.view<DeviceType>().data();
+      } else {
+        k_buf_send_fix.modify<DeviceType>();
+        k_buf_send_fix.sync<LMPHostType>();
+        buf_send_fix = k_buf_send_fix.h_view.data();
+        buf_recv_fix = k_buf_recv_fix.h_view.data();
+      }
+
+      if (recvnum[iswap]) {
+        MPI_Irecv(buf_recv_fix,nsize*recvnum[iswap],MPI_DOUBLE,
+                  recvproc[iswap],0,world,&request);
+      }
+      if (sendnum[iswap])
+        MPI_Send(buf_send_fix,n,MPI_DOUBLE,sendproc[iswap],0,world);
+      if (recvnum[iswap]) MPI_Wait(&request,MPI_STATUS_IGNORE);
+
+      if (!lmp->kokkos->gpu_aware_flag) {
+        k_buf_recv_fix.modify<LMPHostType>();
+        k_buf_recv_fix.sync<DeviceType>();
+      }
+    } else k_buf_recv_fix = k_buf_send_fix;
+
+    // unpack buffer
+
+    fixKKBase->unpack_forward_comm_fix_kokkos(recvnum[iswap],firstrecv[iswap],k_buf_recv_fix);
+    DeviceType().fence();
+  }
+}
+
+/* ---------------------------------------------------------------------- */
+
 void CommKokkos::reverse_comm_fix(Fix *fix, int size)
 {
   k_sendlist.sync<LMPHostType>();
@@ -456,6 +530,12 @@ void CommKokkos::grow_buf_pair(int n) {
   k_buf_recv_pair.resize(max_buf_pair);
 }
 
+void CommKokkos::grow_buf_fix(int n) {
+  max_buf_fix = n * BUFFACTOR;
+  k_buf_send_fix.resize(max_buf_fix);
+  k_buf_recv_fix.resize(max_buf_fix);
+}
+
 void CommKokkos::reverse_comm_pair(Pair *pair)
 {
   k_sendlist.sync<LMPHostType>();
@@ -491,8 +571,8 @@ void CommKokkos::exchange()
     if(!exchange_comm_classic) {
       static int print = 1;
       if(print && comm->me==0) {
-        error->warning(FLERR,"Fixes cannot yet send data in Kokkos communication, "
+        error->warning(FLERR,"Fixes cannot yet send exchange data in Kokkos communication, "
-                      "switching to classic communication");
+                      "switching to classic exchange/border communication");
       }
       print = 0;
       exchange_comm_classic = true;
@@ -742,7 +822,7 @@ void CommKokkos::borders()
          (ghost_velocity && ((AtomVecKokkos*)atom->avec)->no_border_vel_flag)) {
       if (print && comm->me==0) {
         error->warning(FLERR,"Required border comm not yet implemented in Kokkos communication, "
-                      "switching to classic communication");
+                      "switching to classic exchange/border communication");
       }
       print = 0;
       exchange_comm_classic = true;

src/KOKKOS/comm_kokkos.h

@@ -51,6 +51,7 @@ class CommKokkos : public CommBrick {
   template<class DeviceType> void forward_comm_device(int dummy);
   template<class DeviceType> void reverse_comm_device();
   template<class DeviceType> void forward_comm_pair_device(Pair *pair);
+  template<class DeviceType> void forward_comm_fix_device(Fix *fix, int size=0);
   template<class DeviceType> void exchange_device();
   template<class DeviceType> void borders_device();
 
@@ -73,10 +74,11 @@ class CommKokkos : public CommBrick {
   DAT::tdual_int_1d k_sendnum_scan;
   int totalsend;
 
-  int max_buf_pair;
+  int max_buf_pair,max_buf_fix;
-  DAT::tdual_xfloat_1d k_buf_send_pair;
+  DAT::tdual_xfloat_1d k_buf_send_pair, k_buf_send_fix;
-  DAT::tdual_xfloat_1d k_buf_recv_pair;
+  DAT::tdual_xfloat_1d k_buf_recv_pair, k_buf_recv_fix;
   void grow_buf_pair(int);
+  void grow_buf_fix(int);
 
   void grow_send(int, int);
   void grow_recv(int);

src/KOKKOS/fix_qeq_reax_kokkos.cpp

@@ -47,6 +47,7 @@ FixQEqReaxKokkos(LAMMPS *lmp, int narg, char **arg) :
   FixQEqReax(lmp, narg, arg)
 {
   kokkosable = 1;
+  forward_comm_device = 1;
   atomKK = (AtomKokkos *) atom;
   execution_space = ExecutionSpaceFromDevice<DeviceType>::space;
 
@@ -262,19 +263,15 @@ void FixQEqReaxKokkos<DeviceType>::pre_force(int /*vflag*/)
 
   // comm->forward_comm_fix(this); //Dist_vector( s );
   pack_flag = 2;
-  k_s.template modify<DeviceType>();
-  k_s.template sync<LMPHostType>();
-  comm->forward_comm_fix(this);
-  k_s.template modify<LMPHostType>();
   k_s.template sync<DeviceType>();
+  comm->forward_comm_fix(this);
+  k_s.template modify<DeviceType>();
 
   // comm->forward_comm_fix(this); //Dist_vector( t );
   pack_flag = 3;
-  k_t.template modify<DeviceType>();
-  k_t.template sync<LMPHostType>();
-  comm->forward_comm_fix(this);
-  k_t.template modify<LMPHostType>();
   k_t.template sync<DeviceType>();
+  comm->forward_comm_fix(this);
+  k_t.template modify<DeviceType>();
 
   need_dup = lmp->kokkos->need_dup<DeviceType>();
 
@@ -784,11 +781,9 @@ void FixQEqReaxKokkos<DeviceType>::cg_solve1()
 
     // comm->forward_comm_fix(this); //Dist_vector( d );
     pack_flag = 1;
-    k_d.template modify<DeviceType>();
-    k_d.template sync<LMPHostType>();
-    comm->forward_comm_fix(this);
-    k_d.template modify<LMPHostType>();
     k_d.template sync<DeviceType>();
+    comm->forward_comm_fix(this);
+    k_d.template modify<DeviceType>();
 
     // sparse_matvec( &H, d, q );
     FixQEqReaxKokkosSparse22Functor<DeviceType> sparse22_functor(this);
@@ -922,11 +917,9 @@ void FixQEqReaxKokkos<DeviceType>::cg_solve2()
 
     // comm->forward_comm_fix(this); //Dist_vector( d );
     pack_flag = 1;
-    k_d.template modify<DeviceType>();
-    k_d.template sync<LMPHostType>();
-    comm->forward_comm_fix(this);
-    k_d.template modify<LMPHostType>();
     k_d.template sync<DeviceType>();
+    comm->forward_comm_fix(this);
+    k_d.template modify<DeviceType>();
 
     // sparse_matvec( &H, d, q );
     FixQEqReaxKokkosSparse22Functor<DeviceType> sparse22_functor(this);
@@ -1027,11 +1020,9 @@ void FixQEqReaxKokkos<DeviceType>::calculate_q()
 
   pack_flag = 4;
   //comm->forward_comm_fix( this ); //Dist_vector( atom->q );
-  atomKK->k_q.modify<DeviceType>();
-  atomKK->k_q.sync<LMPHostType>();
-  comm->forward_comm_fix(this);
-  atomKK->k_q.modify<LMPHostType>();
   atomKK->k_q.sync<DeviceType>();
+  comm->forward_comm_fix(this);
+  atomKK->k_q.modify<DeviceType>();
 
 }
 
@@ -1361,6 +1352,59 @@ void FixQEqReaxKokkos<DeviceType>::calculate_q_item(int ii) const
 
 /* ---------------------------------------------------------------------- */
 
+template<class DeviceType>
+int FixQEqReaxKokkos<DeviceType>::pack_forward_comm_fix_kokkos(int n, DAT::tdual_int_2d k_sendlist,
+                                                        int iswap_in, DAT::tdual_xfloat_1d &k_buf,
+                                                        int /*pbc_flag*/, int * /*pbc*/)
+{
+  d_sendlist = k_sendlist.view<DeviceType>();
+  iswap = iswap_in;
+  d_buf = k_buf.view<DeviceType>();
+  Kokkos::parallel_for(Kokkos::RangePolicy<DeviceType, TagFixQEqReaxPackForwardComm>(0,n),*this);
+  return n;
+}
+
+template<class DeviceType>
+KOKKOS_INLINE_FUNCTION
+void FixQEqReaxKokkos<DeviceType>::operator()(TagFixQEqReaxPackForwardComm, const int &i) const {
+  int j = d_sendlist(iswap, i);
+
+  if (pack_flag == 1)
+    d_buf[i] = d_d[j];
+  else if( pack_flag == 2 )
+    d_buf[i] = d_s[j];
+  else if( pack_flag == 3 )
+    d_buf[i] = d_t[j];
+  else if( pack_flag == 4 )
+    d_buf[i] = q[j];
+}
+
+/* ---------------------------------------------------------------------- */
+
+template<class DeviceType>
+void FixQEqReaxKokkos<DeviceType>::unpack_forward_comm_fix_kokkos(int n, int first_in, DAT::tdual_xfloat_1d &buf)
+{
+  first = first_in;
+  d_buf = buf.view<DeviceType>();
+  Kokkos::parallel_for(Kokkos::RangePolicy<DeviceType, TagFixQEqReaxUnpackForwardComm>(0,n),*this);
+}
+
+template<class DeviceType>
+KOKKOS_INLINE_FUNCTION
+void FixQEqReaxKokkos<DeviceType>::operator()(TagFixQEqReaxUnpackForwardComm, const int &i) const {
+  if (pack_flag == 1)
+    d_d[i + first] = d_buf[i];
+  else if( pack_flag == 2)
+    d_s[i + first] = d_buf[i];
+  else if( pack_flag == 3)
+    d_t[i + first] = d_buf[i];
+  else if( pack_flag == 4)
+    q[i + first] = d_buf[i];
+
+}
+
+/* ---------------------------------------------------------------------- */
+
 template<class DeviceType>
 int FixQEqReaxKokkos<DeviceType>::pack_forward_comm(int n, int *list, double *buf,
                                                     int /*pbc_flag*/, int * /*pbc*/)

src/KOKKOS/fix_qeq_reax_kokkos.h

@@ -26,6 +26,7 @@ FixStyle(qeq/reax/kk/host,FixQEqReaxKokkos<LMPHostType>)
 #include "kokkos_type.h"
 #include "neigh_list.h"
 #include "neigh_list_kokkos.h"
+#include "kokkos_base.h"
 
 namespace LAMMPS_NS {
 
@@ -33,9 +34,11 @@ struct TagSparseMatvec1 {};
 struct TagSparseMatvec2 {};
 struct TagSparseMatvec3 {};
 struct TagZeroQGhosts{};
+struct TagFixQEqReaxPackForwardComm {};
+struct TagFixQEqReaxUnpackForwardComm {};
 
 template<class DeviceType>
-class FixQEqReaxKokkos : public FixQEqReax {
+class FixQEqReaxKokkos : public FixQEqReax, public KokkosBase {
  public:
   typedef DeviceType device_type;
   typedef ArrayTypes<DeviceType> AT;
@@ -136,6 +139,12 @@ class FixQEqReaxKokkos : public FixQEqReax {
   KOKKOS_INLINE_FUNCTION
   double calculate_H_k(const F_FLOAT &r, const F_FLOAT &shld) const;
 
+  KOKKOS_INLINE_FUNCTION
+  void operator()(TagFixQEqReaxPackForwardComm, const int&) const;
+
+  KOKKOS_INLINE_FUNCTION
+  void operator()(TagFixQEqReaxUnpackForwardComm, const int&) const;
+
   struct params_qeq{
     KOKKOS_INLINE_FUNCTION
     params_qeq(){chi=0;eta=0;gamma=0;};
@@ -144,6 +153,9 @@ class FixQEqReaxKokkos : public FixQEqReax {
     F_FLOAT chi, eta, gamma;
   };
 
+  int pack_forward_comm_fix_kokkos(int, DAT::tdual_int_2d, int, DAT::tdual_xfloat_1d&,
+                       int, int *);
+  void unpack_forward_comm_fix_kokkos(int, int, DAT::tdual_xfloat_1d&);
   virtual int pack_forward_comm(int, int *, double *, int, int *);
   virtual void unpack_forward_comm(int, int, double *);
   int pack_reverse_comm(int, int, double *);
@@ -203,6 +215,11 @@ class FixQEqReaxKokkos : public FixQEqReax {
   Kokkos::Experimental::ScatterView<F_FLOAT*, typename AT::t_ffloat_1d::array_layout, typename KKDevice<DeviceType>::value, Kokkos::Experimental::ScatterSum, Kokkos::Experimental::ScatterDuplicated> dup_o;
   Kokkos::Experimental::ScatterView<F_FLOAT*, typename AT::t_ffloat_1d::array_layout, typename KKDevice<DeviceType>::value, Kokkos::Experimental::ScatterSum, Kokkos::Experimental::ScatterNonDuplicated> ndup_o;
 
+  int iswap;
+  int first;
+  typename AT::t_int_2d d_sendlist;
+  typename AT::t_xfloat_1d_um d_buf;
+
   void init_shielding_k();
   void init_hist();
   void allocate_matrix();
@@ -216,11 +233,6 @@ class FixQEqReaxKokkos : public FixQEqReax {
   int count, isuccess;
   double alpha, beta, delta, cutsq;
 
-  int iswap;
-  int first;
-  typename AT::t_int_2d d_sendlist;
-  typename AT::t_xfloat_1d_um v_buf;
-
   void grow_arrays(int);
   void copy_arrays(int, int, int);
   int pack_exchange(int, double *);

src/KOKKOS/fix_shake_kokkos.h

@@ -0,0 +1,210 @@
+/* -*- 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(shake/kk,FixShakeKokkos<LMPDeviceType>)
+FixStyle(shake/kk/device,FixShakeKokkos<LMPDeviceType>)
+FixStyle(shake/kk/host,FixShakeKokkos<LMPHostType>)
+
+#else
+
+#ifndef LMP_FIX_SHAKE_KOKKOS_H
+#define LMP_FIX_SHAKE_KOKKOS_H
+
+#include "fix_shake.h"
+#include "kokkos_type.h"
+#include "kokkos_base.h"
+
+namespace LAMMPS_NS {
+
+template<int NEIGHFLAG, int EVFLAG>
+struct TagFixShakePostForce{};
+
+template<int PBC_FLAG>
+struct TagFixShakePackForwardComm{};
+
+struct TagFixShakeUnpackForwardComm{};
+
+template<class DeviceType>
+class FixShakeKokkos : public FixShake, public KokkosBase {
+
+ //friend class FixEHEX;
+
+ public:
+  typedef DeviceType device_type;
+  typedef EV_FLOAT value_type;
+  typedef ArrayTypes<DeviceType> AT;
+
+  FixShakeKokkos(class LAMMPS *, int, char **);
+  virtual ~FixShakeKokkos();
+  void init();
+  void pre_neighbor();
+  void post_force(int);
+
+  void grow_arrays(int);
+  void copy_arrays(int, int, int);
+  void set_arrays(int);
+  void update_arrays(int, int);
+
+  int pack_exchange(int, double *);
+  int unpack_exchange(int, double *);
+  int pack_forward_comm_fix_kokkos(int, DAT::tdual_int_2d, int, DAT::tdual_xfloat_1d&,
+                       int, int *);
+  void unpack_forward_comm_fix_kokkos(int, int, DAT::tdual_xfloat_1d&);
+  int pack_forward_comm(int, int *, double *, int, int *);
+  void unpack_forward_comm(int, int, double *);
+
+  void shake_end_of_step(int vflag);
+  void correct_coordinates(int vflag);
+
+  int dof(int);
+
+  void unconstrained_update();
+
+  template<int NEIGHFLAG, int EVFLAG>
+  KOKKOS_INLINE_FUNCTION
+  void operator()(TagFixShakePostForce<NEIGHFLAG,EVFLAG>, const int&, EV_FLOAT&) const;
+
+  template<int NEIGHFLAG, int EVFLAG>
+  KOKKOS_INLINE_FUNCTION
+  void operator()(TagFixShakePostForce<NEIGHFLAG,EVFLAG>, const int&) const;
+
+  template<int PBC_FLAG>
+  KOKKOS_INLINE_FUNCTION
+  void operator()(TagFixShakePackForwardComm<PBC_FLAG>, const int&) const;
+
+  KOKKOS_INLINE_FUNCTION
+  void operator()(TagFixShakeUnpackForwardComm, const int&) const;
+
+ protected:
+
+  typename AT::t_x_array d_x;
+  typename AT::t_v_array d_v;
+  typename AT::t_f_array d_f;
+  typename AT::t_float_1d d_rmass;
+  typename AT::t_float_1d d_mass;
+  typename AT::t_tagint_1d_randomread d_tag;
+  typename AT::t_int_1d d_type;
+  typename AT::t_int_1d d_mask;
+
+  DAT::tdual_efloat_1d k_eatom;
+  typename AT::t_efloat_1d d_eatom;
+
+  DAT::tdual_virial_array k_vatom;
+  typename AT::t_virial_array d_vatom;
+
+  DAT::tdual_float_1d k_bond_distance; // constraint distances
+  typename AT::t_float_1d d_bond_distance;
+  DAT::tdual_float_1d k_angle_distance;
+  typename AT::t_float_1d d_angle_distance;
+
+                                         // atom-based arrays
+  DAT::tdual_int_1d k_shake_flag;
+  typename AT::t_int_1d d_shake_flag; // 0 if atom not in SHAKE cluster
+                                         // 1 = size 3 angle cluster
+                                         // 2,3,4 = size of bond-only cluster
+  DAT::tdual_tagint_2d k_shake_atom;
+  typename AT::t_tagint_2d d_shake_atom; // global IDs of atoms in cluster
+                                         // central atom is 1st
+                                         // lowest global ID is 1st for size 2
+  DAT::tdual_int_2d k_shake_type;
+  typename AT::t_int_2d d_shake_type; // bondtype of each bond in cluster
+                                         // for angle cluster, 3rd value
+                                         //   is angletype
+  DAT::tdual_x_array k_xshake;
+  typename AT::t_x_array d_xshake; // unconstrained atom coords
+
+  DAT::tdual_int_1d k_list;
+  typename AT::t_int_1d d_list; // list of clusters to SHAKE
+
+  DAT::tdual_int_scalar k_error_flag;
+  DAT::tdual_int_scalar k_nlist;
+
+
+  template<int NEIGHFLAG, int EVFLAG>
+  KOKKOS_INLINE_FUNCTION
+  void shake(int, EV_FLOAT&) const;
+
+  template<int NEIGHFLAG, int EVFLAG>
+  KOKKOS_INLINE_FUNCTION
+  void shake3(int, EV_FLOAT&) const;
+
+  template<int NEIGHFLAG, int EVFLAG>
+  KOKKOS_INLINE_FUNCTION
+  void shake4(int, EV_FLOAT&) const;
+
+  template<int NEIGHFLAG, int EVFLAG>
+  KOKKOS_INLINE_FUNCTION
+  void shake3angle(int, EV_FLOAT&) const;
+
+  typedef typename KKDevice<DeviceType>::value KKDeviceType;
+  Kokkos::Experimental::ScatterView<F_FLOAT*[3], typename DAT::t_f_array::array_layout,typename KKDevice<DeviceType>::value,typename Kokkos::Experimental::ScatterSum,Kokkos::Experimental::ScatterDuplicated> dup_f;
+  Kokkos::Experimental::ScatterView<E_FLOAT*, typename DAT::t_efloat_1d::array_layout,typename KKDevice<DeviceType>::value,typename Kokkos::Experimental::ScatterSum,Kokkos::Experimental::ScatterDuplicated> dup_eatom;
+  Kokkos::Experimental::ScatterView<F_FLOAT*[6], typename DAT::t_virial_array::array_layout,typename KKDevice<DeviceType>::value,typename Kokkos::Experimental::ScatterSum,Kokkos::Experimental::ScatterDuplicated> dup_vatom;
+
+  Kokkos::Experimental::ScatterView<F_FLOAT*[3], typename DAT::t_f_array::array_layout,typename KKDevice<DeviceType>::value,typename Kokkos::Experimental::ScatterSum,Kokkos::Experimental::ScatterNonDuplicated> ndup_f;
+  Kokkos::Experimental::ScatterView<E_FLOAT*, typename DAT::t_efloat_1d::array_layout,typename KKDevice<DeviceType>::value,typename Kokkos::Experimental::ScatterSum,Kokkos::Experimental::ScatterNonDuplicated> ndup_eatom;
+  Kokkos::Experimental::ScatterView<F_FLOAT*[6], typename DAT::t_virial_array::array_layout,typename KKDevice<DeviceType>::value,typename Kokkos::Experimental::ScatterSum,Kokkos::Experimental::ScatterNonDuplicated> ndup_vatom;
+
+  int neighflag,need_dup;
+
+  typename AT::t_int_1d d_scalars;
+  HAT::t_int_1d h_scalars;
+  typename AT::t_int_scalar d_error_flag;
+  typename AT::t_int_scalar d_nlist;
+  HAT::t_int_scalar h_error_flag;
+  HAT::t_int_scalar h_nlist;
+
+  template<int NEIGHFLAG>
+  KOKKOS_INLINE_FUNCTION
+  void v_tally(EV_FLOAT&, int, int *, double, double *) const;
+
+  DAT::tdual_int_1d k_map_array;
+  typename AT::t_int_1d_randomread map_array;
+
+  int iswap;
+  int first;
+  typename AT::t_int_2d d_sendlist;
+  typename AT::t_xfloat_1d_um d_buf;
+  X_FLOAT dx,dy,dz;
+
+  int *shake_flag_tmp;
+  tagint **shake_atom_tmp;
+  int **shake_type_tmp;
+
+  // copied from Domain
+
+  KOKKOS_INLINE_FUNCTION
+  void minimum_image(double *) const;
+
+  KOKKOS_INLINE_FUNCTION
+  void minimum_image_once(double *) const;
+
+  void update_domain_variables();
+
+  int triclinic;
+  int xperiodic,yperiodic,zperiodic;
+  X_FLOAT xprd_half,yprd_half,zprd_half;
+  X_FLOAT xprd,yprd,zprd;
+  X_FLOAT xy,xz,yz;
+};
+
+}
+
+#endif
+#endif
+
+/* ERROR/WARNING messages:
+
+*/

src/KOKKOS/kokkos_base.h

@@ -28,6 +28,13 @@ class KokkosBase {
                                        int, int *) {return 0;};
   virtual void unpack_forward_comm_kokkos(int, int, DAT::tdual_xfloat_1d &) {}
 
+  // Fix
+  virtual int pack_forward_comm_fix_kokkos(int, DAT::tdual_int_2d,
+                                           int, DAT::tdual_xfloat_1d &,
+                                           int, int *) {return 0;};
+  virtual void unpack_forward_comm_fix_kokkos(int, int, DAT::tdual_xfloat_1d &) {}
+
+
   // Region
   virtual void match_all_kokkos(int, DAT::tdual_int_1d) {}
 };

src/RIGID/fix_shake.cpp

@@ -240,6 +240,8 @@ FixShake::FixShake(LAMMPS *lmp, int narg, char **arg) :
 
 FixShake::~FixShake()
 {
+  if (copymode) return;
+
   // unregister callbacks to this fix from Atom class
 
   atom->delete_callback(id,Atom::GROW);

src/RIGID/fix_shake.h

@@ -34,14 +34,14 @@ class FixShake : public Fix {
   virtual int setmask();
   virtual void init();
   void setup(int);
-  void pre_neighbor();
+  virtual void pre_neighbor();
   virtual void post_force(int);
   virtual void post_force_respa(int, int, int);
 
   virtual double memory_usage();
   virtual void grow_arrays(int);
   virtual void copy_arrays(int, int, int);
-  void set_arrays(int);
+  virtual void set_arrays(int);
   virtual void update_arrays(int, int);
   void set_molecule(int, tagint, int, double *, double *, double *);
 
@@ -54,7 +54,7 @@ class FixShake : public Fix {
   virtual void correct_coordinates(int vflag);
   virtual void correct_velocities();
 
-  int dof(int);
+  virtual int dof(int);
   virtual void reset_dt();
   void *extract(const char *, int &);
 
@@ -126,7 +126,7 @@ class FixShake : public Fix {
   void shake_info(int *, tagint **, int **);
 
   int masscheck(double);
-  void unconstrained_update();
+  virtual void unconstrained_update();
   void unconstrained_update_respa(int);
   void shake(int);
   void shake3(int);

src/fix.cpp

@@ -111,6 +111,7 @@ Fix::Fix(LAMMPS *lmp, int /*narg*/, char **arg) :
   datamask_modify = ALL_MASK;
 
   kokkosable = 0;
+  forward_comm_device = 0;
   copymode = 0;
 }
 

src/fix.h

@@ -107,6 +107,7 @@ class Fix : protected Pointers {
   // KOKKOS host/device flag and data masks
 
   int kokkosable;                // 1 if Kokkos fix
+  int forward_comm_device;                // 1 if forward comm on Device
   ExecutionSpace execution_space;
   unsigned int datamask_read,datamask_modify;