[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:
Ke Wen
2025-07-14 08:05:13 -07:00
committed by PyTorch MergeBot
parent bc65253369
commit 2043f6911e

View File

@ -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()