mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
[BE] Rename libnvshmem_extension to libtorch_nvshmem (#158234)
`libnvshmem_extension.so` creates an illusion that it is a shared library from NVSHMEM. But indeed it is built from torch source code, for symmetric tensor infrastructure and operations, though leveraging NVSHMEM APIs. Thus this PR renames `libnvshmem_extension.so` to `libtorch_nvshmem.so`. Pull Request resolved: https://github.com/pytorch/pytorch/pull/158234 Approved by: https://github.com/albanD
This commit is contained in:
@ -1027,24 +1027,24 @@ elseif(USE_CUDA)
|
||||
|
||||
# Linking with nvshmem requires the source binary to be built with -rdc
|
||||
# which is not viable for libtorch_cuda. So we isolate the linking of
|
||||
# nvshmem in nvshmem_extension.
|
||||
add_library(nvshmem_extension SHARED
|
||||
# nvshmem in torch_nvshmem.
|
||||
add_library(torch_nvshmem SHARED
|
||||
"${TORCH_SRC_DIR}/csrc/distributed/c10d/cuda/utils.cpp"
|
||||
"${TORCH_SRC_DIR}/csrc/distributed/c10d/symm_mem/nvshmem_extension.cu"
|
||||
"${TORCH_SRC_DIR}/csrc/distributed/c10d/symm_mem/NVSHMEMSymmetricMemory.cu"
|
||||
"${TORCH_SRC_DIR}/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryUtils.cpp"
|
||||
)
|
||||
set_target_properties(nvshmem_extension PROPERTIES CUDA_SEPARABLE_COMPILATION ON)
|
||||
target_compile_options(nvshmem_extension PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-rdc=true>)
|
||||
target_compile_options(nvshmem_extension PRIVATE "-U__CUDA_NO_HALF_OPERATORS__")
|
||||
target_link_libraries(nvshmem_extension PRIVATE
|
||||
set_target_properties(torch_nvshmem PROPERTIES CUDA_SEPARABLE_COMPILATION ON)
|
||||
target_compile_options(torch_nvshmem PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-rdc=true>)
|
||||
target_compile_options(torch_nvshmem PRIVATE "-U__CUDA_NO_HALF_OPERATORS__")
|
||||
target_link_libraries(torch_nvshmem PRIVATE
|
||||
${NVSHMEM_HOST_LIB}
|
||||
${NVSHMEM_DEVICE_LIB}
|
||||
)
|
||||
target_compile_definitions(torch_cuda PUBLIC USE_NVSHMEM)
|
||||
target_compile_definitions(nvshmem_extension PUBLIC USE_NVSHMEM)
|
||||
target_link_libraries(torch_cuda PRIVATE nvshmem_extension)
|
||||
install(TARGETS nvshmem_extension EXPORT Caffe2Targets DESTINATION lib)
|
||||
target_compile_definitions(torch_nvshmem PUBLIC USE_NVSHMEM)
|
||||
target_link_libraries(torch_cuda PRIVATE torch_nvshmem)
|
||||
install(TARGETS torch_nvshmem EXPORT Caffe2Targets DESTINATION lib)
|
||||
else()
|
||||
message(STATUS "NVSHMEM not found, not building with NVSHMEM support.")
|
||||
endif()
|
||||
|
Reference in New Issue
Block a user