Merge pull request lammps#4399 from stanmoore1/kk_bug_fixes

Fix issues in KOKKOS package

cmake/Modules/Packages/KOKKOS.cmake

@@ -117,7 +117,6 @@ set(KOKKOS_PKG_SOURCES ${KOKKOS_PKG_SOURCES_DIR}/kokkos.cpp
                        ${KOKKOS_PKG_SOURCES_DIR}/atom_vec_kokkos.cpp
                        ${KOKKOS_PKG_SOURCES_DIR}/comm_kokkos.cpp
                        ${KOKKOS_PKG_SOURCES_DIR}/comm_tiled_kokkos.cpp
-                       ${KOKKOS_PKG_SOURCES_DIR}/group_kokkos.cpp
                        ${KOKKOS_PKG_SOURCES_DIR}/min_kokkos.cpp
                        ${KOKKOS_PKG_SOURCES_DIR}/min_linesearch_kokkos.cpp
                        ${KOKKOS_PKG_SOURCES_DIR}/neighbor_kokkos.cpp

src/KOKKOS/Install.sh

@@ -230,7 +230,6 @@ action fix_wall_region_kokkos.cpp
 action fix_wall_region_kokkos.h
 action grid3d_kokkos.cpp fft3d.h
 action grid3d_kokkos.h fft3d.h
-action group_kokkos.cpp
 action group_kokkos.h
 action improper_class2_kokkos.cpp improper_class2.cpp
 action improper_class2_kokkos.h improper_class2.h

src/KOKKOS/angle_harmonic_kokkos.cpp

@@ -72,14 +72,14 @@ void AngleHarmonicKokkos<DeviceType>::compute(int eflag_in, int vflag_in)
   // reallocate per-atom arrays if necessary
 
   if (eflag_atom) {
-    if(k_eatom.extent(0) < maxeatom) {
+    if ((int)k_eatom.extent(0) < maxeatom) {
     memoryKK->destroy_kokkos(k_eatom,eatom);
     memoryKK->create_kokkos(k_eatom,eatom,maxeatom,"angle:eatom");
     d_eatom = k_eatom.template view<DeviceType>();
     } else Kokkos::deep_copy(d_eatom,0.0);
   }
   if (vflag_atom) {
-    if(k_vatom.extent(0) < maxvatom) {
+    if ((int)k_vatom.extent(0) < maxvatom) {
     memoryKK->destroy_kokkos(k_vatom,vatom);
     memoryKK->create_kokkos(k_vatom,vatom,maxvatom,"angle:vatom");
     d_vatom = k_vatom.template view<DeviceType>();

src/KOKKOS/angle_hybrid_kokkos.cpp

@@ -76,7 +76,7 @@ void AngleHybridKokkos::compute(int eflag, int vflag)
 
     Kokkos::parallel_for(nanglelist_orig,LAMMPS_LAMBDA(int i) {
       const int m = d_map[d_anglelist_orig(i,3)];
-      if (m >= 0) Kokkos::atomic_increment(&d_nanglelist[m]);
+      if (m >= 0) Kokkos::atomic_inc(&d_nanglelist[m]);
     });
 
     k_nanglelist.modify_device();
@@ -87,7 +87,7 @@ void AngleHybridKokkos::compute(int eflag, int vflag)
       if (h_nanglelist[m] > maxangle_all)
         maxangle_all = h_nanglelist[m] + EXTRA;
 
-    if (k_anglelist.d_view.extent(1) < maxangle_all)
+    if ((int)k_anglelist.d_view.extent(1) < maxangle_all)
       MemKK::realloc_kokkos(k_anglelist, "angle_hybrid:anglelist", nstyles, maxangle_all, 4);
     auto d_anglelist = k_anglelist.d_view;
 

src/KOKKOS/bond_harmonic_kokkos.cpp

@@ -67,14 +67,14 @@ void BondHarmonicKokkos<DeviceType>::compute(int eflag_in, int vflag_in)
   // reallocate per-atom arrays if necessary
 
   if (eflag_atom) {
-    if (k_eatom.extent(0) < maxeatom) {
+    if ((int)k_eatom.extent(0) < maxeatom) {
       memoryKK->destroy_kokkos(k_eatom,eatom);
       memoryKK->create_kokkos(k_eatom,eatom,maxeatom,"improper:eatom");
       d_eatom = k_eatom.template view<KKDeviceType>();
     } else Kokkos::deep_copy(d_eatom,0.0);
   }
   if (vflag_atom) {
-    if (k_vatom.extent(0) < maxvatom) {
+    if ((int)k_vatom.extent(0) < maxvatom) {
       memoryKK->destroy_kokkos(k_vatom,vatom);
       memoryKK->create_kokkos(k_vatom,vatom,maxvatom,"improper:vatom");
       d_vatom = k_vatom.template view<KKDeviceType>();

src/KOKKOS/bond_hybrid_kokkos.cpp

@@ -76,7 +76,7 @@ void BondHybridKokkos::compute(int eflag, int vflag)
 
     Kokkos::parallel_for(nbondlist_orig,LAMMPS_LAMBDA(int i) {
       const int m = d_map[d_bondlist_orig(i,2)];
-      if (m >= 0) Kokkos::atomic_increment(&d_nbondlist[m]);
+      if (m >= 0) Kokkos::atomic_inc(&d_nbondlist[m]);
     });
 
     k_nbondlist.modify_device();
@@ -87,7 +87,7 @@ void BondHybridKokkos::compute(int eflag, int vflag)
       if (h_nbondlist[m] > maxbond_all)
         maxbond_all = h_nbondlist[m] + EXTRA;
 
-    if (k_bondlist.d_view.extent(1) < maxbond_all)
+    if ((int)k_bondlist.d_view.extent(1) < maxbond_all)
       MemKK::realloc_kokkos(k_bondlist, "bond_hybrid:bondlist", nstyles, maxbond_all, 3);
     auto d_bondlist = k_bondlist.d_view;
 

src/KOKKOS/comm_tiled_kokkos.cpp

@@ -37,6 +37,8 @@ static constexpr int BUFEXTRA = 1000;
 CommTiledKokkos::CommTiledKokkos(LAMMPS *_lmp) : CommTiled(_lmp)
 {
   sendlist = nullptr;
+  maxsendlist = nullptr;
+  nprocmaxtot = 0;
 }
 
 /* ---------------------------------------------------------------------- */
@@ -49,14 +51,18 @@ CommTiledKokkos::CommTiledKokkos(LAMMPS *_lmp) : CommTiled(_lmp)
 CommTiledKokkos::CommTiledKokkos(LAMMPS *_lmp, Comm *oldcomm) : CommTiled(_lmp,oldcomm)
 {
   sendlist = nullptr;
+  maxsendlist = nullptr;
+  nprocmaxtot = 0;
 }
 
 /* ---------------------------------------------------------------------- */
 
 CommTiledKokkos::~CommTiledKokkos()
 {
   memoryKK->destroy_kokkos(k_sendlist,sendlist);
+  memory->destroy(maxsendlist);
   sendlist = nullptr;
+  maxsendlist = nullptr;
   buf_send = nullptr;
   buf_recv = nullptr;
 }
@@ -657,12 +663,11 @@ void CommTiledKokkos::grow_list(int iswap, int iwhich, int n)
   k_sendlist.sync<LMPHostType>();
   k_sendlist.modify<LMPHostType>();
 
-  if (size > (int)k_sendlist.extent(2)) {
-    memoryKK->grow_kokkos(k_sendlist,sendlist,maxswap,maxsend,size,"comm:sendlist");
+  memoryKK->grow_kokkos(k_sendlist,sendlist,maxswap,nprocmaxtot,size,"comm:sendlist");
 
-    for (int i = 0; i < maxswap; i++)
-      maxsendlist[iswap][iwhich] = size;
-  }
+  for (int i = 0; i < maxswap; i++)
+    for (int j = 0; j < nprocmaxtot; j++)
+      maxsendlist[i][j] = size;
 }
 
 /* ----------------------------------------------------------------------
@@ -692,24 +697,23 @@ void CommTiledKokkos::grow_swap_send(int i, int n, int /*nold*/)
   memory->destroy(sendbox_multiold[i]);
   memory->create(sendbox_multiold[i],n,atom->ntypes+1,6,"comm:sendbox_multiold");
 
-  delete [] maxsendlist[i];
-  maxsendlist[i] = new int[n];
-
-  for (int j = 0; j < n; j++)
-    maxsendlist[i][j] = BUFMIN;
-
-  if (sendlist && !k_sendlist.d_view.data()) {
-    for (int ii = 0; ii < maxswap; ii++) {
-      if (sendlist[ii]) {
-        for (int jj = 0; jj < nprocmax[ii]; jj++)
-          memory->destroy(sendlist[ii][jj]);
-        delete [] sendlist[ii];
-      }
-    }
+  if (sendlist && !k_sendlist.h_view.data()) {
     delete [] sendlist;
+    delete [] maxsendlist;
+
+    sendlist = nullptr;
+    maxsendlist = nullptr;
   } else {
     memoryKK->destroy_kokkos(k_sendlist,sendlist);
+    memory->destroy(maxsendlist);
   }
 
-  memoryKK->create_kokkos(k_sendlist,sendlist,maxswap,n,BUFMIN,"comm:sendlist");
+  nprocmaxtot = MAX(nprocmaxtot,n);
+
+  memoryKK->create_kokkos(k_sendlist,sendlist,maxswap,nprocmaxtot,BUFMIN,"comm:sendlist");
+  memory->create(maxsendlist,maxswap,nprocmaxtot,"comm:maxsendlist");
+
+  for (int i = 0; i < maxswap; i++)
+    for (int j = 0; j < nprocmaxtot; j++)
+      maxsendlist[i][j] = BUFMIN;
 }

src/KOKKOS/comm_tiled_kokkos.h

@@ -64,18 +64,17 @@ class CommTiledKokkos : public CommTiled {
   template<class DeviceType> void reverse_comm_device();
 
  protected:
+  int nprocmaxtot;
 
   DAT::tdual_int_3d k_sendlist;
-  //DAT::tdual_int_scalar k_total_send;
   DAT::tdual_xfloat_2d k_buf_send,k_buf_recv;
-  //DAT::tdual_int_scalar k_count;
 
-  void grow_send(int, int) override;
-  void grow_recv(int, int flag = 0) override;
+  void grow_send(int, int) override;             // reallocate send buffer
+  void grow_recv(int, int flag = 0) override;    // free/allocate recv buffer
   void grow_send_kokkos(int, int, ExecutionSpace space = Host);
   void grow_recv_kokkos(int, int, ExecutionSpace space = Host);
-  void grow_list(int, int, int) override;
-  void grow_swap_send(int, int, int) override;     // grow swap arrays for send and recv
+  void grow_list(int, int, int) override;        // reallocate sendlist for one swap/proc
+  void grow_swap_send(int, int, int) override;   // grow swap arrays for send and recv
 };
 
 }    // namespace LAMMPS_NS

src/KOKKOS/dihedral_harmonic_kokkos.cpp

@@ -75,14 +75,14 @@ void DihedralHarmonicKokkos<DeviceType>::compute(int eflag_in, int vflag_in)
   // reallocate per-atom arrays if necessary
 
   if (eflag_atom) {
-    if(k_eatom.extent(0) < maxeatom) {
+    if ((int)k_eatom.extent(0) < maxeatom) {
     memoryKK->destroy_kokkos(k_eatom,eatom);
     memoryKK->create_kokkos(k_eatom,eatom,maxeatom,"dihedral:eatom");
     d_eatom = k_eatom.view<DeviceType>();
     } else Kokkos::deep_copy(d_eatom,0.0);
   }
   if (vflag_atom) {
-    if(k_vatom.extent(0) < maxvatom) {
+    if ((int)k_vatom.extent(0) < maxvatom) {
     memoryKK->destroy_kokkos(k_vatom,vatom);
     memoryKK->create_kokkos(k_vatom,vatom,maxvatom,"dihedral:vatom");
     d_vatom = k_vatom.view<DeviceType>();

src/KOKKOS/dihedral_hybrid_kokkos.cpp

@@ -76,7 +76,7 @@ void DihedralHybridKokkos::compute(int eflag, int vflag)
 
     Kokkos::parallel_for(ndihedrallist_orig,LAMMPS_LAMBDA(int i) {
       const int m = d_map[d_dihedrallist_orig(i,4)];
-      if (m >= 0) Kokkos::atomic_increment(&d_ndihedrallist[m]);
+      if (m >= 0) Kokkos::atomic_inc(&d_ndihedrallist[m]);
     });
 
     k_ndihedrallist.modify_device();
@@ -87,7 +87,7 @@ void DihedralHybridKokkos::compute(int eflag, int vflag)
       if (h_ndihedrallist[m] > maxdihedral_all)
         maxdihedral_all = h_ndihedrallist[m] + EXTRA;
 
-    if (k_dihedrallist.d_view.extent(1) < maxdihedral_all)
+    if ((int)k_dihedrallist.d_view.extent(1) < maxdihedral_all)
       MemKK::realloc_kokkos(k_dihedrallist, "dihedral_hybrid:dihedrallist", nstyles, maxdihedral_all, 5);
     auto d_dihedrallist = k_dihedrallist.d_view;
 

src/KOKKOS/fix_cmap_kokkos.cpp

@@ -690,7 +690,7 @@ int FixCMAPKokkos<DeviceType>::pack_exchange_kokkos(
 
   copymode = 1;
 
-  Kokkos::parallel_scan(nsend, KOKKOS_LAMBDA(const int &mysend, int &offset, const bool &final) {
+  Kokkos::parallel_scan(Kokkos::RangePolicy<DeviceType>(0,nsend), KOKKOS_LAMBDA(const int &mysend, int &offset, const bool &final) {
 
     const int i = d_exchange_sendlist(mysend);
 
@@ -782,7 +782,7 @@ void FixCMAPKokkos<DeviceType>::unpack_exchange_kokkos(
 
   copymode = 1;
 
-  Kokkos::parallel_for(nrecv, KOKKOS_LAMBDA(const int &i) {
+  Kokkos::parallel_for(Kokkos::RangePolicy<DeviceType>(0,nrecv), KOKKOS_LAMBDA(const int &i) {
     int index = d_indices(i);
     if (index > -1) {
       int m = d_ubuf(d_buf(i)).i;