403 lines
14 KiB
C++
403 lines
14 KiB
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)
|
|
------------------------------------------------------------------------- */
|
|
|
|
#ifdef MLIAP_PYTHON
|
|
|
|
#include "mliap_unified_kokkos.h"
|
|
#include <Python.h>
|
|
|
|
#include "error.h"
|
|
#include "lmppython.h"
|
|
#include "memory.h"
|
|
#include "mliap_data.h"
|
|
#include "mliap_unified_couple_kokkos.h"
|
|
#include "pair_mliap.h"
|
|
#include "python_compat.h"
|
|
#include "utils.h"
|
|
|
|
using namespace LAMMPS_NS;
|
|
|
|
/* ---------------------------------------------------------------------- */
|
|
|
|
template <class DeviceType>
|
|
MLIAPDummyDescriptorKokkos<DeviceType>::MLIAPDummyDescriptorKokkos(LAMMPS *_lmp) :
|
|
Pointers(_lmp), MLIAPDummyDescriptor(_lmp), MLIAPDescriptorKokkos<DeviceType>(lmp, this) {}
|
|
|
|
template <class DeviceType>
|
|
MLIAPDummyDescriptorKokkos<DeviceType>::~MLIAPDummyDescriptorKokkos()
|
|
{
|
|
// done in base class
|
|
// Py_DECREF(unified_interface);
|
|
}
|
|
|
|
/* ----------------------------------------------------------------------
|
|
invoke compute_descriptors from Cython interface
|
|
---------------------------------------------------------------------- */
|
|
|
|
template <class DeviceType>
|
|
void MLIAPDummyDescriptorKokkos<DeviceType>::compute_descriptors(class MLIAPData *data)
|
|
{
|
|
PyGILState_STATE gstate = PyGILState_Ensure();
|
|
auto *kokkos_data = dynamic_cast<MLIAPDataKokkos<DeviceType>*>(data);
|
|
MLIAPDataKokkosDevice raw_data(*kokkos_data);
|
|
compute_descriptors_python_kokkos(unified_interface, &raw_data);
|
|
if (PyErr_Occurred()) {
|
|
PyErr_Print();
|
|
PyErr_Clear();
|
|
PyGILState_Release(gstate);
|
|
lmp->error->all(FLERR, "Running mliappy unified compute_descriptors failure.");
|
|
}
|
|
PyGILState_Release(gstate);
|
|
}
|
|
|
|
/* ----------------------------------------------------------------------
|
|
invoke compute_forces from Cython interface
|
|
---------------------------------------------------------------------- */
|
|
|
|
template <class DeviceType>
|
|
void MLIAPDummyDescriptorKokkos<DeviceType>::compute_forces(class MLIAPData *data)
|
|
{
|
|
PyGILState_STATE gstate = PyGILState_Ensure();
|
|
auto *kokkos_data = dynamic_cast<MLIAPDataKokkos<DeviceType>*>(data);
|
|
MLIAPDataKokkosDevice raw_data(*kokkos_data);
|
|
compute_forces_python_kokkos(unified_interface, &raw_data);
|
|
if (PyErr_Occurred()) {
|
|
PyErr_Print();
|
|
PyErr_Clear();
|
|
PyGILState_Release(gstate);
|
|
lmp->error->all(FLERR, "Running mliappy unified compute_forces failure.");
|
|
}
|
|
PyGILState_Release(gstate);
|
|
}
|
|
|
|
// not implemented
|
|
template <class DeviceType>
|
|
void MLIAPDummyDescriptorKokkos<DeviceType>::compute_force_gradients(class MLIAPData *)
|
|
{
|
|
error->all(FLERR, "compute_force_gradients not implemented");
|
|
}
|
|
|
|
// not implemented
|
|
template <class DeviceType>
|
|
void MLIAPDummyDescriptorKokkos<DeviceType>::compute_descriptor_gradients(class MLIAPData *)
|
|
{
|
|
error->all(FLERR, "compute_descriptor_gradients not implemented");
|
|
}
|
|
|
|
template <class DeviceType>
|
|
void MLIAPDummyDescriptorKokkos<DeviceType>::init()
|
|
{
|
|
memory->create(radelem, nelements, "mliap_dummy_descriptor:radelem");
|
|
for (int ielem = 0; ielem < nelements; ielem++) { radelem[ielem] = 1; }
|
|
|
|
double cut;
|
|
cutmax = 0.0;
|
|
memory->create(cutsq, nelements, nelements, "mliap/descriptor/dummy:cutsq");
|
|
memory->create(cutghost, nelements, nelements, "mliap/descriptor/dummy:cutghost");
|
|
for (int ielem = 0; ielem < nelements; ielem++) {
|
|
// rcutfac set from python, is global cutoff for all elements
|
|
cut = 2.0 * radelem[ielem] * rcutfac;
|
|
if (cut > cutmax) cutmax = cut;
|
|
cutsq[ielem][ielem] = cut * cut;
|
|
cutghost[ielem][ielem] = cut * cut;
|
|
for (int jelem = ielem + 1; jelem < nelements; jelem++) {
|
|
cut = (radelem[ielem] + radelem[jelem]) * rcutfac;
|
|
cutsq[ielem][jelem] = cutsq[jelem][ielem] = cut * cut;
|
|
cutghost[ielem][jelem] = cutghost[jelem][ielem] = cut * cut;
|
|
}
|
|
}
|
|
}
|
|
|
|
template <class DeviceType>
|
|
void MLIAPDummyDescriptorKokkos<DeviceType>::set_elements(char **elems, int nelems)
|
|
{
|
|
nelements = nelems;
|
|
elements = new char *[nelems];
|
|
for (int i = 0; i < nelems; i++) { elements[i] = utils::strdup(elems[i]); }
|
|
}
|
|
|
|
/* ---------------------------------------------------------------------- */
|
|
|
|
template <class DeviceType>
|
|
MLIAPDummyModelKokkos<DeviceType>::MLIAPDummyModelKokkos(LAMMPS *lmp, char *coefffilename) :
|
|
MLIAPDummyModel(lmp,coefffilename),
|
|
MLIAPModelKokkos<DeviceType>(lmp, this)
|
|
{
|
|
nonlinearflag = 1;
|
|
}
|
|
|
|
template <class DeviceType>
|
|
MLIAPDummyModelKokkos<DeviceType>::~MLIAPDummyModelKokkos()
|
|
{
|
|
// manually decrement borrowed reference from Python
|
|
Py_DECREF(unified_interface);
|
|
}
|
|
|
|
template <class DeviceType>
|
|
int MLIAPDummyModelKokkos<DeviceType>::get_nparams()
|
|
{
|
|
return nparams;
|
|
}
|
|
|
|
template <class DeviceType>
|
|
int MLIAPDummyModelKokkos<DeviceType>::get_gamma_nnz(class MLIAPData *)
|
|
{
|
|
// TODO: get_gamma_nnz
|
|
return 0;
|
|
}
|
|
|
|
/* ----------------------------------------------------------------------
|
|
invoke compute_gradients from Cython interface
|
|
---------------------------------------------------------------------- */
|
|
|
|
template <class DeviceType>
|
|
void MLIAPDummyModelKokkos<DeviceType>::compute_gradients(class MLIAPData *data)
|
|
{
|
|
PyGILState_STATE gstate = PyGILState_Ensure();
|
|
auto *kokkos_data = dynamic_cast<MLIAPDataKokkos<DeviceType>*>(data);
|
|
MLIAPDataKokkosDevice raw_data(*kokkos_data);
|
|
compute_gradients_python_kokkos(unified_interface, &raw_data);
|
|
if (PyErr_Occurred()) {
|
|
PyErr_Print();
|
|
PyErr_Clear();
|
|
PyGILState_Release(gstate);
|
|
MLIAPModelKokkos<DeviceType>::error->all(FLERR, "Running mliappy unified compute_gradients failure.");
|
|
}
|
|
PyGILState_Release(gstate);
|
|
}
|
|
|
|
// not implemented
|
|
template <class DeviceType>
|
|
void MLIAPDummyModelKokkos<DeviceType>::compute_gradgrads(class MLIAPData *)
|
|
{
|
|
MLIAPModelKokkos<DeviceType>::error->all(FLERR, "compute_gradgrads not implemented");
|
|
}
|
|
|
|
// not implemented
|
|
template <class DeviceType>
|
|
void MLIAPDummyModelKokkos<DeviceType>::compute_force_gradients(class MLIAPData *)
|
|
{
|
|
MLIAPModelKokkos<DeviceType>::error->all(FLERR, "compute_force_gradients not implemented");
|
|
}
|
|
|
|
/* ----------------------------------------------------------------------
|
|
memory usage unclear due to Cython/Python implementation
|
|
---------------------------------------------------------------------- */
|
|
|
|
template <class DeviceType>
|
|
double MLIAPDummyModelKokkos<DeviceType>::memory_usage()
|
|
{
|
|
// TODO: implement memory usage in Cython(?)
|
|
return 0;
|
|
}
|
|
|
|
// not implemented
|
|
template <class DeviceType>
|
|
void MLIAPDummyModelKokkos<DeviceType>::read_coeffs(char *)
|
|
{
|
|
MLIAPModelKokkos<DeviceType>::error->all(FLERR, "read_coeffs not implemented");
|
|
}
|
|
|
|
/* ----------------------------------------------------------------------
|
|
build the unified interface object, connect to dummy model and descriptor
|
|
---------------------------------------------------------------------- */
|
|
|
|
template <class DeviceType>
|
|
MLIAPBuildUnifiedKokkos_t<DeviceType> LAMMPS_NS::build_unified(char *unified_fname, MLIAPDataKokkos<DeviceType> *data, LAMMPS *lmp,
|
|
char *coefffilename)
|
|
{
|
|
lmp->python->init();
|
|
PyGILState_STATE gstate = PyGILState_Ensure();
|
|
|
|
PyObject *pyMain = PyImport_AddModule("__main__");
|
|
|
|
if (!pyMain) {
|
|
PyGILState_Release(gstate);
|
|
lmp->error->all(FLERR, "Could not initialize embedded Python");
|
|
}
|
|
|
|
PyImport_ImportModule("mliap_unified_couple_kokkos");
|
|
|
|
if (PyErr_Occurred()) {
|
|
PyErr_Print();
|
|
PyErr_Clear();
|
|
PyGILState_Release(gstate);
|
|
lmp->error->all(FLERR, "Loading mliappy unified module failure.");
|
|
}
|
|
|
|
// Connect dummy model, dummy descriptor, data to Python unified
|
|
MLIAPDummyModelKokkos<DeviceType> *model = new MLIAPDummyModelKokkos<DeviceType>(lmp, coefffilename);
|
|
MLIAPDummyDescriptorKokkos<DeviceType> *descriptor = new MLIAPDummyDescriptorKokkos<DeviceType>(lmp);
|
|
|
|
PyObject *unified_interface = mliap_unified_connect_kokkos(unified_fname, model, descriptor);
|
|
if (PyErr_Occurred()) {
|
|
PyErr_Print();
|
|
PyErr_Clear();
|
|
PyGILState_Release(gstate);
|
|
lmp->error->all(FLERR, "Running mliappy unified module failure.");
|
|
}
|
|
|
|
// Borrowed references must be manually incremented
|
|
model->unified_interface = unified_interface;
|
|
Py_INCREF(unified_interface);
|
|
descriptor->unified_interface = unified_interface;
|
|
Py_INCREF(unified_interface);
|
|
|
|
PyGILState_Release(gstate);
|
|
|
|
MLIAPBuildUnifiedKokkos_t<DeviceType> build = {data, descriptor, model};
|
|
return build;
|
|
}
|
|
|
|
/* ----------------------------------------------------------------------
|
|
set energy for ij atom pairs
|
|
---------------------------------------------------------------------- */
|
|
|
|
void LAMMPS_NS::update_pair_energy(MLIAPDataKokkosDevice *data, double *eij)
|
|
{
|
|
auto d_eatoms = data->eatoms;
|
|
auto d_pair_i= data->pair_i;
|
|
const auto nlocal = data->nlocal;
|
|
Kokkos::parallel_for(nlocal, KOKKOS_LAMBDA(int ii){
|
|
d_eatoms[ii] = 0;
|
|
});
|
|
|
|
Kokkos::parallel_reduce(data->npairs, KOKKOS_LAMBDA(int ii, double &local_sum){
|
|
int i = d_pair_i[ii];
|
|
double e = 0.5 * eij[ii];
|
|
|
|
// must not count any contribution where i is not a local atom
|
|
if (i < nlocal) {
|
|
Kokkos::atomic_add(&d_eatoms[i], e);
|
|
local_sum += e;
|
|
}
|
|
},*data->energy);
|
|
}
|
|
|
|
/* ----------------------------------------------------------------------
|
|
set forces for ij atom pairs
|
|
---------------------------------------------------------------------- */
|
|
|
|
void LAMMPS_NS::update_pair_forces(MLIAPDataKokkosDevice *data, double *fij)
|
|
{
|
|
const auto nlocal = data->nlocal;
|
|
auto *f = data->f;
|
|
auto pair_i = data->pair_i;
|
|
auto j_atoms = data->jatoms;
|
|
auto vflag = data->vflag;
|
|
auto rij = data->rij;
|
|
int vflag_global=data->pairmliap->vflag_global, vflag_atom=data->pairmliap->vflag_atom;
|
|
if (vflag_atom) {
|
|
data->pairmliap->k_vatom.template modify<LMPHostType>();
|
|
data->pairmliap->k_vatom.template sync<LMPDeviceType>();
|
|
}
|
|
auto d_vatom = data->pairmliap->k_vatom.template view<LMPDeviceType>();
|
|
|
|
Kokkos::View<double[6], LMPDeviceType> virial("virial");
|
|
|
|
Kokkos::parallel_for(data->npairs,KOKKOS_LAMBDA (int ii) {
|
|
|
|
int ii3 = ii * 3;
|
|
int i = pair_i[ii];
|
|
int j = j_atoms[ii];
|
|
// must not count any contribution where i is not a local atom
|
|
Kokkos::atomic_add(&f[i*3+0], fij[ii3+0]);
|
|
Kokkos::atomic_add(&f[i*3+1], fij[ii3+1]);
|
|
Kokkos::atomic_add(&f[i*3+2], fij[ii3+2]);
|
|
Kokkos::atomic_add(&f[j*3+0],-fij[ii3+0]);
|
|
Kokkos::atomic_add(&f[j*3+1],-fij[ii3+1]);
|
|
Kokkos::atomic_add(&f[j*3+2],-fij[ii3+2]);
|
|
if (vflag) {
|
|
double v[6];
|
|
v[0] = -rij[ii3+0]*fij[ii3+0];
|
|
v[1] = -rij[ii3+1]*fij[ii3+1];
|
|
v[2] = -rij[ii3+2]*fij[ii3+2];
|
|
v[3] = -rij[ii3+0]*fij[ii3+1];
|
|
v[4] = -rij[ii3+0]*fij[ii3+2];
|
|
v[5] = -rij[ii3+1]*fij[ii3+2];
|
|
if (vflag_global) {
|
|
Kokkos::atomic_add(&virial[0], v[0]);
|
|
Kokkos::atomic_add(&virial[1], v[1]);
|
|
Kokkos::atomic_add(&virial[2], v[2]);
|
|
Kokkos::atomic_add(&virial[3], v[3]);
|
|
Kokkos::atomic_add(&virial[4], v[4]);
|
|
Kokkos::atomic_add(&virial[5], v[5]);
|
|
}
|
|
if (vflag_atom ) {
|
|
Kokkos::atomic_add(&d_vatom(i,0), 0.5*v[0]);
|
|
Kokkos::atomic_add(&d_vatom(i,1), 0.5*v[1]);
|
|
Kokkos::atomic_add(&d_vatom(i,2), 0.5*v[2]);
|
|
Kokkos::atomic_add(&d_vatom(i,3), 0.5*v[3]);
|
|
Kokkos::atomic_add(&d_vatom(i,4), 0.5*v[4]);
|
|
Kokkos::atomic_add(&d_vatom(i,5), 0.5*v[5]);
|
|
|
|
Kokkos::atomic_add(&d_vatom(j,0), 0.5*v[0]);
|
|
Kokkos::atomic_add(&d_vatom(j,1), 0.5*v[1]);
|
|
Kokkos::atomic_add(&d_vatom(j,2), 0.5*v[2]);
|
|
Kokkos::atomic_add(&d_vatom(j,3), 0.5*v[3]);
|
|
Kokkos::atomic_add(&d_vatom(j,4), 0.5*v[4]);
|
|
Kokkos::atomic_add(&d_vatom(j,5), 0.5*v[5]);
|
|
}
|
|
}
|
|
});
|
|
|
|
if (vflag) {
|
|
if (vflag_global) {
|
|
Kokkos::View<double[6], LMPHostType> h_virial("h_virial");
|
|
Kokkos::deep_copy(h_virial,virial);
|
|
for (int i=0;i<6;++i)
|
|
data->pairmliap->virial[i]+=h_virial[i];
|
|
}
|
|
if (vflag_atom) {
|
|
data->pairmliap->k_vatom.template modify<LMPDeviceType>();
|
|
data->pairmliap->k_vatom.template sync<LMPHostType>();
|
|
}
|
|
}
|
|
}
|
|
|
|
/* ----------------------------------------------------------------------
|
|
set energy for i indexed atoms
|
|
---------------------------------------------------------------------- */
|
|
|
|
void LAMMPS_NS::update_atom_energy(MLIAPDataKokkosDevice *data, double *ei)
|
|
{
|
|
auto d_eatoms = data->eatoms;
|
|
const auto nlocal = data->nlocal;
|
|
|
|
Kokkos::parallel_reduce(nlocal, KOKKOS_LAMBDA(int i, double &local_sum){
|
|
double e = ei[i];
|
|
|
|
d_eatoms[i] = e;
|
|
local_sum += e;
|
|
},*data->energy);
|
|
}
|
|
|
|
namespace LAMMPS_NS {
|
|
template class MLIAPDummyModelKokkos<LMPDeviceType>;
|
|
template class MLIAPDummyDescriptorKokkos<LMPDeviceType>;
|
|
template MLIAPBuildUnifiedKokkos_t<LMPDeviceType> LAMMPS_NS::build_unified(char *unified_fname, MLIAPDataKokkos<LMPDeviceType> *data, LAMMPS *lmp,
|
|
char *coefffilename);
|
|
#ifdef LMP_KOKKOS_GPU
|
|
template class MLIAPDummyModelKokkos<LMPHostType>;
|
|
template class MLIAPDummyDescriptorKokkos<LMPHostType>;
|
|
template MLIAPBuildUnifiedKokkos_t<LMPHostType> LAMMPS_NS::build_unified(char *unified_fname, MLIAPDataKokkos<LMPHostType> *data, LAMMPS *lmp,
|
|
char *coefffilename);
|
|
#endif
|
|
}
|
|
#endif
|
|
|