Skip to content

Commit

Permalink
feat(pt): allow PT OP CXXABI different from TF (#3891)
Browse files Browse the repository at this point in the history
- 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>
  • Loading branch information
njzjz authored Jun 21, 2024
1 parent 82e685d commit 7a0ec5d
Show file tree
Hide file tree
Showing 7 changed files with 147 additions and 72 deletions.
2 changes: 2 additions & 0 deletions .github/workflows/test_cuda.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
3 changes: 0 additions & 3 deletions doc/install/install-from-source.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.

:::

::::
Expand Down
29 changes: 25 additions & 4 deletions source/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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})
Expand Down
2 changes: 1 addition & 1 deletion source/api_cc/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
90 changes: 57 additions & 33 deletions source/lib/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -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)
Expand Down
82 changes: 52 additions & 30 deletions source/lib/src/gpu/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down Expand Up @@ -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
Expand All @@ -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)
11 changes: 10 additions & 1 deletion source/op/pt/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down

0 comments on commit 7a0ec5d

Please sign in to comment.