diff --git a/doc/development/create-a-model-pt.md b/doc/development/create-a-model-pt.md index 9df4fdc055..2ab0dee18f 100644 --- a/doc/development/create-a-model-pt.md +++ b/doc/development/create-a-model-pt.md @@ -184,6 +184,15 @@ If you place `SomeDescript` and `descrpt_some_args` into different modules, you After you install your new package, you can now use `dp train` to run your new model. +### Package customized C++ OPs + +You may need to use customized PyTorch C++ OPs in the new model. +Follow [PyTorch documentation](https://pytorch.org/tutorials/advanced/torch_script_custom_ops.html) to create one library. + +When using your customized C++ OPs in the Python interface, use {py:meth}`torch.ops.load_library` to load the OP library in the module defined in `entry_points`. + +When using your customized C++ OPs in the C++ library, define the environment variable {envvar}`DP_PLUGIN_PATH` to load the OP library. + ## Unit tests ### Universal tests diff --git a/doc/development/create-a-model-tf.md b/doc/development/create-a-model-tf.md index a108514e28..9ab3525bb5 100644 --- a/doc/development/create-a-model-tf.md +++ b/doc/development/create-a-model-tf.md @@ -71,3 +71,12 @@ where `deepmd_some_descrtpt` is the module of your codes. It is equivalent to `f If you place `SomeDescript` and `descrpt_some_args` into different modules, you are also expected to add `descrpt_some_args` to `entry_points`. After you install your new package, you can now use `dp train` to run your new model. + +### Package customized C++ OPs + +You may need to use customized TensorFlow C++ OPs in the new model. +Follow [TensorFlow documentation](https://www.tensorflow.org/guide/create_op) to create one library. + +When using your customized C++ OPs in the Python interface, use {py:meth}`tf.load_op_library` to load the OP library in the module defined in `entry_points`. + +When using your customized C++ OPs in the C++ library, define the environment variable {envvar}`DP_PLUGIN_PATH` to load the OP library. diff --git a/doc/env.md b/doc/env.md index 7dfd2d8827..65a50ff163 100644 --- a/doc/env.md +++ b/doc/env.md @@ -77,3 +77,15 @@ Default backend. See [PyTorch documentation](https://pytorch.org/docs/stable/data.html) for details. ::: + +## C++ interface only + +These environment variables also apply to third-party programs using the C++ interface, such as [LAMMPS](./third-party/lammps-command.md). + +:::{envvar} DP_PLUGIN_PATH + +**Type**: List of paths, split by `:` on Unix and `;` on Windows + +List of customized OP plugin libraries to load, such as `/path/to/plugin1.so:/path/to/plugin2.so` on Linux and `/path/to/plugin1.dll;/path/to/plugin2.dll` on Windows. + +::: diff --git a/source/api_cc/src/common.cc b/source/api_cc/src/common.cc index 12309df6a0..baa257d60e 100644 --- a/source/api_cc/src/common.cc +++ b/source/api_cc/src/common.cc @@ -377,12 +377,10 @@ void deepmd::get_env_nthreads(int& num_intra_nthreads, } } -static inline void _load_single_op_library(std::string library_name) { +static inline void _load_library_path(std::string dso_path) { #if defined(_WIN32) - std::string dso_path = library_name + ".dll"; void* dso_handle = LoadLibrary(dso_path.c_str()); #else - std::string dso_path = "lib" + library_name + ".so"; void* dso_handle = dlopen(dso_path.c_str(), RTLD_NOW | RTLD_LOCAL); #endif if (!dso_handle) { @@ -392,6 +390,15 @@ static inline void _load_single_op_library(std::string library_name) { } } +static inline void _load_single_op_library(std::string library_name) { +#if defined(_WIN32) + std::string dso_path = library_name + ".dll"; +#else + std::string dso_path = "lib" + library_name + ".so"; +#endif + _load_library_path(dso_path); +} + void deepmd::load_op_library() { #ifdef BUILD_TENSORFLOW _load_single_op_library("deepmd_op"); @@ -399,6 +406,23 @@ void deepmd::load_op_library() { #ifdef BUILD_PYTORCH _load_single_op_library("deepmd_op_pt"); #endif + // load customized plugins + const char* env_customized_plugins = std::getenv("DP_PLUGIN_PATH"); + if (env_customized_plugins) { +#if !defined(_WIN32) + // note: ":" is a string and ':' is a char + std::string pathvarsep = ":"; +#else + std::string pathvarsep = ";"; +#endif + std::string plugin_path(env_customized_plugins); + std::vector plugin_paths = split(plugin_path, pathvarsep); + for (const auto& plugin : plugin_paths) { + std::cerr << "Loading customized plugin defined in DP_PLUGIN_PATH: " + << plugin << std::endl; + _load_library_path(plugin); + } + } } std::string deepmd::name_prefix(const std::string& scope) {