flydsl.yaml

name: Fly DSL test

on:
  push:
    branches:
      - main
  pull_request:
    branches:
      - main
    types: [opened, synchronize, reopened, labeled]
  workflow_dispatch:
    inputs:
      prepare_runner:
        description: 'Runner for prepare-mlir (compile-only; no GPU needed)'
        type: choice
        default: build-only-flydsl
        options:
          - build-only-flydsl
          - linux-flydsl-mi325-1
          - linux-flydsl-mi355-1
          - linux-flydsl-navi-2

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

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

env:
  DOCKER_IMAGE: "rocm/pytorch:rocm7.2.4_ubuntu24.04_py3.12_pytorch_release_2.10.0"
  LLVM_BUILD_PROFILE: "amd-minimal"
  # Bump when the container or LLVM build environment changes incompatibly.
  MLIR_CACHE_VERSION: "manylinux228-rocm7.2-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:
      # Must match the 'test' matrix below.
      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') }}
    # Compile-only, so it runs on the CPU builder instead of holding a GPU
    # runner. These are ARC scale sets: runs-on names exactly one, and the
    # dispatch input is how to retarget when it is busy.
    runs-on: ${{ inputs.prepare_runner || 'build-only-flydsl' }}
    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

      # A pin bump leaves the base source uncompilable against the PR's MLIR, so
      # the baseline wheel needs the base commit's own install. Every step gated
      # on pin_changed is skipped for a PR that leaves the LLVM inputs alone.
      - name: Resolve base LLVM pin
        id: base-pin
        # `run:` blocks are `bash -e`, and nothing decided here is worth failing
        # prepare-mlir over: with no outputs written, the gated steps skip.
        continue-on-error: true
        # Some self-hosted runners rewrite GitHub URLs to a git cache, and this
        # is the only fetch here that runs on the host, not in the container.
        env:
          GIT_CONFIG_GLOBAL: ${{ runner.temp }}/flydsl-gitconfig
          GIT_CONFIG_NOSYSTEM: "1"
        run: |
          set -uo pipefail
          # Self-hosted workspaces are reused; drop what the last run left.
          rm -rf base-pin
          rm -f mlir_install_base.tgz
          mkdir -p base-pin/thirdparty base-pin/scripts
          if ! git -C flydsl-test fetch "https://github.com/${BASE_REPO_NAME}.git" \
                 "${BASE_COMMIT_SHA}" --no-tags --depth=1; then
            echo "Could not fetch base ${BASE_REPO_NAME}@${BASE_COMMIT_SHA}; treating LLVM pin as unchanged."
            echo "pin_changed=false" >>"${GITHUB_OUTPUT}"
            exit 0
          fi
          for f in thirdparty/llvm-build-info.json thirdparty/llvm-rocdl-lld-argv0.patch scripts/build_llvm.sh; do
            if ! git -C flydsl-test show "FETCH_HEAD:${f}" >"base-pin/${f}"; then
              echo "Base commit has no ${f}; treating LLVM pin as unchanged."
              echo "pin_changed=false" >>"${GITHUB_OUTPUT}"
              exit 0
            fi
          done
          changed=false
          for f in thirdparty/llvm-build-info.json thirdparty/llvm-rocdl-lld-argv0.patch scripts/build_llvm.sh; do
            if ! diff -q "flydsl-test/${f}" "base-pin/${f}" >/dev/null; then
              echo "LLVM input differs from base: ${f}"
              changed=true
            fi
          done
          # sed, not python3: every other python call here runs in the container
          # and the runners promise no host interpreter. Scoped to "upstream" so a
          # second entry carrying its own llvm_hash cannot be picked up.
          base_hash="$(sed -n '/"upstream"/,/}/ s/.*"llvm_hash"[[:space:]]*:[[:space:]]*"\([0-9a-f]\{40\}\)".*/\1/p' \
            base-pin/thirdparty/llvm-build-info.json)"
          if [ -z "${base_hash}" ]; then
            echo "Could not read the base LLVM hash; treating LLVM pin as unchanged."
            echo "pin_changed=false" >>"${GITHUB_OUTPUT}"
            exit 0
          fi
          echo "pin_changed=${changed}" >>"${GITHUB_OUTPUT}"
          echo "base_llvm_hash=${base_hash}" >>"${GITHUB_OUTPUT}"
          echo "Base LLVM pin: ${base_hash}"

      # One script for both keys so they cannot drift; content-derived, which is
      # what makes the base commit's key computable from the extracted files.
      - name: Compute MLIR cache keys
        id: keys
        run: |
          set -euo pipefail
          self_key="$(bash flydsl-test/scripts/ci_mlir_cache_key.sh flydsl-test)"
          echo "self=${self_key}" >>"${GITHUB_OUTPUT}"
          echo "Shared MLIR cache key: ${self_key}"
          if [ "${{ steps.base-pin.outputs.pin_changed }}" = "true" ]; then
            base_key="$(bash flydsl-test/scripts/ci_mlir_cache_key.sh base-pin)"
            echo "base=${base_key}" >>"${GITHUB_OUTPUT}"
            echo "Baseline MLIR cache key: ${base_key}"
          fi

      # Restored before the shared entry and under the same `path`, then moved
      # aside: actions/cache derives its version from the path list, so an entry
      # saved as `mlir_install.tgz` is only findable under that name. No
      # restore-keys either - a prefix fallback would return an install built
      # from another pin.
      #
      # This is the entry every PR that leaves LLVM alone restores on every run,
      # so it stays warm. A PR that also changes MLIR_CACHE_VERSION or
      # LLVM_BUILD_PROFILE misses, since both feed the key from the PR side.
      - name: Restore baseline MLIR cache
        id: base-mlir-cache
        if: steps.base-pin.outputs.pin_changed == 'true'
        continue-on-error: true
        timeout-minutes: 25
        uses: actions/cache/restore@v4
        with:
          path: mlir_install.tgz
          key: ${{ steps.keys.outputs.base }}

      # This install is ~1.4 GB and these runners are spread across regions, so
      # the transfer is slow enough (~1.6 MB/s observed) to be interrupted:
      #
      #   Received 1010827264 of 1396566653 (72.4%), 1.6 MBs/sec
      #   ##[warning]Failed to restore: Server failed to authenticate the
      #             request ... Authorization header ... signature.
      #   Cache not found for input keys: mlir-install-...-4087a417...
      #
      # The entry was there and matched; the blob SAS auth failed mid-download,
      # which actions/cache reports as a plain miss. One retry costs nothing on
      # the normal path and is the difference between having a baseline and
      # silently not having one.
      - name: Retry baseline MLIR cache restore
        id: base-mlir-cache-retry
        if: >-
          steps.base-pin.outputs.pin_changed == 'true' &&
          steps.base-mlir-cache.outputs.cache-hit != 'true'
        continue-on-error: true
        timeout-minutes: 25
        uses: actions/cache/restore@v4
        with:
          path: mlir_install.tgz
          key: ${{ steps.keys.outputs.base }}

      - name: Stash baseline MLIR tarball
        if: >-
          steps.base-mlir-cache.outputs.cache-hit == 'true' ||
          steps.base-mlir-cache-retry.outputs.cache-hit == 'true'
        continue-on-error: true
        run: |
          set -uo pipefail
          if [ ! -s mlir_install.tgz ]; then
            echo "::warning title=Benchmark baseline unavailable::Baseline MLIR cache reported a hit but produced no tarball."
            rm -f mlir_install.tgz
            exit 0
          fi
          mv mlir_install.tgz mlir_install_base.tgz

      - name: Restore shared MLIR cache
        id: mlir-cache
        uses: actions/cache/restore@v4
        with:
          path: mlir_install.tgz
          key: ${{ steps.keys.outputs.self }}

      # Nothing publishes this tag; the fallback builds it the first time a
      # runner needs it. Same step as build-whl.yaml.
      - name: Pull or build manylinux Docker image
        run: |
          ROCM_VERSION=7.2
          IMAGE="ghcr.io/rocm/flydsl-manylinux_2_28:rocm${ROCM_VERSION}-cp310-cp314"
          if ! docker pull "${IMAGE}"; then
            docker build \
              --build-arg ROCM_VERSION="${ROCM_VERSION}" \
              -t "${IMAGE}" \
              - < flydsl-test/.github/workflows/Dockerfile.manylinux_2_28
          fi
          echo "MANYLINUX_IMAGE=${IMAGE}" >> $GITHUB_ENV

      - 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 \
            ${MANYLINUX_IMAGE}
        env:
          GITHUB_WORKSPACE: ${{ github.workspace }}

      # The image ships cmake, ninja, gcc, patchelf, git and ROCm at /opt/rocm.
      - name: Install MLIR build dependencies
        run: |
          docker exec flydsl_mlir_cache bash -c '
            for d in /opt/python/cp*; do
              case "${d}" in *t) continue ;; esac
              py="${d}/bin/python"
              [ -x "${py}" ] || continue
              ver=$("${py}" -c "import sys; print(f\"{sys.version_info.major}.{sys.version_info.minor}\")")
              ln -sf "${py}" "/usr/local/bin/python${ver}"
            done
            ln -sf /usr/local/bin/python3.12 /usr/local/bin/python3
          '
          docker exec flydsl_mlir_cache bash -c "python3 -m pip install -q nanobind"
          docker exec flydsl_mlir_cache bash -c "git config --global --add safe.directory '*'"
          sdk_root="$(docker exec flydsl_mlir_cache bash -c 'if command -v rocm-sdk >/dev/null 2>&1; then rocm-sdk path --root; elif [ -d /opt/rocm ]; then printf "/opt/rocm\n"; fi')"
          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

      # Checked against the VCSRevision.h the install ships: a wrong-pin baseline
      # is worse than none, since it yields a plausible number nobody questions.
      #
      # Only a cache hit produces a baseline. Building the old pin here would
      # cost a second llvm-project clone and full LLVM build on a GPU runner, per
      # push, for an advisory number. Revisit if the logs show misses are common.
      - name: Unpack and verify baseline MLIR
        id: base-mlir
        if: steps.base-pin.outputs.pin_changed == 'true'
        continue-on-error: true
        run: |
          set -uo pipefail
          if [ ! -s ./mlir_install_base.tgz ]; then
            echo "::notice title=Benchmark baseline skipped::No cached MLIR install for the base commit's LLVM (${{ steps.base-pin.outputs.base_llvm_hash }}), so the vs-main comparison is skipped; vs latest tag is unaffected. This resolves itself once main has run under that pin."
            exit 0
          fi
          docker cp ./mlir_install_base.tgz flydsl_mlir_cache:/tmp/mlir_install_base.tgz
          docker exec flydsl_mlir_cache bash -c "
            set -e
            rm -rf /llvm-project/mlir_install_base
            mkdir -p /tmp/mlir_base_extract
            tar -xzf /tmp/mlir_install_base.tgz -C /tmp/mlir_base_extract
            mv /tmp/mlir_base_extract/mlir_install /llvm-project/mlir_install_base
            rmdir /tmp/mlir_base_extract
            test -d /llvm-project/mlir_install_base/lib/cmake/mlir
          " || { echo "::warning title=Benchmark baseline unavailable::Baseline MLIR install could not be unpacked."; exit 0; }
          want="${{ steps.base-pin.outputs.base_llvm_hash }}"
          got="$(docker exec flydsl_mlir_cache sed -n 's/.*LLVM_REVISION R"(\([0-9a-f]\{40\}\))".*/\1/p' \
                   /llvm-project/mlir_install_base/include/llvm/Support/VCSRevision.h 2>/dev/null)"
          if [ -z "${got}" ]; then
            echo "Baseline MLIR carries no VC revision; relying on the content-derived cache key alone."
          elif [ "${got}" != "${want}" ]; then
            echo "::warning title=Benchmark baseline unavailable::Baseline MLIR is LLVM ${got}, expected ${want}; refusing to use it."
            docker exec flydsl_mlir_cache rm -rf /llvm-project/mlir_install_base
            exit 0
          else
            echo "Baseline MLIR verified at LLVM ${got}"
          fi
          echo "ready=true" >>"${GITHUB_OUTPUT}"

      - name: Build current and base FlyDSL wheels
        env:
          BASE_MLIR_PATH: ${{ steps.base-mlir.outputs.ready == 'true' && '/llvm-project/mlir_install_base' || '' }}
          BASELINE_LLVM_HASH: ${{ steps.base-pin.outputs.base_llvm_hash }}
        run: |
          docker exec \
            -e BASE_REPO_NAME \
            -e BASE_COMMIT_SHA \
            -e ROCM_PATH \
            -e BASE_MLIR_PATH \
            flydsl_mlir_cache bash -c '
              export MLIR_PATH=/llvm-project/mlir_install
              if command -v rocm-sdk >/dev/null 2>&1; then
                rocm_bin="$(rocm-sdk path --bin)"
                rocm_cmake="$(rocm-sdk path --cmake)"
              else
                rocm_bin="${ROCM_PATH}/bin"
                rocm_cmake="${ROCM_PATH}/lib/cmake"
              fi
              export PATH="${rocm_bin}:$PATH"
              export CMAKE_PREFIX_PATH="${rocm_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)
            # Mark which LLVM the wheel was really built against. A bump needing
            # no source adaptation still builds from the shared MLIR; labelling
            # that "old LLVM" points at a codegen difference that is not there.
            if [ -n "${BASE_MLIR_PATH}" ] && [ -n "${BASELINE_LLVM_HASH}" ]; then
              printf '%s\n' "${BASELINE_LLVM_HASH}" >flydsl-ci-wheels/base/BASELINE_LLVM
            fi
          else
            echo "::warning title=Baseline wheel unavailable::Uploading PR wheel without an exact base benchmark."
          fi
          (cd flydsl-ci-wheels/tools && sha256sum -c SHA256SUMS)

      # The second install is the whole disk cost, and nothing below needs it.
      - name: Drop baseline MLIR install
        if: always() && steps.base-pin.outputs.pin_changed == 'true'
        continue-on-error: true
        run: |
          rm -f ./mlir_install_base.tgz
          docker exec flydsl_mlir_cache bash -c "rm -rf /llvm-project/mlir_install_base /tmp/mlir_install_base.tgz" || true

      - 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:
      # Must match the 'test-skip' matrix above, and must keep emitting the
      # checks required by the 'main' ruleset. linux-flydsl-mi35x-1 is absent
      # on purpose: it is the same gfx950 target as linux-flydsl-mi355-1.
      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 "if command -v rocm-sdk >/dev/null 2>&1; then rocm-sdk init; fi"
          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 bash -c 'if command -v rocm-sdk >/dev/null 2>&1; then rocm-sdk path --root; elif [ -d /opt/rocm ]; then printf "/opt/rocm\n"; fi')"
          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"
            # From prepare-mlir's marker, not from the pinned hashes: this job
            # cannot see which MLIR the baseline wheel was compiled with.
            display_label="main@${commit:0:8}"
            baseline_llvm="$(cat /flydsl-ci-wheels/base/BASELINE_LLVM 2>/dev/null || true)"
            if [ -n "${baseline_llvm}" ]; then
              display_label="${display_label} (llvm ${baseline_llvm:0:8})"
              echo "::notice title=Benchmark baseline uses a different LLVM::Baseline was built against LLVM ${baseline_llvm:0:8}, the PR against its own pin; deltas include LLVM codegen changes."
            fi
            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 "${display_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 linux-flydsl-mi325-8 and linux-flydsl-mi355-8 independently
  # (linux-flydsl-mi35x-8 is omitted: same gfx950 target as mi355-8).
  # 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 "if command -v rocm-sdk >/dev/null 2>&1; then rocm-sdk init; fi"
          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 bash -c 'if command -v rocm-sdk >/dev/null 2>&1; then rocm-sdk path --root; elif [ -d /opt/rocm ]; then printf "/opt/rocm\n"; fi')"
          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/A4W4 v4_pro accuracy against the torch f32 oracle. The test self-skips on gfx942.
      - name: Run multi-GPU MegaMoEV2 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