diff --git a/src/KOKKOS/mliap_unified_couple_kokkos.pyx b/src/KOKKOS/mliap_unified_couple_kokkos.pyx index 37326263d3..5e274d7f36 100644 --- a/src/KOKKOS/mliap_unified_couple_kokkos.pyx +++ b/src/KOKKOS/mliap_unified_couple_kokkos.pyx @@ -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( 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, 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, 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( 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 = &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 diff --git a/src/KOKKOS/mliap_unified_kokkos.cpp b/src/KOKKOS/mliap_unified_kokkos.cpp index bfb9193df6..1fdf039473 100644 --- a/src/KOKKOS/mliap_unified_kokkos.cpp +++ b/src/KOKKOS/mliap_unified_kokkos.cpp @@ -269,7 +269,6 @@ MLIAPBuildUnifiedKokkos_t 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(); + data->pairmliap->k_vatom.template sync(); + } auto d_vatom = data->pairmliap->k_vatom.template view(); + Kokkos::View 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; template class MLIAPDummyDescriptorKokkos; template MLIAPBuildUnifiedKokkos_t LAMMPS_NS::build_unified(char *unified_fname, MLIAPDataKokkos *data, LAMMPS *lmp, char *coefffilename); -//template void LAMMPS_NS::update_pair_energy(MLIAPDataKokkos *data, double *eij); -//template void LAMMPS_NS::update_pair_forces(MLIAPDataKokkos *data, double *fij); #ifdef LMP_KOKKOS_GPU template class MLIAPDummyModelKokkos; template class MLIAPDummyDescriptorKokkos; template MLIAPBuildUnifiedKokkos_t LAMMPS_NS::build_unified(char *unified_fname, MLIAPDataKokkos *data, LAMMPS *lmp, char *coefffilename); -//template void LAMMPS_NS::update_pair_energy(MLIAPDataKokkos *data, double *eij); -//template void LAMMPS_NS::update_pair_forces(MLIAPDataKokkos *data, double *fij); #endif } #endif diff --git a/src/KOKKOS/pair_mliap_kokkos.cpp b/src/KOKKOS/pair_mliap_kokkos.cpp index 558901c6ff..d19d81e314 100644 --- a/src/KOKKOS/pair_mliap_kokkos.cpp +++ b/src/KOKKOS/pair_mliap_kokkos.cpp @@ -74,7 +74,6 @@ void PairMLIAPKokkos::compute(int eflag, int vflag) int is_kokkos_descriptor = (dynamic_cast*>(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::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::e_tally(MLIAPData* data) if (eflag_global) eng_vdwl += data->energy; if (eflag_atom) { MLIAPDataKokkos *k_data = static_cast*>(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(); auto d_eatoms = k_data->k_eatoms.template view(); auto d_eatom = k_eatom.template view(); diff --git a/src/PYTHON/python_impl.cpp b/src/PYTHON/python_impl.cpp index 9177381a35..57f8ea1cf0 100644 --- a/src/PYTHON/python_impl.cpp +++ b/src/PYTHON/python_impl.cpp @@ -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();