diff --git a/.github/scripts/aiter_prebuild_upload.sh b/.github/scripts/aiter_prebuild_upload.sh index 473ef1c75..73195b490 100755 --- a/.github/scripts/aiter_prebuild_upload.sh +++ b/.github/scripts/aiter_prebuild_upload.sh @@ -12,13 +12,9 @@ set -euo pipefail # Derive ROCm version and aiter commit -> cache key ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" -if [ -n "${ROCM_PATH:-}" ]; then - true # Use provided ROCM_PATH -elif [ -d "/opt/rocm/core" ]; then - ROCM_PATH="/opt/rocm/core" -else - ROCM_PATH="/opt/rocm" -fi +DIR="${ROOT_DIR}/ci" +. "${ROOT_DIR}/ci/_utils.sh" +ROCM_PATH="$(resolve_rocm_path)" export ROCM_PATH ROCM_VER=`head -n1 "${ROCM_PATH}/.info/version" | cut -d. -f1` diff --git a/build_tools/rocm_utils.cmake b/build_tools/rocm_utils.cmake index 526c03937..61c6f9ea1 100644 --- a/build_tools/rocm_utils.cmake +++ b/build_tools/rocm_utils.cmake @@ -7,6 +7,7 @@ include_guard(GLOBAL) if(NOT "$ENV{ROCM_PATH}" STREQUAL "") set(ROCM_PATH "$ENV{ROCM_PATH}") else() + set(_ROCM_SDK_ROOT "") find_program(ROCM_SDK_CLI rocm-sdk) if(ROCM_SDK_CLI) execute_process( @@ -14,18 +15,16 @@ else() OUTPUT_VARIABLE _ROCM_SDK_ROOT OUTPUT_STRIP_TRAILING_WHITESPACE ERROR_QUIET - RESULT_VARIABLE _ROCM_SDK_RESULT ) - if(_ROCM_SDK_RESULT EQUAL 0 AND EXISTS "${_ROCM_SDK_ROOT}/bin/hipcc") - set(ROCM_PATH "${_ROCM_SDK_ROOT}") - endif() endif() - if(NOT DEFINED ROCM_PATH) - if(EXISTS "/opt/rocm/core") - set(ROCM_PATH "/opt/rocm/core") - else() - set(ROCM_PATH "/opt/rocm") - endif() + if(NOT _ROCM_SDK_ROOT STREQUAL "" AND EXISTS "${_ROCM_SDK_ROOT}/bin/hipcc") + set(ROCM_PATH "${_ROCM_SDK_ROOT}") + elseif(EXISTS "/opt/rocm/core") + set(ROCM_PATH "/opt/rocm/core") + elseif(EXISTS "/opt/rocm") + set(ROCM_PATH "/opt/rocm") + else() + message(FATAL_ERROR "Could not find ROCm installation") endif() endif() diff --git a/ci/_utils.sh b/ci/_utils.sh index 3e6b088ac..f2478fc76 100644 --- a/ci/_utils.sh +++ b/ci/_utils.sh @@ -2,6 +2,32 @@ # # See LICENSE for license information. +# ROCM_PATH resolution +resolve_rocm_path() { + if [ -n "${ROCM_PATH:-}" ]; then + echo "$ROCM_PATH" + return 0 + fi + if command -v rocm-sdk >/dev/null 2>&1; then + local _root + _root="$(rocm-sdk path --root)" + if [ -n "$_root" ] && [ -f "${_root}/bin/hipcc" ]; then + echo "$_root" + return 0 + fi + fi + if [ -d "/opt/rocm/core" ]; then + echo /opt/rocm/core + return 0 + fi + if [ -d "/opt/rocm" ]; then + echo /opt/rocm + return 0 + fi + echo "Could not find ROCm installation" >&2 + exit 1 +} + REALPATH=realpath realpath $DIR >/dev/null 2>/dev/null test $? -ne 0 && REALPATH=echo @@ -287,15 +313,9 @@ check_test_filter() { start_message() { echo "Started with TEST_LEVEL=$TEST_LEVEL sGPU='$TEST_SGPU' mGPU='$TEST_MGPU' at `date`" - if [ -n "$ROCM_PATH" ]; then - _rocm_path="$ROCM_PATH" - elif [ -d "/opt/rocm/core" ]; then - _rocm_path="/opt/rocm/core" - else - _rocm_path="/opt/rocm" - fi - _rocm_path=`$REALPATH "$_rocm_path"` - test -d "$_rocm_path" && echo "ROCm: $_rocm_path" || echo "ROCm path not found" + _rocm_path=$(resolve_rocm_path) + _rocm_path=`$REALPATH "$_rocm_path" 2>/dev/null || echo "$_rocm_path"` + echo "ROCM PATH: $_rocm_path" python3 --version } diff --git a/ci/pytorch.sh b/ci/pytorch.sh index 9f1aee749..38cef08ce 100755 --- a/ci/pytorch.sh +++ b/ci/pytorch.sh @@ -12,7 +12,7 @@ TEST_DIR=${TE_PATH}tests/pytorch #: ${TEST_WORKERS:=4} install_prerequisites() { - pip install 'numpy>=1.22.4' pandas safetensors pytest-timeout + pip install 'numpy>=1.22.4' pandas safetensors pyyaml pytest-timeout rc=$? if [ $rc -ne 0 ]; then script_error "Failed to install test prerequisites" diff --git a/tests/cpp/CMakeLists.txt b/tests/cpp/CMakeLists.txt index 3386fd0b7..2041b8bf9 100644 --- a/tests/cpp/CMakeLists.txt +++ b/tests/cpp/CMakeLists.txt @@ -10,12 +10,7 @@ option(USE_ROCM "Use ROCm" ON) set(USE_CUDA OFF) if (USE_ROCM) - #Check if ROCm exists - if( ((EXISTS "/opt/rocm/") OR (EXISTS $ENV{ROCM_PATH})) ) - message("Using AMD Platform") - else() - message(FATAL_ERROR "Could not ROCm installation") - endif() + include("${CMAKE_CURRENT_SOURCE_DIR}/../../build_tools/rocm_utils.cmake") else() #Check if CUDA exists if (EXISTS "$ENV{CUDA_HOME}/bin/nvcc") @@ -33,7 +28,6 @@ if(USE_ROCM) # Disable Asserts In Code (Can't use asserts on HIP stack.) add_definitions(-DNDEBUG) add_definitions(-DUSE_ROCM) - include("${CMAKE_CURRENT_SOURCE_DIR}/../../build_tools/rocm_utils.cmake") else() if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES) if (CUDAToolkit_VERSION VERSION_GREATER_EQUAL 12.8) diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index f2a014184..e646b72c3 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -13,13 +13,7 @@ option(USE_HIPKITTENS_GEMM "Use HipKittens blockwise FP8 and MXFP8 GEMM kernels" set(USE_CUDA OFF) if (USE_ROCM) - #Check if ROCm exists - if( ((EXISTS "/opt/rocm/") OR (EXISTS $ENV{ROCM_PATH})) ) - message("Using AMD Platform") - else() - message(FATAL_ERROR "Could not find ROCm installation") - endif() - + include("${CMAKE_CURRENT_SOURCE_DIR}/../../build_tools/rocm_utils.cmake") else() #Check if CUDA exists if (EXISTS "$ENV{CUDA_HOME}/bin/nvcc") @@ -136,7 +130,6 @@ set(CUTLASS_TOOLS_INCLUDE_DIR else() set(CMAKE_CXX_STANDARD 17) #define GPU ARCH before enabling HIP to avoid potential failures with arch detection - include("${CMAKE_CURRENT_SOURCE_DIR}/../../build_tools/rocm_utils.cmake") project(transformer_engine LANGUAGES HIP CXX) # If using ROCm, ninja must be installed to build TE.