setup PyTorch C++ interface build environement (#3169)

See #3120. - CMake: add `ENABLE_TENSORFLOW` and `ENABLE_PYTORCH`. `BUILD_TENSORFLOW` will be enabled when `TENSORFLOW_ROOT` is not empty or `USE_TF_PYTHON_LIBS` is on. - api_cc: add `BUILD_TENSORFLOW` and `BUILD_PYTORCH` defination. Move several functions from `common.h` to `commonTF.h` to prevent exposing them to header files. - CI: download libtorch in the build/test CC actions. --------- Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
deepmodeling · Jan 24, 2024 · dd53e07 · dd53e07
1 parent 68fb16d
commit dd53e07
Show file tree

Hide file tree

Showing 21 changed files with 298 additions and 154 deletions.
diff --git a/.github/workflows/build_cc.yml b/.github/workflows/build_cc.yml
@@ -27,6 +27,10 @@ jobs:
         cache: 'pip'
     - uses: lukka/get-cmake@latest
     - run: python -m pip install tensorflow
+    - name: Download libtorch
+      run: |
+         wget https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.1.2%2Bcpu.zip -O libtorch.zip
+         unzip libtorch.zip
     - run: |
          wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.0-1_all.deb \
          && sudo dpkg -i cuda-keyring_1.0-1_all.deb \
@@ -48,13 +52,17 @@ jobs:
          && sudo apt-get update \
          && sudo apt-get install -y rocm-dev hipcub-dev
       if: matrix.variant == 'rocm'
-    - run: source/install/build_cc.sh
+    - run: |
+        export CMAKE_PREFIX_PATH=$GITHUB_WORKSPACE/libtorch
+        source/install/build_cc.sh
       env:
         DP_VARIANT: ${{ matrix.dp_variant }}
         DOWNLOAD_TENSORFLOW: "FALSE"
         CMAKE_GENERATOR: Ninja
       if: matrix.variant != 'clang'
-    - run: source/install/build_cc.sh
+    - run: |
+        export CMAKE_PREFIX_PATH=$GITHUB_WORKSPACE/libtorch
+        source/install/build_cc.sh
       env:
         DP_VARIANT: cpu
         DOWNLOAD_TENSORFLOW: "FALSE"

diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml
@@ -37,6 +37,8 @@ jobs:
          && sudo apt-get update \
          && sudo apt-get -y install cuda-cudart-dev-12-2 cuda-nvcc-12-2
         python -m pip install tensorflow
+        wget https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.1.2%2Bcpu.zip -O libtorch.zip
+        unzip libtorch.zip
       env:
         DEBIAN_FRONTEND: noninteractive
     # Initializes the CodeQL tools for scanning.
@@ -46,7 +48,9 @@ jobs:
         languages: ${{ matrix.language }}
         queries: security-extended,security-and-quality
     - name: "Run, Build Application using script"
-      run: source/install/build_cc.sh
+      run: |
+        export CMAKE_PREFIX_PATH=$GITHUB_WORKSPACE/libtorch
+        source/install/build_cc.sh
       env:
         DP_VARIANT: cuda
         DOWNLOAD_TENSORFLOW: "FALSE"

diff --git a/.github/workflows/test_cc.yml b/.github/workflows/test_cc.yml
@@ -18,7 +18,13 @@ jobs:
         mpi: mpich
     - uses: lukka/get-cmake@latest
     - run: python -m pip install tensorflow
-    - run: source/install/test_cc_local.sh
+    - name: Download libtorch
+      run: |
+         wget https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.1.2%2Bcpu.zip -O libtorch.zip
+         unzip libtorch.zip
+    - run: |
+         export CMAKE_PREFIX_PATH=$GITHUB_WORKSPACE/libtorch
+         source/install/test_cc_local.sh
       env:
         OMP_NUM_THREADS: 1
         TF_INTRA_OP_PARALLELISM_THREADS: 1

diff --git a/doc/conf.py b/doc/conf.py
@@ -213,7 +213,10 @@ def setup(app):
 exhale_projects_args = {
     "cc": {
         "containmentFolder": "./API_CC",
-        "exhaleDoxygenStdin": "INPUT = ../source/api_cc/include/",
+        "exhaleDoxygenStdin": """INPUT = ../source/api_cc/include/
+                                 PREDEFINED += BUILD_TENSORFLOW
+                                               BUILD_PYTORCH
+        """,
         "rootFileTitle": "C++ API",
         "rootFileName": "api_cc.rst",
     },

diff --git a/source/CMakeLists.txt b/source/CMakeLists.txt
@@ -2,6 +2,8 @@
 cmake_minimum_required(VERSION 3.16)
 project(DeePMD)
 
+option(ENABLE_TENSORFLOW "Enable TensorFlow interface" OFF)
+option(ENABLE_PYTORCH "Enable PyTorch interface" OFF)
 option(BUILD_TESTING "Build test and enable converage" OFF)
 set(DEEPMD_C_ROOT
     ""
@@ -131,6 +133,7 @@ if(INSTALL_TENSORFLOW)
   set(USE_TF_PYTHON_LIBS TRUE)
 endif(INSTALL_TENSORFLOW)
 if(USE_TF_PYTHON_LIBS)
+  set(ENABLE_TENSORFLOW TRUE)
   if(NOT "$ENV{CIBUILDWHEEL}" STREQUAL "1")
     find_package(
       Python
@@ -141,11 +144,31 @@ if(USE_TF_PYTHON_LIBS)
     set(PYTHON_INCLUDE_DIRS ${PYTHON_INCLUDE_DIR})
   endif()
 endif(USE_TF_PYTHON_LIBS)
+if(TENSORFLOW_ROOT)
+  set(ENABLE_TENSORFLOW TRUE)
+endif()
 
 # find tensorflow, I need tf abi info
-if(NOT DEEPMD_C_ROOT)
+if(ENABLE_TENSORFLOW AND NOT DEEPMD_C_ROOT)
   find_package(tensorflow REQUIRED)
 endif()
+if(ENABLE_PYTORCH AND NOT DEEPMD_C_ROOT)
+  find_package(Torch REQUIRED)
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")
+endif()
+# log enabled backends
+if(NOT DEEPMD_C_ROOT)
+  message(STATUS "Enabled backends:")
+  if(ENABLE_TENSORFLOW)
+    message(STATUS "- TensorFlow")
+  endif()
+  if(ENABLE_PYTORCH)
+    message(STATUS "- PyTorch")
+  endif()
+  if(NOT ENABLE_TENSORFLOW AND NOT ENABLE_PYTORCH)
+    message(FATAL_ERROR "No backend is enabled.")
+  endif()
+endif()
 
 # find threads
 find_package(Threads)
@@ -233,10 +256,13 @@ if(DEEPMD_C_ROOT)
   # use variable for TF path to set deepmd_c path
   set(TensorFlow_LIBRARY_PATH "${DEEPMD_C_ROOT}/lib")
   set(TENSORFLOW_INCLUDE_DIRS "${DEEPMD_C_ROOT}/include")
+  set(TORCH_LIBRARIES "${DEEPMD_C_ROOT}/lib/libtorch.so")
 endif()
 
 if(NOT DEEPMD_C_ROOT)
-  add_subdirectory(op/)
+  if(ENABLE_TENSORFLOW)
+    add_subdirectory(op/)
+  endif()
   add_subdirectory(lib/)
 endif()
 if(BUILD_PY_IF)

diff --git a/source/api_cc/CMakeLists.txt b/source/api_cc/CMakeLists.txt
@@ -11,8 +11,16 @@ add_library(${libname} SHARED ${LIB_SRC})
 
 # link: libdeepmd libdeepmd_op libtensorflow_cc libtensorflow_framework
 target_link_libraries(${libname} PUBLIC ${LIB_DEEPMD})
-target_link_libraries(${libname} PRIVATE TensorFlow::tensorflow_cc
-                                         TensorFlow::tensorflow_framework)
+if(ENABLE_TENSORFLOW)
+  target_link_libraries(${libname} PRIVATE TensorFlow::tensorflow_cc
+                                           TensorFlow::tensorflow_framework)
+  target_compile_definitions(${libname} PRIVATE BUILD_TENSORFLOW)
+endif()
+if(ENABLE_PYTORCH)
+  target_link_libraries(${libname} PRIVATE "${TORCH_LIBRARIES}")
+  target_compile_definitions(${libname} PRIVATE BUILD_PYTORCH)
+endif()
+
 target_include_directories(
   ${libname}
   PUBLIC $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>
@@ -55,3 +63,14 @@ ${CMAKE_INSTALL_PREFIX}/lib/${CMAKE_SHARED_LIBRARY_PREFIX}${libname}${LOW_PREC_V
     add_subdirectory(tests)
   endif()
 endif(BUILD_PY_IF)
+
+if(BUILD_TESTING)
+  # A compilation test to make sure api_cc can compile without any backend
+  add_library(deepmd_cc_test_no_backend SHARED ${LIB_SRC})
+  target_link_libraries(deepmd_cc_test_no_backend PUBLIC ${LIB_DEEPMD})
+  target_include_directories(
+    deepmd_cc_test_no_backend
+    PUBLIC $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>
+           $<BUILD_INTERFACE:${CMAKE_CURRENT_BINARY_DIR}>
+           $<INSTALL_INTERFACE:include>)
+endif()
diff --git a/source/api_cc/include/DataModifierTF.h b/source/api_cc/include/DataModifierTF.h
@@ -3,6 +3,7 @@
 
 #include "DataModifier.h"
 #include "common.h"
+#include "commonTF.h"
 
 namespace deepmd {
 /**

diff --git a/source/api_cc/include/DeepPotTF.h b/source/api_cc/include/DeepPotTF.h
@@ -3,6 +3,7 @@
 
 #include "DeepPot.h"
 #include "common.h"
+#include "commonTF.h"
 #include "neighbor_list.h"
 
 namespace deepmd {

diff --git a/source/api_cc/include/DeepTensorTF.h b/source/api_cc/include/DeepTensorTF.h
@@ -3,6 +3,7 @@
 
 #include "DeepTensor.h"
 #include "common.h"
+#include "commonTF.h"
 #include "neighbor_list.h"
 
 namespace deepmd {

diff --git a/source/api_cc/include/common.h b/source/api_cc/include/common.h
@@ -10,12 +10,6 @@
 #include "neighbor_list.h"
 #include "version.h"
 
-#ifdef TF_PRIVATE
-#include "tf_private.h"
-#else
-#include "tf_public.h"
-#endif
-
 namespace deepmd {
 
 typedef double ENERGYTYPE;
@@ -175,143 +169,8 @@ struct tf_exception : public deepmd::deepmd_exception {
       : deepmd::deepmd_exception(std::string("TensorFlow Error: ") + msg){};
 };
 
-/**
- * @brief Check TensorFlow status. Exit if not OK.
- * @param[in] status TensorFlow status.
- **/
-void check_status(const tensorflow::Status& status);
-
 std::string name_prefix(const std::string& name_scope);
 
-/**
- * @brief Get the value of a tensor.
- * @param[in] session TensorFlow session.
- * @param[in] name The name of the tensor.
- * @param[in] scope The scope of the tensor.
- * @return The value of the tensor.
- **/
-template <typename VT>
-VT session_get_scalar(tensorflow::Session* session,
-                      const std::string name,
-                      const std::string scope = "");
-
-/**
- * @brief Get the vector of a tensor.
- * @param[out] o_vec The output vector.
- * @param[in] session TensorFlow session.
- * @param[in] name The name of the tensor.
- * @param[in] scope The scope of the tensor.
- **/
-template <typename VT>
-void session_get_vector(std::vector<VT>& o_vec,
-                        tensorflow::Session* session,
-                        const std::string name_,
-                        const std::string scope = "");
-
-/**
- * @brief Get the type of a tensor.
- * @param[in] session TensorFlow session.
- * @param[in] name The name of the tensor.
- * @param[in] scope The scope of the tensor.
- * @return The type of the tensor as int.
- **/
-int session_get_dtype(tensorflow::Session* session,
-                      const std::string name,
-                      const std::string scope = "");
-
-/**
- * @brief Get input tensors.
- * @param[out] input_tensors Input tensors.
- * @param[in] dcoord_ Coordinates of atoms.
- * @param[in] ntypes Number of atom types.
- * @param[in] datype_ Atom types.
- * @param[in] dbox Box matrix.
- * @param[in] cell_size Cell size.
- * @param[in] fparam_ Frame parameters.
- * @param[in] aparam_ Atom parameters.
- * @param[in] atommap Atom map.
- * @param[in] scope The scope of the tensors.
- * @param[in] aparam_nall Whether the atomic dimesion of atomic parameters is
- * nall.
- */
-template <typename MODELTYPE, typename VALUETYPE>
-int session_input_tensors(
-    std::vector<std::pair<std::string, tensorflow::Tensor>>& input_tensors,
-    const std::vector<VALUETYPE>& dcoord_,
-    const int& ntypes,
-    const std::vector<int>& datype_,
-    const std::vector<VALUETYPE>& dbox,
-    const double& cell_size,
-    const std::vector<VALUETYPE>& fparam_,
-    const std::vector<VALUETYPE>& aparam_,
-    const deepmd::AtomMap& atommap,
-    const std::string scope = "",
-    const bool aparam_nall = false);
-
-/**
- * @brief Get input tensors.
- * @param[out] input_tensors Input tensors.
- * @param[in] dcoord_ Coordinates of atoms.
- * @param[in] ntypes Number of atom types.
- * @param[in] datype_ Atom types.
- * @param[in] dlist Neighbor list.
- * @param[in] fparam_ Frame parameters.
- * @param[in] aparam_ Atom parameters.
- * @param[in] atommap Atom map.
- * @param[in] nghost Number of ghost atoms.
- * @param[in] ago Update the internal neighbour list if ago is 0.
- * @param[in] scope The scope of the tensors.
- * @param[in] aparam_nall Whether the atomic dimesion of atomic parameters is
- * nall.
- */
-template <typename MODELTYPE, typename VALUETYPE>
-int session_input_tensors(
-    std::vector<std::pair<std::string, tensorflow::Tensor>>& input_tensors,
-    const std::vector<VALUETYPE>& dcoord_,
-    const int& ntypes,
-    const std::vector<int>& datype_,
-    const std::vector<VALUETYPE>& dbox,
-    InputNlist& dlist,
-    const std::vector<VALUETYPE>& fparam_,
-    const std::vector<VALUETYPE>& aparam_,
-    const deepmd::AtomMap& atommap,
-    const int nghost,
-    const int ago,
-    const std::string scope = "",
-    const bool aparam_nall = false);
-
-/**
- * @brief Get input tensors for mixed type.
- * @param[out] input_tensors Input tensors.
- * @param[in] nframes Number of frames.
- * @param[in] dcoord_ Coordinates of atoms.
- * @param[in] ntypes Number of atom types.
- * @param[in] datype_ Atom types.
- * @param[in] dlist Neighbor list.
- * @param[in] fparam_ Frame parameters.
- * @param[in] aparam_ Atom parameters.
- * @param[in] atommap Atom map.
- * @param[in] nghost Number of ghost atoms.
- * @param[in] ago Update the internal neighbour list if ago is 0.
- * @param[in] scope The scope of the tensors.
- * @param[in] aparam_nall Whether the atomic dimesion of atomic parameters is
- * nall.
- */
-template <typename MODELTYPE, typename VALUETYPE>
-int session_input_tensors_mixed_type(
-    std::vector<std::pair<std::string, tensorflow::Tensor>>& input_tensors,
-    const int& nframes,
-    const std::vector<VALUETYPE>& dcoord_,
-    const int& ntypes,
-    const std::vector<int>& datype_,
-    const std::vector<VALUETYPE>& dbox,
-    const double& cell_size,
-    const std::vector<VALUETYPE>& fparam_,
-    const std::vector<VALUETYPE>& aparam_,
-    const deepmd::AtomMap& atommap,
-    const std::string scope = "",
-    const bool aparam_nall = false);
-
 /**
  * @brief Read model file to a string.
  * @param[in] model Path to the model.