diff --git a/CMakeLists.txt b/CMakeLists.txt index d5941243d66db..a87dc617c0827 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -248,6 +248,10 @@ option(USE_COLORIZE_OUTPUT "Colorize output during compilation" ON) option(USE_ASAN "Use Address+Undefined Sanitizers" OFF) option(USE_LSAN "Use Leak Sanitizer" OFF) option(USE_TSAN "Use Thread Sanitizer" OFF) +option( + THEROCK_PYTORCH_PORTABLE_RPATH + "Use TheRock wheel-relative ROCm runtime search paths" + OFF) # Track whether USE_CUDA was explicitly set by the user (before option() is called) # If USE_CUDA is already defined in cache, it means user explicitly set it diff --git a/cmake/Dependencies.cmake b/cmake/Dependencies.cmake index e31ff49aeba0b..70b8aea767ee5 100644 --- a/cmake/Dependencies.cmake +++ b/cmake/Dependencies.cmake @@ -13,7 +13,39 @@ set(CMAKE_BUILD_WITH_INSTALL_RPATH FALSE) set(CMAKE_INSTALL_RPATH "${_rpath_portable_origin}") # Automatically add all linked folders that are NOT in the build directory to # the rpath (per library?) -set(CMAKE_INSTALL_RPATH_USE_LINK_PATH TRUE) +if(THEROCK_PYTORCH_PORTABLE_RPATH) + if(NOT CMAKE_CXX_COMPILER_ID MATCHES "Clang") + message(FATAL_ERROR + "THEROCK_PYTORCH_PORTABLE_RPATH requires a Clang C++ compiler") + endif() + + execute_process( + COMMAND "${CMAKE_CXX_COMPILER}" -print-resource-dir + RESULT_VARIABLE _therock_clang_resource_result + OUTPUT_VARIABLE _therock_clang_resource_dir + OUTPUT_STRIP_TRAILING_WHITESPACE) + if(NOT _therock_clang_resource_result EQUAL 0 OR + NOT _therock_clang_resource_dir) + message(FATAL_ERROR "Could not determine the Clang resource directory") + endif() + get_filename_component( + _therock_clang_resource_version + "${_therock_clang_resource_dir}" + NAME) + + # torch/lib is two levels below the site-packages directory containing the + # separately installed TheRock runtime wheels. + set(CMAKE_INSTALL_RPATH + "${_rpath_portable_origin}" + "${_rpath_portable_origin}/../../_rocm_sdk_core/lib" + "${_rpath_portable_origin}/../../_rocm_sdk_core/lib/host-math/lib" + "${_rpath_portable_origin}/../../_rocm_sdk_core/lib/llvm/lib" + "${_rpath_portable_origin}/../../_rocm_sdk_core/lib/llvm/lib/clang/${_therock_clang_resource_version}/lib/linux" + "${_rpath_portable_origin}/../../_rocm_sdk_libraries/lib") + set(CMAKE_INSTALL_RPATH_USE_LINK_PATH FALSE) +else() + set(CMAKE_INSTALL_RPATH_USE_LINK_PATH TRUE) +endif() # UBSAN triggers when compiling protobuf, so we need to disable it. set(UBSAN_FLAG "-fsanitize=undefined") diff --git a/cmake/Summary.cmake b/cmake/Summary.cmake index 2d68d81c6b367..519872ce39840 100644 --- a/cmake/Summary.cmake +++ b/cmake/Summary.cmake @@ -69,6 +69,7 @@ function(caffe2_print_configuration_summary) message(STATUS " USE_ASAN : ${USE_ASAN}") message(STATUS " USE_LSAN : ${USE_LSAN}") message(STATUS " USE_TSAN : ${USE_TSAN}") + message(STATUS " THEROCK_PYTORCH_PORTABLE_RPATH : ${THEROCK_PYTORCH_PORTABLE_RPATH}") message(STATUS " USE_CPP_CODE_COVERAGE : ${USE_CPP_CODE_COVERAGE}") message(STATUS " USE_CUDA : ${USE_CUDA}") if(${USE_CUDA}) diff --git a/torch/CMakeLists.txt b/torch/CMakeLists.txt index 471fe0c9c1f0b..5a644c704cdfe 100644 --- a/torch/CMakeLists.txt +++ b/torch/CMakeLists.txt @@ -324,7 +324,9 @@ endif() if(CAFFE2_USE_MKL AND BUILD_LIBTORCHLESS) # Use the RPATH of the linked libraries - set(CMAKE_INSTALL_RPATH_USE_LINK_PATH TRUE) + if(NOT THEROCK_PYTORCH_PORTABLE_RPATH) + set(CMAKE_INSTALL_RPATH_USE_LINK_PATH TRUE) + endif() # we need to explicitly link caffe2::mkl in order to have the # correct RPATH in torch_python for the split build target_link_libraries(torch_python PRIVATE caffe2::mkl)