Add atom_map hash option for Kokkos package

src/KOKKOS/Install.sh

@@ -55,6 +55,7 @@ action angle_harmonic_kokkos.cpp angle_harmonic.cpp
 action angle_harmonic_kokkos.h angle_harmonic.h
 action atom_kokkos.cpp
 action atom_kokkos.h
+action atom_map_kokkos.cpp
 action atom_vec_angle_kokkos.cpp atom_vec_angle.cpp
 action atom_vec_angle_kokkos.h atom_vec_angle.h
 action atom_vec_atomic_kokkos.cpp

src/KOKKOS/atom_kokkos.cpp

@@ -29,7 +29,9 @@ using namespace LAMMPS_NS;
 
 /* ---------------------------------------------------------------------- */
 
-AtomKokkos::AtomKokkos(LAMMPS *lmp) : Atom(lmp) {}
+AtomKokkos::AtomKokkos(LAMMPS *lmp) : Atom(lmp) {
+  k_error_flag = DAT::tdual_int_scalar("atom:error_flag");
+}
 
 /* ---------------------------------------------------------------------- */
 
@@ -77,6 +79,8 @@ AtomKokkos::~AtomKokkos()
   memoryKK->destroy_kokkos(k_improper_atom3, improper_atom3);
   memoryKK->destroy_kokkos(k_improper_atom4, improper_atom4);
 
+  map_delete();
+
   // SPIN package
 
   memoryKK->destroy_kokkos(k_sp, sp);

src/KOKKOS/atom_kokkos.h

@@ -14,6 +14,7 @@
 
 #include "atom.h"               // IWYU pragma: export
 #include "kokkos_type.h"
+#include <Kokkos_UnorderedMap.hpp>
 
 #ifndef LMP_ATOM_KOKKOS_H
 #define LMP_ATOM_KOKKOS_H
