feat(pt): allow PT OP CXXABI different from TF (#3891)

- Build PT OP libraries with compatible CXXABI if PT has a different CXX
ABI with TF;
- Enable PT OP in test_cuda workflow;
- Update documentation.

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

- **Documentation**
- Removed outdated instructions related to setting environment variables
for enabling customized C++ OPs in PyTorch.

- **Chores**
- Updated build configuration to handle PyTorch CXX11 ABI compatibility
with TensorFlow.
- Refactored library creation processes for better handling of CUDA and
ROCm toolkits.
- Improved build scripts to dynamically adjust compile definitions and
installation paths based on different build configurations.

- **CI/CD**
- Enhanced the continuous integration workflow to include PyTorch
variable assignments and settings for testing with CUDA.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

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

.github/workflows/test_cuda.yml

@@ -49,11 +49,13 @@ jobs:
     - run: python -m pip install -U uv
     - run: python -m uv pip install --system "tensorflow>=2.15.0rc0" "torch>=2.2.0"
     - run: |
+        export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')
         export TENSORFLOW_ROOT=$(python -c 'import importlib,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
         python -m uv pip install --system -v -e .[gpu,test,lmp,cu12,torch] mpi4py
       env:
         DP_VARIANT: cuda
         DP_ENABLE_NATIVE_OPTIMIZATION: 1
+        DP_ENABLE_PYTORCH: 1
     - run: dp --version
     - run: python -m pytest source/tests --durations=0
       env:

doc/install/install-from-source.md

@@ -123,9 +123,6 @@ Note that TensorFlow may have specific requirements for the compiler version to
 You can set the environment variable `export DP_ENABLE_PYTORCH=1` to enable customized C++ OPs in the PyTorch backend.
 Note that PyTorch may have specific requirements for the compiler version to support the C++ standard version and [`_GLIBCXX_USE_CXX11_ABI`](https://gcc.gnu.org/onlinedocs/libstdc++/manual/using_dual_abi.html) used by PyTorch.
 
-The customized C++ OPs are not enabled by default because TensorFlow and PyTorch packages from the PyPI use different `_GLIBCXX_USE_CXX11_ABI` flags.
-We recommend conda-forge packages in this case.
-
 :::
 
 ::::

source/CMakeLists.txt

@@ -186,13 +186,34 @@ if(ENABLE_PYTORCH AND NOT DEEPMD_C_ROOT)
   string(REGEX MATCH "_GLIBCXX_USE_CXX11_ABI=([0-9]+)" CXXABI_PT_MATCH
                "${TORCH_CXX_FLAGS}")
   if(CXXABI_PT_MATCH)
+    set(OP_CXX_ABI_PT ${CMAKE_MATCH_1})
     message(STATUS "PyTorch CXX11 ABI: ${CMAKE_MATCH_1}")
     if(DEFINED OP_CXX_ABI)
       if(NOT ${CMAKE_MATCH_1} EQUAL ${OP_CXX_ABI})
-        message(
-          FATAL_ERROR
-            "PyTorch CXX11 ABI mismatch TensorFlow: ${CMAKE_MATCH_1} != ${OP_CXX_ABI}"
-        )
+        if(NOT BUILD_PY_IF)
+          message(
+            FATAL_ERROR
+              "PyTorch CXX11 ABI mismatch TensorFlow: ${CMAKE_MATCH_1} != ${OP_CXX_ABI}"
+          )
+        else()
+          if(NOT BUILD_CPP_IF)
+            message(
+              STATUS
+                "PyTorch CXX11 ABI mismatch TensorFlow: ${CMAKE_MATCH_1} != ${OP_CXX_ABI}. "
+                "Try to build libraries with both ABIs.")
+          else()
+            message(
+              WARNING
+                "PyTorch CXX11 ABI mismatch TensorFlow: ${CMAKE_MATCH_1} != ${OP_CXX_ABI}. "
+                "PyTorch C++ OP will be built but PyTorch support for C++ libraries will be disabled. "
+                "Note that we don't officially support building C++ libraries in the Python package, "
+                "except for the wheels we officially release.")
+          endif()
+          set(DEEPMD_BUILD_COMPAT_CXXABI ON)
+          set(OP_CXX_ABI_COMPAT ${OP_CXX_ABI_PT})
+        endif()
+      else()
+        set(DEEPMD_BUILD_COMPAT_CXXABI OFF)
       endif()
     else()
       set(OP_CXX_ABI ${CMAKE_MATCH_1})

source/api_cc/CMakeLists.txt

@@ -16,7 +16,7 @@ if(ENABLE_TENSORFLOW)
                                            TensorFlow::tensorflow_framework)
   target_compile_definitions(${libname} PRIVATE BUILD_TENSORFLOW)
 endif()
-if(ENABLE_PYTORCH)
+if(ENABLE_PYTORCH AND "${OP_CXX_ABI_PT}" EQUAL "${OP_CXX_ABI}")
   target_link_libraries(${libname} PRIVATE "${TORCH_LIBRARIES}")
   target_compile_definitions(${libname} PRIVATE BUILD_PYTORCH)
 endif()

source/lib/CMakeLists.txt

@@ -1,54 +1,78 @@
-# libmd
 set(libname ${LIB_DEEPMD})
 
 file(GLOB LIB_SRC src/*.cc src/*.cpp)
 file(GLOB INC_SRC include/*.h ${CMAKE_CURRENT_BINARY_DIR}/version.h)
 
-add_library(${libname} SHARED ${LIB_SRC})
-target_include_directories(
-  ${libname} PUBLIC $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>
-                    $<INSTALL_INTERFACE:include>)
-
 if(USE_CUDA_TOOLKIT)
   add_definitions("-DGOOGLE_CUDA")
   add_subdirectory(src/gpu)
-  set(EXTRA_LIBS ${EXTRA_LIBS} deepmd_op_cuda)
-  target_link_libraries(${libname} PUBLIC deepmd_dyn_cudart)
-  target_link_libraries(${libname} INTERFACE ${EXTRA_LIBS})
-  # gpu_cuda.h
-  target_include_directories(
-    ${libname} PUBLIC $<BUILD_INTERFACE:${CUDAToolkit_INCLUDE_DIRS}>
-                      $<INSTALL_INTERFACE:include>)
 endif()
-
 if(USE_ROCM_TOOLKIT)
   add_definitions("-DTENSORFLOW_USE_ROCM")
   add_subdirectory(src/gpu)
-  set(EXTRA_LIBS ${EXTRA_LIBS} deepmd_op_rocm)
-  # to define __HIP_PLATFORM_AMD__ in hip_runtime.h
-  target_link_libraries(${libname} PUBLIC hip::host)
-  target_link_libraries(${libname} INTERFACE ${EXTRA_LIBS})
-  # gpu_rocm.h
-  target_include_directories(
-    ${libname} PUBLIC $<BUILD_INTERFACE:${HIP_INCLUDE_DIRS}>
-                      $<INSTALL_INTERFACE:include>)
 endif()
 
-set_target_properties(${libname} PROPERTIES INSTALL_RPATH $ORIGIN)
+function(create_library _suffix)
+  set(libname_suffix "${libname}${_suffix}")
+  add_library(${libname_suffix} SHARED ${LIB_SRC})
+  target_include_directories(
+    ${libname_suffix}
+    PUBLIC $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>
+           $<INSTALL_INTERFACE:include>)
+
+  if(USE_CUDA_TOOLKIT)
+    target_link_libraries(${libname_suffix} PUBLIC deepmd_dyn_cudart)
+    target_link_libraries(${libname_suffix} INTERFACE deepmd_op_cuda${_suffix})
+    target_link_libraries(${libname_suffix} INTERFACE ${EXTRA_LIBS})
+    # gpu_cuda.h
+    target_include_directories(
+      ${libname_suffix} PUBLIC $<BUILD_INTERFACE:${CUDAToolkit_INCLUDE_DIRS}>
+                               $<INSTALL_INTERFACE:include>)
+  endif()
 
-if(CMAKE_TESTING_ENABLED)
-  target_link_libraries(${libname} PRIVATE coverage_config)
+  if(USE_ROCM_TOOLKIT)
+    # to define __HIP_PLATFORM_AMD__ in hip_runtime.h
+    target_link_libraries(${libname_suffix} PUBLIC hip::host)
+    target_link_libraries(${libname_suffix} INTERFACE deepmd_op_rocm${_suffix})
+    target_link_libraries(${libname_suffix} INTERFACE ${EXTRA_LIBS})
+    # gpu_rocm.h
+    target_include_directories(
+      ${libname_suffix} PUBLIC $<BUILD_INTERFACE:${HIP_INCLUDE_DIRS}>
+                               $<INSTALL_INTERFACE:include>)
+  endif()
+
+  set_target_properties(${libname_suffix} PROPERTIES INSTALL_RPATH $ORIGIN)
+
+  if(CMAKE_TESTING_ENABLED)
+    target_link_libraries(${libname_suffix} PRIVATE coverage_config)
+  endif()
+
+  if(BUILD_PY_IF)
+    install(TARGETS ${libname_suffix} DESTINATION deepmd/lib/)
+  else(BUILD_PY_IF)
+    install(
+      TARGETS ${libname_suffix}
+      EXPORT ${CMAKE_PROJECT_NAME}Targets
+      DESTINATION lib/)
+  endif(BUILD_PY_IF)
+endfunction()
+
+remove_definitions(-D_GLIBCXX_USE_CXX11_ABI=${OP_CXX_ABI})
+create_library("")
+target_compile_definitions(
+  ${libname}
+  PUBLIC "$<$<COMPILE_LANGUAGE:CXX>:_GLIBCXX_USE_CXX11_ABI=${OP_CXX_ABI}>")
+if(DEEPMD_BUILD_COMPAT_CXXABI)
+  create_library("_compat_cxxabi")
+  target_compile_definitions(
+    ${libname}_compat_cxxabi
+    PUBLIC
+      "$<$<COMPILE_LANGUAGE:CXX>:_GLIBCXX_USE_CXX11_ABI=${OP_CXX_ABI_COMPAT}>")
 endif()
 
-if(BUILD_PY_IF)
-  install(TARGETS ${libname} DESTINATION deepmd/lib/)
-else(BUILD_PY_IF)
-  install(
-    TARGETS ${libname}
-    EXPORT ${CMAKE_PROJECT_NAME}Targets
-    DESTINATION lib/)
+if(NOT BUILD_PY_IF)
   install(FILES ${INC_SRC} DESTINATION include/deepmd)
-endif(BUILD_PY_IF)
+endif()
 
 if(BUILD_CPP_IF AND CMAKE_TESTING_ENABLED)
   add_subdirectory(tests)

source/lib/src/gpu/CMakeLists.txt

@@ -10,10 +10,6 @@ if(USE_CUDA_TOOLKIT)
   endif()
   enable_language(CUDA)
   set(CMAKE_CUDA_STANDARD 11)
-  if(DEFINED OP_CXX_ABI)
-    add_compile_definitions(
-      "$<$<COMPILE_LANGUAGE:CUDA>:_GLIBCXX_USE_CXX11_ABI=${OP_CXX_ABI}>")
-  endif()
 
   find_package(CUDAToolkit REQUIRED)
 
@@ -55,9 +51,6 @@ if(USE_CUDA_TOOLKIT)
 
   file(GLOB SOURCE_FILES "*.cu")
 
-  add_library(${GPU_LIB_NAME} SHARED ${SOURCE_FILES})
-  target_link_libraries(${GPU_LIB_NAME} PRIVATE deepmd_dyn_cudart)
-
 elseif(USE_ROCM_TOOLKIT)
 
   # required cmake version
@@ -83,31 +76,60 @@ elseif(USE_ROCM_TOOLKIT)
 
   file(GLOB SOURCE_FILES "*.cu")
 
-  add_library(${GPU_LIB_NAME} SHARED ${SOURCE_FILES})
   set_source_files_properties(${SOURCE_FILES} PROPERTIES LANGUAGE HIP)
-  # -fpic
-  set_property(TARGET ${GPU_LIB_NAME} PROPERTY POSITION_INDEPENDENT_CODE ON)
-  target_link_libraries(${GPU_LIB_NAME} PRIVATE hip::hipcub)
-
 endif()
 
-target_include_directories(
+function(create_gpu_lib _suffix)
+  set(GPU_LIB_NAME_SUFFIX ${GPU_LIB_NAME}${_suffix})
+  add_library(${GPU_LIB_NAME_SUFFIX} SHARED ${SOURCE_FILES})
+
+  if(USE_CUDA_TOOLKIT)
+    target_link_libraries(${GPU_LIB_NAME_SUFFIX} PRIVATE deepmd_dyn_cudart)
+
+  elseif(USE_ROCM_TOOLKIT)
+    # -fpic
+    set_property(TARGET ${GPU_LIB_NAME_SUFFIX}
+                 PROPERTY POSITION_INDEPENDENT_CODE ON)
+    target_link_libraries(${GPU_LIB_NAME_SUFFIX} PRIVATE hip::hipcub)
+  endif()
+
+  target_include_directories(
+    ${GPU_LIB_NAME_SUFFIX}
+    PUBLIC $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/../../include/>
+           $<INSTALL_INTERFACE:include>)
+  target_precompile_headers(${GPU_LIB_NAME_SUFFIX} PUBLIC [["device.h"]])
+  if(APPLE)
+    set_target_properties(${GPU_LIB_NAME_SUFFIX} PROPERTIES INSTALL_RPATH
+                                                            @loader_path)
+  else()
+    set_target_properties(${GPU_LIB_NAME_SUFFIX} PROPERTIES INSTALL_RPATH
+                                                            "$ORIGIN")
+  endif()
+
+  if(BUILD_CPP_IF AND NOT BUILD_PY_IF)
+    install(
+      TARGETS ${GPU_LIB_NAME_SUFFIX}
+      EXPORT ${CMAKE_PROJECT_NAME}Targets
+      DESTINATION lib/)
+  endif(BUILD_CPP_IF AND NOT BUILD_PY_IF)
+  if(BUILD_PY_IF)
+    install(TARGETS ${GPU_LIB_NAME_SUFFIX} DESTINATION deepmd/lib/)
+  endif(BUILD_PY_IF)
+endfunction()
+
+remove_definitions(-D_GLIBCXX_USE_CXX11_ABI=${OP_CXX_ABI})
+create_gpu_lib("")
+target_compile_definitions(
   ${GPU_LIB_NAME}
-  PUBLIC $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/../../include/>
-         $<INSTALL_INTERFACE:include>)
-target_precompile_headers(${GPU_LIB_NAME} PUBLIC [["device.h"]])
-if(APPLE)
-  set_target_properties(${GPU_LIB_NAME} PROPERTIES INSTALL_RPATH @loader_path)
-else()
-  set_target_properties(${GPU_LIB_NAME} PROPERTIES INSTALL_RPATH "$ORIGIN")
+  PUBLIC "$<$<COMPILE_LANGUAGE:CXX>:_GLIBCXX_USE_CXX11_ABI=${OP_CXX_ABI}>"
+         "$<$<COMPILE_LANGUAGE:CUDA>:_GLIBCXX_USE_CXX11_ABI=${OP_CXX_ABI}>"
+         "$<$<COMPILE_LANGUAGE:HIP>:_GLIBCXX_USE_CXX11_ABI=${OP_CXX_ABI}>")
+if(DEEPMD_BUILD_COMPAT_CXXABI)
+  create_gpu_lib("_compat_cxxabi")
+  target_compile_definitions(
+    ${GPU_LIB_NAME}_compat_cxxabi
+    PUBLIC
+      "$<$<COMPILE_LANGUAGE:CXX>:_GLIBCXX_USE_CXX11_ABI=${OP_CXX_ABI_COMPAT}>"
+      "$<$<COMPILE_LANGUAGE:CUDA>:_GLIBCXX_USE_CXX11_ABI=${OP_CXX_ABI_COMPAT}>"
+      "$<$<COMPILE_LANGUAGE:HIP>:_GLIBCXX_USE_CXX11_ABI=${OP_CXX_ABI_COMPAT}>")
 endif()
-
-if(BUILD_CPP_IF AND NOT BUILD_PY_IF)
-  install(
-    TARGETS ${GPU_LIB_NAME}
-    EXPORT ${CMAKE_PROJECT_NAME}Targets
-    DESTINATION lib/)
-endif(BUILD_CPP_IF AND NOT BUILD_PY_IF)
-if(BUILD_PY_IF)
-  install(TARGETS ${GPU_LIB_NAME} DESTINATION deepmd/lib/)
-endif(BUILD_PY_IF)

source/op/pt/CMakeLists.txt

@@ -2,7 +2,16 @@ file(GLOB OP_SRC print_summary.cc comm.cc)
 
 add_library(deepmd_op_pt MODULE ${OP_SRC})
 # link: libdeepmd libtorch
-target_link_libraries(deepmd_op_pt PRIVATE ${TORCH_LIBRARIES} ${LIB_DEEPMD})
+target_link_libraries(deepmd_op_pt PRIVATE ${TORCH_LIBRARIES})
+if(${OP_CXX_ABI_PT} EQUAL ${OP_CXX_ABI})
+  target_link_libraries(deepmd_op_pt PRIVATE ${LIB_DEEPMD})
+else()
+  target_link_libraries(deepmd_op_pt PRIVATE ${LIB_DEEPMD}_compat_cxxabi)
+endif()
+remove_definitions(-D_GLIBCXX_USE_CXX11_ABI=${OP_CXX_ABI})
+target_compile_definitions(
+  deepmd_op_pt
+  PUBLIC "$<$<COMPILE_LANGUAGE:CXX>:_GLIBCXX_USE_CXX11_ABI=${OP_CXX_ABI_PT}>")
 if(APPLE)
   set_target_properties(deepmd_op_pt PROPERTIES INSTALL_RPATH "@loader_path")
 else()