Added Kokkos-enabled version of compute temp/deform.
This commit is contained in:
Emily Kahl
285
src/KOKKOS/compute_temp_deform_kokkos.cpp
Normal file
285
src/KOKKOS/compute_temp_deform_kokkos.cpp
Normal file
@ -0,0 +1,285 @@
|
||||
// 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.
|
||||
------------------------------------------------------------------------- */
|
||||
|
||||
/* ----------------------------------------------------------------------
|
||||
Contributing authors: Emily Kahl (UQ)
|
||||
------------------------------------------------------------------------- */
|
||||
|
||||
#include "compute_temp_deform_kokkos.h"
|
||||
|
||||
#include "atom_kokkos.h"
|
||||
#include "atom_masks.h"
|
||||
#include "comm.h"
|
||||
#include "error.h"
|
||||
#include "force.h"
|
||||
#include "update.h"
|
||||
#include "memory_kokkos.h"
|
||||
|
||||
#include <cstring>
|
||||
|
||||
using namespace LAMMPS_NS;
|
||||
|
||||
/* ---------------------------------------------------------------------- */
|
||||
|
||||
template<class DeviceType>
|
||||
ComputeTempDeformKokkos<DeviceType>::ComputeTempDeformKokkos(LAMMPS *lmp, int narg, char **arg) :
|
||||
ComputeTempDeform(lmp, narg, arg)
|
||||
{
|
||||
kokkosable = 1;
|
||||
atomKK = (AtomKokkos *) atom;
|
||||
domainKK = (DomainKokkos *) domain;
|
||||
execution_space = ExecutionSpaceFromDevice<DeviceType>::space;
|
||||
|
||||
datamask_read = V_MASK | MASK_MASK | RMASS_MASK | TYPE_MASK;
|
||||
datamask_modify = EMPTY_MASK;
|
||||
|
||||
maxbias = 0;
|
||||
}
|
||||
|
||||
template<class DeviceType>
|
||||
ComputeTempDeformKokkos<DeviceType>::~ComputeTempDeformKokkos()
|
||||
{
|
||||
|
||||
|
||||
}
|
||||
|
||||
/* ---------------------------------------------------------------------- */
|
||||
|
||||
template<class DeviceType>
|
||||
double ComputeTempDeformKokkos<DeviceType>::compute_scalar()
|
||||
{
|
||||
atomKK->sync(execution_space,datamask_read);
|
||||
atomKK->k_mass.sync<DeviceType>();
|
||||
|
||||
invoked_scalar = update->ntimestep;
|
||||
|
||||
v = atomKK->k_v.view<DeviceType>();
|
||||
x = atomKK->k_x.view<DeviceType>();
|
||||
if (atomKK->rmass)
|
||||
rmass = atomKK->k_rmass.view<DeviceType>();
|
||||
else
|
||||
mass = atomKK->k_mass.view<DeviceType>();
|
||||
type = atomKK->k_type.view<DeviceType>();
|
||||
mask = atomKK->k_mask.view<DeviceType>();
|
||||
int nlocal = atom->nlocal;
|
||||
|
||||
double t = 0.0;
|
||||
CTEMP t_kk;
|
||||
|
||||
/**************** EVK ****************/
|
||||
// Convert from box coords to lamda coords
|
||||
domainKK->x2lamda(nlocal);
|
||||
|
||||
copymode = 1;
|
||||
if (atomKK->rmass)
|
||||
Kokkos::parallel_reduce(Kokkos::RangePolicy<DeviceType, TagComputeTempDeformScalar<1> >(0,nlocal),*this,t_kk);
|
||||
else
|
||||
Kokkos::parallel_reduce(Kokkos::RangePolicy<DeviceType, TagComputeTempDeformScalar<0> >(0,nlocal),*this,t_kk);
|
||||
copymode = 0;
|
||||
|
||||
/**************** EVK ****************/
|
||||
// Convert back to box coords
|
||||
domainKK->lamda2x(nlocal);
|
||||
|
||||
t = t_kk.t0; // could make this more efficient
|
||||
|
||||
MPI_Allreduce(&t,&scalar,1,MPI_DOUBLE,MPI_SUM,world);
|
||||
if (dynamic) dof_compute();
|
||||
if (dof < 0.0 && natoms_temp > 0.0)
|
||||
error->all(FLERR,"Temperature compute degrees of freedom < 0");
|
||||
scalar *= tfactor;
|
||||
|
||||
return scalar;
|
||||
}
|
||||
|
||||
template<class DeviceType>
|
||||
template<int RMASS>
|
||||
KOKKOS_INLINE_FUNCTION
|
||||
void ComputeTempDeformKokkos<DeviceType>::operator()(TagComputeTempDeformScalar<RMASS>, const int &i, CTEMP& t_kk) const {
|
||||
|
||||
double *h_rate = domainKK->h_rate;
|
||||
double *h_ratelo = domainKK->h_ratelo;
|
||||
double vstream[3],vthermal[3];
|
||||
|
||||
vstream[0] = h_rate[0]*x(i,0) + h_rate[5]*x(i,1) + h_rate[4]*x(i,2) + h_ratelo[0];
|
||||
vstream[1] = h_rate[1]*x(i,1) + h_rate[3]*x(i,2) + h_ratelo[1];
|
||||
vstream[2] = h_rate[2]*x(i,2) + h_ratelo[2];
|
||||
vthermal[0] = v(i,0) - vstream[0];
|
||||
vthermal[1] = v(i,1) - vstream[1];
|
||||
vthermal[2] = v(i,2) - vstream[2];
|
||||
if (RMASS) {
|
||||
if (mask[i] & groupbit)
|
||||
t_kk.t0 += (vthermal[0]*vthermal[0] + vthermal[1]*vthermal[1] + vthermal[2]*vthermal[2]) * rmass[i];
|
||||
} else {
|
||||
if (mask[i] & groupbit)
|
||||
t_kk.t0 += (vthermal[0]*vthermal[0] + vthermal[1]*vthermal[1] + vthermal[2]*vthermal[2]) * mass[type[i]];
|
||||
}
|
||||
}
|
||||
|
||||
/* ---------------------------------------------------------------------- */
|
||||
template<class DeviceType>
|
||||
void ComputeTempDeformKokkos<DeviceType>::compute_vector()
|
||||
{
|
||||
atomKK->sync(execution_space,datamask_read);
|
||||
|
||||
int i;
|
||||
|
||||
invoked_vector = update->ntimestep;
|
||||
|
||||
v = atomKK->k_v.view<DeviceType>();
|
||||
x = atomKK->k_x.view<DeviceType>();
|
||||
if (atomKK->rmass)
|
||||
rmass = atomKK->k_rmass.view<DeviceType>();
|
||||
else
|
||||
mass = atomKK->k_mass.view<DeviceType>();
|
||||
type = atomKK->k_type.view<DeviceType>();
|
||||
mask = atomKK->k_mask.view<DeviceType>();
|
||||
int nlocal = atom->nlocal;
|
||||
|
||||
double t[6];
|
||||
for (i = 0; i < 6; i++) t[i] = 0.0;
|
||||
CTEMP t_kk;
|
||||
|
||||
/**************** EVK ****************/
|
||||
// Convert from box coords to lamda coords
|
||||
domainKK->x2lamda(nlocal);
|
||||
|
||||
copymode = 1;
|
||||
if (atomKK->rmass)
|
||||
Kokkos::parallel_reduce(Kokkos::RangePolicy<DeviceType, TagComputeTempDeformVector<1> >(0,nlocal),*this,t_kk);
|
||||
else
|
||||
Kokkos::parallel_reduce(Kokkos::RangePolicy<DeviceType, TagComputeTempDeformVector<0> >(0,nlocal),*this,t_kk);
|
||||
copymode = 0;
|
||||
|
||||
/**************** EVK ****************/
|
||||
// Convert back to box coords
|
||||
domainKK->lamda2x(nlocal);
|
||||
|
||||
t[0] = t_kk.t0;
|
||||
t[1] = t_kk.t1;
|
||||
t[2] = t_kk.t2;
|
||||
t[3] = t_kk.t3;
|
||||
t[4] = t_kk.t4;
|
||||
t[5] = t_kk.t5;
|
||||
|
||||
MPI_Allreduce(t,vector,6,MPI_DOUBLE,MPI_SUM,world);
|
||||
for (i = 0; i < 6; i++) vector[i] *= force->mvv2e;
|
||||
}
|
||||
|
||||
template<class DeviceType>
|
||||
template<int RMASS>
|
||||
KOKKOS_INLINE_FUNCTION
|
||||
void ComputeTempDeformKokkos<DeviceType>::operator()(TagComputeTempDeformVector<RMASS>, const int &i, CTEMP& t_kk) const {
|
||||
|
||||
double *h_rate = domainKK->h_rate;
|
||||
double *h_ratelo = domainKK->h_ratelo;
|
||||
double vstream[3],vthermal[3];
|
||||
|
||||
vstream[0] = h_rate[0]*x(i,0) + h_rate[5]*x(i,1) + h_rate[4]*x(i,2) + h_ratelo[0];
|
||||
vstream[1] = h_rate[1]*x(i,1) + h_rate[3]*x(i,2) + h_ratelo[1];
|
||||
vstream[2] = h_rate[2]*x(i,2) + h_ratelo[2];
|
||||
vthermal[0] = v(i,0) - vstream[0];
|
||||
vthermal[1] = v(i,1) - vstream[1];
|
||||
vthermal[2] = v(i,2) - vstream[2];
|
||||
|
||||
if (mask[i] & groupbit) {
|
||||
F_FLOAT massone = 0.0;
|
||||
if (RMASS) massone = rmass[i];
|
||||
else massone = mass[type[i]];
|
||||
t_kk.t0 += massone * vthermal[0]*vthermal[0];
|
||||
t_kk.t1 += massone * vthermal[1]*vthermal[1];
|
||||
t_kk.t2 += massone * vthermal[2]*vthermal[2];
|
||||
t_kk.t3 += massone * vthermal[0]*vthermal[1];
|
||||
t_kk.t4 += massone * vthermal[0]*vthermal[2];
|
||||
t_kk.t5 += massone * vthermal[1]*vthermal[2];
|
||||
}
|
||||
}
|
||||
|
||||
/* ---------------------------------------------------------------------- */
|
||||
template<class DeviceType>
|
||||
void ComputeTempDeformKokkos<DeviceType>::remove_bias_all()
|
||||
{
|
||||
atomKK->sync(execution_space,datamask_read);
|
||||
v = atomKK->k_v.view<DeviceType>();
|
||||
x = atomKK->k_x.view<DeviceType>();
|
||||
mask = atomKK->k_mask.view<DeviceType>();
|
||||
int nlocal = atom->nlocal;
|
||||
|
||||
if (atom->nmax > maxbias) {
|
||||
//memoryKK->destroy_kokkos(vbiasall);
|
||||
maxbias = atom->nmax;
|
||||
//memoryKK->create_kokkos(vbiasall,maxbias,"temp/deform/kk:vbiasall");
|
||||
vbiasall = typename ArrayTypes<DeviceType>::t_v_array("temp/deform/kk:vbiasall", maxbias);
|
||||
}
|
||||
|
||||
/**************** EVK ****************/
|
||||
// Convert from box coords to lamda coords
|
||||
domainKK->x2lamda(nlocal);
|
||||
|
||||
h_rate = domain->h_rate;
|
||||
h_ratelo = domain->h_ratelo;
|
||||
|
||||
copymode = 1;
|
||||
Kokkos::parallel_for(Kokkos::RangePolicy<DeviceType, TagComputeTempDeformRemoveBias >(0,nlocal),*this);
|
||||
copymode = 0;
|
||||
|
||||
/**************** EVK ****************/
|
||||
// Convert back to box coords
|
||||
domainKK->lamda2x(nlocal);
|
||||
}
|
||||
|
||||
template<class DeviceType>
|
||||
KOKKOS_INLINE_FUNCTION
|
||||
void ComputeTempDeformKokkos<DeviceType>::operator()(TagComputeTempDeformRemoveBias, const int &i) const {
|
||||
if (mask[i] & groupbit) {
|
||||
vbiasall(i,0) = h_rate[0]*x(i,0) + h_rate[5]*x(i,1) + h_rate[4]*x(i,2) + h_ratelo[0];
|
||||
vbiasall(i,1) = h_rate[1]*x(i,1) + h_rate[3]*x(i,2) + h_ratelo[1];
|
||||
vbiasall(i,2) = h_rate[2]*x(i,2) + h_ratelo[2];
|
||||
v(i,0) -= vbiasall(i,0);
|
||||
v(i,1) -= vbiasall(i,1);
|
||||
v(i,2) -= vbiasall(i,2);
|
||||
}
|
||||
}
|
||||
|
||||
/* ---------------------------------------------------------------------- */
|
||||
template<class DeviceType>
|
||||
void ComputeTempDeformKokkos<DeviceType>::restore_bias_all()
|
||||
{
|
||||
atomKK->sync(execution_space,datamask_read);
|
||||
v = atomKK->k_v.view<DeviceType>();
|
||||
x = atomKK->k_x.view<DeviceType>();
|
||||
mask = atomKK->k_mask.view<DeviceType>();
|
||||
int nlocal = atom->nlocal;
|
||||
|
||||
copymode = 1;
|
||||
Kokkos::parallel_for(Kokkos::RangePolicy<DeviceType, TagComputeTempDeformRestoreBias >(0,nlocal),*this);
|
||||
copymode = 0;
|
||||
}
|
||||
|
||||
template<class DeviceType>
|
||||
KOKKOS_INLINE_FUNCTION
|
||||
void ComputeTempDeformKokkos<DeviceType>::operator()(TagComputeTempDeformRestoreBias, const int &i) const {
|
||||
if (mask[i] & groupbit) {
|
||||
v(i,0) += vbiasall(i,0);
|
||||
v(i,1) += vbiasall(i,1);
|
||||
v(i,2) += vbiasall(i,2);
|
||||
}
|
||||
}
|
||||
|
||||
namespace LAMMPS_NS {
|
||||
template class ComputeTempDeformKokkos<LMPDeviceType>;
|
||||
#ifdef LMP_KOKKOS_GPU
|
||||
template class ComputeTempDeformKokkos<LMPHostType>;
|
||||
#endif
|
||||
}
|
||||
127
src/KOKKOS/compute_temp_deform_kokkos.h
Normal file
127
src/KOKKOS/compute_temp_deform_kokkos.h
Normal file
@ -0,0 +1,127 @@
|
||||
// clang-format off
|
||||
/* -*- c++ -*- ----------------------------------------------------------
|
||||
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.
|
||||
------------------------------------------------------------------------- */
|
||||
|
||||
#ifdef COMPUTE_CLASS
|
||||
// clang-format off
|
||||
ComputeStyle(temp/deform/kk,ComputeTempDeformKokkos<LMPDeviceType>);
|
||||
ComputeStyle(temp/deform/kk/device,ComputeTempDeformKokkos<LMPDeviceType>);
|
||||
ComputeStyle(temp/deform/kk/host,ComputeTempDeformKokkos<LMPHostType>);
|
||||
// clang-format on
|
||||
#else
|
||||
|
||||
#ifndef LMP_COMPUTE_TEMP_DEFORM_KOKKOS_H
|
||||
#define LMP_COMPUTE_TEMP_DEFORM_KOKKOS_H
|
||||
|
||||
#include "compute_temp_deform.h"
|
||||
#include "kokkos_type.h"
|
||||
#include "domain_kokkos.h"
|
||||
#include "kokkos_few.h"
|
||||
|
||||
namespace LAMMPS_NS {
|
||||
|
||||
template<int RMASS>
|
||||
struct TagComputeTempDeformScalar{};
|
||||
|
||||
template<int RMASS>
|
||||
struct TagComputeTempDeformVector{};
|
||||
|
||||
struct TagComputeTempDeformRemoveBias{};
|
||||
|
||||
struct TagComputeTempDeformRestoreBias{};
|
||||
|
||||
template<class DeviceType>
|
||||
class ComputeTempDeformKokkos: public ComputeTempDeform {
|
||||
public:
|
||||
struct s_CTEMP {
|
||||
double t0, t1, t2, t3, t4, t5;
|
||||
KOKKOS_INLINE_FUNCTION
|
||||
s_CTEMP() {
|
||||
t0 = t1 = t2 = t3 = t4 = t5 = 0.0;
|
||||
}
|
||||
KOKKOS_INLINE_FUNCTION
|
||||
s_CTEMP& operator+=(const s_CTEMP &rhs) {
|
||||
t0 += rhs.t0;
|
||||
t1 += rhs.t1;
|
||||
t2 += rhs.t2;
|
||||
t3 += rhs.t3;
|
||||
t4 += rhs.t4;
|
||||
t5 += rhs.t5;
|
||||
return *this;
|
||||
}
|
||||
|
||||
KOKKOS_INLINE_FUNCTION
|
||||
void operator+=(const volatile s_CTEMP &rhs) volatile {
|
||||
t0 += rhs.t0;
|
||||
t1 += rhs.t1;
|
||||
t2 += rhs.t2;
|
||||
t3 += rhs.t3;
|
||||
t4 += rhs.t4;
|
||||
t5 += rhs.t5;
|
||||
}
|
||||
};
|
||||
|
||||
typedef s_CTEMP CTEMP;
|
||||
typedef CTEMP value_type;
|
||||
typedef DeviceType device_type;
|
||||
typedef ArrayTypes<DeviceType> AT;
|
||||
|
||||
ComputeTempDeformKokkos(class LAMMPS *, int, char **);
|
||||
~ComputeTempDeformKokkos();
|
||||
double compute_scalar();
|
||||
void compute_vector();
|
||||
void remove_bias_all();
|
||||
void restore_bias_all();
|
||||
|
||||
template<int RMASS>
|
||||
KOKKOS_INLINE_FUNCTION
|
||||
void operator()(TagComputeTempDeformScalar<RMASS>, const int&, CTEMP&) const;
|
||||
|
||||
template<int RMASS>
|
||||
KOKKOS_INLINE_FUNCTION
|
||||
void operator()(TagComputeTempDeformVector<RMASS>, const int&, CTEMP&) const;
|
||||
|
||||
KOKKOS_INLINE_FUNCTION
|
||||
void operator()(TagComputeTempDeformRemoveBias, const int &i) const;
|
||||
|
||||
KOKKOS_INLINE_FUNCTION
|
||||
void operator()(TagComputeTempDeformRestoreBias, const int &i) const;
|
||||
|
||||
protected:
|
||||
typename ArrayTypes<DeviceType>::t_x_array_randomread x;
|
||||
typename ArrayTypes<DeviceType>::t_v_array v;
|
||||
typename ArrayTypes<DeviceType>::t_v_array vbiasall;
|
||||
typename ArrayTypes<DeviceType>::t_float_1d_randomread rmass;
|
||||
typename ArrayTypes<DeviceType>::t_float_1d_randomread mass;
|
||||
typename ArrayTypes<DeviceType>::t_int_1d_randomread type;
|
||||
typename ArrayTypes<DeviceType>::t_int_1d_randomread mask;
|
||||
|
||||
class DomainKokkos *domainKK;
|
||||
|
||||
Few<double, 6> h_rate, h_ratelo;
|
||||
|
||||
};
|
||||
|
||||
}
|
||||
|
||||
#endif
|
||||
#endif
|
||||
|
||||
/* ERROR/WARNING messages:
|
||||
|
||||
E: Temperature compute degrees of freedom < 0
|
||||
|
||||
This should not happen if you are calculating the temperature
|
||||
on a valid set of atoms.
|
||||
|
||||
*/
|
||||
@ -28,6 +28,16 @@ ComputeStyle(temp/kk/host,ComputeTempKokkos<LMPHostType>);
|
||||
|
||||
namespace LAMMPS_NS {
|
||||
|
||||
template<int RMASS>
|
||||
struct TagComputeTempScalar{};
|
||||
|
||||
template<int RMASS>
|
||||
struct TagComputeTempVector{};
|
||||
|
||||
template<class DeviceType>
|
||||
class ComputeTempKokkos : public ComputeTemp {
|
||||
public:
|
||||
|
||||
struct s_CTEMP {
|
||||
double t0, t1, t2, t3, t4, t5;
|
||||
KOKKOS_INLINE_FUNCTION
|
||||
@ -55,23 +65,14 @@ namespace LAMMPS_NS {
|
||||
t5 += rhs.t5;
|
||||
}
|
||||
};
|
||||
|
||||
typedef s_CTEMP CTEMP;
|
||||
|
||||
template<int RMASS>
|
||||
struct TagComputeTempScalar{};
|
||||
|
||||
template<int RMASS>
|
||||
struct TagComputeTempVector{};
|
||||
|
||||
template<class DeviceType>
|
||||
class ComputeTempKokkos : public ComputeTemp {
|
||||
public:
|
||||
typedef DeviceType device_type;
|
||||
typedef CTEMP value_type;
|
||||
typedef ArrayTypes<DeviceType> AT;
|
||||
|
||||
ComputeTempKokkos(class LAMMPS *, int, char **);
|
||||
virtual ~ComputeTempKokkos() {}
|
||||
virtual ~ComputeTempKokkos() {};
|
||||
double compute_scalar();
|
||||
void compute_vector();
|
||||
|
||||
|
||||
@ -55,10 +55,13 @@ ComputeTempDeform::ComputeTempDeform(LAMMPS *lmp, int narg, char **arg) :
|
||||
/* ---------------------------------------------------------------------- */
|
||||
|
||||
ComputeTempDeform::~ComputeTempDeform()
|
||||
{
|
||||
if (!copymode)
|
||||
{
|
||||
memory->destroy(vbiasall);
|
||||
delete [] vector;
|
||||
}
|
||||
}
|
||||
|
||||
/* ---------------------------------------------------------------------- */
|
||||
|
||||
@ -69,7 +72,7 @@ void ComputeTempDeform::init()
|
||||
// check fix deform remap settings
|
||||
|
||||
for (i = 0; i < modify->nfix; i++)
|
||||
if (strcmp(modify->fix[i]->style,"deform") == 0) {
|
||||
if (strncmp(modify->fix[i]->style,"deform", 6) == 0) {
|
||||
if (((FixDeform *) modify->fix[i])->remapflag == Domain::X_REMAP &&
|
||||
comm->me == 0)
|
||||
error->warning(FLERR,"Using compute temp/deform with inconsistent "
|
||||
|
||||
Reference in New Issue
Block a user