diff --git a/lib/Runtime/ROCm/CMakeLists.txt b/lib/Runtime/ROCm/CMakeLists.txt index c4f118c16..5aed46f4c 100644 --- a/lib/Runtime/ROCm/CMakeLists.txt +++ b/lib/Runtime/ROCm/CMakeLists.txt @@ -3,9 +3,40 @@ # # FlyJitRuntime — thin ROCm runtime with GPU module caching. -file(GLOB _rocm_paths LIST_DIRECTORIES true "/opt/rocm*") -list(SORT _rocm_paths ORDER DESCENDING) -find_package(hip REQUIRED CONFIG PATHS ${_rocm_paths}) +# Prefer an explicitly selected ROCm, then the SDK in the active environment. +# NO_DEFAULT_PATH excludes default search locations; a cached hip_DIR still wins. +if(DEFINED ENV{ROCM_PATH} AND NOT "$ENV{ROCM_PATH}" STREQUAL "") + find_package(hip QUIET CONFIG PATHS "$ENV{ROCM_PATH}" NO_DEFAULT_PATH) +endif() + +if(NOT hip_FOUND) + execute_process(COMMAND rocm-sdk path --root + OUTPUT_VARIABLE _rocm_sdk_root + RESULT_VARIABLE _rocm_sdk_result + OUTPUT_STRIP_TRAILING_WHITESPACE ERROR_QUIET) + if(_rocm_sdk_result STREQUAL "0" AND _rocm_sdk_root) + # HIP's config uses ROCM_PATH to find its dependencies. A stale or empty + # value must not redirect those lookups away from the selected SDK. + set(_rocm_path_was_set FALSE) + if(DEFINED ENV{ROCM_PATH}) + set(_rocm_path_was_set TRUE) + set(_rocm_original_path "$ENV{ROCM_PATH}") + endif() + set(ENV{ROCM_PATH} "${_rocm_sdk_root}") + find_package(hip QUIET CONFIG PATHS "${_rocm_sdk_root}" NO_DEFAULT_PATH) + if(_rocm_path_was_set) + set(ENV{ROCM_PATH} "${_rocm_original_path}") + else() + unset(ENV{ROCM_PATH}) + endif() + endif() +endif() + +if(NOT hip_FOUND) + file(GLOB _rocm_paths LIST_DIRECTORIES true "/opt/rocm*") + list(SORT _rocm_paths ORDER DESCENDING) + find_package(hip REQUIRED CONFIG PATHS ${_rocm_paths}) +endif() add_library(FlyJitRuntime SHARED FlyRocmRuntimeWrappers.cpp) target_include_directories(FlyJitRuntime PRIVATE