diff --git a/aten/src/ATen/cuda/detail/LazyNVRTC.cpp b/aten/src/ATen/cuda/detail/LazyNVRTC.cpp index c61f253a3e9..6d321da859f 100644 --- a/aten/src/ATen/cuda/detail/LazyNVRTC.cpp +++ b/aten/src/ATen/cuda/detail/LazyNVRTC.cpp @@ -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; } diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt index fc67f7da26d..a2461892577 100644 --- a/caffe2/CMakeLists.txt +++ b/caffe2/CMakeLists.txt @@ -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) diff --git a/cmake/public/cuda.cmake b/cmake/public/cuda.cmake index 39a48748dfd..62ff4efc5e1 100644 --- a/cmake/public/cuda.cmake +++ b/cmake/public/cuda.cmake @@ -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 diff --git a/torch/CMakeLists.txt b/torch/CMakeLists.txt index f194a544e07..f15d32b5928 100644 --- a/torch/CMakeLists.txt +++ b/torch/CMakeLists.txt @@ -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__