rocm_wheels.yml

name: rocm-wheels

on:
  push:
    branches:
      - develop
  pull_request:
    paths:
      - "packaging/compute_wheel_version.sh"
      - ".github/workflows/wheel*"
      - ".github/actions/setup-windows-runner/action.yml"
      - "setup.py"
      - "requirements*.txt"
  workflow_dispatch:

jobs:
  target_determinator:
    runs-on: ubuntu-latest
    outputs:
      matrix: ${{ steps.set-matrix.outputs.matrix }}
    steps:
    - id: set-matrix
      shell: python
      run: |
        import os
        import json
        environ = os.environ

        PY_VERSIONS = ['3.11']
        TORCH_VERSIONS = ['2.4.0']

        include = []
        for os in ['ubuntu-alola']:
          for python in PY_VERSIONS:
            for torch_version in TORCH_VERSIONS:
              for toolkit_type, toolkit_short_versions in {'rocm': ["6.0", "6.1"]}.items():
                for toolkit_short_version in toolkit_short_versions:
                  include.append(dict(
                    os=os,
                    python=python,
                    torch_version=torch_version,
                    toolkit_type=toolkit_type,
                    toolkit_short_version=toolkit_short_version,
                  ))
                  print(include[-1])
        matrix = {'include': include}
        print(json.dumps(matrix))
        with open(environ["GITHUB_OUTPUT"], "a") as fd:
          fd.write("matrix="+json.dumps(matrix))

  build:
    needs: target_determinator
    strategy:
      fail-fast: false
      matrix: ${{ fromJson(needs.target_determinator.outputs.matrix) }}

    uses: ./.github/workflows/wheels_build.yml
    if: github.repository == 'rocm/xformers' || github.event_name == 'pull_request'
    with:
      os: ${{ matrix.os }}
      python: ${{ matrix.python }}
      torch_version: ${{ matrix.torch_version }}
      toolkit_type: ${{ matrix.toolkit_type }}
      toolkit_short_version: ${{ matrix.toolkit_short_version }}
      artifact_tag: ${{ github.run_id }}

  clean:
    runs-on: self-hosted
    if: ${{ always() }}
    needs: [build]
    steps:
      - name: Remove dangling Docker images
        run: |
          docker images -q -f dangling=true | xargs --no-run-if-empty docker rmi