name: Build Wheels & Release

on:
  workflow_dispatch:
    inputs:
      pypi:
        description: 'Upload wheels to PyPI? 1 = yes, 0 = no'
        default: '1'
        required: true
        type: string

permissions:
  contents: write

jobs:
  build_wheels:
    name: ${{ matrix.os }} Python ${{ matrix.pyver }} CUDA ${{ matrix.cuda }}
    runs-on: ${{ matrix.os }}
    strategy:
      matrix:
        os: [ubuntu-20.04, windows-latest]
        pyver: ["3.8", "3.9", "3.10", "3.11"]
        cuda: ["11.7.0", "11.8.0", "12.1.1"]
    defaults:
      run:
        shell: pwsh
    env:
      CUDAVER: ${{ matrix.cuda }}
      PYVER: ${{ matrix.pyver }}

    steps:
      - name: Free Disk Space
        uses: jlumbroso/free-disk-space@v1.2.0
        if: runner.os == 'Linux'
        with:
          tool-cache: false
          android: true
          dotnet: true
          haskell: true
          large-packages: false
          swap-storage: false
          
      - uses: actions/checkout@v3
      - uses: actions/setup-python@v3
        with:
          python-version: ${{ matrix.pyver }}
        
      - name: Setup Mamba
        uses: conda-incubator/setup-miniconda@v2.2.0
        with:
          activate-environment: "build"
          python-version: ${{ matrix.pyver }}
          miniforge-variant: Mambaforge
          miniforge-version: latest
          use-mamba: true
          add-pip-as-python-dependency: true
          auto-activate-base: false
          
      - name: Install Dependencies
        run: |
          $cudaVersion = $env:CUDAVER
          $cudaVersionPytorch = $env:CUDAVER.Remove($env:CUDAVER.LastIndexOf('.')).Replace('.','')
          
          $cudaChannels = ''
          $cudaNum = [int]$cudaVersion.substring($cudaVersion.LastIndexOf('.')+1)
          while ($cudaNum -ge 0) { $cudaChannels += '-c nvidia/label/cuda-' + $cudaVersion.Remove($cudaVersion.LastIndexOf('.')+1) + $cudaNum + ' '; $cudaNum-- }
          mamba install -y 'cuda' $cudaChannels.TrimEnd().Split()
          
          if ([version]$env:CUDAVER -gt [version]'11.8.0') {$torchver = "torch==2.1.0.* --pre --index-url https://download.pytorch.org/whl/nightly/cu$cudaVersionPytorch"} else {$torchver = "torch --index-url https://download.pytorch.org/whl/cu$cudaVersionPytorch"}
          python -m pip install $torchver.split()
          
          python -m pip install build wheel safetensors sentencepiece ninja
        
      - name: Build Wheel
        id: build-wheel
        run: |
          if ($(Get-Content 'setup.py' -raw) -match 'version = "(\d+\.(?:\d+\.?)*)" \+ \(') {Write-Output $('::notice file=build-wheels-release.yml,line=53,title=Package Version::Detected package version is: {0}' -f $Matches[1]); Write-Output "PACKAGE_VERSION=$($Matches[1])" >> "$env:GITHUB_OUTPUT"} else {Write-Output '::error file=build-wheels-release.yml,line=41::Could not parse version from setup.py! You must upload wheels manually!'; Write-Output "PACKAGE_VERSION=None" >> "$env:GITHUB_OUTPUT"}
          
          $env:CUDA_PATH = $env:CONDA_PREFIX
          $env:CUDA_HOME = $env:CONDA_PREFIX
          
          $cudaVersion = $env:CUDAVER
          $cudaVersionPytorch = $env:CUDAVER.Remove($env:CUDAVER.LastIndexOf('.')).Replace('.','')
          $BUILDTAG = "+cu$cudaVersionPytorch"
          
          if ($IsLinux) {$env:LD_LIBRARY_PATH = $env:CONDA_PREFIX + '/lib:' + $env:LD_LIBRARY_PATH}
          $env:TORCH_CUDA_ARCH_LIST = if ([version]$env:CUDAVER -lt [version]'11.8') {'6.0 6.1 7.0 7.5 8.0 8.6+PTX'} else {'6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX'}
          python -m build -n --wheel -C--build-option=egg_info "-C--build-option=--tag-build=$BUILDTAG"
          
          if ($IsLinux -and $env:PYVER -eq '3.11' -and $env:CUDAVER -eq '11.8.0') {$env:EXLLAMA_NOCOMPILE=1; python -m build -n}
        
      - uses: actions/upload-artifact@v3
        with:
          name: 'wheels'
          path: ./dist/*
          
      - name: Upload files to a GitHub release
        if: steps.build-wheel.outputs.PACKAGE_VERSION != 'None'
        uses: svenstaro/upload-release-action@2.6.1
        with:
          file: ./dist/*.whl
          tag: ${{ steps.build-wheel.outputs.PACKAGE_VERSION }}
          file_glob: true
          overwrite: true
          release_name: ${{ steps.build-wheel.outputs.PACKAGE_VERSION }}
      
  publish-to-pypi:
    name: Publish Python distribution to PyPI
    if: inputs.pypi == '1'
    needs: ['build_wheels','build_rocm']
    runs-on: ubuntu-latest
    
    environment:
      name: pypi
      url: https://pypi.org/p/exllamav2
    permissions:
      id-token: write  # IMPORTANT: mandatory for trusted publishing
      
    steps:
      - name: Download all the wheels
        uses: actions/download-artifact@v3
        with:
          name: wheels
          path: dist/
      - name: Publish distribution to PyPI
        uses: pypa/gh-action-pypi-publish@v1.8.10