[CI] Build the benchmark baseline wheel against the base commit's LLVM

prepare-mlir builds one shared MLIR install from the PR's pin and uses it
for both wheels. An LLVM pin bump also carries the source adaptation that
the new pin requires, so the base commit cannot compile against it and the
baseline wheel is never produced. The run ends at "No usable main benchmark
baseline found" (run 32954898547, PR #1051), losing the vs-main comparison
for exactly the PRs whose performance impact is least predictable.

When the LLVM inputs differ from the base commit's, restore the MLIR install
belonging to the base pin and build the baseline wheel against that.

- ci_mlir_cache_key.sh derives the cache key from file contents, making it
  computable for the base commit. It replaces hashFiles() and serves both
  call sites, so the two keys cannot drift.
- The baseline entry is restored before the shared one and under the same
  path. actions/cache derives its version from the path list, so a renamed
  restore would miss whatever the key said.
- The unpacked install is checked against its VCSRevision.h. A wrong-pin
  baseline is worse than none: it yields a plausible number nobody queries.
- prepare-mlir marks the wheel with the LLVM it was really built against,
  and the table is labelled from that marker. Deriving it from the pinned
  hashes would mislabel a bump that needed no source adaptation.

A cache miss means no baseline. The wanted entry is the one every non-bump
PR restores on every run, so a miss is the exception.

Changing the key formula invalidates the cache, so every PR cold-builds LLVM
until main's next push saves an entry under the new one.

Author: Phil-amd

.github/workflows/flydsl.yaml

@@ -143,12 +143,109 @@ jobs:
           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'
+        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'
+        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: 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') }}
+          key: ${{ steps.keys.outputs.self }}
 
       - name: Start MLIR build container
         run: |
@@ -200,12 +297,56 @@ jobs:
           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
               export PATH="$(rocm-sdk path --bin):$PATH"
@@ -217,11 +358,25 @@ jobs:
           (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' &&
@@ -436,6 +591,14 @@ jobs:
           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"
@@ -460,7 +623,7 @@ jobs:
               status=$?
               if [ "${status}" -eq 0 ] && [ -s "${csv}" ]; then
                 cp "${csv}" /tmp/bench_main.csv
-                echo "${label}" >/tmp/bench_main_label
+                echo "${display_label}" >/tmp/bench_main_label
               else
                 echo "Exact main benchmark baseline failed; not retrying older commits."
               fi

scripts/build_ci_wheels.sh

@@ -77,19 +77,23 @@ MLIR
 build_wheel() {
   local source_dir="$1"
   local label="$2"
+  # The baseline wheel needs the base commit's own MLIR: on a pin bump the base
+  # source predates the API change the new pin requires. The default keeps every
+  # other wheel on the shared install.
+  local mlir_path="${3:-${MLIR_PATH}}"
   local destination="${OUTPUT_DIR}/${label}"
   local wheels=()
   local fly_opt="${source_dir}/build-fly/build_py${PYTHON_SUFFIX}/bin/fly-opt"
 
-  echo "Building ${label} wheel from ${source_dir}"
+  echo "Building ${label} wheel from ${source_dir} (MLIR: ${mlir_path})"
   if ! (
     cd "${source_dir}"
     env \
       PYTHON_VERSIONS="${PYTHON_VERSION}" \
       "${PYTHON_BIN_ENV}=${PYTHON_BIN}" \
       VENV_ROOT="${VENV_ROOT}" \
       ALLOW_ANY_GLIBC=1 \
-      MLIR_PATH="${MLIR_PATH}" \
+      MLIR_PATH="${mlir_path}" \
       bash scripts/build_wheels.sh
   ); then
     return 1
@@ -146,7 +150,7 @@ build_base_wheel() {
   fi
 
   trap cleanup_base_worktree EXIT
-  if ! build_wheel "${BASE_WORKTREE}" base; then
+  if ! build_wheel "${BASE_WORKTREE}" base "${BASE_MLIR_PATH:-${MLIR_PATH}}"; then
     cleanup_base_worktree
     trap - EXIT
     return 1

scripts/ci_mlir_cache_key.sh

@@ -0,0 +1,42 @@
+#!/usr/bin/env bash
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2025 FlyDSL Project Contributors
+#
+# Print the actions/cache key for the MLIR install that the given tree builds.
+#
+# Derived from file *contents*, not paths, so it is computable for any commit
+# whose LLVM inputs have been extracted anywhere on disk - that is what lets a PR
+# look up the install belonging to its base commit's pin. Both workflow call
+# sites use this script; inlining the formula elsewhere would let the two keys
+# drift and return an install built from the wrong pin.
+#
+# Total by construction: a missing input contributes a marker, not an error. It
+# replaced a hashFiles() expression, which tolerated missing files too, and a PR
+# that deletes one must still get a usable key rather than a failed job.
+set -euo pipefail
+
+TREE_ROOT="${1:?usage: ci_mlir_cache_key.sh <tree-root>}"
+
+INPUTS=(
+  thirdparty/llvm-build-info.json
+  thirdparty/llvm-rocdl-lld-argv0.patch
+  scripts/build_llvm.sh
+)
+
+# Per-file digests, so moving bytes across a file boundary changes the key.
+digests=""
+for input in "${INPUTS[@]}"; do
+  if [[ -f "${TREE_ROOT}/${input}" ]]; then
+    digests+="${input}:$(sha256sum <"${TREE_ROOT}/${input}" | cut -d' ' -f1)"$'\n'
+  else
+    digests+="${input}:absent"$'\n'
+  fi
+done
+digest="$(printf '%s' "${digests}" | sha256sum | cut -c1-40)"
+
+printf 'mlir-install-%s-%s-%s-%s-%s\n' \
+  "${RUNNER_OS:?RUNNER_OS must be set}" \
+  "${RUNNER_ARCH:?RUNNER_ARCH must be set}" \
+  "${MLIR_CACHE_VERSION:?MLIR_CACHE_VERSION must be set}" \
+  "${LLVM_BUILD_PROFILE:?LLVM_BUILD_PROFILE must be set}" \
+  "${digest}"