open-standards/.github/workflows/wai-quantum-tensor-program.yml

name: wai-quantum portable tensor program

# The gate as a portable tensor program (quantum_fabric::stablehlo, StableHLO
# as MLIR text): every value a 64-bit integer that never overflows, so a
# compiler for any device must give the exact rows. Two independent
# open-source tensor compilers, both Apache-2.0 and pinned, build it here:
# IREE (3.12.0), whose runtime compares the outputs with the exact ones the
# crate writes, and XLA (through JAX 0.11.2, the compiler tensor processors
# take this program through), whose outputs ci/tensor-program/xla_run.py
# compares. ci/tensor-program/check.sh holds every run to the outcome it must
# have, and fails a run that ends without stating one.
#
# - cpu (the container runner): both compilers on the x86-64 CPU, the program
#   and its three mutants (each must fail where it should); builds for three
#   other GPU architectures (compiled, not run);
# - gpu (the hardware host): every Vulkan device the host has (its discrete
#   and integrated GPUs and the software one), and the discrete GPU through
#   CUDA with each compiler;
# - amd (framework-128gb): its integrated GPU through HIP, built for the GPU's
#   own architecture, and every Vulkan device the host has;
# - metal (the macOS runner): the platform's GPU through Metal, and XLA on the
#   arm64 CPU.
# Each leg past the first also runs one mutant on each device, so the
# comparison there is shown to catch a difference.
#
# A second program, the limb module, is the gate's products as one bf16
# matrix multiply accumulated in f32 (quantum_fabric::stablehlo::limb_module):
# what a compiler maps to a matrix unit. It must be exact on eight-bit limbs
# and differ on nine-bit ones, wherever the compiler builds it exactly: the
# CPUs under both compilers, and the discrete GPU through CUDA under both.

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

on:
  pull_request:
    paths:
      - 'wai/wai-quantum/**'
      - 'rust-toolchain.toml'
      - '.github/workflows/wai-quantum-tensor-program.yml'
  schedule:
    - cron: '40 5 * * *'

env:
  IREE_VERSION: 3.12.0
  JAX_VERSION: 0.11.2

