Enable PyTorch backend for DeePMD-kit v3 (#5253)

* enable PyTorch backend for DeePMD-kit v3

* update CMakeLists.txt

Author: YuLiu98

CMakeLists.txt

@@ -436,22 +436,6 @@ if(ENABLE_FLOAT_FFTW)
 endif()
 
 if(ENABLE_DEEPKS)
-  # Torch uses outdated components to detect CUDA arch, causing failure on
-  # latest CUDA kits. Set CMake variable TORCH_CUDA_ARCH_LIST in the form of
-  # "major.minor" if required.
-  find_package(Torch REQUIRED)
-  if(NOT Torch_VERSION VERSION_LESS "2.1.0")
-    set_if_higher(CMAKE_CXX_STANDARD 17)
-  elseif(NOT Torch_VERSION VERSION_LESS "1.5.0")
-    set_if_higher(CMAKE_CXX_STANDARD 14)
-  endif()
-  include_directories(${TORCH_INCLUDE_DIRS})
-  if(MKL_FOUND)
-    list(PREPEND math_libs ${TORCH_LIBRARIES})
-  else()
-    list(APPEND math_libs ${TORCH_LIBRARIES})
-  endif()
-  add_compile_options(${TORCH_CXX_FLAGS})
   target_link_libraries(${ABACUS_BIN_NAME} deepks)
 
   find_path(libnpy_SOURCE_DIR npy.hpp HINTS ${libnpy_INCLUDE_DIR})
@@ -470,6 +454,25 @@ if(ENABLE_DEEPKS)
   add_compile_definitions(__DEEPKS)
 endif()
 
+# Torch uses outdated components to detect CUDA arch, causing failure on
+# latest CUDA kits. Set CMake variable TORCH_CUDA_ARCH_LIST in the form of
+# "major.minor" if required.
+if(ENABLE_DEEPKS OR DEFINED Torch_DIR)
+  find_package(Torch REQUIRED)
+  if(NOT Torch_VERSION VERSION_LESS "2.1.0")
+    set_if_higher(CMAKE_CXX_STANDARD 17)
+  elseif(NOT Torch_VERSION VERSION_LESS "1.5.0")
+    set_if_higher(CMAKE_CXX_STANDARD 14)
+  endif()
+  include_directories(${TORCH_INCLUDE_DIRS})
+  if(MKL_FOUND)
+    list(PREPEND math_libs ${TORCH_LIBRARIES})
+  else()
+    list(APPEND math_libs ${TORCH_LIBRARIES})
+  endif()
+  add_compile_options(${TORCH_CXX_FLAGS})
+endif()
+
 if (ENABLE_CNPY)
   find_path(cnpy_SOURCE_DIR
     cnpy.h
@@ -590,13 +593,14 @@ if(DEFINED DeePMD_DIR)
     add_compile_definitions(__DPMDC)
   else()
     target_link_libraries(${ABACUS_BIN_NAME} DeePMD::deepmd_cc)
-    if(NOT DEFINED TensorFlow_DIR)
-      set(TensorFlow_DIR ${DeePMD_DIR})
-    endif()
-    find_package(TensorFlow REQUIRED)
-    if(TensorFlow_FOUND)
-      target_link_libraries(${ABACUS_BIN_NAME} TensorFlow::tensorflow_cc)
-    endif()
+  endif()
+endif()
+
+if(DEFINED TensorFlow_DIR)
+  find_package(TensorFlow REQUIRED)
+  include_directories(${TensorFlow_DIR}/include)
+  if(TensorFlow_FOUND)
+    target_link_libraries(${ABACUS_BIN_NAME} TensorFlow::tensorflow_cc)
   endif()
 endif()
 

docs/advanced/install.md

@@ -38,13 +38,22 @@ cmake -B build -DENABLE_DEEPKS=1 -DTorch_DIR=~/libtorch/share/cmake/Torch/ -Dlib
 If the Deep Potential model is employed in Molecule Dynamics calculations, the following prerequisites and steps are needed:
 
 - [DeePMD-kit](https://github.com/deepmodeling/deepmd-kit)
-- [TensorFlow](https://www.tensorflow.org/)
+- [TensorFlow](https://www.tensorflow.org/) (optional)
+- [LibTorch](https://pytorch.org/) (optional)
 
+In the simplest case, the `tensorflow_cc` and `torch` libraries are in the same directory as the `deepmd_c`/`deepmd_cc` libraries, then
 ```bash
-cmake -B build -DDeePMD_DIR=~/deepmd-kit -DTensorFlow_DIR=~/tensorflow
+cmake -B build -DDeePMD_DIR=/dir_to_deepmd-kit
 ```
+DeePMD-kit supports TensorFlow backend but its libraries are placed at another directory, then
 
-> `deepmd_c`/`deepmd_cc` and `tensorflow_cc` libraries would be called according to `DeePMD_DIR` and `TensorFlow_DIR`, which is showed in detail in [this page](https://github.com/deepmodeling/deepmd-kit/blob/master/doc/inference/cxx.md). If `TensorFlow_DIR` is not defined, it will be the same as `DeePMD_DIR`. Note that `tensorflow_cc` is not required if `deepmd_c` is found.
+```bash
+cmake -B build -DDeePMD_DIR=/dir_to_deepmd-kit -DTensorFlow_DIR=/dir_to_tensorflow
+```
+Similarly, DeePMD-kit supports PyTorch backend but its libraries are placed at another directory, then
+```bash
+cmake -B build -DDeePMD_DIR=/dir_to_deepmd-kit -DTorch_DIR=/dir_to_pytorch
+```
 
 ## Build with LibRI and LibComm
 
@@ -282,15 +291,21 @@ directly.
 
 > Note: This part is only required if you want to load a trained DeeP Potential and run molecular dynamics with that. To train the DeeP Potential with DP-GEN, no extra prerequisite is needed and please refer to [this page](http://abacus.deepmodeling.com/en/latest/advanced/interface/dpgen.html) for ABACUS interface with DP-GEN.
 
-To compile ABACUS with DeePMD-kit, you need to define `DeePMD_DIR` and `TensorFlow_DIR` in the file `Makefile.vars` or use
+To compile ABACUS with DeePMD-kit, you need to define `DeePMD_DIR` and `TensorFlow_DIR` (TensorFlow Backend, optional) and/or `LIBTORCH_DIR` (PyTorch Backend, optional) in the file `Makefile.vars`.
 
+Or the `tensorflow_cc` and `torch` libraries are in the same directory as the `deepmd_c`/`deepmd_cc` libraries, then
 ```makefile
-make DeePMD_DIR=~/deepmd-kit TensorFlow_DIR=~/tensorflow
+make DeePMD_DIR=/dir_to_deepmd-kit
 ```
+DeePMD-kit supports TensorFlow backend but its libraries are placed at another directory, then
 
-directly.
-
-> `deepmd_c`/`deepmd_cc` and `tensorflow_cc` libraries would be called according to `DeePMD_DIR` and `TensorFlow_DIR`, which is showed in detail in [this page](https://github.com/deepmodeling/deepmd-kit/blob/master/doc/inference/cxx.md). If `TensorFlow_DIR` is not defined, it will be the same as `DeePMD_DIR`. Note that `tensorflow_cc` is not required if `deepmd_c` is found.
+```makefile
+make DeePMD_DIR=/dir_to_deepmd-kit TensorFlow_DIR=/dir_to_tensorflow
+```
+Similarly, DeePMD-kit supports PyTorch backend but its libraries are placed at another directory, then
+```makefile
+make DeePMD_DIR=/dir_to_deepmd-kit Torch_DIR=/dir_to_pytorch
+```
 
 ### Add LibRI Support
 To use new EXX, you need two libraries: [LibRI](https://github.com/abacusmodeling/LibRI) and [LibComm](https://github.com/abacusmodeling/LibComm) and need to define `LIBRI_DIR` and `LIBCOMM_DIR` in the file `Makefile.vars` or use

source/Makefile

@@ -135,16 +135,17 @@ ifdef LIBRI_DIR
     HONG += -D__EXX -DEXX_H_COMM=2 -DEXX_DM=3 -DTEST_EXX_LCAO=0 -DTEST_EXX_RADIAL=1
 endif
 
+ifdef LIBNPY_DIR
+    CNPY_INCLUDE_DIR = -I${LIBNPY_DIR}/include
+    HONG += -D__DEEPKS
+    INCLUDES += $(CNPY_INCLUDE_DIR)
+endif
+
 ifdef LIBTORCH_DIR
- ifdef LIBNPY_DIR
-      LIBTORCH_INCLUDE_DIR =  -I${LIBTORCH_DIR}/include -I${LIBTORCH_DIR}/include/torch/csrc/api/include
-      LIBTORCH_LIB_DIR= ${LIBTORCH_DIR}/lib
-      LIBTORCH_LIB = -L${LIBTORCH_LIB_DIR} -ltorch -lc10 -Wl,-rpath,${LIBTORCH_LIB_DIR} -Wl,--no-as-needed,"${LIBTORCH_LIB_DIR}/libtorch_cpu.so" -Wl,--as-needed ${LIBTORCH_LIB_DIR}/libc10.so -lpthread -Wl,--no-as-needed,"${LIBTORCH_LIB_DIR}/libtorch.so" -Wl,--as-needed 
-      CNPY_INCLUDE_DIR = -I${LIBNPY_DIR}/include
-      HONG += -D__DEEPKS
-      INCLUDES += $(LIBTORCH_INCLUDE_DIR) $(CNPY_INCLUDE_DIR)
-      LIBS += $(LIBTORCH_LIB)
-  endif
+    LIBTORCH_INCLUDE_DIR =  -I${LIBTORCH_DIR}/include -I${LIBTORCH_DIR}/include/torch/csrc/api/include
+    INCLUDES += $(LIBTORCH_INCLUDE_DIR)
+    LIBTORCH_LIB_DIR= ${LIBTORCH_DIR}/lib
+    LIBS += -L${LIBTORCH_LIB_DIR} -ltorch -lc10 -Wl,-rpath,${LIBTORCH_LIB_DIR} -Wl,--no-as-needed,"${LIBTORCH_LIB_DIR}/libtorch_cpu.so" -Wl,--as-needed ${LIBTORCH_LIB_DIR}/libc10.so -lpthread -Wl,--no-as-needed,"${LIBTORCH_LIB_DIR}/libtorch.so" -Wl,--as-needed 
 endif
 
 ifdef DeePMD_DIR
@@ -162,20 +163,19 @@ ifdef DeePMD_DIR
         LIBS += -L${DeeP_LIB_DIR} -ldeepmd_c -Wl,-rpath=${DeeP_LIB_DIR}
     else
         LIBS += -L${DeeP_LIB_DIR} -ldeepmd_cc -Wl,-rpath=${DeeP_LIB_DIR}
-
-        #==========================
-        # TensorFlow package
-        #==========================
-        ifndef TensorFlow_DIR
-            TensorFlow_DIR = ${DeePMD_DIR}
-        endif
-        TensorFlow_LIB_DIR = ${TensorFlow_DIR}/lib
-        LIBS += -L${TensorFlow_LIB_DIR} -ltensorflow_cc -Wl,-rpath=${TensorFlow_LIB_DIR}
-        TensorFlow_INCLUDE_DIR = ${TensorFlow_DIR}/include
-        INCLUDES += -I${TensorFlow_INCLUDE_DIR}
     endif
 endif
 
+ifdef TensorFlow_DIR
+    #==========================
+    # TensorFlow package
+    #==========================
+    TensorFlow_LIB_DIR = ${TensorFlow_DIR}/lib
+    LIBS += -L${TensorFlow_LIB_DIR} -ltensorflow_cc -Wl,-rpath=${TensorFlow_LIB_DIR}
+    TensorFlow_INCLUDE_DIR = ${TensorFlow_DIR}/include
+    INCLUDES += -I${TensorFlow_INCLUDE_DIR}
+endif
+
 ifdef PEXSI_DIR
     OBJS_ABACUS += ${OBJS_HSOLVER_PEXSI}
     INCLUDES += -I${PEXSI_DIR}/include -I${PARMETIS_DIR}/include -I${DSUPERLU_DIR}/include

source/Makefile.vars

@@ -57,7 +57,7 @@ CEREAL_DIR    = /usr/local/include/cereal
 ##-------------------  OPTIONAL LIBS  ---------------------------------
 ## To use DEEPKS: set LIBTORCH_DIR and LIBNPY_DIR
 ## To use LIBXC:  set LIBXC_DIR which contains include and lib/libxc.a (>5.1.7)
-## To use DeePMD: set DeePMD_DIR and TensorFlow_DIR
+## To use DeePMD: set DeePMD_DIR LIBTORCH_DIR and TensorFlow_DIR
 ## To use LibRI:  set LIBRI_DIR and LIBCOMM_DIR
 ## To use PEXSI: set PEXSI_DIR DSUPERLU_DIR and PARMETIS_DIR
 ##---------------------------------------------------------------------