dynamically load op library in C++ interface (#1384)

* dynamically load op library in C++ interface C++ interface will dynamically load OP libraries, just like Python interface, so it no longer needs linking. * version cannot be NULL... * set LD_LIBRARY_PATH * install runUnitTest * remove CMAKE_LINK_WHAT_YOU_USE It is not necessary anymore. * also remove CMAKE_LINK_WHAT_YOU_USE in api_cc/tests * change the type of op library to MODULE * add the absolute path of library directory to cc rpath * Revert "add the absolute path of library directory to cc rpath" This reverts commit fabdac9. * add `-Wl,--disable-new-dtags` * Revert "add `-Wl,--disable-new-dtags`" This reverts commit ecccb57. * dlopen from dp lib but not TF
deepmodeling · Jan 21, 2022 · 7068698 · 7068698
1 parent 057e6ab
commit 7068698
Show file tree

Hide file tree

Showing 13 changed files with 61 additions and 21 deletions.
diff --git a/.gitignore b/.gitignore
@@ -29,6 +29,8 @@ _templates
 API_CC
 doc/api_py/
 dp/
+dp_test/
+dp_test_cc/
 build_lammps/
 .idea/
 build_tests/

diff --git a/source/CMakeLists.txt b/source/CMakeLists.txt
@@ -1,8 +1,5 @@
 cmake_minimum_required(VERSION 3.7)
 project(DeePMD)
-if (CMAKE_COMPILER_IS_GNUCXX)
-  set(CMAKE_LINK_WHAT_YOU_USE TRUE)
-endif ()
 
 # build cpp or python interfaces
 if (NOT DEFINED BUILD_CPP_IF) 

diff --git a/source/api_cc/CMakeLists.txt b/source/api_cc/CMakeLists.txt
@@ -17,7 +17,6 @@ add_library(${libname} SHARED ${LIB_SRC})
 
 # link: libdeepmd libdeepmd_op libtensorflow_cc libtensorflow_framework
 target_link_libraries (${libname} PUBLIC ${LIB_DEEPMD}	${TensorFlow_LIBRARY} ${TensorFlowFramework_LIBRARY})
-target_link_libraries (${libname} PRIVATE	${LIB_DEEPMD_OP})
 target_include_directories(${libname} PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include ${CMAKE_CURRENT_BINARY_DIR} ${TensorFlow_INCLUDE_DIRS})
 
 set_target_properties(

diff --git a/source/api_cc/include/common.h b/source/api_cc/include/common.h
@@ -116,6 +116,12 @@ void
 get_env_nthreads(int & num_intra_nthreads,
 		 int & num_inter_nthreads);
 
+/**
+ * @brief Dynamically load OP library. This should be called before loading graphs.
+ */
+void
+load_op_library();
+
 /** @struct deepmd::deepmd_exception
  **/
 

diff --git a/source/api_cc/src/DataModifier.cc b/source/api_cc/src/DataModifier.cc
@@ -33,6 +33,7 @@ init (const std::string & model,
   get_env_nthreads(num_intra_nthreads, num_inter_nthreads);
   options.config.set_inter_op_parallelism_threads(num_inter_nthreads);
   options.config.set_intra_op_parallelism_threads(num_intra_nthreads);
+  deepmd::load_op_library();
   deepmd::check_status(NewSession(options, &session));
   deepmd::check_status(ReadBinaryProto(Env::Default(), model, &graph_def));
   deepmd::check_status(session->Create(graph_def));  

diff --git a/source/api_cc/src/DeepPot.cc b/source/api_cc/src/DeepPot.cc
@@ -189,6 +189,7 @@ init (const std::string & model, const int & gpu_rank, const std::string & file_
   get_env_nthreads(num_intra_nthreads, num_inter_nthreads);
   options.config.set_inter_op_parallelism_threads(num_inter_nthreads);
   options.config.set_intra_op_parallelism_threads(num_intra_nthreads);
+  deepmd::load_op_library();
 
   if(file_content.size() == 0)
     check_status (ReadBinaryProto(Env::Default(), model, &graph_def));

diff --git a/source/api_cc/src/DeepTensor.cc b/source/api_cc/src/DeepTensor.cc
@@ -33,6 +33,7 @@ init (const std::string & model,
   get_env_nthreads(num_intra_nthreads, num_inter_nthreads);
   options.config.set_inter_op_parallelism_threads(num_inter_nthreads);
   options.config.set_intra_op_parallelism_threads(num_intra_nthreads);
+  deepmd::load_op_library();
   deepmd::check_status (NewSession(options, &session));
   deepmd::check_status (ReadBinaryProto(Env::Default(), model, &graph_def));
   deepmd::check_status (session->Create(graph_def));  

diff --git a/source/api_cc/src/common.cc b/source/api_cc/src/common.cc
@@ -1,6 +1,7 @@
 #include "common.h"
 #include "AtomMap.h"
 #include "device.h"
+#include <dlfcn.h>
 
 using namespace tensorflow;
 
@@ -228,6 +229,18 @@ get_env_nthreads(int & num_intra_nthreads,
   }
 }
 
+void
+deepmd::
+load_op_library()
+{
+  tensorflow::Env* env = tensorflow::Env::Default();
+  std::string dso_path = env->FormatLibraryFileName("deepmd_op", "");
+  void* dso_handle = dlopen(dso_path.c_str(), RTLD_NOW | RTLD_LOCAL);
+  if (!dso_handle) {
+    throw deepmd::deepmd_exception(dso_path + " is not found! You can add the library directory to LD_LIBRARY_PATH");
+  }
+}
+
 std::string
 deepmd::
 name_prefix(const std::string & scope)

diff --git a/source/api_cc/tests/CMakeLists.txt b/source/api_cc/tests/CMakeLists.txt
@@ -1,6 +1,5 @@
 cmake_minimum_required(VERSION 3.9)
 project(deepmd_api_test)
-set(CMAKE_LINK_WHAT_YOU_USE TRUE)
 
 if (NOT DEFINED BUILD_CPP_IF) 
   set(BUILD_CPP_IF TRUE)
@@ -99,11 +98,11 @@ else()
 endif()
 
 if (USE_CUDA_TOOLKIT)
-  target_link_libraries(runUnitTests gtest gtest_main ${libname} ${apiname} ${opname} pthread ${TensorFlow_LIBRARY} rt deepmd_op_cuda coverage_config)
+  target_link_libraries(runUnitTests gtest gtest_main ${libname} ${apiname} pthread ${TensorFlow_LIBRARY} rt deepmd_op_cuda coverage_config)
 elseif(USE_ROCM_TOOLKIT)
-  target_link_libraries(runUnitTests gtest gtest_main ${libname} ${apiname} ${opname} pthread ${TensorFlow_LIBRARY} rt deepmd_op_rocm coverage_config)
+  target_link_libraries(runUnitTests gtest gtest_main ${libname} ${apiname} pthread ${TensorFlow_LIBRARY} rt deepmd_op_rocm coverage_config)
 else()
-  target_link_libraries(runUnitTests gtest gtest_main ${libname} ${apiname} ${opname} pthread ${TensorFlow_LIBRARY} rt coverage_config)
+  target_link_libraries(runUnitTests gtest gtest_main ${libname} ${apiname} pthread ${TensorFlow_LIBRARY} rt coverage_config)
 endif()
 
 find_package(Protobuf)
@@ -134,3 +133,16 @@ if(NOT GTEST_LIBRARIES)
 else ()
   include_directories(${GTEST_INCLUDE_DIRS})
 endif ()
+
+set_target_properties(
+  runUnitTests
+  PROPERTIES 
+  INSTALL_RPATH "$ORIGIN/../lib"
+)
+set_target_properties(
+  ${apiname}
+  PROPERTIES 
+  INSTALL_RPATH "$ORIGIN;${TensorFlow_LIBRARY_PATH}"
+)
+install(TARGETS runUnitTests DESTINATION bin/)
+install(TARGETS ${libname} ${apiname} ${opname} DESTINATION lib/)
diff --git a/source/install/test_cc.sh b/source/install/test_cc.sh
@@ -7,29 +7,32 @@ NPROC=$(nproc --all)
 
 #------------------
 
+INSTALL_PREFIX=${SCRIPT_PATH}/../../dp_test
 BUILD_TMP_DIR=${SCRIPT_PATH}/../build_tests
 mkdir -p ${BUILD_TMP_DIR}
 cd ${BUILD_TMP_DIR}
-cmake ../lib/tests
+cmake ../lib/tests -DCMAKE_INSTALL_PREFIX=${INSTALL_PREFIX}
 make -j${NPROC}
+make install
 
 #------------------
-${BUILD_TMP_DIR}/runUnitTests
+${INSTALL_PREFIX}/bin/runUnitTests
 
 
 #------------------
 
 BUILD_TMP_DIR=${SCRIPT_PATH}/../build_cc_tests
-INSTALL_PREFIX=${SCRIPT_PATH}/../../dp
+INSTALL_PREFIX=${SCRIPT_PATH}/../../dp_test_cc
 mkdir -p ${BUILD_TMP_DIR}
 mkdir -p ${INSTALL_PREFIX}
 cd ${BUILD_TMP_DIR}
 cmake -DINSTALL_TENSORFLOW=TRUE -DCMAKE_INSTALL_PREFIX=${INSTALL_PREFIX} ../api_cc/tests
 make -j${NPROC}
+make install
 
 #------------------
 cd ${SCRIPT_PATH}/../api_cc/tests
-${BUILD_TMP_DIR}/runUnitTests
+${INSTALL_PREFIX}/bin/runUnitTests
 
 #------------------
 # upload to codecov

diff --git a/source/install/test_cc_local.sh b/source/install/test_cc_local.sh
@@ -6,27 +6,29 @@ SCRIPT_PATH=$(dirname $(realpath -s $0))
 NPROC=$(nproc --all)
 
 #------------------
-
+INSTALL_PREFIX=${SCRIPT_PATH}/../../dp_test
 BUILD_TMP_DIR=${SCRIPT_PATH}/../build_tests
 mkdir -p ${BUILD_TMP_DIR}
 cd ${BUILD_TMP_DIR}
-cmake ../lib/tests
+cmake ../lib/tests -DCMAKE_INSTALL_PREFIX=${INSTALL_PREFIX}
 make -j${NPROC}
+make install
 
 #------------------
-${BUILD_TMP_DIR}/runUnitTests
+${INSTALL_PREFIX}/bin/runUnitTests
 
 
 #------------------
 
 echo "try to find tensorflow in ${tensorflow_root}"
 BUILD_TMP_DIR=${SCRIPT_PATH}/../build_cc_tests
-INSTALL_PREFIX=${SCRIPT_PATH}/../../dp
+INSTALL_PREFIX=${SCRIPT_PATH}/../../dp_test_cc
 mkdir -p ${BUILD_TMP_DIR}
 cd ${BUILD_TMP_DIR}
-cmake -DINSTALL_TENSORFLOW=FALSE -DTENSORFLOW_ROOT=${tensorflow_root} ../api_cc/tests
+cmake -DINSTALL_TENSORFLOW=FALSE -DTENSORFLOW_ROOT=${tensorflow_root} ../api_cc/tests -DCMAKE_INSTALL_PREFIX=${INSTALL_PREFIX}
 make -j${NPROC}
+make install
 
 #------------------
 cd ${SCRIPT_PATH}/../api_cc/tests
-${BUILD_TMP_DIR}/runUnitTests
+${INSTALL_PREFIX}/bin/runUnitTests
diff --git a/source/lib/tests/CMakeLists.txt b/source/lib/tests/CMakeLists.txt
@@ -106,3 +106,6 @@ if(NOT GTEST_LIBRARIES)
 else ()
   include_directories(${GTEST_INCLUDE_DIRS})
 endif ()
+
+install(TARGETS runUnitTests DESTINATION bin/)
+install(TARGETS runUnitTests DESTINATION lib/)
diff --git a/source/op/CMakeLists.txt b/source/op/CMakeLists.txt
@@ -8,7 +8,7 @@ file(GLOB OP_GRADS_SRC custom_op.cc prod_force_grad.cc prod_force_grad_multi_dev
 file(GLOB OP_PY *.py)
 
 if (BUILD_CPP_IF) 
-  add_library(${LIB_DEEPMD_OP} SHARED ${OP_SRC})
+  add_library(${LIB_DEEPMD_OP} MODULE ${OP_SRC})
   # link: libdeepmd libtensorflow_cc libtensorflow_framework
   target_link_libraries (${LIB_DEEPMD_OP} PUBLIC ${TensorFlow_LIBRARY} ${TensorFlowFramework_LIBRARY})
   target_link_libraries (${LIB_DEEPMD_OP} PRIVATE ${LIB_DEEPMD})
@@ -17,8 +17,8 @@ if (BUILD_CPP_IF)
 endif (BUILD_CPP_IF)
 
 if (BUILD_PY_IF)
-  add_library(op_abi SHARED ${OP_SRC} ${OP_LIB})
-  add_library(op_grads SHARED ${OP_GRADS_SRC})
+  add_library(op_abi MODULE ${OP_SRC} ${OP_LIB})
+  add_library(op_grads MODULE ${OP_GRADS_SRC})
 
   message(STATUS ${TensorFlowFramework_LIBRARY})
   # link: libdeepmd libtensorflow_framework