Add gfx950 ROCm CI for AOTI

.ci/scripts/test-rocm-aoti.sh

@@ -0,0 +1,158 @@
+#!/usr/bin/env bash
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+set -euo pipefail
+
+ROCM_VERSION="${ROCM_VERSION:-7.1}"
+ROCM_PATH="${ROCM_PATH:-/opt/rocm}"
+PYTORCH_ROCM_INDEX="${PYTORCH_ROCM_INDEX:-https://download.pytorch.org/whl/test/rocm${ROCM_VERSION}}"
+TORCHAO_ROCM_WHEEL_BASE="${TORCHAO_ROCM_WHEEL_BASE:-https://download.pytorch.org/whl/nightly/rocm${ROCM_VERSION}}"
+ROCM_CI_TMP_ROOT="${RUNNER_TEMP:-/tmp}"
+mkdir -p "${ROCM_CI_TMP_ROOT}" 2>/dev/null || ROCM_CI_TMP_ROOT=/tmp
+ROCM_CI_TMPDIR="$(mktemp -d "${ROCM_CI_TMP_ROOT}/executorch-rocm-ci.XXXXXX")"
+trap 'rm -rf "${ROCM_CI_TMPDIR}"' EXIT
+
+export ROCM_PATH
+export HIP_VISIBLE_DEVICES=0
+export CUDA_VISIBLE_DEVICES=0
+export TORCHINDUCTOR_CACHE_DIR="${ROCM_CI_TMPDIR}/inductor-cache"
+export TORCHINDUCTOR_COMPILE_THREADS=1
+
+read -r TORCH_VERSION TORCHAO_VERSION < <(
+  python - <<'PY'
+from install_requirements import TORCHAO_NIGHTLY_VERSION
+from torch_pin import TORCH_VERSION
+
+print(TORCH_VERSION, TORCHAO_NIGHTLY_VERSION)
+PY
+)
+# TorchAO ROCm wheels are not exposed by the per-version pip index.
+TORCHAO_WHEEL="${TORCHAO_ROCM_WHEEL_BASE}/torchao-${TORCHAO_VERSION}"
+TORCHAO_WHEEL+="%2Brocm${ROCM_VERSION}-cp310-abi3-manylinux_2_28_x86_64.whl"
+python -m pip install "torch==${TORCH_VERSION}" \
+  --index-url "${PYTORCH_ROCM_INDEX}"
+python -m pip install -r requirements-dev.txt \
+  "${TORCHAO_WHEEL}"
+EXECUTORCH_BUILD_MINIMAL=1 \
+  python -m pip install --editable . --no-build-isolation
+
+if ! command -v conda >/dev/null; then
+  echo "The ROCm CI image must provide conda for its C++ runtime libraries"
+  exit 1
+fi
+conda install -y -c conda-forge 'libstdcxx-ng>=12'
+
+python - <<'PY'
+import torch
+import torchao
+import triton
+
+assert torch.version.hip is not None, "PyTorch is not a ROCm build"
+assert torch.version.cuda is None, "PyTorch unexpectedly reports a CUDA runtime"
+assert torch.cuda.is_available(), "No AMD GPU is visible through PyTorch"
+assert "+rocm" in torchao.__version__, "TorchAO is not a ROCm build"
+
+device = torch.cuda.get_device_properties(0)
+arch = device.gcnArchName.split(":", 1)[0]
+assert arch == "gfx950", f"Expected gfx950, got {device.gcnArchName}"
+assert device.warp_size == 64, f"Expected wave64, got {device.warp_size}"
+assert arch in torch.cuda.get_arch_list(), (
+    f"{arch} is not supported by this PyTorch build: {torch.cuda.get_arch_list()}"
+)
+
+print("PyTorch:", torch.__version__)
+print("TorchAO:", torchao.__version__)
+print("ROCm:", torch.version.hip)
+print("Triton:", triton.__version__)
+print("Device:", device.name, device.gcnArchName)
+print("Architectures:", torch.cuda.get_arch_list())
+PY
+
+if command -v rocminfo >/dev/null; then
+  rocminfo | sed -n '1,160p'
+fi
+if command -v rocm-smi >/dev/null; then
+  rocm-smi --showproductname --showmeminfo vram --showuse --showtemp || true
+elif command -v amd-smi >/dev/null; then
+  amd-smi static --gpu all || true
+fi
+
+TORCH_CMAKE_PREFIX="$(python -c 'import torch; print(torch.utils.cmake_prefix_path)')"
+export CMAKE_PREFIX_PATH="${TORCH_CMAKE_PREFIX}"
+
+cmake --preset llm-release-rocm -DEXECUTORCH_BUILD_TESTS=ON
+cmake --build cmake-out-rocm-llm \
+  --target executor_runner test_cuda_allocator test_cuda_mutable_state \
+  --parallel "$(nproc)"
+
+PYTHON_PREFIX="$(python -c 'import sys; print(sys.prefix)')"
+export LD_LIBRARY_PATH="${PYTHON_PREFIX}/lib:${PWD}/cmake-out-rocm-llm/backends/cuda:${PWD}/cmake-out-rocm-llm/extension/cuda:${ROCM_PATH}/lib:${LD_LIBRARY_PATH:-}"
+
+ctest --test-dir cmake-out-rocm-llm \
+  -R 'test_cuda_(allocator|mutable_state)' \
+  -V
+
+# Targeted builds leave the shim in the build tree rather than install lib/.
+for ROCM_CI_BINARY in \
+  cmake-out-rocm-llm/backends/cuda/libaoti_cuda_shims.so \
+  cmake-out-rocm-llm/executor_runner; do
+  ROCM_CI_LINKS="$(ldd "${ROCM_CI_BINARY}")"
+  printf '%s\n' "${ROCM_CI_LINKS}"
+  if grep -q 'not found' <<<"${ROCM_CI_LINKS}"; then
+    echo "Missing runtime dependency in ${ROCM_CI_BINARY}"
+    exit 1
+  fi
+  if grep -q 'libcudart' <<<"${ROCM_CI_LINKS}"; then
+    echo "Unexpected CUDA runtime dependency in ${ROCM_CI_BINARY}"
+    exit 1
+  fi
+  if ! grep -q 'libamdhip64' <<<"${ROCM_CI_LINKS}"; then
+    echo "HIP runtime dependency missing from ${ROCM_CI_BINARY}"
+    exit 1
+  fi
+done
+
+if ! grep -q '^EXECUTORCH_BUILD_CUDA:BOOL=OFF$' \
+  cmake-out-rocm-llm/CMakeCache.txt; then
+  echo "ROCm CI unexpectedly enabled the CUDA build"
+  exit 1
+fi
+if ! grep -q '^EXECUTORCH_BUILD_ROCM:BOOL=ON$' \
+  cmake-out-rocm-llm/CMakeCache.txt; then
+  echo "ROCm CI did not enable the ROCm build"
+  exit 1
+fi
+if ! grep -q 'find_dependency(hip CONFIG)' \
+  cmake-out-rocm-llm/executorch-backend-dependencies.cmake; then
+  echo "ROCm package metadata does not declare its HIP dependency"
+  exit 1
+fi
+if grep -q 'find_dependency(CUDAToolkit)' \
+  cmake-out-rocm-llm/executorch-backend-dependencies.cmake; then
+  echo "ROCm package metadata unexpectedly depends on CUDAToolkit"
+  exit 1
+fi
+
+export EXECUTORCH_EXECUTOR_RUNNER="${PWD}/cmake-out-rocm-llm/executor_runner"
+ROCM_EXCLUDED_TESTS=(
+  # This focused job does not install the optional flash-linear-attention package.
+  --ignore=backends/cuda/tests/test_chunk_gated_delta_rule.py
+  # Defer specialized SDPA suites until their CI cost is measured.
+  --ignore=backends/cuda/tests/test_tq4_sdpa.py
+  --ignore=backends/cuda/tests/test_triton_sdpa.py
+  --ignore=backends/cuda/tests/test_triton_sdpa_splitk.py
+  # This test process core-dumps on ROCm instead of reporting a failure.
+  --ignore=backends/cuda/tests/test_triton_sdpa_nan.py
+)
+python -m pytest -v -o 'addopts=' \
+  "${ROCM_EXCLUDED_TESTS[@]}" \
+  backends/cuda/tests \
+  backends/cuda/passes/tests
+
+ROCM_POINTWISE_DIR="${ROCM_CI_TMPDIR}/pointwise"
+python -m examples.cuda.scripts.export_amd_pointwise \
+  --output-dir "${ROCM_POINTWISE_DIR}"

