Fix libnvrtc discoverability in package patched by auditwheel (#52184)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/52184

`auditwheel` inserts first 8 symbols of sha256 checksum of the library before relocating into the wheel package. This change adds logic for computing the same short sha sum and embedding it into LazyNVRTC as alternative name for libnvrt.so

Fixes https://github.com/pytorch/pytorch/issues/52075

Test Plan: Imported from OSS

Reviewed By: seemethere

Differential Revision: D26417403

Pulled By: malfet

fbshipit-source-id: e366dd22e95e219979f6c2fa39acb11585b34c72
This commit is contained in:
Nikita Shulga 2021-02-13 19:32:22 -08:00 committed by Facebook GitHub Bot
parent 357e5baf7e
commit de4c9ecc35
4 changed files with 28 additions and 3 deletions

View file

@ -23,10 +23,17 @@ at::DynamicLibrary& getNVRTCLibrary() {
constexpr auto minor = ( CUDA_VERSION / 10 ) % 10;
#if defined(_WIN32)
auto libname = std::string("nvrtc64_") + std::to_string(major) + std::to_string(minor) + "_0.dll";
std::string alt_libname;
#else
static auto libname = std::string("libnvrtc.so.") + std::to_string(major) + "." + std::to_string(minor);
static auto lib_version = std::to_string(major) + "." + std::to_string(minor);
static auto libname = std::string("libnvrtc.so.") + lib_version;
#ifdef NVRTC_SHORTHASH
static auto alt_libname = std::string("libnvrtc-") + C10_STRINGIZE(NVRTC_SHORTHASH) + ".so." + lib_version;
#else
std::string alt_libname;
#endif
static at::DynamicLibrary lib(libname.c_str());
#endif
static at::DynamicLibrary lib(libname.c_str(), alt_libname.empty() ? nullptr : alt_libname.c_str());
return lib;
}

View file

@ -595,6 +595,10 @@ if(NOT INTERN_BUILD_MOBILE OR NOT BUILD_CAFFE2_MOBILE)
list(APPEND Caffe2_GPU_SRCS
${TORCH_SRC_DIR}/csrc/cuda/nccl.cpp)
endif()
set_source_files_properties(
${TORCH_ROOT}/aten/src/ATen/cuda/detail/LazyNVRTC.cpp
PROPERTIES COMPILE_DEFINITIONS "NVRTC_SHORTHASH=${CUDA_NVRTC_SHORTHASH}"
)
endif()
if(USE_ROCM)

View file

@ -188,6 +188,20 @@ find_library(CUDA_CUDA_LIB cuda
find_library(CUDA_NVRTC_LIB nvrtc
PATHS ${CUDA_TOOLKIT_ROOT_DIR}
PATH_SUFFIXES lib lib64 lib/x64)
if(CUDA_NVRTC_LIB AND NOT CUDA_NVRTC_SHORTHASH)
execute_process(
COMMAND "${PYTHON_EXECUTABLE}" -c
"import hashlib;hash=hashlib.sha256();hash.update(open('${CUDA_NVRTC_LIB}','rb').read());print(hash.hexdigest()[:8])"
RESULT_VARIABLE _retval
OUTPUT_VARIABLE CUDA_NVRTC_SHORTHASH)
if(NOT _retval EQUAL 0)
message(WARNING "Failed to compute shorthash for libnvrtc.so")
set(CUDA_NVRTC_SHORTHASH "XXXXXXXX")
else()
string(STRIP "${CUDA_NVRTC_SHORTHASH}" CUDA_NVRTC_SHORTHASH)
message(STATUS "${CUDA_NVRTC_LIB} shorthash is ${CUDA_NVRTC_SHORTHASH}")
endif()
endif()
# Create new style imported libraries.
# Several of these libraries have a hardcoded path if CAFFE2_STATIC_LINK_CUDA

View file

@ -162,7 +162,7 @@ endif()
# In the most recent CMake versions, a new 'TRANSFORM' subcommand of 'list' allows much of the boilerplate of defining the lists
# of type stub files to be omitted.
# For comptability with older CMake versions, we omit it for now, but leave it as a comment in case comptability with the older
# For compatibility with older CMake versions, we omit it for now, but leave it as a comment in case compatibility with the older
# CMake versions is eventually dropped.
# set(Modules
# __init__