kernel-bench.yml

name: Triton Kernel Benchmarking
on:
  pull_request:
    types: [opened, synchronize]
  workflow_dispatch:
  schedule:
  # Runs at 12:00 PM UTC, which is 5:00 AM PST
   - cron: '0 12 * * *'

concurrency:
  # A PR number if a pull request and otherwise the commit hash. This cancels
  # queued and in-progress runs for the same PR (presubmit) or commit
  # (postsubmit). The workflow name is prepended to avoid conflicts between
  # different workflows.
  group: ${{ github.workflow }}-${{ github.event.number || github.sha }}
  cancel-in-progress: true

jobs:
  kernel-bench:
    runs-on: amdgpu-mi250-x86-64
    env:
      KERNEL_VENV_DIR: ${{ github.workspace }}/triton_bench
    steps:
      - name: Checkout Repo
        uses: actions/checkout@v4

      - name: Setup Python
        uses: actions/setup-python@v5
        with:
          python-version: "3.11"

      - name : Set Test File
        run: |
          python3.11 ci-tools/parse.py --test "${{ github.event.pull_request.body }}"
          echo "TEST_FILE=$TEST_FILE" >> $GITHUB_ENV

      - name: Setup Triton Python Venv
        run: |
          python3.11 -m venv ${KERNEL_VENV_DIR}
          source ${KERNEL_VENV_DIR}/bin/activate
          pip install --upgrade pip
          pip install --pre pytorch-triton-rocm==3.1.0+cf34004b8a torch==2.6.0.dev20241023+rocm6.2 --index-url https://download.pytorch.org/whl/nightly/rocm6.2
          pip install pytest-xdist
          pip install -r requirements.txt

      - name: Run Kernel Benchmarking
        run: |
          source ${KERNEL_VENV_DIR}/bin/activate
          pip freeze
          pytest kernels/${TEST_FILE} -s -n 4
      
      - uses: actions/upload-artifact@master
        with:
          name: perf-artifacts
          path: ./perf-artifacts/**/*.csv