cmake: add opt-in portable RPATHs for TheRock ROCm wheels

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

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")

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})

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)