GitHub workflow to build Triton wheels (#4431)

New workflow to build PyTorch wheels that will replace eventually the
existing "Nightly wheels" workflow. Currently only wheels for Windows
are supported to address user requests for pre-built Triton for Windows.

---------

Signed-off-by: Pavel Chekin <[email protected]>

Author: pbchekin

.github/workflows/wheels-pytorch.yml

@@ -6,7 +6,6 @@ on:
     branches:
       - main
     paths:
-      - .github/workflows/wheels-pytorch.yml
       - .github/pins/pytorch.txt
 
 permissions: read-all

.github/workflows/wheels-triton.yml

@@ -0,0 +1,98 @@
+name: Triton wheels
+
+on:
+  workflow_dispatch:
+
+permissions: read-all
+
+env:
+  NEW_WORKSPACE: C:\gh${{ github.run_id }}
+
+jobs:
+  windows:
+    runs-on:
+      - windows
+      - b580
+    strategy:
+      matrix:
+        python:
+          - "3.9"
+          - "3.10"
+          - "3.11"
+          - "3.12"
+          - "3.13"
+      fail-fast: false
+      max-parallel: 2
+    steps:
+      - name: Checkout Triton repository
+        uses: actions/checkout@v4
+
+      - name: Setup Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Clean up old workspaces
+        shell: bash
+        run: |
+          rm -rf /c/gh*
+
+      # Copy workspace to a temporary location with a shorter name.
+      - name: Copy workspace
+        run: |
+          Copy-Item -Path ${{ github.workspace }} -Destination ${{ env.NEW_WORKSPACE }} -Recurse
+
+      - name: Create venv
+        run: |
+          cd ${{ env.NEW_WORKSPACE }}
+          python -m venv .venv
+
+      - name: Install build dependencies
+        run: |
+          cd ${{ env.NEW_WORKSPACE }}
+          .venv\Scripts\activate.ps1
+          pip install cibuildwheel
+
+      - name: Set Triton version
+        shell: bash
+        run: |
+          cd /c/gh${{ github.run_id }}
+          version="$(./scripts/triton-version.sh)"
+          version="$version.post$GITHUB_RUN_NUMBER+git$(git rev-parse --short=7 HEAD)"
+          ./scripts/triton-version.sh "$version"
+          echo "Triton version: $(./scripts/triton-version.sh)"
+
+      - name: Build Triton wheels
+        run: |
+          cd ${{ env.NEW_WORKSPACE }}
+          .venv\Scripts\activate.ps1
+          Invoke-BatchFile "C:\Program Files\Microsoft Visual Studio\2022\Community\VC\Auxiliary\Build\vcvarsall.bat" x64
+
+          $env:CIBW_BUILD_VERBOSITY = "1"
+          $env:CIBW_SKIP = "cp{35,36,37}-*"
+          $env:CIBW_BUILD = "cp" + "${{ matrix.python}}".replace(".", "") + "*-win_amd64"
+          $env:CIBW_CACHE_PATH = "${{ env.NEW_WORKSPACE }}/cibw"
+
+          python -m cibuildwheel --output-dir wheelhouse
+
+      - name: Artifact name
+        run: |
+          (Get-ChildItem -Path "${{ env.NEW_WORKSPACE }}\wheelhouse\*.whl").Name | Tee-Object -Variable WHEELS_NAME
+          echo "WHEELS_NAME=$WHEELS_NAME" | Out-File -FilePath $env:GITHUB_ENV -Encoding utf8 -Append
+
+      - name: Upload wheels
+        uses: actions/upload-artifact@v4
+        with:
+          name: ${{ env.WHEELS_NAME }}
+          path: ${{ env.NEW_WORKSPACE }}\wheelhouse\*.whl
+
+      - name: Clean up workspace
+        if: ${{ always() }}
+        run: |
+          Remove-Item -LiteralPath ${{ env.NEW_WORKSPACE }} -Force -Recurse -ErrorAction Ignore
+
+      - name: Clean up temporary files
+        if: ${{ always() }}
+        shell: bash
+        run: |
+          rm -rf rm -rf /tmp/triton-* /tmp/tmp*

scripts/triton-version.sh

@@ -0,0 +1,30 @@
+#!/usr/bin/env bash
+
+# Script to get and set Triton version.
+# Currently Triton does not have a single file to set a version. There are two files:
+# 1. In `setup.py`, the version is set to `{version}+git{sha}` and this version is used for the
+#    Python package/wheel.
+# 2. In `python/triton/__init__.py`, the version is set to `{version}` and used in runtime for
+#    `triton.__version__`.
+# When building a wheel, the both files need to be updated to use the same version.
+
+set -euo pipefail
+
+PROJECT_ROOT="$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && cd .. && pwd )"
+
+get_version() {
+    # Use setup.py, as the most reliable source
+    sed -n -E "s/^TRITON_VERSION = \"([^\"]+)\".*/\1/p" "$PROJECT_ROOT/setup.py"
+}
+
+set_version() {
+    local version="$1"
+    sed -i -E "s/^TRITON_VERSION = \"([^\"]+)\".*/TRITON_VERSION = \"$version\"/" "$PROJECT_ROOT/setup.py"
+    sed -i -E "s/^__version__ = '[^']+'/__version__ = '$version'/" "$PROJECT_ROOT/python/triton/__init__.py"
+}
+
+if (( $# == 1 )); then
+    set_version "$1"
+else
+    get_version
+fi