flydsl.yaml

name: Fly DSL test

on:
  push:
    branches:
      - main
  pull_request:
    branches:
      - main
    types: [opened, synchronize, reopened, labeled]
  workflow_dispatch:

permissions:
  contents: read
  actions: read
  pull-requests: read

concurrency:
  group: ${{ github.workflow }}-${{ github.ref }}
  cancel-in-progress: true

env:
  DOCKER_IMAGE: "rocm/pytorch:rocm7.14_ubuntu24.04_py3.12_pytorch_release_2.12.0@sha256:c38eeda81d85f00fbe35d3d50ce42ce59c524e87d810624f4eb5c52fddb3b9ad"
  LLVM_BUILD_PROFILE: "amd-minimal"
  # Bump when the container or LLVM build environment changes incompatibly.
  MLIR_CACHE_VERSION: "rocm7.14-ubuntu24.04-py3.12-v1"
  GITHUB_REPO_NAME: ${{ github.event.pull_request.head.repo.full_name || github.repository }}
  GITHUB_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.event.head_commit.id || github.sha }}
  # PRs compare against their exact base; main pushes compare against the previous commit.
  BASE_REPO_NAME: ${{ github.event.pull_request.base.repo.full_name || github.repository }}
  BASE_COMMIT_SHA: ${{ github.event.pull_request.base.sha || github.event.before || 'main' }}

