ICODE-ADC/lammps
4
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
0e1e8161efd06c147a01afda73d71fd5aef0deb4
lammps/src/KOKKOS/mliap_data_kokkos.cpp

365 lines
16 KiB
C++
Raw Blame History

// clang-format off
/* -*- c++ -*- ----------------------------------------------------------
LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
https://www.lammps.org/, Sandia National Laboratories
LAMMPS Development team: developers@lammps.org
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.
------------------------------------------------------------------------- */
/* ----------------------------------------------------------------------
Contributing author: Matt Bettencourt (NVIDIA)
------------------------------------------------------------------------- */
#include "mliap_data_kokkos.h"
#include "atom_kokkos.h"
#include "kokkos_type.h"
#include "pair_mliap_kokkos.h"
#include "atom_masks.h"
#include "mliap_descriptor.h"
#include "lammps.h"
#include "kokkos.h"
/* ---------------------------------------------------------------------- */
namespace LAMMPS_NS {
template<class DeviceType>
MLIAPDataKokkos<DeviceType>::MLIAPDataKokkos(LAMMPS *lmp_in, int gradgradflag_in, int *map_in,
class MLIAPModel* model_in,
class MLIAPDescriptor* descriptor_in,
class PairMLIAPKokkos<DeviceType>* pairmliap_in) :
MLIAPData(lmp_in, gradgradflag_in, map_in, model_in, descriptor_in, pairmliap_in),
k_pairmliap(pairmliap_in),
lmp(lmp_in)
{
execution_space = ExecutionSpaceFromDevice<DeviceType>::space;
}
/* ---------------------------------------------------------------------- */
template<class DeviceType>
MLIAPDataKokkos<DeviceType>::~MLIAPDataKokkos() {
memoryKK->destroy_kokkos(k_gradforce,gradforce);
memoryKK->destroy_kokkos(k_betas,betas);
memoryKK->destroy_kokkos(k_descriptors,descriptors);
memoryKK->destroy_kokkos(k_eatoms,eatoms);
memoryKK->destroy_kokkos(k_gamma_row_index,gamma_row_index);
memoryKK->destroy_kokkos(k_gamma_col_index,gamma_col_index);
memoryKK->destroy_kokkos(k_gamma,gamma);
memoryKK->destroy_kokkos(k_iatoms,iatoms);
memoryKK->destroy_kokkos(k_ielems,ielems);
memoryKK->destroy_kokkos(k_numneighs,numneighs);
memoryKK->destroy_kokkos(k_jatoms,jatoms);
memoryKK->destroy_kokkos(k_pair_i,pair_i);
memoryKK->destroy_kokkos(k_jelems,jelems);
memoryKK->destroy_kokkos(k_elems,elems);
memoryKK->destroy_kokkos(k_ij);
memoryKK->destroy_kokkos(k_rij,rij);
memoryKK->destroy_kokkos(k_graddesc,graddesc);
}
/* ---------------------------------------------------------------------- */
template<class DeviceType>
void MLIAPDataKokkos<DeviceType>::generate_neighdata(class NeighList *list_in, int eflag_in, int vflag_in) {
list = list_in;
// grow nmax gradforce array if necessary
if (atom->nmax > nmax) {
nmax = atom->nmax;
memoryKK->destroy_kokkos(k_gradforce,gradforce);
memoryKK->create_kokkos(k_gradforce, gradforce, nmax, size_gradforce, "mliap_data:gradforce");
memoryKK->destroy_kokkos(k_elems,elems);
memoryKK->create_kokkos(k_elems, elems, nmax, "mliap_data:elems"); }
// clear gradforce array
int nall = atom->nlocal + atom->nghost;
ntotal = nall;
auto d_gradforce = k_gradforce.template view<DeviceType>();
Kokkos::deep_copy(d_gradforce, 0.);
auto d_elems = k_elems.template view<DeviceType>();
Kokkos::deep_copy(d_elems, 0.);
// grow arrays if necessary
nlistatoms = list->inum;
if (nlistatoms_max < nlistatoms) {
memoryKK->destroy_kokkos(k_betas,betas);
memoryKK->create_kokkos(k_betas, betas, nlistatoms, ndescriptors, "mliap_data:betas");
memoryKK->destroy_kokkos(k_descriptors,descriptors);
memoryKK->create_kokkos(k_descriptors, descriptors, nlistatoms, ndescriptors, "mliap_data:descriptors");
memoryKK->destroy_kokkos(k_eatoms,eatoms);
memoryKK->create_kokkos(k_eatoms, eatoms, nlistatoms, "mliap_data:eatoms");
nlistatoms_max = nlistatoms;
}
// grow gamma arrays if necessary
if (gradgradflag == 1) {
if (natomgamma_max < nlistatoms) {
memoryKK->destroy_kokkos(k_gamma_row_index,gamma_row_index);
memoryKK->create_kokkos(k_gamma_row_index, gamma_row_index, nlistatoms, gamma_nnz, "mliap_data:gamma_row_index");
memoryKK->destroy_kokkos(k_gamma_col_index,gamma_col_index);
memoryKK->create_kokkos(k_gamma_col_index, gamma_col_index, nlistatoms, gamma_nnz, "mliap_data:gamma_col_index");
memoryKK->destroy_kokkos(k_gamma,gamma);
memoryKK->create_kokkos(k_gamma, gamma, nlistatoms, gamma_nnz, "mliap_data:");
natomgamma_max = nlistatoms;
}
}
atomKK->sync(execution_space,X_MASK | V_MASK | F_MASK | TYPE_MASK);
k_pairmliap->x = atomKK->k_x.view<DeviceType>();
k_pairmliap->v = atomKK->k_v.view<DeviceType>();
k_pairmliap->f = atomKK->k_f.view<DeviceType>();
grow_neigharrays();
// Use the ielems memory for prefix scan and set it at the end to the i type
auto d_iatoms = k_iatoms.template view<DeviceType>();
auto d_ielems = k_ielems.template view<DeviceType>();
auto d_ij = k_ij.template view<DeviceType>();
auto d_numneighs = k_numneighs.template view<DeviceType>();
auto d_jatoms = k_jatoms.template view<DeviceType>();
auto d_pair_i= k_pair_i.template view<DeviceType>();
auto d_jelems= k_jelems.template view<DeviceType>();
auto d_rij= k_rij.template view<DeviceType>();
NeighListKokkos<DeviceType>* k_list = static_cast<NeighListKokkos<DeviceType>*>(list);
auto d_numneigh = k_list->d_numneigh;
auto d_neighbors = k_list->d_neighbors;
auto d_ilist = k_list->d_ilist;
auto d_cutsq=k_pairmliap->k_cutsq.template view<DeviceType>();
AtomKokkos *atomKK = (AtomKokkos *) atom;
auto x = atomKK->k_x.view<DeviceType>();
auto type = atomKK->k_type.view<DeviceType>();
auto map=k_pairmliap->k_map.template view<DeviceType>();
Kokkos::parallel_scan(nlistatoms, KOKKOS_LAMBDA (int ii, int &update, const bool final) {
if (final)
d_ij(ii) = update;
update += d_numneighs(ii);
});
Kokkos::parallel_for(nlistatoms, KOKKOS_LAMBDA (int ii) {
int ij = d_ij(ii);
const int i = d_ilist[ii];
const double xtmp = x(i, 0);
const double ytmp = x(i, 1);
const double ztmp = x(i, 2);
const int itype = type(i);
const int ielem = map(itype);
const int jnum = d_numneigh(i);
for (int jj = 0; jj < jnum; jj++) {
int j = d_neighbors(i,jj);
j &= NEIGHMASK;
const double delx = x(j,0) - xtmp;
const double dely = x(j,1) - ytmp;
const double delz = x(j,2) - ztmp;
const double rsq = delx*delx + dely*dely + delz*delz;
int jtype = type(j);
const int jelem = map(jtype);
if (rsq < d_cutsq(itype,jtype)) {
d_jatoms(ij) = j;
d_pair_i(ij) = i;
d_jelems(ij) = jelem;
d_rij(ij, 0) = delx;
d_rij(ij, 1) = dely;
d_rij(ij, 2) = delz;
ij++;
}
}
d_iatoms[ii] = i;
d_ielems[ii] = ielem;
});
Kokkos::parallel_for(nmax, KOKKOS_LAMBDA (int i) {
const int itype = type(i);
d_elems(i) = map(itype);
});
modified(execution_space, NUMNEIGHS_MASK | IATOMS_MASK | IELEMS_MASK | ELEMS_MASK | JATOMS_MASK | PAIR_I_MASK | JELEMS_MASK | RIJ_MASK | IJ_MASK );
eflag = eflag_in;
vflag = vflag_in;
}
/* ---------------------------------------------------------------------- */
template<class DeviceType>
void MLIAPDataKokkos<DeviceType>::grow_neigharrays() {
AtomKokkos *atomKK = (AtomKokkos *) atom;
f = atom->f;
f_device = atomKK->k_f.view<DeviceType>().data();
// grow neighbor arrays if necessary
if (natomneigh_max < nlistatoms) {
natomneigh_max = nlistatoms;
memoryKK->destroy_kokkos(k_iatoms,iatoms);
memoryKK->create_kokkos(k_iatoms, iatoms, natomneigh_max, "mliap_data:iatoms");
memoryKK->destroy_kokkos(k_ielems,ielems);
memoryKK->create_kokkos(k_ielems, ielems, natomneigh_max, "mliap_data:ielems");
memoryKK->destroy_kokkos(k_ij);
memoryKK->create_kokkos(k_ij, natomneigh_max, "mliap_data:ij");
memoryKK->destroy_kokkos(k_numneighs,numneighs);
memoryKK->create_kokkos(k_numneighs, numneighs, natomneigh_max, "mliap_data:numneighs");
}
NeighListKokkos<DeviceType>* k_list = static_cast<NeighListKokkos<DeviceType>*>(list);
auto d_numneigh = k_list->d_numneigh;
auto d_neighbors = k_list->d_neighbors;
auto d_ilist = k_list->d_ilist;
auto x = atomKK->k_x.view<DeviceType>();
auto type = atomKK->k_type.view<DeviceType>();
auto d_cutsq=k_pairmliap->k_cutsq.template view<DeviceType>();
auto h_cutsq=k_pairmliap->k_cutsq.template view<LMPHostType>();
auto d_numneighs = k_numneighs.template view<DeviceType>();
Kokkos::parallel_reduce(nlistatoms, KOKKOS_LAMBDA (int ii, int &contrib) {
const int i = d_ilist[ii];
int count=0;
const double xtmp = x(i, 0);
const double ytmp = x(i, 1);
const double ztmp = x(i, 2);
const int itype = type(i);
const int jnum = d_numneigh(i);
for (int jj = 0; jj < jnum; jj++) {
int j = d_neighbors(i,jj);
j &= NEIGHMASK;
const double delx = x(j,0) - xtmp;
const double dely = x(j,1) - ytmp;
const double delz = x(j,2) - ztmp;
const double rsq = delx*delx + dely*dely + delz*delz;
int jtype = type(j);
if (rsq < d_cutsq(itype,jtype))
count++;
}
d_numneighs(ii) = count;
contrib += count;
}, npairs);
modified(execution_space, NUMNEIGHS_MASK);
if (nneigh_max < npairs) {
memoryKK->destroy_kokkos(k_jatoms,jatoms);
memoryKK->create_kokkos(k_jatoms, jatoms, npairs, "mliap_data:jatoms");
memoryKK->destroy_kokkos(k_pair_i,pair_i);
memoryKK->create_kokkos(k_pair_i, pair_i, npairs, "mliap_data:pair_i");
memoryKK->destroy_kokkos(k_jelems,jelems);
memoryKK->create_kokkos(k_jelems, jelems, npairs, "mliap_data:jelems");
memoryKK->destroy_kokkos(k_rij,rij);
memoryKK->create_kokkos(k_rij, rij, npairs, 3, "mliap_data:rij");
if (gradgradflag == 0){
memoryKK->destroy_kokkos(k_graddesc,graddesc);
memoryKK->create_kokkos(k_graddesc, graddesc, npairs, ndescriptors,3, "mliap_data:graddesc");
}
nneigh_max = npairs;
}
}
/* ---------------------------------------------------------------------- */
template<class DeviceType>
void MLIAPDataKokkos<DeviceType>::modified(ExecutionSpace space, unsigned int mask, bool ignore_auto_sync) {
if (space == Device) {
if (mask & IATOMS_MASK ) k_iatoms .modify<LMPDeviceType>();
if (mask & IELEMS_MASK ) k_ielems .modify<LMPDeviceType>();
if (mask & JATOMS_MASK ) k_jatoms .modify<LMPDeviceType>();
if (mask & PAIR_I_MASK ) k_pair_i .modify<LMPDeviceType>();
if (mask & JELEMS_MASK ) k_jelems .modify<LMPDeviceType>();
if (mask & ELEMS_MASK ) k_elems .modify<LMPDeviceType>();
if (mask & IJ_MASK ) k_ij .modify<LMPDeviceType>();
if (mask & BETAS_MASK ) k_betas .modify<LMPDeviceType>();
if (mask & DESCRIPTORS_MASK ) k_descriptors .modify<LMPDeviceType>();
if (mask & EATOMS_MASK ) k_eatoms .modify<LMPDeviceType>();
if (mask & RIJ_MASK ) k_rij .modify<LMPDeviceType>();
if (mask & GRADFORCE_MASK ) k_gradforce .modify<LMPDeviceType>();
if (mask & GRADDESC_MASK ) k_graddesc .modify<LMPDeviceType>();
if (mask & NUMNEIGHS_MASK ) k_numneighs .modify<LMPDeviceType>();
if (mask & GAMMA_MASK_MASK ) k_gamma .modify<LMPDeviceType>();
if (mask & GAMMA_ROW_MASK ) k_gamma_row_index.modify<LMPDeviceType>();
if (mask & GAMMA_COL_MASK ) k_gamma_col_index.modify<LMPDeviceType>();
if (lmp->kokkos->auto_sync && !ignore_auto_sync) sync(Host, mask, true);
} else {
if (mask & IATOMS_MASK ) k_iatoms .modify<LMPHostType>();
if (mask & IELEMS_MASK ) k_ielems .modify<LMPHostType>();
if (mask & JATOMS_MASK ) k_jatoms .modify<LMPHostType>();
if (mask & PAIR_I_MASK ) k_pair_i .modify<LMPHostType>();
if (mask & JELEMS_MASK ) k_jelems .modify<LMPHostType>();
if (mask & ELEMS_MASK ) k_elems .modify<LMPHostType>();
if (mask & IJ_MASK ) k_ij .modify<LMPHostType>();
if (mask & BETAS_MASK ) k_betas .modify<LMPHostType>();
if (mask & DESCRIPTORS_MASK ) k_descriptors .modify<LMPHostType>();
if (mask & EATOMS_MASK ) k_eatoms .modify<LMPHostType>();
if (mask & RIJ_MASK ) k_rij .modify<LMPHostType>();
if (mask & GRADFORCE_MASK ) k_gradforce .modify<LMPHostType>();
if (mask & GRADDESC_MASK ) k_graddesc .modify<LMPHostType>();
if (mask & NUMNEIGHS_MASK ) k_numneighs .modify<LMPHostType>();
if (mask & GAMMA_MASK_MASK ) k_gamma .modify<LMPHostType>();
if (mask & GAMMA_ROW_MASK ) k_gamma_row_index.modify<LMPHostType>();
if (mask & GAMMA_COL_MASK ) k_gamma_col_index.modify<LMPHostType>();
if (lmp->kokkos->auto_sync && !ignore_auto_sync) sync(Device, mask, true);
}
}
/* ---------------------------------------------------------------------- */
template<class DeviceType>
void MLIAPDataKokkos<DeviceType>::sync(ExecutionSpace space, unsigned int mask, bool ignore_auto_sync) {
if (space == Device) {
if (lmp->kokkos->auto_sync && !ignore_auto_sync) modified(Host, mask, true);
if (mask & IATOMS_MASK ) k_iatoms .sync<LMPDeviceType>();
if (mask & IELEMS_MASK ) k_ielems .sync<LMPDeviceType>();
if (mask & JATOMS_MASK ) k_jatoms .sync<LMPDeviceType>();
if (mask & PAIR_I_MASK ) k_pair_i .sync<LMPDeviceType>();
if (mask & JELEMS_MASK ) k_jelems .sync<LMPDeviceType>();
if (mask & ELEMS_MASK ) k_elems .sync<LMPDeviceType>();
if (mask & IJ_MASK ) k_ij .sync<LMPDeviceType>();
if (mask & BETAS_MASK ) k_betas .sync<LMPDeviceType>();
if (mask & DESCRIPTORS_MASK ) k_descriptors .sync<LMPDeviceType>();
if (mask & EATOMS_MASK ) k_eatoms .sync<LMPDeviceType>();
if (mask & RIJ_MASK ) k_rij .sync<LMPDeviceType>();
if (mask & GRADFORCE_MASK ) k_gradforce .sync<LMPDeviceType>();
if (mask & GRADDESC_MASK ) k_graddesc .sync<LMPDeviceType>();
if (mask & NUMNEIGHS_MASK ) k_numneighs .sync<LMPDeviceType>();
if (mask & GAMMA_MASK_MASK ) k_gamma .sync<LMPDeviceType>();
if (mask & GAMMA_ROW_MASK ) k_gamma_row_index.sync<LMPDeviceType>();
if (mask & GAMMA_COL_MASK ) k_gamma_col_index.sync<LMPDeviceType>();
} else {
if (lmp->kokkos->auto_sync && !ignore_auto_sync) modified(Device, mask, true);
if (mask & IATOMS_MASK ) k_iatoms .sync<LMPHostType>();
if (mask & IELEMS_MASK ) k_ielems .sync<LMPHostType>();
if (mask & JATOMS_MASK ) k_jatoms .sync<LMPHostType>();
if (mask & PAIR_I_MASK ) k_pair_i .sync<LMPHostType>();
if (mask & JELEMS_MASK ) k_jelems .sync<LMPHostType>();
if (mask & ELEMS_MASK ) k_elems .sync<LMPHostType>();
if (mask & IJ_MASK ) k_ij .sync<LMPHostType>();
if (mask & BETAS_MASK ) k_betas .sync<LMPHostType>();
if (mask & DESCRIPTORS_MASK ) k_descriptors .sync<LMPHostType>();
if (mask & EATOMS_MASK ) k_eatoms .sync<LMPHostType>();
if (mask & RIJ_MASK ) k_rij .sync<LMPHostType>();
if (mask & GRADFORCE_MASK ) k_gradforce .sync<LMPHostType>();
if (mask & GRADDESC_MASK ) k_graddesc .sync<LMPHostType>();
if (mask & NUMNEIGHS_MASK ) k_numneighs .sync<LMPHostType>();
if (mask & GAMMA_MASK_MASK ) k_gamma .sync<LMPHostType>();
if (mask & GAMMA_ROW_MASK ) k_gamma_row_index.sync<LMPHostType>();
if (mask & GAMMA_COL_MASK ) k_gamma_col_index.sync<LMPHostType>();
}
}
/* ---------------------------------------------------------------------- */
template class MLIAPDataKokkos<LMPDeviceType>;
#ifdef LMP_KOKKOS_GPU
template class MLIAPDataKokkos<LMPHostType>;
#endif
}// namespace
Reference in New Issue View Git Blame Copy Permalink