feat: load customized OP library in the C++ interface (#4073)

Fix #4041.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit


- **New Features**
- Added documentation for integrating customized C++ operations (OPs) in
both PyTorch and TensorFlow models.
- Introduced guidelines on using the `DP_PLUGIN_PATH` environment
variable for loading custom OP libraries.
- Clarified the usage of environment variables in the context of a C++
interface, improving usability for third-party programs.

- **Documentation**
- Enhanced documentation to provide clearer instructions for model
development with custom C++ functionality.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>

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

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.

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.
+
+:::

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,13 +390,39 @@ 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");
 #endif
 #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<std::string> 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) {