jobs:
  # Does this pull request's own diff touch the paths below? This forge matches a
  # workflow's `paths:` against the push, so a merge of main into a branch fires the
  # workflow whatever the pull request changes; every job here waits on this answer.
  # A push to main, a schedule, or anything the script cannot determine answers
  # `true` (scripts/pr-touches-paths.py).
  relevant:
    name: does this change touch it
    runs-on: linux-arm64
    outputs:
      run: ${{ steps.gate.outputs.run }}
    steps:
      - uses: actions/checkout@v5
        with:
          fetch-depth: 0
      - id: gate
        run: |
          set -euo pipefail
          r=$(python3 scripts/pr-touches-paths.py .github/workflows/wai-quantum-tensor-program.yml \
                "${{ github.event.pull_request.base.sha }}" \
                "${{ github.event.pull_request.head.sha }}")
          echo "this change touches wai-quantum-tensor-program: $r"
          echo "run=$r" >> "$GITHUB_OUTPUT"

  cpu:
    name: tensor program · CPU, mutants, GPU builds
    needs: relevant
    if: needs.relevant.outputs.run == 'true'
    runs-on: linux
    timeout-minutes: 60
    container:
      image: debian@sha256:913f6706df59a68922d1dd08f78c2476560a8d367897200a6005b00e5f67c2d5
      options: --memory=6g --memory-swap=6g
    steps:
      - name: Tools
        shell: bash
        run: |
          set -euo pipefail
          apt-get update -qq
          apt-get install -y -qq --no-install-recommends ca-certificates curl git build-essential python3 python3-venv >/dev/null
      # No Node for a checkout action: the forge's archive of the exact commit,
      # its embedded commit id checked.
      - name: Fetch the source
        shell: bash
        run: |
          set -euo pipefail
          want="${{ github.event.pull_request.head.sha || github.sha }}"
          rm -rf src src.tar.gz src.tar && mkdir src
          curl -sSf -o src.tar.gz "http://100.64.224.0:3400/Transaction-Science/open-standards/archive/$want.tar.gz"
          gzip -dc src.tar.gz > src.tar
          got="$(git get-tar-commit-id < src.tar)"
          [ "$got" = "$want" ] || { echo "::error::the archive is of $got, not $want"; exit 1; }
          tar -xf src.tar -C src --strip-components=1
          rm src.tar.gz src.tar
      - name: The program and its mutants on the CPU, both compilers; the GPU builds
        shell: bash
        working-directory: src/wai/wai-quantum
        run: |
          set -euo pipefail
          cache=/ci-cache/wai-quantum-tensor-program
          mkdir -p "$cache"
          exec 9>"$cache/.lock"
          flock 9
          venv="$cache/compilers-iree-$IREE_VERSION-jax-$JAX_VERSION"
          [ -x "$venv/bin/iree-compile" ] || { rm -rf "$venv"; python3 -m venv "$venv"; "$venv/bin/pip" install -q "iree-base-compiler==$IREE_VERSION" "iree-base-runtime==$IREE_VERSION" "jax==$JAX_VERSION"; }
          export RUSTUP_HOME="$cache/rustup" CARGO_HOME="$cache/cargo" CARGO_TARGET_DIR="$cache/target"
          export PATH="$venv/bin:$CARGO_HOME/bin:$PATH" XLA_PYTHON="$venv/bin/python"
          if ! command -v rustup >/dev/null; then
            curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --no-modify-path --profile minimal --default-toolchain none
          fi
          rustup toolchain install
          cargo test --locked --features quantum_fabric --lib -- quantum_fabric::stablehlo:: 2>&1 | tee "$RUNNER_TEMP/unit.log"
          grep -qE '^test result: ok\. [1-9][0-9]* passed; 0 failed' "$RUNNER_TEMP/unit.log" || { echo "::error::no unit test ran"; exit 1; }
          check="bash $PWD/ci/tensor-program/check.sh"
          d="$RUNNER_TEMP/program"
          cargo run --locked --release -q --example stablehlo_gate --features quantum_fabric -- "$d"
          cd "$d"
          # The program exact; m1, m2 (the rounding, the shift) both functions
          # wrong; m3 (row 0's last sign) the rows wrong, the probe exact.
          for g in 'gate:==' 'gate_m1:!!' 'gate_m2:!!' 'gate_m3:!='; do
            m="${g%%:*}" want="${g#*:}"
            iree-compile $m.mlir --iree-hal-target-device=local --iree-hal-local-target-device-backends=llvm-cpu --iree-llvmcpu-target-cpu=host -o $m.cpu.vmfb
            $check iree local-task $m.cpu.vmfb "$want"
            $check xla cpu $m.mlir "$want"
          done
          # The products on a matrix unit: one bf16 multiply accumulated in f32,
          # exact on eight-bit limbs; nine-bit limbs, which bf16 cannot hold,
          # must differ (quantum_fabric::stablehlo::limb_module).
          iree-compile limbs.mlir --iree-hal-target-device=local --iree-hal-local-target-device-backends=llvm-cpu --iree-llvmcpu-target-cpu=host -o limbs.cpu.vmfb
          $check iree local-task limbs.cpu.vmfb '=!'
          $check xla cpu limbs.mlir '=!'
          for t in gfx942 gfx1100 gfx1201; do
            iree-compile gate.mlir --iree-hal-target-backends=rocm --iree-rocm-target=$t -o gate.$t.vmfb
            echo "built for $t: $(stat -c %s gate.$t.vmfb) bytes"
          done

  gpu:
    name: tensor program · GPUs (jetson-hub)
    needs: relevant
    if: needs.relevant.outputs.run == 'true'
    runs-on: linux-hw
    timeout-minutes: 60
    env:
      # The Vulkan devices this host has: its discrete and integrated GPUs and
      # the software one. Every one must give the exact rows.
      VULKAN_DEVICES: '3'
    steps:
      - name: Fetch the source
        shell: bash
        run: |
          set -euo pipefail
          want="${{ github.event.pull_request.head.sha || github.sha }}"
          rm -rf src && mkdir src && cd src
          git init -q
          # Fetch the commit itself, not the ref: a later commit on the ref that this
          # workflow's paths do not match would otherwise be fetched instead.
          git fetch -q --depth 1 http://100.64.224.0:3400/Transaction-Science/open-standards.git "$want"
          git checkout -q FETCH_HEAD
          got="$(git rev-parse HEAD)"
          [ "$got" = "$want" ] || { echo "::error::fetched $got, not $want"; exit 1; }
      - name: The program on every Vulkan device, and through CUDA with each compiler
        shell: bash
        working-directory: src/wai/wai-quantum
        run: |
          set -euo pipefail
          export PATH="$HOME/.cargo/bin:$PATH"
          iree="$HOME/iree-$IREE_VERSION"
          [ -x "$iree/bin/iree-compile" ] || { rm -rf "$iree"; python3 -m venv "$iree"; "$iree/bin/pip" install -q "iree-base-compiler==$IREE_VERSION" "iree-base-runtime==$IREE_VERSION"; }
          xla="$HOME/xla-$JAX_VERSION-cuda13"
          [ -x "$xla/bin/python" ] && "$xla/bin/python" -I -c 'import jax' 2>/dev/null || { rm -rf "$xla"; python3 -m venv "$xla"; "$xla/bin/pip" install -q "jax[cuda13]==$JAX_VERSION"; }
          export PATH="$iree/bin:$PATH" XLA_PYTHON="$xla/bin/python"
          check="bash $PWD/ci/tensor-program/check.sh"
          d="$RUNNER_TEMP/program"
          CARGO_TARGET_DIR="$HOME/hw-target/wai-quantum-tensor-program" nice -n 10 cargo run --locked --release -q --example stablehlo_gate --features quantum_fabric -- "$d"
          cd "$d"
          # A Vulkan capability profile with 64-bit integers; the default has none.
          for m in gate gate_m3; do
            iree-compile $m.mlir --iree-hal-target-device=vulkan --iree-vulkan-target=ampere -o $m.vulkan.vmfb
            iree-compile $m.mlir --iree-hal-target-device=cuda --iree-cuda-target=sm_89 -o $m.cuda.vmfb
          done
          devices="$(iree-run-module --list_devices 2>/dev/null | grep -oE 'vulkan://[^ ]+' | sort -u)"
          count="$(grep -c . <<<"$devices")"
          [ "$count" = "$VULKAN_DEVICES" ] || { echo "::error::$count Vulkan devices, not $VULKAN_DEVICES"; iree-run-module --list_devices; exit 1; }
          for dev in $devices; do
            $check iree "$dev" gate.vulkan.vmfb '=='
            $check iree "$dev" gate_m3.vulkan.vmfb '!='
          done
          $check iree cuda gate.cuda.vmfb '=='
          $check iree cuda gate_m3.cuda.vmfb '!='
          $check xla cuda gate.mlir '=='
          $check xla cuda gate_m3.mlir '!='
          # The products on the GPU's matrix path, through both compilers.
          iree-compile limbs.mlir --iree-hal-target-device=cuda --iree-cuda-target=sm_89 -o limbs.cuda.vmfb
          $check iree cuda limbs.cuda.vmfb '=!'
          $check xla cuda limbs.mlir '=!'

  amd:
    name: tensor program · integrated GPU (framework-128gb)
    needs: relevant
    if: needs.relevant.outputs.run == 'true'
    runs-on: linux-amd
    timeout-minutes: 60
    env:
      # The Vulkan devices this host has: its integrated GPU and the software
      # one. Every one must give the exact rows.
      VULKAN_DEVICES: '2'
    steps:
      - name: Fetch the source
        shell: bash
        run: |
          set -euo pipefail
          want="${{ github.event.pull_request.head.sha || github.sha }}"
          rm -rf src && mkdir src && cd src
          git init -q
          # Fetch the commit itself, not the ref. The host names its own route to
          # the forge (WAI_FORGE_URL).
          git fetch -q --depth 1 "${WAI_FORGE_URL:-http://100.64.224.0:3400}/Transaction-Science/open-standards.git" "$want"
          git checkout -q FETCH_HEAD
          got="$(git rev-parse HEAD)"
          [ "$got" = "$want" ] || { echo "::error::fetched $got, not $want"; exit 1; }
      - name: The program through HIP and on every Vulkan device
        shell: bash
        working-directory: src/wai/wai-quantum
        run: |
          set -euo pipefail
          export PATH="$HOME/.cargo/bin:$PATH"
          iree="$HOME/iree-$IREE_VERSION"
          [ -x "$iree/bin/iree-compile" ] || { rm -rf "$iree"; python3 -m venv "$iree"; "$iree/bin/pip" install -q "iree-base-compiler==$IREE_VERSION" "iree-base-runtime==$IREE_VERSION"; }
          export PATH="$iree/bin:$PATH"
          check="bash $PWD/ci/tensor-program/check.sh"
          d="$RUNNER_TEMP/program"
          CARGO_TARGET_DIR="$HOME/hw-target/wai-quantum-tensor-program" nice -n 10 cargo run --locked --release -q --example stablehlo_gate --features quantum_fabric -- "$d"
          cd "$d"
          # The GPU's own architecture for both: the HIP build's and the Vulkan
          # capability profile's (which has 64-bit integers).
          arch="${WAI_QUANTUM_AMD_GPU_ARCH:-gfx1151}"
          for m in gate gate_m3; do
            iree-compile $m.mlir --iree-hal-target-backends=rocm --iree-rocm-target="$arch" -o $m.hip.vmfb
            iree-compile $m.mlir --iree-hal-target-device=vulkan --iree-vulkan-target="$arch" -o $m.vulkan.vmfb
          done
          $check iree hip gate.hip.vmfb '=='
          $check iree hip gate_m3.hip.vmfb '!='
          devices="$(iree-run-module --list_devices 2>/dev/null | grep -oE 'vulkan://[^ ]+' | sort -u)"
          count="$(grep -c . <<<"$devices")"
          [ "$count" = "$VULKAN_DEVICES" ] || { echo "::error::$count Vulkan devices, not $VULKAN_DEVICES"; iree-run-module --list_devices; exit 1; }
          for dev in $devices; do
            $check iree "$dev" gate.vulkan.vmfb '=='
            $check iree "$dev" gate_m3.vulkan.vmfb '!='
          done
          # The products through this compiler's HIP build: shown, not
          # required. Its bf16 multiply for this GPU returns values that are
          # not integers (measured 2026-10-15), so the coefficients differ; the
          # step says what it found, and a build that becomes exact shows here.
          iree-compile limbs.mlir --iree-hal-target-backends=rocm --iree-rocm-target="$arch" -o limbs.hip.vmfb
          $check iree hip limbs.hip.vmfb '=!' || echo "(the HIP build's limb products are not exact; not required)"

  metal:
    name: tensor program · Metal and the arm64 CPU (macos)
    needs: relevant
    if: needs.relevant.outputs.run == 'true'
    runs-on: macos
    timeout-minutes: 60
    steps:
      - uses: actions/checkout@v5
      # launchd's PATH omits cargo on the macOS host runner.
      - name: Put cargo on PATH
        shell: bash
        run: echo "$HOME/.cargo/bin" >> "$GITHUB_PATH"
      - name: The program on the platform's GPU and, through XLA, on its CPU
        shell: bash
        working-directory: wai/wai-quantum
        run: |
          set -euo pipefail
          # The system Python is older than the compilers' packages support.
          py=/opt/homebrew/bin/python3.12
          [ -x "$py" ] || py="$(command -v python3.12 || command -v python3)"
          iree="$HOME/iree-$IREE_VERSION"
          [ -x "$iree/bin/iree-compile" ] || { rm -rf "$iree"; "$py" -m venv "$iree"; "$iree/bin/pip" install -q "iree-base-compiler==$IREE_VERSION" "iree-base-runtime==$IREE_VERSION"; }
          xla="$HOME/xla-$JAX_VERSION"
          [ -x "$xla/bin/python" ] && "$xla/bin/python" -I -c 'import jax' 2>/dev/null || { rm -rf "$xla"; "$py" -m venv "$xla"; "$xla/bin/pip" install -q "jax==$JAX_VERSION"; }
          export PATH="$iree/bin:$PATH" XLA_PYTHON="$xla/bin/python"
          check="bash $PWD/ci/tensor-program/check.sh"
          d="$RUNNER_TEMP/program"
          nice -n 10 cargo run --locked --release -q --example stablehlo_gate --features quantum_fabric -- "$d"
          cd "$d"
          # Shader source in the module, compiled by the system at load: no
          # developer toolchain needed.
          for m in gate gate_m3; do
            iree-compile $m.mlir --iree-hal-target-device=metal --iree-metal-compile-to-metallib=false -o $m.metal.vmfb
          done
          $check iree metal gate.metal.vmfb '=='
          $check iree metal gate_m3.metal.vmfb '!='
          $check xla cpu gate.mlir '=='
          $check xla cpu gate_m3.mlir '!='
          # The products on the arm64 CPU through XLA. (The Metal build of
          # this compiler version does not lower bf16.)
          $check xla cpu limbs.mlir '=!'