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

Merge pull request #3661 from bathmatt/hippy-fixes

Browse Source 
Hippy fixes
This commit is contained in:
Axel Kohlmeyer
2023-03-01 20:36:37 -05:00
committed by GitHub
parent d44ba14030 cad3c5ba95
commit c2fe2b4360
4 changed files with 59 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
src/KOKKOS/mliap_unified_couple_kokkos.pyx
View File

@ -115,7 +115,6 @@ cdef create_array(device, void *pointer, shape,is_int):
size=1
for i in shape:
size = size*i
if ( device == 1):
mem = cupy.cuda.UnownedMemory(ptr=int( <uintptr_t> pointer), owner=None, size=size)
memptr = cupy.cuda.MemoryPointer(mem, 0)
@ -146,10 +145,18 @@ cdef class MLIAPDataPy:
self.data = NULL
def update_pair_energy_cpu(self, eij):
cdef double[:] eij_arr = eij
cdef double[:] eij_arr
try:
eij_arr = eij
except:
eij_arr = eij.detach().numpy().astype(np.double)
update_pair_energy(self.data, &eij_arr[0])
def update_pair_energy_gpu(self, eij):
cdef uintptr_t ptr = eij.data.ptr
cdef uintptr_t ptr;
try:
ptr = eij.data.ptr
except:
ptr = eij.data_ptr()
update_pair_energy(self.data, <double*>ptr)
def update_pair_energy(self, eij):
if self.data.dev==0:
@ -158,10 +165,18 @@ cdef class MLIAPDataPy:
self.update_pair_energy_gpu(eij)
def update_pair_forces_cpu(self, fij):
cdef double[:, ::1] fij_arr = fij
cdef double[:, ::1] fij_arr
try:
fij_arr = fij
except:
fij_arr = fij.detach().numpy().astype(np.double)
update_pair_forces(self.data, &fij_arr[0][0])
def update_pair_forces_gpu(self, fij):
cdef uintptr_t ptr = fij.data.ptr
cdef uintptr_t ptr
try:
ptr = fij.data.ptr
except:
ptr = fij.data_ptr()
update_pair_forces(self.data, <double*>ptr)
def update_pair_forces(self, fij):
if self.data.dev==0:
@ -172,7 +187,8 @@ cdef class MLIAPDataPy:
def f(self):
if self.data.f is NULL:
return None
return cupy.asarray(<double[:self.ntotal, :3]> self.data.f)
return create_array(self.data.dev, self.data.f, [self.ntotal, 3],False)
@property
def size_gradforce(self):
@ -205,14 +221,11 @@ cdef class MLIAPDataPy:
descriptors_view[:] = value_view
print("This code has not been tested or optimized for the GPU, if you are getting this warning optimize descriptors")
@write_only_property
def eatoms(self, value):
@property
def eatoms(self):
if self.data.eatoms is NULL:
raise ValueError("attempt to set NULL eatoms")
cdef double[:] eatoms_view = <double[:self.nlistatoms]> &self.data.eatoms[0]
cdef double[:] value_view = value
eatoms_view[:] = value_view
print("This code has not been tested or optimized for the GPU, if you are getting this warning optimize eatoms")
return create_array(self.data.dev, self.data.eatoms, [self.nlistatoms],False)
@write_only_property
@ -351,7 +364,7 @@ cdef class MLIAPDataPy:
# Interface between C and Python compute functions
cdef class MLIAPUnifiedInterface:
cdef class MLIAPUnifiedInterfaceKokkos:
cdef MLIAPDummyModel * model
cdef MLIAPDummyDescriptor * descriptor
cdef unified_impl
@ -404,7 +417,7 @@ cdef public object mliap_unified_connect_kokkos(char *fname, MLIAPDummyModel * m
with open(str_fname, 'rb') as pfile:
unified = pickle.load(pfile)
unified_int = MLIAPUnifiedInterface(unified)
unified_int = MLIAPUnifiedInterfaceKokkos(unified)
unified_int.model = model
unified_int.descriptor = descriptor

15
src/KOKKOS/mliap_unified_kokkos.cpp
View File

@ -269,7 +269,6 @@ MLIAPBuildUnifiedKokkos_t<DeviceType> LAMMPS_NS::build_unified(char *unified_fna
void LAMMPS_NS::update_pair_energy(MLIAPDataKokkosDevice *data, double *eij)
{
double e_total = 0.0;
auto d_eatoms = data->eatoms;
auto d_pair_i= data->pair_i;
const auto nlistatoms = data->nlistatoms;
@ -301,8 +300,13 @@ void LAMMPS_NS::update_pair_forces(MLIAPDataKokkosDevice *data, double *fij)
auto j_atoms = data->jatoms;
auto vflag = data->vflag;
auto rij = data->rij;
int vflag_either=data->pairmliap->vflag_either, vflag_global=data->pairmliap->vflag_global, vflag_atom=data->pairmliap->vflag_atom;
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) {
@ -310,7 +314,6 @@ void LAMMPS_NS::update_pair_forces(MLIAPDataKokkosDevice *data, double *fij)
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
if (i < nlistatoms) {
Kokkos::atomic_add(&f[i*3+0], fij[ii3+0]);
@ -335,7 +338,7 @@ void LAMMPS_NS::update_pair_forces(MLIAPDataKokkosDevice *data, double *fij)
Kokkos::atomic_add(&virial[4], v[4]);
Kokkos::atomic_add(&virial[5], v[5]);
}
if (vflag_atom) {
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]);
@ -373,15 +376,11 @@ 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);
//template void LAMMPS_NS::update_pair_energy(MLIAPDataKokkos<LMPDeviceType> *data, double *eij);
//template void LAMMPS_NS::update_pair_forces(MLIAPDataKokkos<LMPDeviceType> *data, double *fij);
#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);
//template void LAMMPS_NS::update_pair_energy(MLIAPDataKokkos<LMPHostType> *data, double *eij);
//template void LAMMPS_NS::update_pair_forces(MLIAPDataKokkos<LMPHostType> *data, double *fij);
#endif
}
#endif