@@ -69,6 +70,32 @@ class AtomKokkos : public Atom {
   AtomKokkos(class LAMMPS *);
   ~AtomKokkos();
 
+  void map_init(int check = 1);
+  void map_clear();
+  void map_set();
+  void map_delete();
+
+  DAT::tdual_int_1d k_sametag;
+  DAT::tdual_int_1d k_map_array;
+  DAT::tdual_int_scalar k_error_flag;
+
+  typedef Kokkos::UnorderedMap<tagint,int,LMPDeviceType> hash_type;
+  typedef Kokkos::DualView<hash_type, LMPDeviceType::array_layout, LMPDeviceType> dual_hash_type;
+  typedef dual_hash_type::t_host::data_type host_hash_type;
+  dual_hash_type k_map_hash;
+
+  template<class DeviceType>
+  KOKKOS_INLINE_FUNCTION
+  static int map_find_hash_kokkos(tagint global, dual_hash_type k_map_hash)
+  {
+    int local = -1;
+    auto d_map_hash = k_map_hash.view<DeviceType>()();
+    auto index = d_map_hash.find(global);
+    if (d_map_hash.valid_at(index))
+      local = d_map_hash.value_at(index);
+    return local;
+  }
+
   virtual void allocate_type_arrays();
   void sync(const ExecutionSpace space, unsigned int mask);
   void modified(const ExecutionSpace space, unsigned int mask);

src/KOKKOS/atom_map_kokkos.cpp

@@ -0,0 +1,287 @@
+// clang-format off
+/* ----------------------------------------------------------------------
+   LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
+   https://www.lammps.org/, 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 "atom_kokkos.h"
+#include "comm.h"
+#include "error.h"
+#include "memory_kokkos.h"
+#include "atom_masks.h"
+
+#include <cmath>
+
+using namespace LAMMPS_NS;
+
+#define EXTRA 1000
+
+/* ----------------------------------------------------------------------
+   allocate and initialize array or hash table for global -> local map
+   for array option:
+     array length = 1 to map_tag_max
+     set entire array to -1 as initial values
+   for hash option:
+     map_nhash = length of hash table
+     map_nbucket = # of hash buckets, prime larger than map_nhash * 2
+       so buckets will only be filled with 0 or 1 atoms on average
+------------------------------------------------------------------------- */
+
+void AtomKokkos::map_init(int check)
+{
+  // check for new map style if max atomID changed (check = 1 = default)
+  // recreate = 1 if must delete old map and create new map
+  // recreate = 0 if can re-use old map w/out realloc and just adjust settings
+  // map_maxarray/map_nhash initially -1, to force recreate even when no atoms
+
+  int recreate = 0;
+  if (check) recreate = map_style_set();
+
+  if (map_style == MAP_ARRAY && map_tag_max > map_maxarray) recreate = 1;
+  else if (map_style == MAP_HASH && nlocal+nghost > map_nhash) recreate = 1;
+
+  // if not recreating:
+  // for array, initialize current map_tag_max values
+  // for hash, set all buckets to empty, put all entries in free list
+
+  if (!recreate) {
+    if (map_style == MAP_ARRAY) {
+      for (int i = 0; i <= map_tag_max; i++) map_array[i] = -1;
+    } else {
+      for (int i = 0; i < map_nbucket; i++) map_bucket[i] = -1;
+      map_nused = 0;
+      map_free = 0;
+      for (int i = 0; i < map_nhash; i++) map_hash[i].next = i+1;
+      if (map_nhash > 0) map_hash[map_nhash-1].next = -1;
+
+      k_map_hash.h_view().clear();
+    }
+
+  // recreating: delete old map and create new one for array or hash
+
+  } else {
+    map_delete();
+
+    if (map_style == MAP_ARRAY) {
+      map_maxarray = map_tag_max;
+      memoryKK->create_kokkos(k_map_array,map_array,map_maxarray+1,"atom:map_array");
+      for (int i = 0; i <= map_tag_max; i++) map_array[i] = -1;
+
+    } else {
+
+      // map_nhash = max # of atoms that can be hashed on this proc
+      // set to max of ave atoms/proc or atoms I can store
+      // multiply by 2, require at least 1000
+      // doubling means hash table will need to be re-init only rarely
+
+      int nper = static_cast<int> (natoms/comm->nprocs);
+      map_nhash = MAX(nper,nmax);
+      map_nhash *= 2;
+      map_nhash = MAX(map_nhash,1000);
+
+      // map_nbucket = prime just larger than map_nhash
+      // next_prime() should be fast enough,
+      //   about 10% of odd integers are prime above 1M
+
+      map_nbucket = next_prime(map_nhash);
+
+      // set all buckets to empty
+      // set hash to map_nhash in length
+      // put all hash entries in free list and point them to each other
+
+      map_bucket = new int[map_nbucket];
+      for (int i = 0; i < map_nbucket; i++) map_bucket[i] = -1;
+
+      map_hash = new HashElem[map_nhash];
+      map_nused = 0;
+      map_free = 0;
+      for (int i = 0; i < map_nhash; i++) map_hash[i].next = i+1;
+      map_hash[map_nhash-1].next = -1;
+
+      k_map_hash = dual_hash_type("atom:map_hash");
+      k_map_hash.h_view() = host_hash_type(map_nhash);
+    }
+  }
+
+  if (map_style == Atom::MAP_ARRAY)
+    k_map_array.modify_host();
+  else if (map_style == Atom::MAP_HASH)
+    k_map_hash.modify_host();
+}
+
+/* ----------------------------------------------------------------------
+   clear global -> local map for all of my own and ghost atoms
+   for hash table option:
+     global ID may not be in table if image atom was already cleared
+------------------------------------------------------------------------- */
+
+void AtomKokkos::map_clear()
+{
+  Atom::map_clear();
+
+  if (map_style == MAP_ARRAY) {
+    k_map_array.modify_host();
+  } else {
+    k_map_hash.h_view().clear();
+    k_map_hash.modify_host();
+  }
+  k_sametag.modify_host();
+}
+
+/* ----------------------------------------------------------------------
+   set global -> local map for all of my own and ghost atoms
+   loop in reverse order so that nearby images take precedence over far ones
+     and owned atoms take precedence over images
+   this enables valid lookups of bond topology atoms
+   for hash table option:
+     if hash table too small, re-init
+     global ID may already be in table if image atom was set
+------------------------------------------------------------------------- */
+
+void AtomKokkos::map_set()
+{
+  int nall = nlocal + nghost;
+
+  atomKK->sync(Host,TAG_MASK);
+
+  k_sametag.sync_host();
+  if (map_style == Atom::MAP_ARRAY)
+    k_map_array.sync_host();
+
+  if (map_style == MAP_ARRAY) {
+
+    // possible reallocation of sametag must come before loop over atoms
+    // since loop sets sametag
+
+    if (nall > max_same) {
+      max_same = nall + EXTRA;
+      memoryKK->destroy_kokkos(k_sametag,sametag);
+      memoryKK->create_kokkos(k_sametag,sametag,max_same,"atom:sametag");
+    }
+
+    for (int i = nall-1; i >= 0 ; i--) {
+      sametag[i] = map_array[tag[i]];
+      map_array[tag[i]] = i;
+    }
+
+  } else {
+
+    // if this proc has more atoms than hash table size, call map_init()
+    //   call with 0 since max atomID in system has not changed
+    // possible reallocation of sametag must come after map_init(),
+    //   b/c map_init() may invoke map_delete(), whacking sametag
+
+    if (nall > map_nhash) map_init(0);
+    if (nall > max_same) {
+      max_same = nall + EXTRA;
+      memoryKK->destroy_kokkos(k_sametag,sametag);
+      memoryKK->create_kokkos(k_sametag,sametag,max_same,"atom:sametag");
+    }
+
+    int previous,ibucket,index;
+    tagint global;
+
+    for (int i = nall-1; i >= 0 ; i--) {
+      sametag[i] = map_find_hash(tag[i]);
+
+      // search for key
+      // if found it, just overwrite local value with index
+
+      previous = -1;
+      global = tag[i];
+      ibucket = global % map_nbucket;
+      index = map_bucket[ibucket];
+      while (index > -1) {
+        if (map_hash[index].global == global) break;
+        previous = index;
+        index = map_hash[index].next;
+      }
+      if (index > -1) {
+        map_hash[index].local = i;
+        continue;
+      }
+
+      // take one entry from free list
+      // add the new global/local pair as entry at end of bucket list
+      // special logic if this entry is 1st in bucket
+
+      index = map_free;
+      map_free = map_hash[map_free].next;
+      if (previous == -1) map_bucket[ibucket] = index;
+      else map_hash[previous].next = index;
+      map_hash[index].global = global;
+      map_hash[index].local = i;
+      map_hash[index].next = -1;
+      map_nused++;
+    }
+
+    // Copy to Kokkos hash
+
+    k_map_hash.h_view().clear();
+    auto h_map_hash = k_map_hash.h_view();
+
+    for (int i = nall-1; i >= 0 ; i--) {
+
+      // search for key
+      // if don't find it, done
+
+      previous = -1;
+      global = tag[i];
+      ibucket = global % map_nbucket;
+      index = map_bucket[ibucket];
+      while (index > -1) {
+        if (map_hash[index].global == global) break;
+        previous = index;
+        index = map_hash[index].next;
+      }
+      if (index == -1) continue;
+
+      int local = map_hash[index].local;
+
+      auto insert_result = h_map_hash.insert(global,local);
+      if (insert_result.failed())
+        error->one(FLERR, "Kokkos::UnorderedMap insertion failed");
+    }
+  }
+
+  k_sametag.modify_host();
+  if (map_style == Atom::MAP_ARRAY)
+    k_map_array.modify_host();
+  else if (map_style == Atom::MAP_HASH)
+    k_map_hash.modify_host();
+}
+
+
+/* ----------------------------------------------------------------------
+   free the array or hash table for global to local mapping
+------------------------------------------------------------------------- */
+
+void AtomKokkos::map_delete()
+{
+  memoryKK->destroy_kokkos(k_sametag,sametag);
+  sametag = nullptr;
+  max_same = 0;
+
+  if (map_style == MAP_ARRAY) {
+    memoryKK->destroy_kokkos(k_map_array,map_array);
+    map_array = nullptr;
+  } else {
+    if (map_nhash) {
+      delete [] map_bucket;
+      delete [] map_hash;
+      map_bucket = nullptr;
+      map_hash = nullptr;
+      k_map_hash = dual_hash_type();
+    }
+    map_nhash = map_nbucket = 0;
+  }
+}
+

src/KOKKOS/comm_kokkos.cpp

@@ -1128,9 +1128,10 @@ void CommKokkos::borders_device() {
   max = MAX(maxforward*rmax,maxreverse*smax);
   if (max > maxrecv) grow_recv_kokkos(max);
 
+  atomKK->modified(exec_space,ALL_MASK);
+
   // reset global->local map
 
-  atomKK->modified(exec_space,ALL_MASK);
   if (map_style != Atom::MAP_NONE) {
     atomKK->sync(Host,TAG_MASK);
     atom->map_set();

src/KOKKOS/fix_shake_kokkos.cpp

@@ -201,6 +201,15 @@ void FixShakeKokkos<DeviceType>::pre_neighbor()
   type = atom->type;
   nlocal = atom->nlocal;
 
+  map_style = atom->map_style;
+  if (map_style == Atom::MAP_ARRAY) {
+    k_map_array = atomKK->k_map_array;
+    k_map_array.template sync<DeviceType>();
+  } else if (map_style == Atom::MAP_HASH) {
+    k_map_hash = atomKK->k_map_hash;
+    k_map_hash.template sync<DeviceType>();
+  }
+
   k_shake_flag.sync<DeviceType>();
   k_shake_atom.sync<DeviceType>();
 
@@ -213,21 +222,17 @@ void FixShakeKokkos<DeviceType>::pre_neighbor()
     d_list = k_list.view<DeviceType>();
   }
 
-  // don't yet have atom_map_kokkos routines, so move data from host to device
+  // Atom Map
 
-  if (atom->map_style != Atom::MAP_ARRAY)
+  map_style = atom->map_style;
-    error->all(FLERR,"Must use atom map style array with Kokkos");
 
-  int* map_array_host = atom->get_map_array();
+  if (map_style == Atom::MAP_ARRAY) {
-  int map_size = atom->get_map_size();
+    k_map_array = atomKK->k_map_array;
-  int map_maxarray = atom->get_map_maxarray();
+    k_map_array.template sync<DeviceType>();
-  if (map_maxarray > (int)k_map_array.extent(0))
+  } else if (map_style == Atom::MAP_HASH) {
-    k_map_array = DAT::tdual_int_1d("NeighBond:map_array",map_maxarray);
+    k_map_hash = atomKK->k_map_hash;
-  for (int i=0; i<map_size; i++)
+    k_map_hash.template sync<DeviceType>();
-    k_map_array.h_view[i] = map_array_host[i];
+  }
-  k_map_array.template modify<LMPHostType>();
-  k_map_array.template sync<DeviceType>();
-  map_array = k_map_array.view<DeviceType>();
 
   // build list of SHAKE clusters I compute
 
@@ -241,14 +246,16 @@ void FixShakeKokkos<DeviceType>::pre_neighbor()
     auto d_list = this->d_list;
     auto d_error_flag = this->d_error_flag;
     auto d_nlist = this->d_nlist;
-    auto map_array = this->map_array;
+    auto map_style = atom->map_style;
+    auto k_map_array = this->k_map_array;
+    auto k_map_hash = this->k_map_hash;
 
     Kokkos::parallel_for(Kokkos::RangePolicy<DeviceType>(0,nlocal),
      LAMMPS_LAMBDA(const int& i) {
       if (d_shake_flag[i]) {
         if (d_shake_flag[i] == 2) {
-          const int atom1 = map_array(d_shake_atom(i,0));
+          const int atom1 = map_kokkos(d_shake_atom(i,0),map_style,k_map_array,k_map_hash);
-          const int atom2 = map_array(d_shake_atom(i,1));
+          const int atom2 = map_kokkos(d_shake_atom(i,1),map_style,k_map_array,k_map_hash);
           if (atom1 == -1 || atom2 == -1) {
             d_error_flag() = 1;
           }
@@ -257,9 +264,9 @@ void FixShakeKokkos<DeviceType>::pre_neighbor()
             d_list[nlist] = i;
           }
         } else if (d_shake_flag[i] % 2 == 1) {
-          const int atom1 = map_array(d_shake_atom(i,0));
+          const int atom1 = map_kokkos(d_shake_atom(i,0),map_style,k_map_array,k_map_hash);
-          const int atom2 = map_array(d_shake_atom(i,1));
+          const int atom2 = map_kokkos(d_shake_atom(i,1),map_style,k_map_array,k_map_hash);
-          const int atom3 = map_array(d_shake_atom(i,2));
+          const int atom3 = map_kokkos(d_shake_atom(i,2),map_style,k_map_array,k_map_hash);
           if (atom1 == -1 || atom2 == -1 || atom3 == -1)
             d_error_flag() = 1;
           if (i <= atom1 && i <= atom2 && i <= atom3) {
@@ -267,10 +274,10 @@ void FixShakeKokkos<DeviceType>::pre_neighbor()
             d_list[nlist] = i;
           }
         } else {
-          const int atom1 = map_array(d_shake_atom(i,0));
+          const int atom1 = map_kokkos(d_shake_atom(i,0),map_style,k_map_array,k_map_hash);
-          const int atom2 = map_array(d_shake_atom(i,1));
+          const int atom2 = map_kokkos(d_shake_atom(i,1),map_style,k_map_array,k_map_hash);
-          const int atom3 = map_array(d_shake_atom(i,2));
+          const int atom3 = map_kokkos(d_shake_atom(i,2),map_style,k_map_array,k_map_hash);
-          const int atom4 = map_array(d_shake_atom(i,3));
+          const int atom4 = map_kokkos(d_shake_atom(i,3),map_style,k_map_array,k_map_hash);
           if (atom1 == -1 || atom2 == -1 || atom3 == -1 || atom4 == -1)
             d_error_flag() = 1;
           if (i <= atom1 && i <= atom2 && i <= atom3 && i <= atom4) {
@@ -307,6 +314,15 @@ void FixShakeKokkos<DeviceType>::post_force(int vflag)
   d_mass = atomKK->k_mass.view<DeviceType>();
   nlocal = atomKK->nlocal;
 
+  map_style = atom->map_style;
+  if (map_style == Atom::MAP_ARRAY) {
+    k_map_array = atomKK->k_map_array;
+    k_map_array.template sync<DeviceType>();
+  } else if (map_style == Atom::MAP_HASH) {
+    k_map_hash = atomKK->k_map_hash;
+    k_map_hash.template sync<DeviceType>();
+  }
+
   if (d_rmass.data())
     atomKK->sync(execution_space,X_MASK|F_MASK|RMASS_MASK);
   else
@@ -586,8 +602,8 @@ void FixShakeKokkos<DeviceType>::shake(int m, EV_FLOAT& ev) const
 
   // local atom IDs and constraint distances
 
-  int i0 = map_array(d_shake_atom(m,0));
+  int i0 = map_kokkos(d_shake_atom(m,0),map_style,k_map_array,k_map_hash);
-  int i1 = map_array(d_shake_atom(m,1));
+  int i1 = map_kokkos(d_shake_atom(m,1),map_style,k_map_array,k_map_hash);
   double bond1 = d_bond_distance[d_shake_type(m,0)];
 
   // r01 = distance vec between atoms, with PBC
@@ -697,9 +713,9 @@ void FixShakeKokkos<DeviceType>::shake3(int m, EV_FLOAT& ev) const
 
   // local atom IDs and constraint distances
 
-  int i0 = map_array(d_shake_atom(m,0));
+  int i0 = map_kokkos(d_shake_atom(m,0),map_style,k_map_array,k_map_hash);
-  int i1 = map_array(d_shake_atom(m,1));
+  int i1 = map_kokkos(d_shake_atom(m,1),map_style,k_map_array,k_map_hash);
-  int i2 = map_array(d_shake_atom(m,2));
+  int i2 = map_kokkos(d_shake_atom(m,2),map_style,k_map_array,k_map_hash);
   double bond1 = d_bond_distance[d_shake_type(m,0)];
   double bond2 = d_bond_distance[d_shake_type(m,1)];
 
@@ -880,10 +896,10 @@ void FixShakeKokkos<DeviceType>::shake4(int m, EV_FLOAT& ev) const
 
   // local atom IDs and constraint distances
 
-  int i0 = map_array(d_shake_atom(m,0));
+  int i0 = map_kokkos(d_shake_atom(m,0),map_style,k_map_array,k_map_hash);
-  int i1 = map_array(d_shake_atom(m,1));
+  int i1 = map_kokkos(d_shake_atom(m,1),map_style,k_map_array,k_map_hash);
-  int i2 = map_array(d_shake_atom(m,2));
+  int i2 = map_kokkos(d_shake_atom(m,2),map_style,k_map_array,k_map_hash);
-  int i3 = map_array(d_shake_atom(m,3));
+  int i3 = map_kokkos(d_shake_atom(m,3),map_style,k_map_array,k_map_hash);
   double bond1 = d_bond_distance[d_shake_type(m,0)];
   double bond2 = d_bond_distance[d_shake_type(m,1)];
   double bond3 = d_bond_distance[d_shake_type(m,2)];
@@ -1142,9 +1158,9 @@ void FixShakeKokkos<DeviceType>::shake3angle(int m, EV_FLOAT& ev) const
 
   // local atom IDs and constraint distances
 
-  int i0 = map_array(d_shake_atom(m,0));
+  int i0 = map_kokkos(d_shake_atom(m,0),map_style,k_map_array,k_map_hash);
-  int i1 = map_array(d_shake_atom(m,1));
+  int i1 = map_kokkos(d_shake_atom(m,1),map_style,k_map_array,k_map_hash);
-  int i2 = map_array(d_shake_atom(m,2));
+  int i2 = map_kokkos(d_shake_atom(m,2),map_style,k_map_array,k_map_hash);
   double bond1 = d_bond_distance[d_shake_type(m,0)];
   double bond2 = d_bond_distance[d_shake_type(m,1)];
   double bond12 = d_angle_distance[d_shake_type(m,2)];
@@ -1905,6 +1921,24 @@ void FixShakeKokkos<DeviceType>::minimum_image_once(double *delta) const
 
 /* ---------------------------------------------------------------------- */
 
+// functions for global to local ID mapping
+// map lookup function inlined for efficiency
+// return -1 if no map defined
+
+template<class DeviceType>
+KOKKOS_INLINE_FUNCTION
+int FixShakeKokkos<DeviceType>::map_kokkos(tagint global, int map_style, DAT::tdual_int_1d k_map_array, dual_hash_type k_map_hash)
+{
+  if (map_style == 1)
+    return k_map_array.view<DeviceType>()(global);
+  else if (map_style == 2)
+    return AtomKokkos::map_find_hash_kokkos<DeviceType>(global,k_map_hash);
+  else
+    return -1;
+}
+
+/* ---------------------------------------------------------------------- */
+
 namespace LAMMPS_NS {
 template class FixShakeKokkos<LMPDeviceType>;
 #ifdef LMP_KOKKOS_GPU

src/KOKKOS/fix_shake_kokkos.h

@@ -26,6 +26,7 @@ FixStyle(shake/kk/host,FixShakeKokkos<LMPHostType>);
 #include "fix_shake.h"
 #include "kokkos_type.h"
 #include "kokkos_base.h"
+#include <Kokkos_UnorderedMap.hpp>
 
 namespace LAMMPS_NS {
 
@@ -172,9 +173,6 @@ class FixShakeKokkos : public FixShake, public KokkosBase {
   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;
@@ -185,6 +183,17 @@ class FixShakeKokkos : public FixShake, public KokkosBase {
   tagint **shake_atom_tmp;
   int **shake_type_tmp;
 
+  int map_style;
+
+  DAT::tdual_int_1d k_map_array;
+
+  typedef Kokkos::UnorderedMap<tagint,int,LMPDeviceType> hash_type;
+  typedef Kokkos::DualView<hash_type, LMPDeviceType::array_layout, LMPDeviceType> dual_hash_type;
+  dual_hash_type k_map_hash;
+
+  KOKKOS_INLINE_FUNCTION
+  static int map_kokkos(tagint, int, DAT::tdual_int_1d, dual_hash_type);
+
   // copied from Domain
 
   KOKKOS_INLINE_FUNCTION

src/KOKKOS/neigh_bond_kokkos.cpp

@@ -207,30 +207,7 @@ void NeighBondKokkos<DeviceType>::build_topology_kk()
 
   lostbond = output->thermo->lostbond;
 
-  // don't yet have atom_map_kokkos routines, so move data from host to device
+  update_class_variables();
-
-  if (atom->map_style != Atom::MAP_ARRAY)
-    error->all(FLERR,"Must use atom map style array with Kokkos");
-
-  int* map_array_host = atom->get_map_array();
-  int map_size = atom->get_map_size();
-  int map_maxarray = atom->get_map_maxarray();
-  if (map_maxarray > (int)k_map_array.extent(0))
-    k_map_array = DAT::tdual_int_1d("NeighBond:map_array",map_maxarray);
-  for (int i=0; i<map_size; i++)
-    k_map_array.h_view[i] = map_array_host[i];
-  k_map_array.template modify<LMPHostType>();
-  k_map_array.template sync<DeviceType>();
-  map_array = k_map_array.view<DeviceType>();
-
-  int* sametag_host = atomKK->sametag;
-  if (nmax > (int)k_sametag.extent(0))
-    k_sametag = DAT::tdual_int_1d("NeighBond:sametag",nmax);
-  for (int i=0; i<nall; i++)
-    k_sametag.h_view[i] = sametag_host[i];
-  k_sametag.template modify<LMPHostType>();
-  k_sametag.template sync<DeviceType>();
-  sametag = k_sametag.view<DeviceType>();
 
   if (force->bond) (this->*bond_build_kk)();
   if (force->angle) (this->*angle_build_kk)();
@@ -306,7 +283,7 @@ template<class DeviceType>
 KOKKOS_INLINE_FUNCTION
 void NeighBondKokkos<DeviceType>::operator()(TagNeighBondBondAll, const int &i, int &nmissing) const {
   for (int m = 0; m < num_bond[i]; m++) {
-    int atom1 = map_array(bond_atom(i,m));
+    int atom1 = map_kokkos(bond_atom(i,m));
     if (atom1 == -1) {
       nmissing++;
       if (lostbond == Thermo::ERROR) return;
@@ -394,7 +371,7 @@ KOKKOS_INLINE_FUNCTION
 void NeighBondKokkos<DeviceType>::operator()(TagNeighBondBondPartial, const int &i, int &nmissing) const {
   for (int m = 0; m < num_bond[i]; m++) {
     if (bond_type(i,m) <= 0) continue;
-    int atom1 = map_array(bond_atom(i,m));
+    int atom1 = map_kokkos(bond_atom(i,m));
     if (atom1 == -1) {
       nmissing++;
       if (lostbond == Thermo::ERROR) return;
@@ -420,7 +397,6 @@ void NeighBondKokkos<DeviceType>::bond_check()
 {
   int flag = 0;
 
-  update_domain_variables();
   atomKK->sync(execution_space, X_MASK);
   k_bondlist.sync<DeviceType>();
 
@@ -507,9 +483,9 @@ template<class DeviceType>
 KOKKOS_INLINE_FUNCTION
 void NeighBondKokkos<DeviceType>::operator()(TagNeighBondAngleAll, const int &i, int &nmissing) const {
   for (int m = 0; m < num_angle[i]; m++) {
-    int atom1 = map_array(angle_atom1(i,m));
+    int atom1 = map_kokkos(angle_atom1(i,m));
-    int atom2 = map_array(angle_atom2(i,m));
+    int atom2 = map_kokkos(angle_atom2(i,m));
-    int atom3 = map_array(angle_atom3(i,m));
+    int atom3 = map_kokkos(angle_atom3(i,m));
     if (atom1 == -1 || atom2 == -1 || atom3 == -1) {
       nmissing++;
       if (lostbond == Thermo::ERROR) return;
@@ -602,9 +578,9 @@ KOKKOS_INLINE_FUNCTION
 void NeighBondKokkos<DeviceType>::operator()(TagNeighBondAnglePartial, const int &i, int &nmissing) const {
   for (int m = 0; m < num_angle[i]; m++) {
     if (angle_type(i,m) <= 0) continue;
-    int atom1 = map_array(angle_atom1(i,m));
+    int atom1 = map_kokkos(angle_atom1(i,m));
-    int atom2 = map_array(angle_atom2(i,m));
+    int atom2 = map_kokkos(angle_atom2(i,m));
-    int atom3 = map_array(angle_atom3(i,m));
+    int atom3 = map_kokkos(angle_atom3(i,m));
     if (atom1 == -1 || atom2 == -1 || atom3 == -1) {
       nmissing++;
       if (lostbond == Thermo::ERROR) return;
@@ -636,7 +612,6 @@ void NeighBondKokkos<DeviceType>::angle_check()
   // check all 3 distances
   // in case angle potential computes any of them
 
-  update_domain_variables();
   atomKK->sync(execution_space, X_MASK);
   k_anglelist.sync<DeviceType>();
 
@@ -735,10 +710,10 @@ template<class DeviceType>
 KOKKOS_INLINE_FUNCTION
 void NeighBondKokkos<DeviceType>::operator()(TagNeighBondDihedralAll, const int &i, int &nmissing) const {
   for (int m = 0; m < num_dihedral[i]; m++) {
-    int atom1 = map_array(dihedral_atom1(i,m));
+    int atom1 = map_kokkos(dihedral_atom1(i,m));
-    int atom2 = map_array(dihedral_atom2(i,m));
+    int atom2 = map_kokkos(dihedral_atom2(i,m));
-    int atom3 = map_array(dihedral_atom3(i,m));
+    int atom3 = map_kokkos(dihedral_atom3(i,m));
-    int atom4 = map_array(dihedral_atom4(i,m));
+    int atom4 = map_kokkos(dihedral_atom4(i,m));
     if (atom1 == -1 || atom2 == -1 || atom3 == -1 || atom4 == -1) {
       nmissing++;
       if (lostbond == Thermo::ERROR) return;
@@ -835,10 +810,10 @@ KOKKOS_INLINE_FUNCTION
 void NeighBondKokkos<DeviceType>::operator()(TagNeighBondDihedralPartial, const int &i, int &nmissing) const {
   for (int m = 0; m < num_dihedral[i]; m++) {
     if (dihedral_type(i,m) <= 0) continue;
-    int atom1 = map_array(dihedral_atom1(i,m));
+    int atom1 = map_kokkos(dihedral_atom1(i,m));
-    int atom2 = map_array(dihedral_atom2(i,m));
+    int atom2 = map_kokkos(dihedral_atom2(i,m));
-    int atom3 = map_array(dihedral_atom3(i,m));
+    int atom3 = map_kokkos(dihedral_atom3(i,m));
-    int atom4 = map_array(dihedral_atom4(i,m));
+    int atom4 = map_kokkos(dihedral_atom4(i,m));
     if (atom1 == -1 || atom2 == -1 || atom3 == -1 || atom4 == -1) {
       nmissing++;
       if (lostbond == Thermo::ERROR) return;
@@ -874,7 +849,6 @@ void NeighBondKokkos<DeviceType>::dihedral_check(int nlist, typename AT::t_int_2
   // check all 6 distances
   // in case dihedral/improper potential computes any of them
 
-  update_domain_variables();
   atomKK->sync(execution_space, X_MASK);
   k_dihedrallist.sync<DeviceType>();
 
@@ -990,10 +964,10 @@ template<class DeviceType>
 KOKKOS_INLINE_FUNCTION
 void NeighBondKokkos<DeviceType>::operator()(TagNeighBondImproperAll, const int &i, int &nmissing) const {
   for (int m = 0; m < num_improper[i]; m++) {
-    int atom1 = map_array(improper_atom1(i,m));
+    int atom1 = map_kokkos(improper_atom1(i,m));
-    int atom2 = map_array(improper_atom2(i,m));
+    int atom2 = map_kokkos(improper_atom2(i,m));
-    int atom3 = map_array(improper_atom3(i,m));
+    int atom3 = map_kokkos(improper_atom3(i,m));
-    int atom4 = map_array(improper_atom4(i,m));
+    int atom4 = map_kokkos(improper_atom4(i,m));
     if (atom1 == -1 || atom2 == -1 || atom3 == -1 || atom4 == -1) {
       nmissing++;
       if (lostbond == Thermo::ERROR) return;
@@ -1090,10 +1064,10 @@ KOKKOS_INLINE_FUNCTION
 void NeighBondKokkos<DeviceType>::operator()(TagNeighBondImproperPartial, const int &i, int &nmissing) const {
   for (int m = 0; m < num_improper[i]; m++) {
     if (improper_type(i,m) <= 0) continue;
-    int atom1 = map_array(improper_atom1(i,m));
+    int atom1 = map_kokkos(improper_atom1(i,m));
-    int atom2 = map_array(improper_atom2(i,m));
+    int atom2 = map_kokkos(improper_atom2(i,m));
-    int atom3 = map_array(improper_atom3(i,m));
+    int atom3 = map_kokkos(improper_atom3(i,m));
-    int atom4 = map_array(improper_atom4(i,m));
+    int atom4 = map_kokkos(improper_atom4(i,m));
     if (atom1 == -1 || atom2 == -1 || atom3 == -1 || atom4 == -1) {
       nmissing++;
       if (lostbond == Thermo::ERROR) return;
@@ -1137,8 +1111,8 @@ int NeighBondKokkos<DeviceType>::closest_image(const int i, int j) const
   X_FLOAT rsqmin = delx*delx + dely*dely + delz*delz;
   X_FLOAT rsq;
 
-  while (sametag[j] >= 0) {
+  while (d_sametag[j] >= 0) {
-    j = sametag[j];
+    j = d_sametag[j];
     delx = xi0 - x(j,0);
     dely = xi1 - x(j,1);
     delz = xi2 - x(j,2);
@@ -1217,9 +1191,29 @@ void NeighBondKokkos<DeviceType>::minimum_image(X_FLOAT &dx, X_FLOAT &dy, X_FLOA
 
 /* ---------------------------------------------------------------------- */
 
+// functions for global to local ID mapping
+// map lookup function inlined for efficiency
+// return -1 if no map defined
+
 template<class DeviceType>
-void NeighBondKokkos<DeviceType>::update_domain_variables()
+KOKKOS_INLINE_FUNCTION
+int NeighBondKokkos<DeviceType>::map_kokkos(tagint global) const
 {
+  if (map_style == 1)
+    return k_map_array.view<DeviceType>()(global);
+  else if (map_style == 2)
+    return AtomKokkos::map_find_hash_kokkos<DeviceType>(global,k_map_hash);
+  else
+    return -1;
+}
+
+/* ---------------------------------------------------------------------- */
+
+template<class DeviceType>
+void NeighBondKokkos<DeviceType>::update_class_variables()
+{
+  // Domain
+
   triclinic = domain->triclinic;
   xperiodic = domain->xperiodic;
   xprd_half = domain->xprd_half;
@@ -1233,6 +1227,22 @@ void NeighBondKokkos<DeviceType>::update_domain_variables()
   xy = domain->xy;
   xz = domain->xz;
   yz = domain->yz;
+
+  // Atom Map
+
+  map_style = atom->map_style;
+
+  k_sametag = atomKK->k_sametag;
+  k_sametag.template sync<DeviceType>();
+  d_sametag = k_sametag.view<DeviceType>();
+
+  if (map_style == Atom::MAP_ARRAY) {
+    k_map_array = atomKK->k_map_array;
+    k_map_array.template sync<DeviceType>();
+  } else if (map_style == Atom::MAP_HASH) {
+    k_map_hash = atomKK->k_map_hash;
+    k_map_hash.template sync<DeviceType>();  
+  }
 }
 
 /* ---------------------------------------------------------------------- */

src/KOKKOS/neigh_bond_kokkos.h

@@ -19,6 +19,7 @@
 #include "kokkos_type.h"
 #include "domain_kokkos.h"
 #include "pointers.h"
+#include <Kokkos_UnorderedMap.hpp>
 
 namespace LAMMPS_NS {
 
@@ -81,13 +82,19 @@ class NeighBondKokkos : protected Pointers  {
   int me,nprocs;
 
  private:
-
+  int map_style;
-
-  DAT::tdual_int_1d k_map_array;
-  typename AT::t_int_1d_randomread map_array;
 
   DAT::tdual_int_1d k_sametag;
-  typename AT::t_int_1d_randomread sametag;
+  typename AT::t_int_1d d_sametag;
+
+  DAT::tdual_int_1d k_map_array;
+
+  typedef Kokkos::UnorderedMap<tagint,int,LMPDeviceType> hash_type;
+  typedef Kokkos::DualView<hash_type, LMPDeviceType::array_layout, LMPDeviceType> dual_hash_type;
+  dual_hash_type k_map_hash;
+
+  KOKKOS_INLINE_FUNCTION
+  int map_kokkos(tagint) const;
 
   typename AT::t_int_2d v_bondlist;
   typename AT::t_int_2d v_anglelist;
@@ -130,7 +137,7 @@ class NeighBondKokkos : protected Pointers  {
   KOKKOS_INLINE_FUNCTION
   void minimum_image(X_FLOAT &dx, X_FLOAT &dy, X_FLOAT &dz) const;
 
-  void update_domain_variables();
+  void update_class_variables();
 
   // topology build functions
 

src/atom.h

@@ -364,13 +364,13 @@ class Atom : protected Pointers {
       return -1;
   };
 
-  void map_init(int check = 1);
+  virtual void map_init(int check = 1);
-  void map_clear();
+  virtual void map_clear();
-  void map_set();
+  virtual void map_set();
-  void map_one(tagint, int);
+  virtual void map_one(tagint, int);
   int map_style_set();
-  void map_delete();
+  virtual void map_delete();
-  int map_find_hash(tagint);
+  virtual int map_find_hash(tagint);
 
  protected:
   // global to local ID mapping

src/atom_map.cpp

@@ -309,7 +309,7 @@ int Atom::map_style_set()
   if (map_user == MAP_ARRAY || map_user == MAP_HASH) {
     map_style = map_user;
   } else {  // map_user == MAP_YES
-    if (map_tag_max > 1000000 && !lmp->kokkos) map_style = MAP_HASH;
+    if (map_tag_max > 1000000) map_style = MAP_HASH;
     else map_style = MAP_ARRAY;
   }