.github/pytorch-probot.yml

@@ -7,6 +7,7 @@ ciflow_push_tags:
 - ciflow/cuda-perf
 - ciflow/metal
 - ciflow/mlx
+- ciflow/rocm
 - ciflow/nightly
 - ciflow/trunk
 - ciflow/binaries

.github/workflows/rocm.yml

@@ -0,0 +1,123 @@
+# Test ExecuTorch AOTI ROCm support.
+# Use gfx950 for functional coverage and reserve scarce gfx1100 RDNA runners
+# for a future focused performance job.
+
+name: Test ROCm AOTI
+
+on:
+  push:
+    branches:
+      - main
+      - release/*
+    tags:
+      - ciflow/rocm/*
+  pull_request:
+    paths:
+      - .github/workflows/rocm.yml
+      - .github/workflows/_ci-run-decision.yml
+      - .github/workflows/_get-changed-files.yml
+      - .ci/scripts/test-rocm-aoti.sh
+      - CMakeLists.txt
+      - CMakePresets.json
+      - install_requirements.py
+      - torch_pin.py
+      - backends/aoti/**
+      - '!backends/aoti/**.md'
+      - backends/cuda/**
+      - '!backends/cuda/**.md'
+      - extension/cuda/**
+      - examples/cuda/**
+      - '!examples/cuda/**.md'
+  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}-${{ github.event_name == 'schedule' }}
+  cancel-in-progress: false
+
+permissions:
+  contents: read
+
+jobs:
+  changed-files:
+    name: Get changed files
+    uses: ./.github/workflows/_get-changed-files.yml
+    with:
+      include-push-diff: true
+
+  run-decision:
+    name: CI run decision
+    uses: ./.github/workflows/_ci-run-decision.yml
+
+  aoti-run-decision:
+    name: AOTI CI run decision
+    needs: [changed-files, run-decision]
+    runs-on: ubuntu-latest
+    outputs:
+      run-aoti: ${{ steps.decide.outputs.run-aoti }}
+    steps:
+      - name: Select AOTI runs
+        id: decide
+        env:
+          CHANGED_FILES: ${{ needs.changed-files.outputs.changed-files }}
+          FULL_RUN: ${{ needs.run-decision.outputs.is-full-run }}
+        run: |
+          set -efu
+          RUN_AOTI=false
+
+          if [[ "$CHANGED_FILES" == "*" || "$FULL_RUN" == "true" ]]; then
+            RUN_AOTI=true
+          fi
+
+          for path in $CHANGED_FILES; do
+            case "$path" in
+              backends/aoti/*.md | backends/cuda/*.md | examples/cuda/*.md)
+                continue
+                ;;
+              .github/workflows/rocm.yml | \
+              .github/workflows/_ci-run-decision.yml | \
+              .github/workflows/_get-changed-files.yml | \
+              .ci/scripts/test-rocm-aoti.sh | \
+              CMakeLists.txt | \
+              CMakePresets.json | \
+              install_requirements.py | \
+              torch_pin.py | \
+              backends/aoti/* | \
+              backends/cuda/* | \
+              extension/cuda/* | \
+              examples/cuda/*)
+                RUN_AOTI=true
+                break
+                ;;
+            esac
+          done
+
+          echo "run-aoti=$RUN_AOTI" >> "$GITHUB_OUTPUT"
+
+  unittest-rocm-gfx950:
+    name: unittest-rocm-gfx950-rocm${{ matrix.rocm-version }}
+    needs: [aoti-run-decision]
+    if: needs.aoti-run-decision.outputs.run-aoti == 'true'
+    strategy:
+      fail-fast: false
+      matrix:
+        rocm-version: ["7.1"]
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+    permissions:
+      id-token: write
+      contents: read
+    with:
+      timeout: 90
+      no-sudo: true
+      runner: linux.rocm.gpu.gfx950.1
+      gpu-arch-type: rocm
+      gpu-arch-version: ${{ matrix.rocm-version }}
+      use-custom-docker-registry: false
+      docker-image: pytorch/manylinux2_28-builder:rocm${{ matrix.rocm-version }}
+      submodules: recursive
+      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+      script: |
+        set -eux
+        eval "$(conda shell.bash hook)"
+        conda create -n executorch-rocm-ci python=3.10 -y
+        conda activate executorch-rocm-ci
+        ROCM_VERSION=${{ matrix.rocm-version }} bash .ci/scripts/test-rocm-aoti.sh

CMakePresets.json

@@ -161,6 +161,25 @@
         "list": ["Linux", "Windows"]
       }
     },
+    {
+      "name": "llm-release-rocm",
+      "displayName": "LLM release build with ROCm",
+      "inherits": [
+        "llm-release"
+      ],
+      "binaryDir": "${sourceDir}/cmake-out-rocm-llm",
+      "cacheVariables": {
+        "CMAKE_INSTALL_PREFIX": "${sourceDir}/cmake-out-rocm-llm",
+        "CMAKE_PREFIX_PATH": "$penv{ROCM_PATH};/opt/rocm;$penv{CMAKE_PREFIX_PATH}",
+        "EXECUTORCH_BUILD_CUDA": "OFF",
+        "EXECUTORCH_BUILD_ROCM": "ON"
+      },
+      "condition": {
+        "lhs": "${hostSystemName}",
+        "type": "equals",
+        "rhs": "Linux"
+      }
+    },
     {
       "name": "llm-release-metal",
       "displayName": "LLM release build with Metal",
@@ -397,6 +416,16 @@
       ],
       "jobs": 0
     },
+    {
+      "name": "llm-release-rocm-install",
+      "displayName": "Build and install LLM extension release artifacts (ROCm)",
+      "configurePreset": "llm-release-rocm",
+      "configuration": "Release",
+      "targets": [
+        "install"
+      ],
+      "jobs": 0
+    },
     {
       "name": "llm-release-metal-install",
       "displayName": "Build and install LLM extension release artifacts (Metal)",
@@ -500,6 +529,20 @@
         }
       ]
     },
+    {
+      "name": "llm-release-rocm",
+      "displayName": "Configure, build and install ExecuTorch LLM extension with ROCm enabled",
+      "steps": [
+        {
+          "type": "configure",
+          "name": "llm-release-rocm"
+        },
+        {
+          "type": "build",
+          "name": "llm-release-rocm-install"
+        }
+      ]
+    },
     {
       "name": "llm-release-metal",
       "displayName": "Configure, build and install ExecuTorch LLM extension with Metal enabled",

backends/cuda/rocm.md

@@ -7,7 +7,7 @@ GPUs as device type `cuda`, so names such as `CudaBackend` and
 `EXECUTORCH_BUILD_ROCM` is off by default, is never auto-enabled, and is
 mutually exclusive with `EXECUTORCH_BUILD_CUDA`. It requires
 `EXECUTORCH_BUILD_EXTENSION_TENSOR`. Execution has been validated only on
-MI300X (`gfx942`).
+MI300X (`gfx942`); CI is configured to exercise `gfx950`.
 
 ## Requirements and limitations
 
@@ -18,23 +18,19 @@ MI300X (`gfx942`).
 - Python pybindings cannot allocate ROCm device memory; use a native runner such
   as `executor_runner`.
 - Installed CMake consumers must be able to find the HIP package.
-- Model runners do not yet link the ROCm backend, and there is no ROCm CI.
+- Model runners do not yet claim ROCm support. The gfx950 CI job covers the
+  backend build, native runtime, AOTI export and execution, and Triton W4 tests.
 
 ## Build
 
 ```bash
-cmake -S . -B cmake-out-rocm \
-  -DCMAKE_BUILD_TYPE=Release \
-  -DCMAKE_PREFIX_PATH="$(python -c 'import torch; print(torch.utils.cmake_prefix_path)');/opt/rocm" \
-  -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \
-  -DEXECUTORCH_BUILD_ROCM=ON
-cmake --build cmake-out-rocm --target aoti_cuda_backend aoti_cuda_shims -j
+cmake --workflow --preset llm-release-rocm
 ```
 
 Verify that the build uses HIP and not the NVIDIA runtime:
 
 ```bash
-ldd cmake-out-rocm/backends/cuda/libaoti_cuda_shims.so | grep -E 'amdhip64|cudart'
+ldd cmake-out-rocm-llm/lib/libaoti_cuda_shims.so | grep -E 'amdhip64|cudart'
 ```
 
 A `libcudart` dependency is a configuration error.