335 lines
9.3 KiB
C++
335 lines
9.3 KiB
C++
// clang-format off
|
|
/* ----------------------------------------------------------------------
|
|
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: Trung Nguyen (U Chicago)
|
|
------------------------------------------------------------------------- */
|
|
|
|
#include "fix_spring_self_kokkos.h"
|
|
|
|
#include "atom_kokkos.h"
|
|
#include "update.h"
|
|
#include "domain_kokkos.h"
|
|
#include "memory_kokkos.h"
|
|
#include "error.h"
|
|
#include "atom_masks.h"
|
|
|
|
using namespace LAMMPS_NS;
|
|
using namespace FixConst;
|
|
|
|
/* ---------------------------------------------------------------------- */
|
|
|
|
template<class DeviceType>
|
|
FixSpringSelfKokkos<DeviceType>::FixSpringSelfKokkos(LAMMPS *lmp, int narg, char **arg) :
|
|
FixSpringSelf(lmp, narg, arg)
|
|
{
|
|
kokkosable = 1;
|
|
exchange_comm_device = 1;
|
|
atomKK = (AtomKokkos *) atom;
|
|
execution_space = ExecutionSpaceFromDevice<DeviceType>::space;
|
|
datamask_read = EMPTY_MASK;
|
|
datamask_modify = EMPTY_MASK;
|
|
|
|
xoriginal_tmp = xoriginal;
|
|
xoriginal = nullptr;
|
|
|
|
int nmax = atom->nmax;
|
|
grow_arrays(nmax);
|
|
|
|
for (int i = 0; i < atom->nlocal; i++) {
|
|
k_xoriginal.h_view(i,0) = xoriginal_tmp[i][0];
|
|
k_xoriginal.h_view(i,1) = xoriginal_tmp[i][1];
|
|
k_xoriginal.h_view(i,2) = xoriginal_tmp[i][2];
|
|
}
|
|
|
|
k_xoriginal.modify_host();
|
|
|
|
d_count = typename AT::t_int_scalar("spring/self:count");
|
|
h_count = Kokkos::create_mirror_view(d_count);
|
|
|
|
memory->destroy(xoriginal_tmp);
|
|
}
|
|
|
|
/* ---------------------------------------------------------------------- */
|
|
|
|
template<class DeviceType>
|
|
FixSpringSelfKokkos<DeviceType>::~FixSpringSelfKokkos()
|
|
{
|
|
if (copymode) return;
|
|
|
|
memoryKK->destroy_kokkos(k_xoriginal,xoriginal);
|
|
xoriginal = nullptr;
|
|
}
|
|
|
|
/* ---------------------------------------------------------------------- */
|
|
|
|
template<class DeviceType>
|
|
void FixSpringSelfKokkos<DeviceType>::init()
|
|
{
|
|
FixSpringSelf::init();
|
|
|
|
if (kstyle != CONSTANT)
|
|
error->all(FLERR, "Fix spring/self/kk does not support variable spring constants (yet)");
|
|
|
|
if (utils::strmatch(update->integrate_style,"^respa"))
|
|
error->all(FLERR,"Cannot (yet) use respa with Kokkos");
|
|
}
|
|
|
|
/* ---------------------------------------------------------------------- */
|
|
|
|
template<class DeviceType>
|
|
void FixSpringSelfKokkos<DeviceType>::post_force(int /*vflag*/)
|
|
{
|
|
atomKK->sync(execution_space, X_MASK | F_MASK | IMAGE_MASK | MASK_MASK);
|
|
|
|
x = atomKK->k_x.view<DeviceType>();
|
|
f = atomKK->k_f.view<DeviceType>();
|
|
image = atomKK->k_image.view<DeviceType>();
|
|
mask = atomKK->k_mask.view<DeviceType>();
|
|
int nlocal = atom->nlocal;
|
|
|
|
double espring_kk;
|
|
|
|
k_xoriginal.modify<LMPHostType>();
|
|
k_xoriginal.sync<DeviceType>();
|
|
|
|
copymode = 1;
|
|
|
|
{
|
|
// local variables for lambda capture
|
|
auto prd = Few<double,3>(domain->prd);
|
|
auto h = Few<double,6>(domain->h);
|
|
auto triclinic = domain->triclinic;
|
|
auto l_k = k;
|
|
auto l_xoriginal = d_xoriginal;
|
|
|
|
auto l_x = x;
|
|
auto l_f = f;
|
|
auto l_mask = mask;
|
|
auto l_image = image;
|
|
auto l_groupbit = groupbit;
|
|
auto l_xflag = xflag;
|
|
auto l_yflag = yflag;
|
|
auto l_zflag = zflag;
|
|
|
|
Kokkos::parallel_reduce(nlocal, LAMMPS_LAMBDA(const int& i, double& espring_kk) {
|
|
if (l_mask[i] & l_groupbit) {
|
|
Few<double,3> x_i;
|
|
x_i[0] = l_x(i,0);
|
|
x_i[1] = l_x(i,1);
|
|
x_i[2] = l_x(i,2);
|
|
auto unwrap = DomainKokkos::unmap(prd,h,triclinic,x_i,l_image(i));
|
|
auto dx = unwrap[0] - l_xoriginal(i, 0);
|
|
auto dy = unwrap[1] - l_xoriginal(i, 1);
|
|
auto dz = unwrap[2] - l_xoriginal(i, 2);
|
|
if (!l_xflag) dx = 0.0;
|
|
if (!l_yflag) dy = 0.0;
|
|
if (!l_zflag) dz = 0.0;
|
|
l_f(i,0) -= l_k*dx;
|
|
l_f(i,1) -= l_k*dy;
|
|
l_f(i,2) -= l_k*dz;
|
|
espring_kk += l_k * (dx*dx + dy*dy + dz*dz);
|
|
}
|
|
},espring_kk);
|
|
}
|
|
|
|
copymode = 0;
|
|
|
|
atomKK->modified(execution_space, F_MASK);
|
|
|
|
espring = 0.5*espring_kk;
|
|
}
|
|
|
|
/* ----------------------------------------------------------------------
|
|
allocate local atom-based arrays
|
|
------------------------------------------------------------------------- */
|
|
|
|
template<class DeviceType>
|
|
void FixSpringSelfKokkos<DeviceType>::grow_arrays(int nmax)
|
|
{
|
|
memoryKK->grow_kokkos(k_xoriginal,xoriginal,nmax,"spring/self:xoriginal");
|
|
d_xoriginal = k_xoriginal.view<DeviceType>();
|
|
}
|
|
|
|
/* ----------------------------------------------------------------------
|
|
copy values within local atom-based arrays
|
|
------------------------------------------------------------------------- */
|
|
|
|
template<class DeviceType>
|
|
void FixSpringSelfKokkos<DeviceType>::copy_arrays(int i, int j, int delflag)
|
|
{
|
|
k_xoriginal.sync_host();
|
|
|
|
FixSpringSelf::copy_arrays(i,j,delflag);
|
|
|
|
k_xoriginal.modify_host();
|
|
}
|
|
|
|
/* ---------------------------------------------------------------------- */
|
|
|
|
template<class DeviceType>
|
|
KOKKOS_INLINE_FUNCTION
|
|
void FixSpringSelfKokkos<DeviceType>::pack_exchange_item(const int &mysend, int &offset, const bool &/*final*/) const
|
|
{
|
|
const int i = d_exchange_sendlist(mysend);
|
|
|
|
int m = nsend + offset;
|
|
d_buf[mysend] = m;
|
|
d_buf[m++] = d_xoriginal(i,0);
|
|
d_buf[m++] = d_xoriginal(i,1);
|
|
d_buf[m++] = d_xoriginal(i,2);
|
|
if (mysend == nsend-1) d_count() = m;
|
|
offset = m - nsend;
|
|
|
|
const int j = d_copylist(mysend);
|
|
if (j > -1) {
|
|
d_xoriginal(i,0) = d_xoriginal(j,0);
|
|
d_xoriginal(i,1) = d_xoriginal(j,1);
|
|
d_xoriginal(i,2) = d_xoriginal(j,2);
|
|
}
|
|
}
|
|
|
|
/* ---------------------------------------------------------------------- */
|
|
|
|
template<class DeviceType>
|
|
int FixSpringSelfKokkos<DeviceType>::pack_exchange_kokkos(
|
|
const int &nsend, DAT::tdual_xfloat_2d &k_buf,
|
|
DAT::tdual_int_1d k_exchange_sendlist, DAT::tdual_int_1d k_copylist,
|
|
ExecutionSpace space)
|
|
{
|
|
|
|
k_buf.sync<DeviceType>();
|
|
k_copylist.sync<DeviceType>();
|
|
k_exchange_sendlist.sync<DeviceType>();
|
|
|
|
d_buf = typename ArrayTypes<DeviceType>::t_xfloat_1d_um(
|
|
k_buf.template view<DeviceType>().data(),
|
|
k_buf.extent(0)*k_buf.extent(1));
|
|
d_copylist = k_copylist.view<DeviceType>();
|
|
d_exchange_sendlist = k_exchange_sendlist.view<DeviceType>();
|
|
this->nsend = nsend;
|
|
|
|
|
|
k_xoriginal.template sync<DeviceType>();
|
|
|
|
Kokkos::deep_copy(d_count,0);
|
|
|
|
copymode = 1;
|
|
|
|
FixSpringSelfKokkosPackExchangeFunctor<DeviceType> pack_exchange_functor(this);
|
|
Kokkos::parallel_scan(nsend,pack_exchange_functor);
|
|
|
|
copymode = 0;
|
|
|
|
k_buf.modify<DeviceType>();
|
|
|
|
if (space == Host) k_buf.sync<LMPHostType>();
|
|
else k_buf.sync<LMPDeviceType>();
|
|
|
|
k_xoriginal.template modify<DeviceType>();
|
|
|
|
Kokkos::deep_copy(h_count,d_count);
|
|
|
|
return h_count();
|
|
}
|
|
|
|
/* ---------------------------------------------------------------------- */
|
|
|
|
template<class DeviceType>
|
|
KOKKOS_INLINE_FUNCTION
|
|
void FixSpringSelfKokkos<DeviceType>::operator()(TagFixSpringSelfUnpackExchange, const int &i) const
|
|
{
|
|
int index = d_indices(i);
|
|
|
|
if (index > -1) {
|
|
int m = d_buf[i];
|
|
if (i >= nrecv1)
|
|
m = nextrarecv1 + d_buf[nextrarecv1 + i - nrecv1];
|
|
|
|
d_xoriginal(index,0) = static_cast<tagint> (d_buf[m++]);
|
|
d_xoriginal(index,1) = static_cast<tagint> (d_buf[m++]);
|
|
d_xoriginal(index,2) = static_cast<tagint> (d_buf[m++]);
|
|
}
|
|
}
|
|
|
|
/* ---------------------------------------------------------------------- */
|
|
|
|
template <class DeviceType>
|
|
void FixSpringSelfKokkos<DeviceType>::unpack_exchange_kokkos(
|
|
DAT::tdual_xfloat_2d &k_buf, DAT::tdual_int_1d &k_indices, int nrecv,
|
|
int nrecv1, int nextrarecv1,
|
|
ExecutionSpace /*space*/)
|
|
{
|
|
k_buf.sync<DeviceType>();
|
|
k_indices.sync<DeviceType>();
|
|
|
|
d_buf = typename ArrayTypes<DeviceType>::t_xfloat_1d_um(
|
|
k_buf.template view<DeviceType>().data(),
|
|
k_buf.extent(0)*k_buf.extent(1));
|
|
d_indices = k_indices.view<DeviceType>();
|
|
|
|
this->nrecv1 = nrecv1;
|
|
this->nextrarecv1 = nextrarecv1;
|
|
|
|
k_xoriginal.template sync<DeviceType>();
|
|
|
|
copymode = 1;
|
|
|
|
Kokkos::parallel_for(Kokkos::RangePolicy<DeviceType,TagFixSpringSelfUnpackExchange>(0,nrecv),*this);
|
|
|
|
copymode = 0;
|
|
|
|
k_xoriginal.template modify<DeviceType>();
|
|
}
|
|
|
|
/* ----------------------------------------------------------------------
|
|
pack values in local atom-based arrays for exchange with another proc
|
|
------------------------------------------------------------------------- */
|
|
|
|
template<class DeviceType>
|
|
int FixSpringSelfKokkos<DeviceType>::pack_exchange(int i, double *buf)
|
|
{
|
|
k_xoriginal.sync_host();
|
|
|
|
int m = FixSpringSelf::pack_exchange(i,buf);
|
|
|
|
k_xoriginal.modify_host();
|
|
|
|
return m;
|
|
}
|
|
|
|
/* ----------------------------------------------------------------------
|
|
unpack values in local atom-based arrays from exchange with another proc
|
|
------------------------------------------------------------------------- */
|
|
|
|
template<class DeviceType>
|
|
int FixSpringSelfKokkos<DeviceType>::unpack_exchange(int nlocal, double *buf)
|
|
{
|
|
k_xoriginal.sync_host();
|
|
|
|
int m = FixSpringSelf::unpack_exchange(nlocal,buf);
|
|
|
|
k_xoriginal.modify_host();
|
|
|
|
return m;
|
|
}
|
|
|
|
namespace LAMMPS_NS {
|
|
template class FixSpringSelfKokkos<LMPDeviceType>;
|
|
#ifdef LMP_KOKKOS_GPU
|
|
template class FixSpringSelfKokkos<LMPHostType>;
|
|
#endif
|
|
}
|
|
|