jobs:
  # ---------------------------------------------------------------------------
  # Detect whether the change touches code. Docs-only / non-code changes (docs,
  # markdown, root docs, images, non-workflow .github files) set code_changed to
  # 'false', and the heavy test / multi-gpu jobs below are skipped at the job
  # level (a skipped required check still counts as passing for branch
  # protection), so no GPU runner is used. workflow_dispatch always forces a
  # full run, and a detection failure fails safe to running the tests.
  # ---------------------------------------------------------------------------
  detect-changes:
    runs-on: ubuntu-latest
    timeout-minutes: 10
    outputs:
      # true when any changed file is code, OR on manual dispatch (force run).
      code_changed: ${{ steps.filter.outputs.code == 'true' || github.event_name == 'workflow_dispatch' }}
    steps:
      - name: Checkout code
        uses: actions/checkout@v4

      - name: Detect code vs non-code changes
        id: filter
        uses: dorny/paths-filter@v3
        with:
          # 'every': a file is code only if it matches all patterns, so
          # docs-only changes yield code=false. The default 'some' would match
          # any file via '**' and make 'code' always true.
          predicate-quantifier: 'every'
          filters: |
            code:
              - '**'
              - '!docs/**'
              - '!**/*.md'
              - '!**/*.rst'
              - '!LICENSE'
              - '!**/*.png'
              - '!**/*.jpg'
              - '!**/*.jpeg'
              - '!**/*.gif'
              - '!**/*.svg'
              - '!**/*.webp'
              - '!.github/ISSUE_TEMPLATE/**'

  check-signal:
    if: ${{ github.event_name != 'pull_request' || github.event.action != 'labeled' || github.event.label.name == 'multi-gpu' }}
    runs-on: ubuntu-latest
    timeout-minutes: 15
    steps:
      - name: Checkout code
        uses: actions/checkout@v4

      - name: Wait for Checks workflow
        run: ./.github/scripts/check_signal.sh
        env:
          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
          GITHUB_SHA: ${{ github.sha }}
          MAX_RETRIES: 5
          RETRY_INTERVAL_SECONDS: 30

  # Reports the required 'test (<runner>)' checks when the GPU 'test' job is
  # skipped for a non-code change (a skipped matrix job emits no per-runner
  # checks, which would leave branch protection waiting forever). Fails them if
  # check-signal failed/was cancelled, since check-signal is not itself required.
  test-skip:
    needs: [check-signal, detect-changes]
    if: >-
      ${{ always()
          && needs.detect-changes.result == 'success'
          && needs.detect-changes.outputs.code_changed != 'true' }}
    name: test
    strategy:
      matrix:
        runners: [
          'linux-flydsl-mi325-1',
          'linux-flydsl-mi355-1',
          'linux-flydsl-navi-2',
        ]
    runs-on: ubuntu-latest
    steps:
      - name: Report skipped GPU tests (green) or fail if pre-checks failed (red)
        env:
          CHECK_SIGNAL_RESULT: ${{ needs.check-signal.result }}
        run: |
          if [ "$CHECK_SIGNAL_RESULT" = "failure" ] || [ "$CHECK_SIGNAL_RESULT" = "cancelled" ]; then
            echo "::error::check-signal ${CHECK_SIGNAL_RESULT}; failing 'test (${{ matrix.runners }})' for this docs-only change."
            exit 1
          fi
          echo "Docs-only change: GPU tests skipped (check-signal=${CHECK_SIGNAL_RESULT}); reporting 'test (${{ matrix.runners }})' as passed."

  # ---------------------------------------------------------------------------
  # Prepare one shared MLIR install, then build current/base FlyDSL wheels once.
  # GPU jobs consume the compact wheels instead of transferring the full MLIR
  # install. Main and trusted same-repository PRs persist rebuilt MLIR caches.
  # ---------------------------------------------------------------------------
  prepare-mlir:
    needs: [check-signal, detect-changes]
    if: >-
      ${{ !cancelled()
          && needs.check-signal.result == 'success'
          && (needs.detect-changes.result != 'success'
              || needs.detect-changes.outputs.code_changed == 'true') }}
    runs-on: linux-flydsl-mi325-1
    timeout-minutes: 150
    steps:
      - name: Checkout code
        uses: actions/checkout@v4
        env:
          GIT_CONFIG_GLOBAL: ${{ runner.temp }}/flydsl-gitconfig
          GIT_CONFIG_NOSYSTEM: "1"
        with:
          repository: ${{ env.GITHUB_REPO_NAME }}
          ref: ${{ env.GITHUB_COMMIT_SHA }}
          path: flydsl-test

      - name: Restore shared MLIR cache
        id: mlir-cache
        uses: actions/cache/restore@v4
        with:
          path: mlir_install.tgz
          key: mlir-install-${{ runner.os }}-${{ runner.arch }}-${{ env.MLIR_CACHE_VERSION }}-${{ env.LLVM_BUILD_PROFILE }}-${{ hashFiles('flydsl-test/thirdparty/llvm-build-info.json', 'flydsl-test/thirdparty/llvm-rocdl-lld-argv0.patch', 'flydsl-test/scripts/build_llvm.sh') }}

      - name: Start MLIR build container
        run: |
          docker ps -aq -f name=flydsl_mlir_cache | xargs -r docker stop | xargs -r docker rm || true
          docker run -dt --network=host --user root \
            -v "${GITHUB_WORKSPACE:-$PWD}/flydsl-test:/flydsl-test" \
            --ipc=host --shm-size 16g \
            --cap-add=SYS_PTRACE \
            --security-opt seccomp=unconfined \
            -w /flydsl-test \
            --name flydsl_mlir_cache \
            ${{ env.DOCKER_IMAGE }}
        env:
          GITHUB_WORKSPACE: ${{ github.workspace }}

      - name: Install MLIR build dependencies
        run: |
          docker exec flydsl_mlir_cache bash -c "apt-get update && apt-get install -y cmake build-essential patchelf"
          docker exec flydsl_mlir_cache bash -c "python3 -m pip install -U pip setuptools wheel"
          docker exec flydsl_mlir_cache bash -c "python3 -m pip install ninja>=1.11.1"
          docker exec flydsl_mlir_cache bash -c "python3 -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ 'rocm[devel]==7.14.0' && rocm-sdk init"
          docker exec flydsl_mlir_cache bash -c "git config --global --add safe.directory /flydsl-test"
          sdk_root="$(docker exec flydsl_mlir_cache rocm-sdk path --root)"
          test -n "${sdk_root}"
          docker exec flydsl_mlir_cache test -d "${sdk_root}"
          printf 'ROCM_PATH=%s\n' "${sdk_root}" >>"${GITHUB_ENV}"

      - name: Use cached MLIR install
        if: steps.mlir-cache.outputs.cache-hit == 'true'
        run: |
          test -s ./mlir_install.tgz
          docker cp ./mlir_install.tgz flydsl_mlir_cache:/tmp/mlir_install.tgz
          docker exec flydsl_mlir_cache bash -c "mkdir -p /llvm-project && tar -xzf /tmp/mlir_install.tgz -C /llvm-project"
          docker exec flydsl_mlir_cache bash -c "test -d /llvm-project/mlir_install/lib/cmake/mlir"

      - name: Build LLVM once
        if: steps.mlir-cache.outputs.cache-hit != 'true'
        # A cold LLVM build measures ~10 min on these runners. Fail fast instead
        # of letting a stalled fetch or build hold the GPU runner until the
        # job-level timeout.
        timeout-minutes: 45
        run: |
          set -ex
          docker exec -e LLVM_BUILD_PROFILE -e ROCM_PATH flydsl_mlir_cache bash -c '
            cd /flydsl-test
            bash scripts/build_llvm.sh
          '
          docker exec flydsl_mlir_cache bash -c "ls -la /llvm-project/mlir_install/lib/cmake/mlir"
          docker cp flydsl_mlir_cache:/llvm-project/mlir_install.tgz ./mlir_install.tgz
          test -s ./mlir_install.tgz

      - name: Build current and base FlyDSL wheels
        run: |
          docker exec \
            -e BASE_REPO_NAME \
            -e BASE_COMMIT_SHA \
            -e ROCM_PATH \
            flydsl_mlir_cache bash -c '
              export MLIR_PATH=/llvm-project/mlir_install
              export PATH="$(rocm-sdk path --bin):$PATH"
              export CMAKE_PREFIX_PATH="$(rocm-sdk path --cmake)${CMAKE_PREFIX_PATH:+:$CMAKE_PREFIX_PATH}"
              cd /flydsl-test
              bash scripts/build_ci_wheels.sh
            '
          docker cp flydsl_mlir_cache:/tmp/flydsl-ci-wheels ./flydsl-ci-wheels
          (cd flydsl-ci-wheels/pr && sha256sum -c SHA256SUMS)
          if [ -d flydsl-ci-wheels/base ]; then
            (cd flydsl-ci-wheels/base && sha256sum -c SHA256SUMS)
          else
            echo "::warning title=Baseline wheel unavailable::Uploading PR wheel without an exact base benchmark."
          fi
          (cd flydsl-ci-wheels/tools && sha256sum -c SHA256SUMS)

      - name: Save shared MLIR cache
        if: >-
          steps.mlir-cache.outputs.cache-hit != 'true' &&
          ((github.ref == 'refs/heads/main' &&
            (github.event_name == 'push' || github.event_name == 'workflow_dispatch')) ||
           (github.event_name == 'pull_request' &&
            github.event.pull_request.head.repo.full_name == github.repository))
        uses: actions/cache/save@v4
        with:
          path: mlir_install.tgz
          key: ${{ steps.mlir-cache.outputs.cache-primary-key }}

      - name: Upload FlyDSL CI wheels for this run
        uses: actions/upload-artifact@v4
        with:
          name: flydsl-ci-wheels-${{ github.run_id }}
          path: flydsl-ci-wheels
          if-no-files-found: error
          retention-days: 1

      - name: Clean up MLIR build container
        if: always()
        run: |
          docker stop flydsl_mlir_cache || true
          docker rm flydsl_mlir_cache || true

  # ---------------------------------------------------------------------------
  # Single-GPU tests: kernels, unit, examples, MLIR FileCheck, benchmarks.
  # Runs on 1-GPU and Navi runners only.
  # ---------------------------------------------------------------------------
  test:
    needs: [check-signal, detect-changes, prepare-mlir]
    # Run only when check-signal passed AND the change touched code. Docs-only
    # changes (code_changed == 'false') skip this job entirely, so no GPU runner
    # is used; a detect-changes failure falls back to running.
    if: >-
      ${{ !cancelled()
          && needs.check-signal.result == 'success'
          && needs.prepare-mlir.result == 'success'
          && (needs.detect-changes.result != 'success'
              || needs.detect-changes.outputs.code_changed == 'true') }}
    env:
      # Some self-hosted runners configure a GitHub -> git-cache URL rewrite.
      # If that cache is unavailable, checkout fails before the tests start.
      GIT_CONFIG_GLOBAL: /dev/null
      GIT_CONFIG_NOSYSTEM: "1"
      # Temporary quarantine for linux-flydsl-navi-2 GPU 1; see #858.
      CI_HIP_VISIBLE_DEVICES: ${{ contains(matrix.runners, 'navi') && '3' || '' }}
    strategy:
      matrix:
        runners: [
          'linux-flydsl-mi325-1',
          'linux-flydsl-mi355-1',
          'linux-flydsl-navi-2',
        ]
      fail-fast: false
    runs-on: ${{ matrix.runners }}
    steps:
      - name: Checkout code
        uses: actions/checkout@v4
        env:
          # Some self-hosted runners configure a GitHub -> git-cache URL rewrite.
          # If that cache is unavailable, checkout fails before the tests start.
          GIT_CONFIG_GLOBAL: ${{ runner.temp }}/flydsl-gitconfig
          GIT_CONFIG_NOSYSTEM: "1"
        with:
          repository: ${{ env.GITHUB_REPO_NAME }}
          ref: ${{ env.GITHUB_COMMIT_SHA }}
          path: flydsl-test

      # download-artifact overlays its destination. Self-hosted workspaces are
      # persistent, so remove wheels left by an older package version first.
      - name: Remove stale FlyDSL CI wheels
        run: rm -rf "${GITHUB_WORKSPACE:?}/flydsl-ci-wheels"

      - name: Download FlyDSL CI wheels
        uses: actions/download-artifact@v4
        with:
          name: flydsl-ci-wheels-${{ github.run_id }}
          path: flydsl-ci-wheels

      - name: Verify FlyDSL CI wheels
        run: |
          shopt -s nullglob
          pr_wheels=(flydsl-ci-wheels/pr/*.whl)
          if [ "${#pr_wheels[@]}" -ne 1 ]; then
            printf 'Expected exactly one PR wheel, found %s:\n' "${#pr_wheels[@]}" >&2
            printf '  %s\n' "${pr_wheels[@]}" >&2
            exit 1
          fi
          (cd flydsl-ci-wheels/pr && sha256sum -c SHA256SUMS)
          if [ -d flydsl-ci-wheels/base ]; then
            base_wheels=(flydsl-ci-wheels/base/*.whl)
            if [ "${#base_wheels[@]}" -ne 1 ]; then
              printf 'Expected exactly one base wheel, found %s:\n' "${#base_wheels[@]}" >&2
              printf '  %s\n' "${base_wheels[@]}" >&2
              exit 1
            fi
            (cd flydsl-ci-wheels/base && sha256sum -c SHA256SUMS)
          else
            echo "::warning title=Baseline wheel unavailable::Exact base benchmark will be skipped."
          fi
          (cd flydsl-ci-wheels/tools && sha256sum -c SHA256SUMS)

      - name: Start CI container
        run: |
          echo "Clean up containers..."
          docker ps -aq -f name=flydsl_test | xargs -r docker stop | xargs -r docker rm || true

          echo "Start CI container..."
          if [ -f "/etc/podinfo/gha-render-devices" ]; then
            DEVICE_FLAG=$(cat /etc/podinfo/gha-render-devices)
          else
            DEVICE_FLAG="--device /dev/dri"
          fi

          GPU_ENV_ARGS=()
          if [ -n "${CI_HIP_VISIBLE_DEVICES}" ]; then
            GPU_ENV_ARGS=(--env "HIP_VISIBLE_DEVICES=${CI_HIP_VISIBLE_DEVICES}")
          fi

          echo "Starting container: flydsl_test:ci"
          docker run -dt --network=host --user root --device=/dev/kfd $DEVICE_FLAG "${GPU_ENV_ARGS[@]}" \
          -v "${GITHUB_WORKSPACE:-$PWD}/flydsl-test:/flydsl-test" \
          -v "${GITHUB_WORKSPACE:-$PWD}/flydsl-ci-wheels:/flydsl-ci-wheels:ro" \
          --ipc=host --group-add video \
          --shm-size 16g \
          --cap-add=SYS_PTRACE \
          --security-opt seccomp=unconfined \
          -w /flydsl-test \
          --name flydsl_test \
          ${{ env.DOCKER_IMAGE }}
        env:
          GITHUB_WORKSPACE: ${{ github.workspace }}

      - name: Install dependencies
        run: |
          docker exec flydsl_test bash -c "apt-get update && apt-get install -y cmake build-essential patchelf"
          docker exec flydsl_test bash -c "python3 -m pip install -U pip setuptools wheel"
          docker exec flydsl_test bash -c "python3 -m pip install ninja>=1.11.1"
          docker exec flydsl_test bash -c "python3 -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ 'rocm[devel]==7.14.0' && rocm-sdk init"
          docker exec flydsl_test bash -c "python3 -m pip install -U 'hypothesis>=6.82.0'"
          docker exec flydsl_test bash -c "git config --global --add safe.directory /flydsl-test && cd /flydsl-test && git log"
          sdk_root="$(docker exec flydsl_test rocm-sdk path --root)"
          test -n "${sdk_root}"
          docker exec flydsl_test test -d "${sdk_root}"
          printf 'ROCM_PATH=%s\n' "${sdk_root}" >>"${GITHUB_ENV}"

      - name: Install FlyDSL CI wheel
        run: |
          docker exec flydsl_test bash -c "python3 -m pip install --force-reinstall --no-deps /flydsl-ci-wheels/pr/*.whl"
          docker exec flydsl_test bash -c "mkdir -p /flydsl-test/build-fly/bin && cp /flydsl-ci-wheels/pr/fly-opt /flydsl-test/build-fly/bin/fly-opt && chmod +x /flydsl-test/build-fly/bin/fly-opt"
          docker exec flydsl_test bash -c "cp /flydsl-ci-wheels/tools/FileCheck /usr/local/bin/FileCheck && chmod +x /usr/local/bin/FileCheck"
          docker exec -e ROCM_PATH flydsl_test bash -c "/flydsl-test/build-fly/bin/fly-opt --version"
          docker exec -e ROCM_PATH flydsl_test bash -c 'python3 -c "import flydsl; from flydsl._mlir.ir import Context; print(\"Installed FlyDSL CI wheel\")"'

      - name: Prepare aiter
        if: ${{ !contains(matrix.runners, 'navi') }}
        run: |
          docker exec flydsl_test bash -c "rm -rf /tmp/aiter && git clone --depth 1 --recursive --shallow-submodules https://github.com/ROCm/aiter.git /tmp/aiter"
          docker exec flydsl_test bash -c "python3 -c \"from pathlib import Path; src = Path('/tmp/aiter/requirements.txt'); dst = Path('/tmp/aiter/requirements-flydsl-ci.txt'); lines = [line for line in src.read_text().splitlines() if line.strip() and not line.strip().startswith('flydsl==')]; dst.write_text('\\n'.join(lines) + '\\n')\" && python3 -m pip install -r /tmp/aiter/requirements-flydsl-ci.txt"
          docker exec flydsl_test bash -c "python3 -c 'import torch; import triton; assert tuple(map(int, triton.__version__.split(\"+\", 1)[0].split(\".\")[:3])) >= (3, 6, 0), triton.__version__; print(\"Using image Triton\", triton.__version__, \"with PyTorch\", torch.__version__)'"

      - name: Run tests
        timeout-minutes: 60
        run: |
          docker exec -e ROCM_PATH flydsl_test bash -c 'export PYTHONPATH=/tmp/aiter:${PYTHONPATH:-} && export AITER_REPO=/tmp/aiter && cd /flydsl-test && RUN_TESTS_FULL=1 bash scripts/run_tests.sh'

      - name: Show tests logs
        if: failure()
        run: |
          log_dir="$(mktemp -d "${RUNNER_TEMP}/flydsl-test-logs.XXXXXX")"
          docker exec flydsl_test bash -c 'cd /tmp && tar czf /tmp/logs.tgz *.log 2>/dev/null || echo "no logs"'
          if docker cp flydsl_test:/tmp/logs.tgz "${log_dir}/logs.tgz"; then
            tar -xzf "${log_dir}/logs.tgz" -C "${log_dir}" || true
            find "${log_dir}" -maxdepth 1 -type f -name '*.log' -print -exec cat {} \;
          else
            echo "logs.tgz not found; skipping log extraction"
          fi

      - name: Run benchmarks
        id: benchmarks
        run: |
          docker exec -i -e ROCM_PATH flydsl_test bash <<'BASH'
          set -e -o pipefail
          export PYTHONPATH=/tmp/aiter:${PYTHONPATH:-}
          export AITER_REPO=/tmp/aiter
          cd /flydsl-test
          BENCH_LOG_DIR=/tmp/flydsl_bench_current bash scripts/run_benchmark.sh 2>&1 | tee /tmp/bench_current.out
          python3 scripts/benchmark_output_to_csv.py /tmp/bench_current.out /tmp/bench_current.csv
          BASH

      - name: Run benchmark baselines
        id: bench-baselines
        timeout-minutes: 45
        continue-on-error: true
        run: |
          docker exec -i \
            -e BASE_REPO_NAME \
            -e BASE_COMMIT_SHA \
            -e ROCM_PATH \
            flydsl_test bash <<'BASH'
          set -u
          cd /flydsl-test
          rm -rf /tmp/flydsl-bench-main /tmp/flydsl-bench-tag
          rm -f /tmp/bench_main.csv /tmp/bench_latest_tag.csv /tmp/bench_main_label /tmp/bench_latest_tag_label
          git worktree prune || true
          base_repo_url="https://github.com/${BASE_REPO_NAME}.git"

          if [ ! -d /flydsl-ci-wheels/base ]; then
            echo "::warning title=Benchmark baseline unavailable::Base wheel was not produced; skipping exact main comparison."
          elif git fetch "${base_repo_url}" "${BASE_COMMIT_SHA}" --no-tags --depth=1; then
            commit="$(git rev-parse FETCH_HEAD)"
            label="main"
            worktree="/tmp/flydsl-bench-main"
            csv="/tmp/bench_main_candidate.csv"
            output="/tmp/bench_main.out"
            log_dir="/tmp/flydsl_bench_${label}"
            echo "Trying exact benchmark baseline ${BASE_REPO_NAME}@${commit}"

            if ! git worktree add --detach "${worktree}" "${commit}"; then
              echo "Failed to create worktree for exact main baseline."
            else
              (
                set -e -o pipefail
                cd "${worktree}"
                python3 -m pip install --force-reinstall --no-deps /flydsl-ci-wheels/base/*.whl 2>&1 | tail -5
                mkdir -p build-fly/bin
                cp /flydsl-ci-wheels/base/fly-opt build-fly/bin/fly-opt
                chmod +x build-fly/bin/fly-opt
                export PYTHONPATH=/tmp/aiter:${PYTHONPATH:-}
                export AITER_REPO=/tmp/aiter
                BENCH_LOG_DIR="${log_dir}" bash scripts/run_benchmark.sh 2>&1 | tee "${output}"
                python3 /flydsl-test/scripts/benchmark_output_to_csv.py "${output}" "${csv}"
              )
              status=$?
              if [ "${status}" -eq 0 ] && [ -s "${csv}" ]; then
                cp "${csv}" /tmp/bench_main.csv
                echo "${label}" >/tmp/bench_main_label
              else
                echo "Exact main benchmark baseline failed; not retrying older commits."
              fi
            fi
          else
            echo "Failed to fetch exact main baseline ${BASE_REPO_NAME}@${BASE_COMMIT_SHA}."
          fi

          if [ ! -s /tmp/bench_main.csv ]; then
            echo "::warning title=Benchmark baseline unavailable::No usable exact main baseline found; skipping main comparison."
          fi

          if git fetch "${base_repo_url}" 'refs/tags/v*:refs/tags/v*' --force --depth=1; then
            latest_tag="$(git tag --list 'v*' --sort=-v:refname | sed -n '1p')"
            if [ -n "${latest_tag}" ]; then
              worktree="/tmp/flydsl-bench-tag"
              csv="/tmp/bench_latest_tag_candidate.csv"
              output="/tmp/bench_latest_tag.out"
              log_dir="/tmp/flydsl_bench_${latest_tag}"
              package_version="${latest_tag#v}"
              echo "Trying benchmark baseline ${latest_tag} from pip package flydsl==${package_version}"

              if git worktree add --detach "${worktree}" "${latest_tag}"; then
                (
                  set -e -o pipefail
                  cd "${worktree}"
                  python3 -m pip install --only-binary=:all: "flydsl==${package_version}" 2>&1 | tail -5
                  export PYTHONPATH=/tmp/aiter:${PYTHONPATH:-}
                  export AITER_REPO=/tmp/aiter
                  BENCH_LOG_DIR="${log_dir}" bash scripts/run_benchmark.sh 2>&1 | tee "${output}"
                  python3 /flydsl-test/scripts/benchmark_output_to_csv.py "${output}" "${csv}"
                )
                status=$?
                if [ "${status}" -eq 0 ] && [ -s "${csv}" ]; then
                  cp "${csv}" /tmp/bench_latest_tag.csv
                  echo "${latest_tag}" >/tmp/bench_latest_tag_label
                else
                  echo "Benchmark baseline ${latest_tag} failed."
                fi
              else
                echo "Failed to create worktree for ${latest_tag}; skipping latest-tag comparison."
              fi
            else
              echo "::warning title=Tag baseline unavailable::No v* tags found; skipping latest-tag comparison."
            fi
          else
            echo "::warning title=Tag baseline unavailable::Failed to fetch tags; skipping latest-tag comparison."
          fi
          BASH

      - name: Check benchmark performance (current vs main)
        if: steps.bench-baselines.outcome != 'skipped'
        timeout-minutes: 5
        run: |
          docker exec flydsl_test bash -c "
            if [ ! -f /tmp/bench_main.csv ]; then
              echo '::warning title=Benchmark comparison skipped::No usable main benchmark baseline found.'
              exit 0
            fi
            cd /flydsl-test
            main_label=\$(cat /tmp/bench_main_label 2>/dev/null || echo main)
            python3 scripts/compare_benchmark.py /tmp/bench_main.csv /tmp/bench_current.csv \
              --baseline-label \"\${main_label}\" --current-label current
          "

      - name: Check benchmark performance (current vs latest tag)
        if: steps.bench-baselines.outcome != 'skipped'
        timeout-minutes: 5
        run: |
          docker exec flydsl_test bash -c "
            if [ ! -f /tmp/bench_latest_tag.csv ]; then
              echo '::warning title=Benchmark comparison skipped::No usable latest-tag benchmark baseline found.'
              exit 0
            fi
            cd /flydsl-test
            latest_tag_label=\$(cat /tmp/bench_latest_tag_label 2>/dev/null || echo latest-tag)
            python3 scripts/compare_benchmark.py /tmp/bench_latest_tag.csv /tmp/bench_current.csv \
              --baseline-label \"\${latest_tag_label}\" --current-label current
          "

      - name: Show benchmarks logs
        if: failure()
        run: |
          log_dir="$(mktemp -d "${RUNNER_TEMP}/flydsl-benchmark-logs.XXXXXX")"
          docker exec flydsl_test bash -c 'cd /tmp && tar czf /tmp/flydsl_bench_logs.tgz flydsl_bench* bench_*.csv 2>/dev/null || echo "no logs"'
          if docker cp flydsl_test:/tmp/flydsl_bench_logs.tgz "${log_dir}/logs.tgz"; then
            tar -xzf "${log_dir}/logs.tgz" -C "${log_dir}" || true
            find "${log_dir}" -type f -name '*.log' -print -exec cat {} \;
            find "${log_dir}" -maxdepth 1 -type f -name 'bench_*.csv' -print -exec cat {} \;
          else
            echo "flydsl_bench_logs.tgz not found; skipping log extraction"
          fi
      
      - name: Clean up
        if: always()
        run: |
          docker stop flydsl_test
          docker rm flydsl_test

  # ---------------------------------------------------------------------------
  # Multi-GPU communication operator tests: ONLY for 8-GPU runners.
  # Runs on BOTH linux-flydsl-mi325-8 AND linux-flydsl-mi355-8 independently.
  # Triggered when PR has label "multi-gpu" (added by a maintainer), or when
  # the workflow is manually dispatched.
  # fail-fast: false ensures both runners always complete even if one fails.
  # ---------------------------------------------------------------------------
  multi-gpu:
    needs: [test, detect-changes, prepare-mlir]
    name: Multi-GPU Communication Operator Tests (${{ matrix.runners }})
    timeout-minutes: 120
    env:
      # Keep checkout independent of runner-local git-cache rewrites.
      GIT_CONFIG_GLOBAL: /dev/null
      GIT_CONFIG_NOSYSTEM: "1"
    # Also gated on code_changed: docs-only changes skip this job (a
    # detect-changes failure falls back to running).
    if: |
      !cancelled() &&
      needs.prepare-mlir.result == 'success' &&
      (needs.detect-changes.result != 'success' ||
       needs.detect-changes.outputs.code_changed == 'true') && (
        (github.event_name == 'pull_request' &&
         contains(github.event.pull_request.labels.*.name, 'multi-gpu')) ||
        github.event_name == 'workflow_dispatch'
      )
    strategy:
      matrix:
        runners: [
          'linux-flydsl-mi325-8',
          'linux-flydsl-mi355-8',
        ]
      fail-fast: false
    runs-on: ${{ matrix.runners }}
    steps:
      - name: Checkout code
        uses: actions/checkout@v4
        env:
          # Keep checkout independent of runner-local git-cache rewrites.
          GIT_CONFIG_GLOBAL: ${{ runner.temp }}/flydsl-gitconfig
          GIT_CONFIG_NOSYSTEM: "1"
        with:
          repository: ${{ env.GITHUB_REPO_NAME }}
          ref: ${{ env.GITHUB_COMMIT_SHA }}
          path: flydsl-test

      # download-artifact overlays its destination. Self-hosted workspaces are
      # persistent, so remove wheels left by an older package version first.
      - name: Remove stale FlyDSL CI wheels
        run: rm -rf "${GITHUB_WORKSPACE:?}/flydsl-ci-wheels"

      - name: Download FlyDSL CI wheels
        uses: actions/download-artifact@v4
        with:
          name: flydsl-ci-wheels-${{ github.run_id }}
          path: flydsl-ci-wheels

      - name: Verify FlyDSL CI wheels
        run: |
          shopt -s nullglob
          pr_wheels=(flydsl-ci-wheels/pr/*.whl)
          if [ "${#pr_wheels[@]}" -ne 1 ]; then
            printf 'Expected exactly one PR wheel, found %s:\n' "${#pr_wheels[@]}" >&2
            printf '  %s\n' "${pr_wheels[@]}" >&2
            exit 1
          fi
          (cd flydsl-ci-wheels/pr && sha256sum -c SHA256SUMS)
          if [ -d flydsl-ci-wheels/base ]; then
            base_wheels=(flydsl-ci-wheels/base/*.whl)
            if [ "${#base_wheels[@]}" -ne 1 ]; then
              printf 'Expected exactly one base wheel, found %s:\n' "${#base_wheels[@]}" >&2
              printf '  %s\n' "${base_wheels[@]}" >&2
              exit 1
            fi
            (cd flydsl-ci-wheels/base && sha256sum -c SHA256SUMS)
          else
            echo "::warning title=Baseline wheel unavailable::Exact base benchmark will be skipped."
          fi
          (cd flydsl-ci-wheels/tools && sha256sum -c SHA256SUMS)

      - name: Start CI container
        run: |
          echo "Clean up containers..."
          docker ps -aq -f name=flydsl_test | xargs -r docker stop | xargs -r docker rm || true

          echo "Start CI container..."
          if [ -f "/etc/podinfo/gha-render-devices" ]; then
            DEVICE_FLAG=$(cat /etc/podinfo/gha-render-devices)
          else
            DEVICE_FLAG="--device /dev/dri"
          fi

          docker run -dt --network=host --user root --device=/dev/kfd $DEVICE_FLAG \
          -v "${GITHUB_WORKSPACE:-$PWD}/flydsl-test:/flydsl-test" \
          -v "${GITHUB_WORKSPACE:-$PWD}/flydsl-ci-wheels:/flydsl-ci-wheels:ro" \
          --ipc=host --group-add video \
          --shm-size 16g \
          --cap-add=SYS_PTRACE \
          --security-opt seccomp=unconfined \
          -w /flydsl-test \
          --name flydsl_test \
          ${{ env.DOCKER_IMAGE }}
        env:
          GITHUB_WORKSPACE: ${{ github.workspace }}

      - name: Install dependencies
        run: |
          docker exec flydsl_test bash -c "apt-get update && apt-get install -y cmake build-essential patchelf"
          docker exec flydsl_test bash -c "python3 -m pip install -U pip setuptools wheel"
          docker exec flydsl_test bash -c "python3 -m pip install ninja>=1.11.1"
          docker exec flydsl_test bash -c "python3 -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ 'rocm[devel]==7.14.0' && rocm-sdk init"
          docker exec flydsl_test bash -c "python3 -m pip install -U 'hypothesis>=6.82.0'"
          docker exec flydsl_test bash -c "git config --global --add safe.directory /flydsl-test && cd /flydsl-test && git log"
          sdk_root="$(docker exec flydsl_test rocm-sdk path --root)"
          test -n "${sdk_root}"
          docker exec flydsl_test test -d "${sdk_root}"
          printf 'ROCM_PATH=%s\n' "${sdk_root}" >>"${GITHUB_ENV}"

      - name: Install FlyDSL CI wheel
        run: |
          docker exec flydsl_test bash -c "python3 -m pip install --force-reinstall --no-deps /flydsl-ci-wheels/pr/*.whl"
          docker exec flydsl_test bash -c "mkdir -p /flydsl-test/build-fly/bin && cp /flydsl-ci-wheels/pr/fly-opt /flydsl-test/build-fly/bin/fly-opt && chmod +x /flydsl-test/build-fly/bin/fly-opt"
          docker exec flydsl_test bash -c "cp /flydsl-ci-wheels/tools/FileCheck /usr/local/bin/FileCheck && chmod +x /usr/local/bin/FileCheck"
          docker exec -e ROCM_PATH flydsl_test bash -c "/flydsl-test/build-fly/bin/fly-opt --version"
          docker exec -e ROCM_PATH flydsl_test bash -c 'python3 -c "import flydsl; from flydsl._mlir.ir import Context; print(\"Installed FlyDSL CI wheel\")"'

      - name: Install mori (for shmem regression test)
        timeout-minutes: 15
        run: |
          docker exec -e ROCM_PATH flydsl_test bash -c "
            apt-get install -y libpci-dev libibverbs-dev libgrpc++1.51 libgrpc29 &&
            python3 -m pip install amd_mori &&
            MORI_PRECOMPILE=1 python3 -c 'import mori'
          "

      - name: Run multi-GPU shmem regression test
        timeout-minutes: 10
        run: |
          docker exec -e ROCM_PATH flydsl_test bash -c "
            cd /flydsl-test
            python3 -m pytest tests/kernels/test_flydsl_shmem.py \
              -m multi_gpu -v --no-header --tb=short
          "

      - name: Run multi-GPU dispatch/combine CI sweep (8-GPU, accuracy + cudagraph perf)
        timeout-minutes: 30
        run: |
          docker exec -e ROCM_PATH flydsl_test bash -c "
            cd /flydsl-test &&
            MORI_SOCKET_IFNAME=lo python tests/kernels/test_profiler_dispatch_combine.py \
              --ci-sweep \
              --skip-verify-spawn \
              --world-size 8 \
              --port 29503 \
              --output-dir /tmp/flydsl_ci_sweep
          "

      # MegaMoEV2 A8W4 v4_pro accuracy against the torch f32 oracle. The test self-skips on gfx942.
      - name: Run multi-GPU MegaMoEV2 A8W4 v4_pro accuracy tests
        timeout-minutes: 45
        run: |
          docker exec -e ROCM_PATH flydsl_test bash -c "
            cd /flydsl-test &&
            python3 -c 'import torch; n = torch.cuda.device_count(); assert n >= 8, f\"requires 8 GPUs, found {n}\"' &&
            python3 -m pytest \
              tests/kernels/test_mega_moe_v2.py::test_mega_moe_8gpu_accuracy \
              -v --no-header --tb=short
          "

      - name: Run multi-GPU allreduce tests
        timeout-minutes: 30
        run: |
          docker exec -e ROCM_PATH flydsl_test bash -c "
            cd /flydsl-test
            python3 -m pytest tests/kernels/test_allreduce.py \
              -m multi_gpu -v --no-header --tb=short
          "

      - name: Run allreduce benchmark (PR)
        timeout-minutes: 30
        run: |
          docker exec -e ROCM_PATH flydsl_test bash -c "
            cd /flydsl-test
            python3 tests/kernels/test_allreduce.py \
              --world_size 8 --iters 51 --warmup 5 \
              --allreduce_impl flydsl --mode cudagraph \
              --shapes '2,7168,fp16;32,8192,fp32;128,8192,fp16;1024,7168,bf16;4096,8192,bf16' \
              --output_csv /tmp/bench_pr.csv
          "

      - name: Build main branch baseline
        id: build-main
        timeout-minutes: 20
        continue-on-error: true
        run: |
          docker exec \
            -e BASE_REPO_NAME \
            -e BASE_COMMIT_SHA \
            flydsl_test bash -c "
            set -e -o pipefail
            if [ ! -d /flydsl-ci-wheels/base ]; then
              echo '::warning title=Benchmark baseline unavailable::Base wheel was not produced.'
              exit 1
            fi
            cd /flydsl-test
            git fetch \"https://github.com/\${BASE_REPO_NAME}.git\" \"\${BASE_COMMIT_SHA}\" --no-tags --depth=1
            git worktree add --detach /tmp/flydsl-main FETCH_HEAD
            cd /tmp/flydsl-main
            python3 -m pip install --force-reinstall --no-deps /flydsl-ci-wheels/base/*.whl 2>&1 | tail -5
            mkdir -p build-fly/bin
            cp /flydsl-ci-wheels/base/fly-opt build-fly/bin/fly-opt
            chmod +x build-fly/bin/fly-opt
          "

      - name: Run allreduce benchmark (main)
        id: bench-main
        if: steps.build-main.outcome == 'success'
        timeout-minutes: 30
        continue-on-error: true
        run: |
          docker exec -e ROCM_PATH flydsl_test bash -c "
            cp /flydsl-test/tests/kernels/test_allreduce.py \
               /tmp/flydsl-main/tests/kernels/test_allreduce.py
            cd /tmp/flydsl-main
            python3 tests/kernels/test_allreduce.py \
              --world_size 8 --iters 51 --warmup 5 \
              --allreduce_impl flydsl --mode cudagraph \
              --shapes '2,7168,fp16;32,8192,fp32;128,8192,fp16;1024,7168,bf16;4096,8192,bf16' \
              --output_csv /tmp/bench_main.csv
          "

      - name: Check performance regression (PR vs main)
        if: steps.bench-main.outcome != 'skipped'
        timeout-minutes: 5
        run: |
          docker exec flydsl_test bash -c "
            cd /flydsl-test
            python3 tests/kernels/compare_allreduce_benchmark.py \
              /tmp/bench_main.csv /tmp/bench_pr.csv
          "

      - name: Show test logs
        if: failure()
        run: |
          log_dir="$(mktemp -d "${RUNNER_TEMP}/flydsl-test-logs.XXXXXX")"
          docker exec flydsl_test bash -c 'cd /tmp && tar czf /tmp/logs.tgz *.log 2>/dev/null || echo "no logs"'
          if docker cp flydsl_test:/tmp/logs.tgz "${log_dir}/logs.tgz"; then
            tar -xzf "${log_dir}/logs.tgz" -C "${log_dir}" || true
            find "${log_dir}" -maxdepth 1 -type f -name '*.log' -print -exec cat {} \;
          else
            echo "logs.tgz not found; skipping log extraction"
          fi

      - name: Clean up
        if: always()
        run: |
          docker stop flydsl_test
          docker rm flydsl_test