ICODE-ADC/lammps
4
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add atom_map hash option for Kokkos package

Browse Source
This commit is contained in:
Stan Gerald Moore
2021-06-28 08:38:31 -06:00
parent ddc596170e
commit a4c2bc13cf
11 changed files with 487 additions and 107 deletions
Download Patch File Download Diff File Expand all files Collapse all files

104
src/KOKKOS/fix_shake_kokkos.cpp
View File

@ -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)
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<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>();
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>();
}
// 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 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<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 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<DeviceType>::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<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 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<DeviceType>::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<DeviceType>::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<DeviceType>::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<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