diff --git a/cmake/Modules/Packages/KOKKOS.cmake b/cmake/Modules/Packages/KOKKOS.cmake index 9b795de3df..8f4313e50e 100644 --- a/cmake/Modules/Packages/KOKKOS.cmake +++ b/cmake/Modules/Packages/KOKKOS.cmake @@ -3,6 +3,7 @@ if(CMAKE_CXX_STANDARD LESS 14) message(FATAL_ERROR "The KOKKOS package requires the C++ standard to be set to at least C++14") endif() + ######################################################################## # consistency checks and Kokkos options/settings required by LAMMPS if(Kokkos_ENABLE_CUDA) @@ -89,7 +90,6 @@ else() endif() add_subdirectory(${LAMMPS_LIB_KOKKOS_SRC_DIR} ${LAMMPS_LIB_KOKKOS_BIN_DIR}) - set(Kokkos_INCLUDE_DIRS ${LAMMPS_LIB_KOKKOS_SRC_DIR}/core/src ${LAMMPS_LIB_KOKKOS_SRC_DIR}/containers/src ${LAMMPS_LIB_KOKKOS_SRC_DIR}/algorithms/src @@ -143,7 +143,23 @@ if(PKG_ML-IAP) list(APPEND KOKKOS_PKG_SOURCES ${KOKKOS_PKG_SOURCES_DIR}/mliap_data_kokkos.cpp ${KOKKOS_PKG_SOURCES_DIR}/mliap_descriptor_so3_kokkos.cpp ${KOKKOS_PKG_SOURCES_DIR}/mliap_model_linear_kokkos.cpp + ${KOKKOS_PKG_SOURCES_DIR}/mliap_model_python_kokkos.cpp ${KOKKOS_PKG_SOURCES_DIR}/mliap_so3_kokkos.cpp) + + # Add KOKKOS version of ML-IAP Python coupling if non-KOKKOS version is included + if(MLIAP_ENABLE_PYTHON AND Cythonize_EXECUTABLE) + file(GLOB MLIAP_KOKKOS_CYTHON_SRC ${LAMMPS_SOURCE_DIR}/KOKKOS/*.pyx) + foreach(MLIAP_CYTHON_FILE ${MLIAP_KOKKOS_CYTHON_SRC}) + get_filename_component(MLIAP_CYTHON_BASE ${MLIAP_CYTHON_FILE} NAME_WE) + add_custom_command(OUTPUT ${MLIAP_BINARY_DIR}/${MLIAP_CYTHON_BASE}.cpp ${MLIAP_BINARY_DIR}/${MLIAP_CYTHON_BASE}.h + COMMAND ${CMAKE_COMMAND} -E copy_if_different ${MLIAP_CYTHON_FILE} ${MLIAP_BINARY_DIR}/${MLIAP_CYTHON_BASE}.pyx + COMMAND ${Cythonize_EXECUTABLE} -3 ${MLIAP_BINARY_DIR}/${MLIAP_CYTHON_BASE}.pyx + WORKING_DIRECTORY ${MLIAP_BINARY_DIR} + MAIN_DEPENDENCY ${MLIAP_CYTHON_FILE} + COMMENT "Generating C++ sources with cythonize...") + list(APPEND KOKKOS_PKG_SOURCES ${MLIAP_BINARY_DIR}/${MLIAP_CYTHON_BASE}.cpp) + endforeach() + endif() endif() if(PKG_PHONON) diff --git a/examples/mliap/mliap_pytorch_Ta06A_kokkos.py b/examples/mliap/mliap_pytorch_Ta06A_kokkos.py new file mode 100644 index 0000000000..da0d3fcaff --- /dev/null +++ b/examples/mliap/mliap_pytorch_Ta06A_kokkos.py @@ -0,0 +1,104 @@ +# Demonstrate how to load a model from the python side. +# This is essentially the same as in.mliap.pytorch.Ta06A +# except that python is the driving program, and lammps +# is in library mode. + +before_loading =\ +"""# Demonstrate MLIAP/PyTorch interface to linear SNAP potential + +# Initialize simulation + +variable nsteps index 100 +variable nrep equal 4 +variable a equal 3.316 +units metal + +# generate the box and atom positions using a BCC lattice + +variable nx equal ${nrep} +variable ny equal ${nrep} +variable nz equal ${nrep} + +boundary p p p + +lattice bcc $a +region box block 0 ${nx} 0 ${ny} 0 ${nz} +create_box 1 box +create_atoms 1 box + +mass 1 180.88 + +# choose potential + +# DATE: 2014-09-05 UNITS: metal CONTRIBUTOR: Aidan Thompson athomps@sandia.gov CITATION: Thompson, Swiler, Trott, Foiles and Tucker, arxiv.org, 1409.3880 (2014) + +# Definition of SNAP potential Ta_Cand06A +# Assumes 1 LAMMPS atom type + +variable zblcutinner equal 4 +variable zblcutouter equal 4.8 +variable zblz equal 73 + +# Specify hybrid with SNAP, ZBL + +pair_style hybrid/overlay & +zbl ${zblcutinner} ${zblcutouter} & +mliap model mliappy LATER & +descriptor sna Ta06A.mliap.descriptor +pair_coeff 1 1 zbl ${zblz} ${zblz} +pair_coeff * * mliap Ta +""" +after_loading =\ +""" + +# Setup output + +compute eatom all pe/atom +compute energy all reduce sum c_eatom + +compute satom all stress/atom NULL +compute str all reduce sum c_satom[1] c_satom[2] c_satom[3] +variable press equal (c_str[1]+c_str[2]+c_str[3])/(3*vol) + +thermo_style custom step temp epair c_energy etotal press v_press +thermo 10 +thermo_modify norm yes + +# Set up NVE run + +timestep 0.5e-3 +neighbor 1.0 bin +neigh_modify once no every 1 delay 0 check yes + +# Run MD + +velocity all create 300.0 4928459 loop geom +fix 1 all nve +run ${nsteps} +""" + +import lammps + +lmp = lammps.lammps(cmdargs=['-k', 'on', 'g', '1', '-sf', 'kk', '-pk', 'kokkos', 'neigh', 'full', 'newton', 'on', '-echo', 'both']) + +# Before defining the pair style, one must do the following: +import lammps.mliap +lammps.mliap.activate_mliappy_kokkos(lmp) +# Otherwise, when running lammps in library mode, +# you will get an error: +# "ERROR: Loading MLIAPPY coupling module failure." + +# Setup the simulation and declare an empty model +# by specifying model filename as "LATER" +lmp.commands_string(before_loading) + +# Define the model however you like. In this example +# we load it from disk: +import torch +model = torch.load('Ta06A.mliap.pytorch.model.pt') + +# Connect the PyTorch model to the mliap pair style. +lammps.mliap.load_model_kokkos(model) + +# run the simulation with the mliap pair style +lmp.commands_string(after_loading) diff --git a/python/lammps/mliap/__init__.py b/python/lammps/mliap/__init__.py index 3e50f43098..9a12b075c4 100644 --- a/python/lammps/mliap/__init__.py +++ b/python/lammps/mliap/__init__.py @@ -31,5 +31,8 @@ if not pylib.Py_IsInitialized(): "in undefined behavior.") else: from .loader import load_model, load_unified, activate_mliappy - + try: + from .loader import load_model_kokkos, activate_mliappy_kokkos + except: + pass del sysconfig, ctypes, library, pylib diff --git a/python/lammps/mliap/loader.py b/python/lammps/mliap/loader.py index 0d44413371..940bd10f1f 100644 --- a/python/lammps/mliap/loader.py +++ b/python/lammps/mliap/loader.py @@ -72,6 +72,24 @@ def activate_mliappy(lmp): except Exception as ee: raise ImportError("Could not load ML-IAP python coupling module.") from ee +def activate_mliappy_kokkos(lmp): + try: + library = lmp.lib + module_names = ["mliap_model_python_couple_kokkos"] + api_version = library.lammps_python_api_version() + + for module_name in module_names: + # Make Machinery + loader = DynamicLoader(module_name,library,api_version) + spec = importlib.util.spec_from_loader(module_name,loader) + + # Do the import + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + except Exception as ee: + raise ImportError("Could not load ML-IAP python coupling module.") from ee + def load_model(model): try: import mliap_model_python_couple @@ -81,6 +99,16 @@ def load_model(model): ) from ie mliap_model_python_couple.load_from_python(model) +def load_model_kokkos(model): + try: + import mliap_model_python_couple_kokkos + except ImportError as ie: + raise ImportError("ML-IAP python module must be activated before loading\n" + "the pair style. Call lammps.mliap.activate_mliappy(lmp)." + ) from ie + mliap_model_python_couple_kokkos.load_from_python(model) + + def load_unified(model): try: import mliap_unified_couple @@ -89,3 +117,4 @@ def load_unified(model): "the pair style. Call lammps.mliap.activate_mliappy(lmp)." ) from ie mliap_unified_couple.load_from_python(model) + diff --git a/python/lammps/mliap/pytorch.py b/python/lammps/mliap/pytorch.py index b02d5ee3dd..fdf04b77d3 100644 --- a/python/lammps/mliap/pytorch.py +++ b/python/lammps/mliap/pytorch.py @@ -89,7 +89,6 @@ class TorchWrapper(torch.nn.Module): """ super().__init__() - self.model = model self.device = device self.dtype = dtype @@ -105,7 +104,7 @@ class TorchWrapper(torch.nn.Module): self.n_descriptors = n_descriptors self.n_elements = n_elements - def forward(self, elems, descriptors, beta, energy): + def forward(self, elems, descriptors, beta, energy,use_gpu_data=False): """ Takes element types and descriptors calculated via lammps and calculates the per atom energies and forces. @@ -130,20 +129,28 @@ class TorchWrapper(torch.nn.Module): ------- None """ - - descriptors = torch.from_numpy(descriptors).to(dtype=self.dtype, device=self.device).requires_grad_(True) - elems = torch.from_numpy(elems).to(dtype=torch.long, device=self.device) - 1 - + descriptors = torch.as_tensor(descriptors,dtype=self.dtype, device=self.device).requires_grad_(True) + elems = torch.as_tensor(elems,dtype=torch.int32, device=self.device) + elems=elems-1 with torch.autograd.enable_grad(): - energy_nn = self.model(descriptors, elems) - if energy_nn.ndim > 1: - energy_nn = energy_nn.flatten() + if (use_gpu_data): + energy_nn = torch.as_tensor(energy,dtype=self.dtype, device=self.device) + energy_nn[:] = self.model(descriptors, elems).flatten() + else: + energy_nn = self.model(descriptors, elems).flatten() + energy[:] = energy_nn.detach().cpu().numpy().astype(np.float64) + #if energy_nn.ndim > 1: + # energy_nn = energy_nn.flatten() + if (use_gpu_data): + beta_nn = torch.as_tensor(beta,dtype=self.dtype, device=self.device) + beta_nn[:] = torch.autograd.grad(energy_nn.sum(), descriptors)[0] + else: beta_nn = torch.autograd.grad(energy_nn.sum(), descriptors)[0] + beta[:] = beta_nn.detach().cpu().numpy().astype(np.float64) - beta[:] = beta_nn.detach().cpu().numpy().astype(np.float64) - energy[:] = energy_nn.detach().cpu().numpy().astype(np.float64) + elems=elems+1 class IgnoreElems(torch.nn.Module): diff --git a/src/Depend.sh b/src/Depend.sh index 90dfdbba7a..cbc907df16 100755 --- a/src/Depend.sh +++ b/src/Depend.sh @@ -151,6 +151,7 @@ fi if (test $1 = "PYTHON") then depend ML-IAP + depend KOKKOS fi if (test $1 = "PHONON") then @@ -163,13 +164,17 @@ if (test $1 = "RIGID") then depend DPD-SMOOTH fi +if (test $1 = "ML-IAP") then + depend KOKKOS +fi + if (test $1 = "ML-PACE") then depend KOKKOS fi if (test $1 = "ML-SNAP") then - depend KOKKOS depend ML-IAP + depend KOKKOS fi if (test $1 = "CG-SPICA") then diff --git a/src/KOKKOS/Install.sh b/src/KOKKOS/Install.sh index 77a9932d90..e7472f4e88 100755 --- a/src/KOKKOS/Install.sh +++ b/src/KOKKOS/Install.sh @@ -201,6 +201,8 @@ action mliap_descriptor_so3_kokkos.cpp mliap_descriptor_so3.cpp action mliap_descriptor_so3_kokkos.h mliap_descriptor_so3.h action mliap_model_linear_kokkos.cpp mliap_model_linear.cpp action mliap_model_linear_kokkos.h mliap_model_linear.h +action mliap_model_python_kokkos.cpp mliap_model_linear.cpp +action mliap_model_python_kokkos.h mliap_model_linear.h action mliap_model_kokkos.h mliap_model.h action mliap_so3_kokkos.cpp mliap_so3.cpp action mliap_so3_kokkos.h mliap_so3.h @@ -361,6 +363,9 @@ action transpose_helper_kokkos.h action verlet_kokkos.cpp action verlet_kokkos.h +# Install cython pyx file only if non-KOKKOS version is present +action mliap_model_python_couple_kokkos.pyx mliap_model_python_couple.pyx + # edit 2 Makefile.package files to include/exclude package info if (test $1 = 1) then @@ -411,3 +416,22 @@ elif (test $1 = 0) then fi fi + +# Python cython stuff. Only need to convert/remove sources. +# Package settings were already done in ML-IAP package Install.sh script. + +if (test $1 = 1) then + if (type cythonize > /dev/null 2>&1 && test -e ../python_impl.cpp) then + cythonize -3 ../mliap_model_python_couple_kokkos.pyx + fi + +elif (test $1 = 0) then + rm -f ../mliap_model_python_couple_kokkos.cpp ../mliap_model_python_couple_kokkos.h + +elif (test $1 = 2) then + if (type cythonize > /dev/null 2>&1 && test -e ../python_impl.cpp) then + cythonize -3 ../mliap_model_python_couple_kokkos.pyx + else + rm -f ../mliap_model_python_couple_kokkos.cpp ../mliap_model_python_couple_kokkos.h + fi +fi diff --git a/src/KOKKOS/mliap_model_python_couple_kokkos.pyx b/src/KOKKOS/mliap_model_python_couple_kokkos.pyx new file mode 100644 index 0000000000..6dec8cdbfe --- /dev/null +++ b/src/KOKKOS/mliap_model_python_couple_kokkos.pyx @@ -0,0 +1,160 @@ +# cython: language_level=3 +# distutils: language = c++ + +cimport cython + +from libc.stdint cimport uintptr_t + +import pickle + +# For converting C arrays to numpy arrays +import numpy +import torch +try: + import cupy +except ImportError: + pass + +# For converting void * to integer for tracking object identity +from libc.stdint cimport uintptr_t + +from libcpp.string cimport string + + +cdef extern from "mliap_data_kokkos.h" namespace "LAMMPS_NS": + cdef cppclass MLIAPDataKokkosDevice: + # Array shapes + int nlistatoms + int ndescriptors + + # Input data + int * ielems # types for all atoms in list + double * descriptors # descriptors for all atoms in list + + # Output data to write to + double * betas # betas for all atoms in list + double * eatoms # energy for all atoms in list + double *energy + int dev +cdef extern from "mliap_model_python_kokkos.h" namespace "LAMMPS_NS": + cdef cppclass MLIAPModelPythonKokkosDevice: + void connect_param_counts() + + +class MLIAPPYKokkosModelNotLinked(Exception): pass + + +LOADED_MODELS = {} + +cdef object c_id(MLIAPModelPythonKokkosDevice * c_model): + """ + Use python-style id of object to keep track of identity. + Note, this is probably not a perfect general strategy but it should work fine with LAMMPS pair styles. + """ + return int( c_model) + +cdef object retrieve(MLIAPModelPythonKokkosDevice * c_model) with gil: + try: + model = LOADED_MODELS[c_id(c_model)] + except KeyError as ke: + raise KeyError("Model has not been loaded.") from ke + if model is None: + raise MLIAPPYKokkosModelNotLinked("Model not linked, connect the model from the python side.") + return model + +cdef public int MLIAPPYKokkos_load_model(MLIAPModelPythonKokkosDevice * c_model, char* fname) with gil: + str_fname = fname.decode('utf-8') # Python 3 only; not Python 2 not supported. + if str_fname == "LATER": + model = None + returnval = 0 + else: + if str_fname.endswith(".pt") or str_fname.endswith('.pth'): + import torch + model = torch.load(str_fname) + else: + with open(str_fname,'rb') as pfile: + model = pickle.load(pfile) + returnval = 1 + LOADED_MODELS[c_id(c_model)] = model + return returnval + +def load_from_python(model): + unloaded_models = [k for k, v in LOADED_MODELS.items() if v is None] + num_models = len(unloaded_models) + cdef MLIAPModelPythonKokkosDevice * lmp_model + + if num_models == 0: + raise ValueError("No model in the waiting area.") + elif num_models > 1: + raise ValueError("Model is amibguous, more than one model in waiting area.") + else: + c_id = unloaded_models[0] + LOADED_MODELS[c_id]=model + lmp_model = c_id + lmp_model.connect_param_counts() + + +cdef public void MLIAPPYKokkos_unload_model(MLIAPModelPythonKokkosDevice * c_model) with gil: + del LOADED_MODELS[c_id(c_model)] + +cdef public int MLIAPPYKokkos_nparams(MLIAPModelPythonKokkosDevice * c_model) with gil: + return int(retrieve(c_model).n_params) + +cdef public int MLIAPPYKokkos_nelements(MLIAPModelPythonKokkosDevice * c_model) with gil: + return int(retrieve(c_model).n_elements) + +cdef public int MLIAPPYKokkos_ndescriptors(MLIAPModelPythonKokkosDevice * c_model) with gil: + return int(retrieve(c_model).n_descriptors) + +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) + type=cupy.double + if (is_int): + type=cupy.int32 + return cupy.ndarray(shape, type, memptr=memptr) + else: + if (len(shape) == 1 ): + if (is_int): + return numpy.asarray(pointer) + else: + return numpy.asarray(pointer) + else: + if (is_int): + return numpy.asarray(pointer) + else: + return numpy.asarray(pointer) + + +cdef public void MLIAPPYKokkos_compute_gradients(MLIAPModelPythonKokkosDevice * c_model, MLIAPDataKokkosDevice * data) with gil: + + dev=data.dev + + torch.cuda.nvtx.range_push("set data fields") + model = retrieve(c_model) + n_d = data.ndescriptors + n_a = data.nlistatoms + + cdef void* ptr = data.ielems + # Make numpy arrays from pointers + elem_cp = create_array(dev, data.ielems, (n_d,), True) + en_cp = create_array(dev, data.eatoms, (n_a,), False) + beta_cp = create_array(dev, data.betas, (n_a, n_d), False) + desc_cp = create_array(dev, data.descriptors, (n_a, n_d), False) + torch.cuda.nvtx.range_pop() + + # Invoke python model on numpy arrays. + torch.cuda.nvtx.range_push("call model") + model(elem_cp,desc_cp,beta_cp,en_cp,dev==1) + torch.cuda.nvtx.range_pop() + + # Get the total energy from the atom energy. + energy = cupy.sum(en_cp) + data.energy[0] = energy + return diff --git a/src/KOKKOS/mliap_model_python_kokkos.cpp b/src/KOKKOS/mliap_model_python_kokkos.cpp new file mode 100644 index 0000000000..6bfc58057c --- /dev/null +++ b/src/KOKKOS/mliap_model_python_kokkos.cpp @@ -0,0 +1,178 @@ +// clang-format off +/* -*- 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_model_python_kokkos.h" + +#include "mliap_data_kokkos.h" +#include "comm.h" +#include "error.h" +#include "utils.h" +#include "mliap_model_python_couple_kokkos.h" +#include "lmppython.h" +#include "python_compat.h" +#include + +using namespace LAMMPS_NS; + +template +MLIAPModelPythonKokkos::~MLIAPModelPythonKokkos() { + auto nontemplated_this = static_cast((void*)this); + if (model_loaded) + MLIAPPYKokkos_unload_model(nontemplated_this); + model_loaded=false; +} + +template +MLIAPModelPythonKokkos::MLIAPModelPythonKokkos(LAMMPS *lmp, char *coefffilename) : + MLIAPModelPython(lmp,coefffilename,true), + MLIAPModelKokkos(lmp, this) +{ + if (!std::is_same::value ) + MLIAPModelKokkos::error->all(FLERR, "MLIAP Kokkos version of the python interface is ONLY available on device"); + + model_loaded = 0; + MLIAPModelKokkos::python->init(); + PyGILState_STATE gstate = PyGILState_Ensure(); + + PyObject *pyMain = PyImport_AddModule("__main__"); + + if (!pyMain) { + PyGILState_Release(gstate); + MLIAPModelKokkos::error->all(FLERR, "Could not initialize embedded Python"); + } + + PyObject *coupling_module = PyImport_ImportModule("mliap_model_python_couple_kokkos"); + + if (!coupling_module) { + PyErr_Print(); + PyErr_Clear(); + PyGILState_Release(gstate); + MLIAPModelKokkos::error->all(FLERR, "Loading MLIAPPYKokkos coupling module failure."); + } + // Recipe from lammps/src/pair_python.cpp : + // add current directory to PYTHONPATH + PyObject *py_path = PySys_GetObject((char *) "path"); + PyList_Append(py_path, PY_STRING_FROM_STRING(".")); + + // if LAMMPS_POTENTIALS environment variable is set, add it to PYTHONPATH as well + const char *potentials_path = getenv("LAMMPS_POTENTIALS"); + if (potentials_path != nullptr) { + PyList_Append(py_path, PY_STRING_FROM_STRING(potentials_path)); + } + PyGILState_Release(gstate); + if (coefffilename) read_coeffs(coefffilename); + + if (coefffilename) MLIAPModelKokkos::set_k_coeffelem(); + + nonlinearflag = 1; +} +/* ---------------------------------------------------------------------- */ + +template +void MLIAPModelPythonKokkos::read_coeffs(char *fname) +{ + PyGILState_STATE gstate = PyGILState_Ensure(); + auto nontemplated_this = static_cast((void*)this); + model_loaded = MLIAPPYKokkos_load_model(nontemplated_this, fname); + if (PyErr_Occurred()) { + PyErr_Print(); + PyErr_Clear(); + PyGILState_Release(gstate); + MLIAPModelKokkos::error->all(FLERR, "Loading python model failure."); + } + PyGILState_Release(gstate); + + if (model_loaded) { + this->connect_param_counts(); + } else { + if (MLIAPModelKokkos::comm->me == 0) + utils::logmesg(MLIAPModelKokkos::lmp, "Loading python model deferred.\n"); + } +} + +/* ---------------------------------------------------------------------- */ +// Finalize loading of the model. +template +void MLIAPModelPythonKokkos::connect_param_counts() +{ + PyGILState_STATE gstate = PyGILState_Ensure(); + auto nontemplated_this = static_cast((void*)this); + nelements = MLIAPPYKokkos_nelements(nontemplated_this); + nparams = MLIAPPYKokkos_nparams(nontemplated_this); + ndescriptors = MLIAPPYKokkos_ndescriptors(nontemplated_this); + + if (PyErr_Occurred()) { + PyErr_Print(); + PyErr_Clear(); + PyGILState_Release(gstate); + MLIAPModelKokkos::error->all(FLERR, "Loading python model failure."); + } + PyGILState_Release(gstate); + model_loaded = 1; + utils::logmesg(MLIAPModelKokkos::lmp, "Loading python model complete.\n"); +} + +/* ---------------------------------------------------------------------- */ + +template +void MLIAPModelPythonKokkos::compute_gradients(class MLIAPData *data) +{ + if (!model_loaded) { MLIAPModelKokkos::error->all(FLERR, "Model not loaded."); } + + PyGILState_STATE gstate = PyGILState_Ensure(); + + auto nontemplated_this = static_cast((void*)this); + auto *kokkos_data = dynamic_cast*>(data); + MLIAPDataKokkosDevice raw_data(*kokkos_data); + MLIAPPYKokkos_compute_gradients(nontemplated_this, &raw_data); + if (PyErr_Occurred()) { + PyErr_Print(); + PyErr_Clear(); + PyGILState_Release(gstate); + MLIAPModelKokkos::error->all(FLERR, "Running python model failure."); + } + PyGILState_Release(gstate); +} + +/* ---------------------------------------------------------------------- */ + +template +void MLIAPModelPythonKokkos::compute_gradgrads(class MLIAPData *data) +{ + MLIAPModelPython::compute_gradgrads(data); +} + +/* ---------------------------------------------------------------------- */ + +template +void MLIAPModelPythonKokkos::compute_force_gradients(class MLIAPData *data) +{ + MLIAPModelPython::compute_force_gradients(data); +} + +/* ---------------------------------------------------------------------- */ + +namespace LAMMPS_NS { +template class MLIAPModelPythonKokkos; +#ifdef LMP_KOKKOS_GPU +template class MLIAPModelPythonKokkos; +#endif +} + +#endif diff --git a/src/KOKKOS/mliap_model_python_kokkos.h b/src/KOKKOS/mliap_model_python_kokkos.h new file mode 100644 index 0000000000..e8c9909b88 --- /dev/null +++ b/src/KOKKOS/mliap_model_python_kokkos.h @@ -0,0 +1,86 @@ +/* -*- 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) + ------------------------------------------------------------------------- */ + +#ifndef LMP_MLIAP_MODEL_PYTHON_KOKKOS_H +#define LMP_MLIAP_MODEL_PYTHON_KOKKOS_H + +#include "mliap_model_python.h" +#include "mliap_model_kokkos.h" +#include "kokkos_type.h" + +namespace LAMMPS_NS { + +template +class MLIAPModelPythonKokkos : public MLIAPModelPython, public MLIAPModelKokkos { + public: + MLIAPModelPythonKokkos(LAMMPS *, char * = nullptr); + ~MLIAPModelPythonKokkos(); + void read_coeffs(char *fname); + + void compute_gradients(class MLIAPData *) override; + void compute_gradgrads(class MLIAPData *) override; + void compute_force_gradients(class MLIAPData *) override; + void connect_param_counts(); +}; +} // namespace LAMMPS_NS + + + + +#include "mliap_data_kokkos.h" + +namespace LAMMPS_NS { +class MLIAPModelPythonKokkosDevice: public MLIAPModelPythonKokkos { +}; + +class MLIAPDataKokkosDevice { +public: + + MLIAPDataKokkosDevice(MLIAPDataKokkos &base) : + ndescriptors(base.ndescriptors), + nlistatoms(base.nlistatoms), + ielems(base.k_ielems.d_view.data()), + descriptors(base.k_descriptors.d_view.data()), + betas(base.k_betas.d_view.data()), + eatoms(base.k_eatoms.d_view.data()), + energy(&base.energy), +#if defined(KOKKOS_ENABLE_CUDA) + dev(1) +#else + dev(0) +#endif + { } + + const int ndescriptors; + const int nlistatoms; + int *ielems; + double *descriptors; + double *betas; + double *eatoms; + double *energy; + int dev; + +#ifdef LMP_KOKKOS_GPU + MLIAPDataKokkosDevice(MLIAPDataKokkos &base) : ndescriptors(-1),nlistatoms(-1) + { + // It cannot get here, but needed for compilation + } +#endif +}; +} + +#endif diff --git a/src/KOKKOS/pair_mliap_kokkos.cpp b/src/KOKKOS/pair_mliap_kokkos.cpp index 5f00c7473a..856492565a 100644 --- a/src/KOKKOS/pair_mliap_kokkos.cpp +++ b/src/KOKKOS/pair_mliap_kokkos.cpp @@ -22,6 +22,9 @@ #include "mliap_data_kokkos.h" #include "mliap_descriptor_so3_kokkos.h" #include "mliap_model_linear_kokkos.h" +#ifdef MLIAP_PYTHON +#include "mliap_model_python_kokkos.h" +#endif #include "error.h" #include "neigh_request.h" #include "lammps.h" @@ -38,6 +41,7 @@ PairMLIAPKokkos::PairMLIAPKokkos(class LAMMPS* l) : PairMLIAP(l) kokkosable = 1; execution_space = ExecutionSpaceFromDevice::space; datamask_modify = 0; + is_child=true; } /* ---------------------------------------------------------------------- */ @@ -50,6 +54,10 @@ PairMLIAPKokkos::~PairMLIAPKokkos() memoryKK->destroy_kokkos(k_setflag, setflag); memoryKK->destroy_kokkos(k_eatom,eatom); memoryKK->destroy_kokkos(k_vatom,vatom); + delete model; + delete descriptor; + model=nullptr; + descriptor=nullptr; allocated = 0; } @@ -143,7 +151,7 @@ void PairMLIAPKokkos::allocate() template void PairMLIAPKokkos::settings(int narg, char ** arg) { - PairMLIAP::settings(narg, arg); + std::vector new_args; int iarg=0; while (iarg < narg) { if (strcmp(arg[iarg],"model") == 0) { @@ -152,8 +160,19 @@ void PairMLIAPKokkos::settings(int narg, char ** arg) delete model; model = new MLIAPModelLinearKokkos(lmp,arg[iarg+2]); iarg += 3; - } else - iarg += 2; + } else if (strcmp(arg[iarg+1],"mliappy") == 0) { +#ifdef MLIAP_PYTHON + if (iarg+3 > narg) utils::missing_cmd_args(FLERR, "pair_style mliap mliappy", error); + delete model; + model = new MLIAPModelPythonKokkos(lmp,arg[iarg+2]); + iarg += 3; +#else + error->all(FLERR,"Using pair_style mliap model mliappy requires ML-IAP with python support"); +#endif + } else { + new_args.push_back(arg[iarg++]); + new_args.push_back(arg[iarg++]); + } } else if (strcmp(arg[iarg],"descriptor") == 0) { if (strcmp(arg[iarg+1],"so3") == 0) { if (iarg+3 > narg) error->all(FLERR,"Illegal pair_style mliap command"); @@ -161,10 +180,12 @@ void PairMLIAPKokkos::settings(int narg, char ** arg) descriptor = new MLIAPDescriptorSO3Kokkos(lmp,arg[iarg+2]); iarg += 3; } else - iarg ++; + new_args.push_back(arg[iarg++]); } else - iarg++; + new_args.push_back(arg[iarg++]); } + PairMLIAP::settings(new_args.size(), new_args.data()); + } /* ---------------------------------------------------------------------- */ diff --git a/src/ML-IAP/mliap_model_python.cpp b/src/ML-IAP/mliap_model_python.cpp index 774d14098d..c06bbdc121 100644 --- a/src/ML-IAP/mliap_model_python.cpp +++ b/src/ML-IAP/mliap_model_python.cpp @@ -33,10 +33,14 @@ using namespace LAMMPS_NS; /* ---------------------------------------------------------------------- */ -MLIAPModelPython::MLIAPModelPython(LAMMPS *lmp, char *coefffilename) : +MLIAPModelPython::MLIAPModelPython(LAMMPS *lmp, char *coefffilename, bool is_child) : MLIAPModel(lmp, coefffilename) { model_loaded = 0; + nonlinearflag = 1; + + if (is_child) + return; python->init(); PyGILState_STATE gstate = PyGILState_Ensure(); @@ -66,17 +70,18 @@ MLIAPModelPython::MLIAPModelPython(LAMMPS *lmp, char *coefffilename) : PyList_Append(py_path, PY_STRING_FROM_STRING(potentials_path)); } PyGILState_Release(gstate); - if (coefffilename) read_coeffs(coefffilename); - nonlinearflag = 1; + } /* ---------------------------------------------------------------------- */ MLIAPModelPython::~MLIAPModelPython() { - MLIAPPY_unload_model(this); + if (model_loaded!=0) + MLIAPPY_unload_model(this); + model_loaded=0; } /* ---------------------------------------------------------------------- @@ -92,7 +97,7 @@ void MLIAPModelPython::read_coeffs(char *fname) { PyGILState_STATE gstate = PyGILState_Ensure(); - int loaded = MLIAPPY_load_model(this, fname); + model_loaded = MLIAPPY_load_model(this, fname); if (PyErr_Occurred()) { PyErr_Print(); PyErr_Clear(); @@ -101,7 +106,7 @@ void MLIAPModelPython::read_coeffs(char *fname) } PyGILState_Release(gstate); - if (loaded) { + if (model_loaded) { this->connect_param_counts(); } else { if (comm->me == 0) utils::logmesg(lmp, "Loading python model deferred.\n"); diff --git a/src/ML-IAP/mliap_model_python.h b/src/ML-IAP/mliap_model_python.h index e7aee37345..4ae499d5b4 100644 --- a/src/ML-IAP/mliap_model_python.h +++ b/src/ML-IAP/mliap_model_python.h @@ -20,7 +20,7 @@ namespace LAMMPS_NS { class MLIAPModelPython : public MLIAPModel { public: - MLIAPModelPython(LAMMPS *, char * = nullptr); + MLIAPModelPython(LAMMPS *, char * = nullptr, bool is_child=false); ~MLIAPModelPython() override; int get_nparams() override; int get_gamma_nnz(class MLIAPData *) override; diff --git a/src/ML-IAP/pair_mliap.cpp b/src/ML-IAP/pair_mliap.cpp index a79a8c0627..28e8a1e493 100644 --- a/src/ML-IAP/pair_mliap.cpp +++ b/src/ML-IAP/pair_mliap.cpp @@ -49,7 +49,10 @@ PairMLIAP::PairMLIAP(LAMMPS *lmp) : restartinfo = 0; one_coeff = 1; manybody_flag = 1; + is_child = false; centroidstressflag = CENTROID_NOTAVAIL; + model=nullptr; + descriptor=nullptr; } /* ---------------------------------------------------------------------- */ @@ -61,7 +64,9 @@ PairMLIAP::~PairMLIAP() delete model; delete descriptor; delete data; - + model=nullptr; + descriptor=nullptr; + data=nullptr; if (allocated) { memory->destroy(setflag); memory->destroy(cutsq); @@ -131,17 +136,16 @@ void PairMLIAP::settings(int narg, char ** arg) { if (narg < 2) utils::missing_cmd_args(FLERR, "pair_style mliap", error); - // set flags for required keywords - - delete model; - model = nullptr; - delete descriptor; - descriptor = nullptr; + // This is needed because the unit test calls settings twice + if (!is_child) { + delete model; + model = nullptr; + delete descriptor; + descriptor = nullptr; + } // process keywords - int iarg = 0; - while (iarg < narg) { if (strcmp(arg[iarg],"model") == 0) { if (iarg+2 > narg) utils::missing_cmd_args(FLERR, "pair_style mliap model", error); diff --git a/src/ML-IAP/pair_mliap.h b/src/ML-IAP/pair_mliap.h index a8327b11c6..9d7e08964d 100644 --- a/src/ML-IAP/pair_mliap.h +++ b/src/ML-IAP/pair_mliap.h @@ -44,6 +44,7 @@ class PairMLIAP : public Pair { class MLIAPModel *model; class MLIAPDescriptor *descriptor; class MLIAPData *data; + bool is_child; }; } // namespace LAMMPS_NS diff --git a/src/lammps.cpp b/src/lammps.cpp index cc32578f22..faf91ff0c2 100644 --- a/src/lammps.cpp +++ b/src/lammps.cpp @@ -442,7 +442,9 @@ LAMMPS::LAMMPS(int narg, char **arg, MPI_Comm communicator) : iarg += 3; while (iarg < narg && arg[iarg][0] != '-') iarg++; - } else error->universe_all(FLERR,"Invalid command-line argument"); + } else { + error->universe_all(FLERR, fmt::format("Invalid command-line argument: {}", arg[iarg]) ); + } } // if no partition command-line switch, universe is one world with all procs