[ROCm] LoadHIP CMake cleanup (#137112)

Should help mitigate issues reported here: https://github.com/pytorch/pytorch/issues/128313

While working on https://github.com/pytorch/pytorch/pull/136700, we realized that some of the ROCm CMake can be streamlined.

This PR does not fix any bugs or provide any new functionality. Strictly clean-up.

The remaining `${ROCM_ROCTX_LIB}` will be removed when we transition to the rocprofiler-sdk (to be done in a separate PR).

Pull Request resolved: https://github.com/pytorch/pytorch/pull/137112
Approved by: https://github.com/jithunnair-amd, https://github.com/jeffdaily

caffe2/CMakeLists.txt

@@ -606,7 +606,7 @@ if(USE_ROCM)
   # caffe2_nvrtc's stubs to driver APIs are useful for HIP.
   # See NOTE [ ATen NVRTC Stub and HIP ]
   add_library(caffe2_nvrtc SHARED ${ATen_NVRTC_STUB_SRCS})
-  target_link_libraries(caffe2_nvrtc ${PYTORCH_HIP_LIBRARIES} ${ROCM_HIPRTC_LIB})
+  target_link_libraries(caffe2_nvrtc hip::amdhip64 hiprtc::hiprtc)
   target_include_directories(caffe2_nvrtc PRIVATE ${CMAKE_BINARY_DIR})
   target_compile_definitions(caffe2_nvrtc PRIVATE USE_ROCM __HIP_PLATFORM_AMD__)
   install(TARGETS caffe2_nvrtc DESTINATION "${TORCH_INSTALL_LIB_DIR}")

cmake/External/rccl.cmake

@@ -7,8 +7,7 @@ if(NOT __NCCL_INCLUDED)
     if(rccl_FOUND)
       message(STATUS "RCCL Found!")
       add_library(__caffe2_nccl INTERFACE)
-      target_link_libraries(__caffe2_nccl INTERFACE ${PYTORCH_RCCL_LIBRARIES})
+      target_link_libraries(__caffe2_nccl INTERFACE roc::rccl)
-      target_include_directories(__caffe2_nccl INTERFACE ${RCCL_INCLUDE_DIRS})
     else()
       message(STATUS "RCCL NOT Found!")
     endif()

cmake/public/LoadHIP.cmake

@@ -44,78 +44,60 @@ endif()
 message("Building PyTorch for GPU arch: ${PYTORCH_ROCM_ARCH}")
 
 # Add HIP to the CMAKE Module Path
+# needed because the find_package call to this module uses the Module mode search
+# https://cmake.org/cmake/help/latest/command/find_package.html#search-modes
 set(CMAKE_MODULE_PATH ${ROCM_PATH}/lib/cmake/hip ${CMAKE_MODULE_PATH})
 
+# Add ROCM_PATH to CMAKE_PREFIX_PATH, needed because the find_package
+# call to individual ROCM components uses the Config mode search
+list(APPEND CMAKE_PREFIX_PATH ${ROCM_PATH})
+
 macro(find_package_and_print_version PACKAGE_NAME)
   find_package("${PACKAGE_NAME}" ${ARGN})
   message("${PACKAGE_NAME} VERSION: ${${PACKAGE_NAME}_VERSION}")
 endmacro()
 
 # Find the HIP Package
-find_package_and_print_version(HIP 1.0)
+# MODULE argument is added for clarity that CMake is searching
+# for FindHIP.cmake in Module mode
+find_package_and_print_version(HIP 1.0 MODULE)
 
 if(HIP_FOUND)
   set(PYTORCH_FOUND_HIP TRUE)
-  set(FOUND_ROCM_VERSION_H FALSE)
-
-  set(PROJECT_RANDOM_BINARY_DIR "${PROJECT_BINARY_DIR}")
-  set(file "${PROJECT_BINARY_DIR}/detect_rocm_version.cc")
 
   # Find ROCM version for checks
-  # ROCM 5.0 and later will have header api for version management
+  if(EXISTS ${ROCM_INCLUDE_DIRS}/rocm-core/rocm_version.h)
-  if(EXISTS ${ROCM_INCLUDE_DIRS}/rocm_version.h)
+    set(ROCM_HEADER_FILE ${ROCM_INCLUDE_DIRS}/rocm-core/rocm_version.h)
-    set(FOUND_ROCM_VERSION_H TRUE)
-    file(WRITE ${file} ""
-      "#include <rocm_version.h>\n"
-      )
-  elseif(EXISTS ${ROCM_INCLUDE_DIRS}/rocm-core/rocm_version.h)
-    set(FOUND_ROCM_VERSION_H TRUE)
-    file(WRITE ${file} ""
-      "#include <rocm-core/rocm_version.h>\n"
-      )
   else()
-    message("********************* rocm_version.h couldnt be found ******************\n")
+    message(FATAL_ERROR "********************* rocm_version.h could not be found ******************\n")
   endif()
 
-  if(FOUND_ROCM_VERSION_H)
+  # Read the ROCM headerfile into a variable
-    file(APPEND ${file} ""
+  file(READ ${ROCM_HEADER_FILE} ROCM_HEADER_CONTENT)
-      "#include <cstdio>\n"
 
-      "#ifndef ROCM_VERSION_PATCH\n"
+  # Below we use a RegEx to find ROCM version numbers.
-      "#define ROCM_VERSION_PATCH 0\n"
+  # Note that CMake does not support \s for blank space. That is
-      "#endif\n"
+  # why in the regular expressions below we have a blank space in
-      "#define STRINGIFYHELPER(x) #x\n"
+  # the square brackets.
-      "#define STRINGIFY(x) STRINGIFYHELPER(x)\n"
+  # There are three steps:
-      "int main() {\n"
+  # 1. Match regular expression
-      "  printf(\"%d.%d.%s\", ROCM_VERSION_MAJOR, ROCM_VERSION_MINOR, STRINGIFY(ROCM_VERSION_PATCH));\n"
+  # 2. Strip the non-numerical part of the string
-      "  return 0;\n"
+  # 3. Strip leading and trailing spaces
-      "}\n"
+  string(REGEX MATCH "ROCM_VERSION_MAJOR[ ]+[0-9]+" TEMP1 ${ROCM_HEADER_CONTENT})
-      )
+  string(REPLACE "ROCM_VERSION_MAJOR" "" TEMP2 ${TEMP1})
+  string(STRIP ${TEMP2} ROCM_VERSION_DEV_MAJOR)
+  string(REGEX MATCH "ROCM_VERSION_MINOR[ ]+[0-9]+" TEMP1 ${ROCM_HEADER_CONTENT})
+  string(REPLACE "ROCM_VERSION_MINOR" "" TEMP2 ${TEMP1})
+  string(STRIP ${TEMP2} ROCM_VERSION_DEV_MINOR)
+  string(REGEX MATCH "ROCM_VERSION_PATCH[ ]+[0-9]+" TEMP1 ${ROCM_HEADER_CONTENT})
+  string(REPLACE "ROCM_VERSION_PATCH" "" TEMP2 ${TEMP1})
+  string(STRIP ${TEMP2} ROCM_VERSION_DEV_PATCH)
 
-    try_run(run_result compile_result ${PROJECT_RANDOM_BINARY_DIR} ${file}
+  # Create ROCM_VERSION_DEV_INT which is later used as a preprocessor macros
-      CMAKE_FLAGS "-DINCLUDE_DIRECTORIES=${ROCM_INCLUDE_DIRS}"
+  set(ROCM_VERSION_DEV "${ROCM_VERSION_DEV_MAJOR}.${ROCM_VERSION_DEV_MINOR}.${ROCM_VERSION_DEV_PATCH}")
-      RUN_OUTPUT_VARIABLE rocm_version_from_header
+  math(EXPR ROCM_VERSION_DEV_INT "(${ROCM_VERSION_DEV_MAJOR}*10000) + (${ROCM_VERSION_DEV_MINOR}*100) + ${ROCM_VERSION_DEV_PATCH}")
-      COMPILE_OUTPUT_VARIABLE output_var
-      )
-    # We expect the compile to be successful if the include directory exists.
-    if(NOT compile_result)
-      message(FATAL_ERROR "Caffe2: Couldn't determine version from header: " ${output_var})
-    endif()
-    message(STATUS "Caffe2: Header version is: " ${rocm_version_from_header})
-    set(ROCM_VERSION_DEV_RAW ${rocm_version_from_header})
-    message("\n***** ROCm version from rocm_version.h ****\n")
-  endif()
-
-  string(REGEX MATCH "^([0-9]+)\.([0-9]+)\.([0-9]+).*$" ROCM_VERSION_DEV_MATCH ${ROCM_VERSION_DEV_RAW})
-
-  if(ROCM_VERSION_DEV_MATCH)
-    set(ROCM_VERSION_DEV_MAJOR ${CMAKE_MATCH_1})
-    set(ROCM_VERSION_DEV_MINOR ${CMAKE_MATCH_2})
-    set(ROCM_VERSION_DEV_PATCH ${CMAKE_MATCH_3})
-    set(ROCM_VERSION_DEV "${ROCM_VERSION_DEV_MAJOR}.${ROCM_VERSION_DEV_MINOR}.${ROCM_VERSION_DEV_PATCH}")
-    math(EXPR ROCM_VERSION_DEV_INT "(${ROCM_VERSION_DEV_MAJOR}*10000) + (${ROCM_VERSION_DEV_MINOR}*100) + ${ROCM_VERSION_DEV_PATCH}")
-  endif()
 
+  message("\n***** ROCm version from rocm_version.h ****\n")
   message("ROCM_VERSION_DEV: ${ROCM_VERSION_DEV}")
   message("ROCM_VERSION_DEV_MAJOR: ${ROCM_VERSION_DEV_MAJOR}")
   message("ROCM_VERSION_DEV_MINOR: ${ROCM_VERSION_DEV_MINOR}")
@@ -127,42 +109,9 @@ if(HIP_FOUND)
   message("HIP_VERSION_MINOR: ${HIP_VERSION_MINOR}")
   message("TORCH_HIP_VERSION: ${TORCH_HIP_VERSION}")
 
-  message("\n***** Library versions from dpkg *****\n")
+  # Find ROCM components using Config mode
-  execute_process(COMMAND dpkg -l COMMAND grep rocm-dev COMMAND awk "{print $2 \" VERSION: \" $3}")
+  # These components will be searced for recursively in ${ROCM_PATH}
-  execute_process(COMMAND dpkg -l COMMAND grep rocm-libs COMMAND awk "{print $2 \" VERSION: \" $3}")
-  execute_process(COMMAND dpkg -l COMMAND grep hsakmt-roct COMMAND awk "{print $2 \" VERSION: \" $3}")
-  execute_process(COMMAND dpkg -l COMMAND grep rocr-dev COMMAND awk "{print $2 \" VERSION: \" $3}")
-  execute_process(COMMAND dpkg -l COMMAND grep -w hcc COMMAND awk "{print $2 \" VERSION: \" $3}")
-  execute_process(COMMAND dpkg -l COMMAND grep hip-base COMMAND awk "{print $2 \" VERSION: \" $3}")
-  execute_process(COMMAND dpkg -l COMMAND grep hip_hcc COMMAND awk "{print $2 \" VERSION: \" $3}")
-
   message("\n***** Library versions from cmake find_package *****\n")
-
-  set(CMAKE_HIP_CLANG_FLAGS_DEBUG ${CMAKE_CXX_FLAGS_DEBUG})
-  set(CMAKE_HIP_CLANG_FLAGS_RELEASE ${CMAKE_CXX_FLAGS_RELEASE})
-  ### Remove setting of Flags when FindHIP.CMake PR #558 is accepted.###
-
-  set(hip_DIR ${ROCM_PATH}/lib/cmake/hip)
-  set(hsa-runtime64_DIR ${ROCM_PATH}/lib/cmake/hsa-runtime64)
-  set(AMDDeviceLibs_DIR ${ROCM_PATH}/lib/cmake/AMDDeviceLibs)
-  set(amd_comgr_DIR ${ROCM_PATH}/lib/cmake/amd_comgr)
-  set(rocrand_DIR ${ROCM_PATH}/lib/cmake/rocrand)
-  set(hiprand_DIR ${ROCM_PATH}/lib/cmake/hiprand)
-  set(rocblas_DIR ${ROCM_PATH}/lib/cmake/rocblas)
-  set(hipblas_DIR ${ROCM_PATH}/lib/cmake/hipblas)
-  set(hipblaslt_DIR ${ROCM_PATH}/lib/cmake/hipblaslt)
-  set(miopen_DIR ${ROCM_PATH}/lib/cmake/miopen)
-  set(rocfft_DIR ${ROCM_PATH}/lib/cmake/rocfft)
-  set(hipfft_DIR ${ROCM_PATH}/lib/cmake/hipfft)
-  set(hipsparse_DIR ${ROCM_PATH}/lib/cmake/hipsparse)
-  set(rccl_DIR ${ROCM_PATH}/lib/cmake/rccl)
-  set(rocprim_DIR ${ROCM_PATH}/lib/cmake/rocprim)
-  set(hipcub_DIR ${ROCM_PATH}/lib/cmake/hipcub)
-  set(rocthrust_DIR ${ROCM_PATH}/lib/cmake/rocthrust)
-  set(hipsolver_DIR ${ROCM_PATH}/lib/cmake/hipsolver)
-  set(hiprtc_DIR ${ROCM_PATH}/lib/cmake/hiprtc)
-
-
   find_package_and_print_version(hip REQUIRED)
   find_package_and_print_version(hsa-runtime64 REQUIRED)
   find_package_and_print_version(amd_comgr REQUIRED)
@@ -181,27 +130,11 @@ if(HIP_FOUND)
   find_package_and_print_version(hipsolver REQUIRED)
   find_package_and_print_version(hiprtc REQUIRED)
 
-
-  find_library(PYTORCH_HIP_LIBRARIES amdhip64 HINTS ${ROCM_PATH}/lib)
-  # TODO: miopen_LIBRARIES should return fullpath to the library file,
-  # however currently it's just the lib name
-  if(TARGET ${miopen_LIBRARIES})
-    set(PYTORCH_MIOPEN_LIBRARIES ${miopen_LIBRARIES})
-  else()
-    find_library(PYTORCH_MIOPEN_LIBRARIES ${miopen_LIBRARIES} HINTS ${ROCM_PATH}/lib)
-  endif()
-  # TODO: rccl_LIBRARIES should return fullpath to the library file,
-  # however currently it's just the lib name
-  if(TARGET ${rccl_LIBRARIES})
-    set(PYTORCH_RCCL_LIBRARIES ${rccl_LIBRARIES})
-  else()
-    find_library(PYTORCH_RCCL_LIBRARIES ${rccl_LIBRARIES} HINTS ${ROCM_PATH}/lib)
-  endif()
-  find_library(ROCM_HIPRTC_LIB hiprtc HINTS ${ROCM_PATH}/lib)
   # roctx is part of roctracer
   find_library(ROCM_ROCTX_LIB roctx64 HINTS ${ROCM_PATH}/lib)
 
   # check whether HIP declares new types
+  set(PROJECT_RANDOM_BINARY_DIR "${PROJECT_BINARY_DIR}")
   set(file "${PROJECT_BINARY_DIR}/hip_new_types.cc")
   file(WRITE ${file} ""
     "#include <hip/library_types.h>\n"

test/cpp/jit/CMakeLists.txt

@@ -143,8 +143,8 @@ if(USE_CUDA)
   endif()
 elseif(USE_ROCM)
   target_link_libraries(test_jit PRIVATE
-    ${ROCM_HIPRTC_LIB}
+    hiprtc::hiprtc
-    ${PYTORCH_HIP_LIBRARIES}
+    hip::amdhip64
     ${TORCH_CUDA_LIBRARIES})
 
   target_compile_definitions(test_jit PRIVATE USE_ROCM)

test/cpp/lazy/CMakeLists.txt

@@ -36,8 +36,8 @@ if(USE_CUDA)
   target_compile_definitions(test_lazy PRIVATE USE_CUDA)
 elseif(USE_ROCM)
   target_link_libraries(test_lazy PRIVATE
-    ${ROCM_HIPRTC_LIB}
+    hiprtc::hiprtc
-    ${PYTORCH_HIP_LIBRARIES}
+    hip::amdhip64
     ${TORCH_CUDA_LIBRARIES})
 
   target_compile_definitions(test_lazy PRIVATE USE_ROCM)

test/cpp/tensorexpr/CMakeLists.txt

@@ -58,14 +58,14 @@ if(USE_CUDA)
   target_compile_definitions(tutorial_tensorexpr PRIVATE USE_CUDA)
 elseif(USE_ROCM)
   target_link_libraries(test_tensorexpr PRIVATE
-    ${ROCM_HIPRTC_LIB}
+    hiprtc::hiprtc
-    ${PYTORCH_HIP_LIBRARIES}
+    hip::amdhip64
     ${TORCH_CUDA_LIBRARIES})
   target_compile_definitions(test_tensorexpr PRIVATE USE_ROCM)
 
   target_link_libraries(tutorial_tensorexpr PRIVATE
-    ${ROCM_HIPRTC_LIB}
+    hiprtc::hiprtc
-    ${PYTORCH_HIP_LIBRARIES}
+    hip::amdhip64
     ${TORCH_CUDA_LIBRARIES})
   target_compile_definitions(tutorial_tensorexpr PRIVATE USE_ROCM)
 endif()