7
src/KOKKOS/pair_mliap_kokkos.cpp
View File

@ -74,7 +74,6 @@ void PairMLIAPKokkos<DeviceType>::compute(int eflag, int vflag)
int is_kokkos_descriptor = (dynamic_cast<MLIAPDescriptorKokkos<DeviceType>*>(descriptor)) != nullptr;
auto model_space = is_kokkos_model ? execution_space : Host;
auto descriptor_space = is_kokkos_descriptor? execution_space : Host;
// consistency checks
if (data->ndescriptors != model->ndescriptors)
error->all(FLERR, "Incompatible model and descriptor descriptor count");
@ -107,12 +106,14 @@ void PairMLIAPKokkos<DeviceType>::compute(int eflag, int vflag)
k_data->sync(model_space, IELEMS_MASK | DESCRIPTORS_MASK);
model->compute_gradients(data);
k_data->modified(model_space, BETAS_MASK);
if (eflag_atom)
if (eflag_atom) {
k_data->modified(model_space, EATOMS_MASK);
}
// calculate force contributions beta_i*dB_i/dR_j
atomKK->sync(descriptor_space,F_MASK);
k_data->sync(descriptor_space, NUMNEIGHS_MASK | IATOMS_MASK | IELEMS_MASK | ELEMS_MASK | BETAS_MASK | JATOMS_MASK | PAIR_I_MASK | JELEMS_MASK | RIJ_MASK );
descriptor->compute_forces(data);
e_tally(data);
@ -293,7 +294,7 @@ void PairMLIAPKokkos<DeviceType>::e_tally(MLIAPData* data)
if (eflag_global) eng_vdwl += data->energy;
if (eflag_atom) {
MLIAPDataKokkos<DeviceType> *k_data = static_cast<MLIAPDataKokkos<DeviceType>*>(data);
k_data->sync(execution_space, IATOMS_MASK | EATOMS_MASK);
k_data->sync(execution_space, IATOMS_MASK | EATOMS_MASK, true);
auto d_iatoms = k_data->k_iatoms.template view<DeviceType>();
auto d_eatoms = k_data->k_eatoms.template view<DeviceType>();
auto d_eatom = k_eatom.template view<DeviceType>();

21
src/PYTHON/python_impl.cpp
View File

@ -35,6 +35,17 @@
// However, that exposes -too many- headers.
#include "mliap_model_python_couple.h"
#include "mliap_unified_couple.h"
#ifdef LMP_KOKKOS
#include "mliap_model_python_kokkos.h"
#include "mliap_unified_kokkos.h"
// The above should somehow really be included in the next file.
// We could get around this with cython --capi-reexport-cincludes
// However, that exposes -too many- headers.
#include "mliap_model_python_couple_kokkos.h"
#include "mliap_unified_couple_kokkos.h"
#endif
#endif
using namespace LAMMPS_NS;
@ -71,6 +82,16 @@ PythonImpl::PythonImpl(LAMMPS *lmp) : Pointers(lmp)
err = PyImport_AppendInittab("mliap_unified_couple", PyInit_mliap_unified_couple);
if (err) error->all(FLERR, "Could not register MLIAPPY unified embedded python module.");
#ifdef LMP_KOKKOS
// Inform python intialization scheme of the mliappy module.
// This -must- happen before python is initialized.
err = PyImport_AppendInittab("mliap_model_python_couple_kokkos", PyInit_mliap_model_python_couple_kokkos);
if (err) error->all(FLERR, "Could not register MLIAPPY embedded python module.");
err = PyImport_AppendInittab("mliap_unified_couple_kokkos", PyInit_mliap_unified_couple_kokkos);
if (err) error->all(FLERR, "Could not register MLIAPPY unified embedded python module.");
#endif
#endif
Py_Initialize();