diff --git a/src/KOKKOS/Install.sh b/src/KOKKOS/Install.sh index 0ac90a53f2..a28b17270a 100755 --- a/src/KOKKOS/Install.sh +++ b/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 diff --git a/src/KOKKOS/atom_kokkos.cpp b/src/KOKKOS/atom_kokkos.cpp index a8527989d7..c8d2268937 100644 --- a/src/KOKKOS/atom_kokkos.cpp +++ b/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); diff --git a/src/KOKKOS/atom_kokkos.h b/src/KOKKOS/atom_kokkos.h index e807180f35..d2629b4441 100644 --- a/src/KOKKOS/atom_kokkos.h +++ b/src/KOKKOS/atom_kokkos.h @@ -14,6 +14,7 @@ #include "atom.h" // IWYU pragma: export #include "kokkos_type.h" +#include #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 hash_type; + typedef Kokkos::DualView dual_hash_type; + typedef dual_hash_type::t_host::data_type host_hash_type; + dual_hash_type k_map_hash; + + template + 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()(); + 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); diff --git a/src/KOKKOS/atom_map_kokkos.cpp b/src/KOKKOS/atom_map_kokkos.cpp new file mode 100644 index 0000000000..4c3d5cd1e9 --- /dev/null +++ b/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 + +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 (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; + } +} + diff --git a/src/KOKKOS/comm_kokkos.cpp b/src/KOKKOS/comm_kokkos.cpp index fd5318e1d1..588ff94d73 100644 --- a/src/KOKKOS/comm_kokkos.cpp +++ b/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(); diff --git a/src/KOKKOS/fix_shake_kokkos.cpp b/src/KOKKOS/fix_shake_kokkos.cpp index 804da9fd22..594c9c60db 100644 --- a/src/KOKKOS/fix_shake_kokkos.cpp +++ b/src/KOKKOS/fix_shake_kokkos.cpp @@ -201,6 +201,15 @@ void FixShakeKokkos::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(); + } else if (map_style == Atom::MAP_HASH) { + k_map_hash = atomKK->k_map_hash; + k_map_hash.template sync(); + } + k_shake_flag.sync(); k_shake_atom.sync(); @@ -213,21 +222,17 @@ void FixShakeKokkos::pre_neighbor() d_list = k_list.view(); } - // don't yet have atom_map_kokkos routines, so move data from host to device + // Atom Map - if (atom->map_style != Atom::MAP_ARRAY) - error->all(FLERR,"Must use atom map style array with Kokkos"); + map_style = atom->map_style; - 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(); - k_map_array.template sync(); - map_array = k_map_array.view(); + if (map_style == Atom::MAP_ARRAY) { + k_map_array = atomKK->k_map_array; + k_map_array.template sync(); + } else if (map_style == Atom::MAP_HASH) { + k_map_hash = atomKK->k_map_hash; + k_map_hash.template sync(); + } // build list of SHAKE clusters I compute @@ -241,14 +246,16 @@ void FixShakeKokkos::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(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 atom2 = map_array(d_shake_atom(i,1)); + const int atom1 = map_kokkos(d_shake_atom(i,0),map_style,k_map_array,k_map_hash); + 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::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 atom2 = map_array(d_shake_atom(i,1)); - const int atom3 = map_array(d_shake_atom(i,2)); + const int atom1 = map_kokkos(d_shake_atom(i,0),map_style,k_map_array,k_map_hash); + const int atom2 = map_kokkos(d_shake_atom(i,1),map_style,k_map_array,k_map_hash); + 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::pre_neighbor() d_list[nlist] = i; } } else { - const int atom1 = map_array(d_shake_atom(i,0)); - const int atom2 = map_array(d_shake_atom(i,1)); - const int atom3 = map_array(d_shake_atom(i,2)); - const int atom4 = map_array(d_shake_atom(i,3)); + const int atom1 = map_kokkos(d_shake_atom(i,0),map_style,k_map_array,k_map_hash); + const int atom2 = map_kokkos(d_shake_atom(i,1),map_style,k_map_array,k_map_hash); + const int atom3 = map_kokkos(d_shake_atom(i,2),map_style,k_map_array,k_map_hash); + 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::post_force(int vflag) d_mass = atomKK->k_mass.view(); 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(); + } else if (map_style == Atom::MAP_HASH) { + k_map_hash = atomKK->k_map_hash; + k_map_hash.template sync(); + } + if (d_rmass.data()) atomKK->sync(execution_space,X_MASK|F_MASK|RMASS_MASK); else @@ -586,8 +602,8 @@ void FixShakeKokkos::shake(int m, EV_FLOAT& ev) const // local atom IDs and constraint distances - int i0 = map_array(d_shake_atom(m,0)); - int i1 = map_array(d_shake_atom(m,1)); + int i0 = map_kokkos(d_shake_atom(m,0),map_style,k_map_array,k_map_hash); + 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::shake3(int m, EV_FLOAT& ev) const // local atom IDs and constraint distances - int i0 = map_array(d_shake_atom(m,0)); - int i1 = map_array(d_shake_atom(m,1)); - int i2 = map_array(d_shake_atom(m,2)); + int i0 = map_kokkos(d_shake_atom(m,0),map_style,k_map_array,k_map_hash); + int i1 = map_kokkos(d_shake_atom(m,1),map_style,k_map_array,k_map_hash); + 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::shake4(int m, EV_FLOAT& ev) const // local atom IDs and constraint distances - int i0 = map_array(d_shake_atom(m,0)); - int i1 = map_array(d_shake_atom(m,1)); - int i2 = map_array(d_shake_atom(m,2)); - int i3 = map_array(d_shake_atom(m,3)); + int i0 = map_kokkos(d_shake_atom(m,0),map_style,k_map_array,k_map_hash); + int i1 = map_kokkos(d_shake_atom(m,1),map_style,k_map_array,k_map_hash); + int i2 = map_kokkos(d_shake_atom(m,2),map_style,k_map_array,k_map_hash); + 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::shake3angle(int m, EV_FLOAT& ev) const // local atom IDs and constraint distances - int i0 = map_array(d_shake_atom(m,0)); - int i1 = map_array(d_shake_atom(m,1)); - int i2 = map_array(d_shake_atom(m,2)); + int i0 = map_kokkos(d_shake_atom(m,0),map_style,k_map_array,k_map_hash); + int i1 = map_kokkos(d_shake_atom(m,1),map_style,k_map_array,k_map_hash); + 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::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 +KOKKOS_INLINE_FUNCTION +int FixShakeKokkos::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()(global); + else if (map_style == 2) + return AtomKokkos::map_find_hash_kokkos(global,k_map_hash); + else + return -1; +} + +/* ---------------------------------------------------------------------- */ + namespace LAMMPS_NS { template class FixShakeKokkos; #ifdef LMP_KOKKOS_GPU diff --git a/src/KOKKOS/fix_shake_kokkos.h b/src/KOKKOS/fix_shake_kokkos.h index b95c042ac1..fbf6e1de88 100644 --- a/src/KOKKOS/fix_shake_kokkos.h +++ b/src/KOKKOS/fix_shake_kokkos.h @@ -26,6 +26,7 @@ FixStyle(shake/kk/host,FixShakeKokkos); #include "fix_shake.h" #include "kokkos_type.h" #include "kokkos_base.h" +#include 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 hash_type; + typedef Kokkos::DualView 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 diff --git a/src/KOKKOS/neigh_bond_kokkos.cpp b/src/KOKKOS/neigh_bond_kokkos.cpp index eed0026af3..c427d18216 100644 --- a/src/KOKKOS/neigh_bond_kokkos.cpp +++ b/src/KOKKOS/neigh_bond_kokkos.cpp @@ -207,30 +207,7 @@ void NeighBondKokkos::build_topology_kk() lostbond = output->thermo->lostbond; - // don't yet have atom_map_kokkos routines, so move data from host to device - - 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(); - k_map_array.template sync(); - map_array = k_map_array.view(); - - 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(); - k_sametag.template sync(); - sametag = k_sametag.view(); + update_class_variables(); if (force->bond) (this->*bond_build_kk)(); if (force->angle) (this->*angle_build_kk)(); @@ -306,7 +283,7 @@ template KOKKOS_INLINE_FUNCTION void NeighBondKokkos::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::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::bond_check() { int flag = 0; - update_domain_variables(); atomKK->sync(execution_space, X_MASK); k_bondlist.sync(); @@ -507,9 +483,9 @@ template KOKKOS_INLINE_FUNCTION void NeighBondKokkos::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 atom2 = map_array(angle_atom2(i,m)); - int atom3 = map_array(angle_atom3(i,m)); + int atom1 = map_kokkos(angle_atom1(i,m)); + int atom2 = map_kokkos(angle_atom2(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::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 atom2 = map_array(angle_atom2(i,m)); - int atom3 = map_array(angle_atom3(i,m)); + int atom1 = map_kokkos(angle_atom1(i,m)); + int atom2 = map_kokkos(angle_atom2(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::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(); @@ -735,10 +710,10 @@ template KOKKOS_INLINE_FUNCTION void NeighBondKokkos::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 atom2 = map_array(dihedral_atom2(i,m)); - int atom3 = map_array(dihedral_atom3(i,m)); - int atom4 = map_array(dihedral_atom4(i,m)); + int atom1 = map_kokkos(dihedral_atom1(i,m)); + int atom2 = map_kokkos(dihedral_atom2(i,m)); + int atom3 = map_kokkos(dihedral_atom3(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::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 atom2 = map_array(dihedral_atom2(i,m)); - int atom3 = map_array(dihedral_atom3(i,m)); - int atom4 = map_array(dihedral_atom4(i,m)); + int atom1 = map_kokkos(dihedral_atom1(i,m)); + int atom2 = map_kokkos(dihedral_atom2(i,m)); + int atom3 = map_kokkos(dihedral_atom3(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::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(); @@ -990,10 +964,10 @@ template KOKKOS_INLINE_FUNCTION void NeighBondKokkos::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 atom2 = map_array(improper_atom2(i,m)); - int atom3 = map_array(improper_atom3(i,m)); - int atom4 = map_array(improper_atom4(i,m)); + int atom1 = map_kokkos(improper_atom1(i,m)); + int atom2 = map_kokkos(improper_atom2(i,m)); + int atom3 = map_kokkos(improper_atom3(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::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 atom2 = map_array(improper_atom2(i,m)); - int atom3 = map_array(improper_atom3(i,m)); - int atom4 = map_array(improper_atom4(i,m)); + int atom1 = map_kokkos(improper_atom1(i,m)); + int atom2 = map_kokkos(improper_atom2(i,m)); + int atom3 = map_kokkos(improper_atom3(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::closest_image(const int i, int j) const X_FLOAT rsqmin = delx*delx + dely*dely + delz*delz; X_FLOAT rsq; - while (sametag[j] >= 0) { - j = sametag[j]; + while (d_sametag[j] >= 0) { + 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::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 -void NeighBondKokkos::update_domain_variables() +KOKKOS_INLINE_FUNCTION +int NeighBondKokkos::map_kokkos(tagint global) const { + if (map_style == 1) + return k_map_array.view()(global); + else if (map_style == 2) + return AtomKokkos::map_find_hash_kokkos(global,k_map_hash); + else + return -1; +} + +/* ---------------------------------------------------------------------- */ + +template +void NeighBondKokkos::update_class_variables() +{ + // Domain + triclinic = domain->triclinic; xperiodic = domain->xperiodic; xprd_half = domain->xprd_half; @@ -1233,6 +1227,22 @@ void NeighBondKokkos::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(); + d_sametag = k_sametag.view(); + + if (map_style == Atom::MAP_ARRAY) { + k_map_array = atomKK->k_map_array; + k_map_array.template sync(); + } else if (map_style == Atom::MAP_HASH) { + k_map_hash = atomKK->k_map_hash; + k_map_hash.template sync(); + } } /* ---------------------------------------------------------------------- */ diff --git a/src/KOKKOS/neigh_bond_kokkos.h b/src/KOKKOS/neigh_bond_kokkos.h index 887ead95f6..defcf990e3 100644 --- a/src/KOKKOS/neigh_bond_kokkos.h +++ b/src/KOKKOS/neigh_bond_kokkos.h @@ -19,6 +19,7 @@ #include "kokkos_type.h" #include "domain_kokkos.h" #include "pointers.h" +#include namespace LAMMPS_NS { @@ -81,13 +82,19 @@ class NeighBondKokkos : protected Pointers { int me,nprocs; private: - - - DAT::tdual_int_1d k_map_array; - typename AT::t_int_1d_randomread map_array; + int map_style; 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 hash_type; + typedef Kokkos::DualView 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 diff --git a/src/atom.h b/src/atom.h index 392e5c5d5c..e922630706 100644 --- a/src/atom.h +++ b/src/atom.h @@ -364,13 +364,13 @@ class Atom : protected Pointers { return -1; }; - void map_init(int check = 1); - void map_clear(); - void map_set(); - void map_one(tagint, int); + virtual void map_init(int check = 1); + virtual void map_clear(); + virtual void map_set(); + virtual void map_one(tagint, int); int map_style_set(); - void map_delete(); - int map_find_hash(tagint); + virtual void map_delete(); + virtual int map_find_hash(tagint); protected: // global to local ID mapping diff --git a/src/atom_map.cpp b/src/atom_map.cpp index 95f09d2639..2eb4a08ea6 100644 --- a/src/atom_map.cpp +++ b/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; }