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

Merge branch 'develop' into collected-small-changes

Browse Source
This commit is contained in:
Axel Kohlmeyer
2023-01-06 14:41:22 -05:00
parent ddc8ed8c2d 9a8455f546
commit f007eaf946
16 changed files with 681 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
cmake/Modules/Packages/KOKKOS.cmake
View File

@ -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)

104
examples/mliap/mliap_pytorch_Ta06A_kokkos.py Normal file
View File

@ -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)

5
python/lammps/mliap/__init__.py
View File

@ -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

29
python/lammps/mliap/loader.py
View File

@ -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)

29
python/lammps/mliap/pytorch.py
View File

@ -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):

7
src/Depend.sh
View File

@ -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

24
src/KOKKOS/Install.sh
View File

@ -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

160
src/KOKKOS/mliap_model_python_couple_kokkos.pyx Normal file
View File

@ -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(<uintptr_t> 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 = <MLIAPModelPythonKokkosDevice *> <uintptr_t> 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( <uintptr_t> 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(<int[:shape[0]]>pointer)
else:
return numpy.asarray(<double[:shape[0]]>pointer)
else:
if (is_int):
return numpy.asarray(<int[:shape[0],:shape[1]]>pointer)
else:
return numpy.asarray(<double[:shape[0],:shape[1]]>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] = <double> energy
return

178
src/KOKKOS/mliap_model_python_kokkos.cpp Normal file
View File

@ -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 <type_traits>
using namespace LAMMPS_NS;
template<class DeviceType>
MLIAPModelPythonKokkos<DeviceType>::~MLIAPModelPythonKokkos() {
auto nontemplated_this = static_cast<MLIAPModelPythonKokkosDevice*>((void*)this);
if (model_loaded)
MLIAPPYKokkos_unload_model(nontemplated_this);
model_loaded=false;
}
template<class DeviceType>
MLIAPModelPythonKokkos<DeviceType>::MLIAPModelPythonKokkos(LAMMPS *lmp, char *coefffilename) :
MLIAPModelPython(lmp,coefffilename,true),
MLIAPModelKokkos<DeviceType>(lmp, this)
{
if (!std::is_same<DeviceType,LMPDeviceType>::value )
MLIAPModelKokkos<DeviceType>::error->all(FLERR, "MLIAP Kokkos version of the python interface is ONLY available on device");
model_loaded = 0;
MLIAPModelKokkos<DeviceType>::python->init();
PyGILState_STATE gstate = PyGILState_Ensure();
PyObject *pyMain = PyImport_AddModule("__main__");
if (!pyMain) {
PyGILState_Release(gstate);
MLIAPModelKokkos<DeviceType>::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<DeviceType>::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<DeviceType>::set_k_coeffelem();
nonlinearflag = 1;
}
/* ---------------------------------------------------------------------- */
template<class DeviceType>
void MLIAPModelPythonKokkos<DeviceType>::read_coeffs(char *fname)
{
PyGILState_STATE gstate = PyGILState_Ensure();
auto nontemplated_this = static_cast<MLIAPModelPythonKokkosDevice*>((void*)this);
model_loaded = MLIAPPYKokkos_load_model(nontemplated_this, fname);
if (PyErr_Occurred()) {
PyErr_Print();
PyErr_Clear();
PyGILState_Release(gstate);
MLIAPModelKokkos<DeviceType>::error->all(FLERR, "Loading python model failure.");
}
PyGILState_Release(gstate);
if (model_loaded) {
this->connect_param_counts();
} else {
if (MLIAPModelKokkos<DeviceType>::comm->me == 0)
utils::logmesg(MLIAPModelKokkos<DeviceType>::lmp, "Loading python model deferred.\n");
}
}
/* ---------------------------------------------------------------------- */
// Finalize loading of the model.
template<class DeviceType>
void MLIAPModelPythonKokkos<DeviceType>::connect_param_counts()
{
PyGILState_STATE gstate = PyGILState_Ensure();
auto nontemplated_this = static_cast<MLIAPModelPythonKokkosDevice*>((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<DeviceType>::error->all(FLERR, "Loading python model failure.");
}
PyGILState_Release(gstate);
model_loaded = 1;
utils::logmesg(MLIAPModelKokkos<DeviceType>::lmp, "Loading python model complete.\n");
}
/* ---------------------------------------------------------------------- */
template<class DeviceType>
void MLIAPModelPythonKokkos<DeviceType>::compute_gradients(class MLIAPData *data)
{
if (!model_loaded) { MLIAPModelKokkos<DeviceType>::error->all(FLERR, "Model not loaded."); }
PyGILState_STATE gstate = PyGILState_Ensure();
auto nontemplated_this = static_cast<MLIAPModelPythonKokkosDevice*>((void*)this);
auto *kokkos_data = dynamic_cast<MLIAPDataKokkos<DeviceType>*>(data);
MLIAPDataKokkosDevice raw_data(*kokkos_data);
MLIAPPYKokkos_compute_gradients(nontemplated_this, &raw_data);
if (PyErr_Occurred()) {
PyErr_Print();
PyErr_Clear();
PyGILState_Release(gstate);
MLIAPModelKokkos<DeviceType>::error->all(FLERR, "Running python model failure.");
}
PyGILState_Release(gstate);
}
/* ---------------------------------------------------------------------- */
template<class DeviceType>
void MLIAPModelPythonKokkos<DeviceType>::compute_gradgrads(class MLIAPData *data)
{
MLIAPModelPython::compute_gradgrads(data);
}
/* ---------------------------------------------------------------------- */
template<class DeviceType>
void MLIAPModelPythonKokkos<DeviceType>::compute_force_gradients(class MLIAPData *data)
{
MLIAPModelPython::compute_force_gradients(data);
}
/* ---------------------------------------------------------------------- */
namespace LAMMPS_NS {
template class MLIAPModelPythonKokkos<LMPDeviceType>;
#ifdef LMP_KOKKOS_GPU
template class MLIAPModelPythonKokkos<LMPHostType>;
#endif
}
#endif

86
src/KOKKOS/mliap_model_python_kokkos.h Normal file
View File

@ -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 DeviceType>
class MLIAPModelPythonKokkos : public MLIAPModelPython, public MLIAPModelKokkos<DeviceType> {
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<LMPDeviceType> {
};
class MLIAPDataKokkosDevice {
public:
MLIAPDataKokkosDevice(MLIAPDataKokkos<LMPDeviceType> &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<LMPHostType> &base) : ndescriptors(-1),nlistatoms(-1)
{
// It cannot get here, but needed for compilation
}
#endif
};
}
#endif

31
src/KOKKOS/pair_mliap_kokkos.cpp
View File

@ -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<DeviceType>::PairMLIAPKokkos(class LAMMPS* l) : PairMLIAP(l)
kokkosable = 1;
execution_space = ExecutionSpaceFromDevice<DeviceType>::space;
datamask_modify = 0;
is_child=true;
}
/* ---------------------------------------------------------------------- */
@ -50,6 +54,10 @@ PairMLIAPKokkos<DeviceType>::~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<DeviceType>::allocate()
template<class DeviceType>
void PairMLIAPKokkos<DeviceType>::settings(int narg, char ** arg)
{
PairMLIAP::settings(narg, arg);
std::vector<char*> new_args;
int iarg=0;
while (iarg < narg) {
if (strcmp(arg[iarg],"model") == 0) {
@ -152,8 +160,19 @@ void PairMLIAPKokkos<DeviceType>::settings(int narg, char ** arg)
delete model;
model = new MLIAPModelLinearKokkos<DeviceType>(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<DeviceType>(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<DeviceType>::settings(int narg, char ** arg)
descriptor = new MLIAPDescriptorSO3Kokkos<DeviceType>(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());
}
/* ---------------------------------------------------------------------- */

17
src/ML-IAP/mliap_model_python.cpp
View File

@ -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");

2
src/ML-IAP/mliap_model_python.h
View File

@ -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;

22
src/ML-IAP/pair_mliap.cpp
View File

@ -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);

1
src/ML-IAP/pair_mliap.h
View File

@ -44,6 +44,7 @@ class PairMLIAP : public Pair {
class MLIAPModel *model;
class MLIAPDescriptor *descriptor;
class MLIAPData *data;
bool is_child;
};
} // namespace LAMMPS_NS

4
src/lammps.cpp
View File

@ -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