name: Build Wheels & Release

on:
  workflow_dispatch:
    inputs:
      release:
        description: 'Release? 1 = yes, 0 = no'
        default: '0'
        required: true
        type: string

permissions:
  contents: write

jobs:
  build_wheels:
    name: ${{ matrix.os }} Python=${{ matrix.pyver }} CUDA=${{ matrix.cuda }} ROCm=${{ matrix.rocm }} Torch=${{ matrix.torch }}
    runs-on: ${{ matrix.os }}
    strategy:
      matrix:
        include:

        # Ubuntu 20.04 CUDA
         - { artname: 'wheel', os: ubuntu-20.04, pyver:  '3.8', cuda: '11.7.0', rocm:    '', torch: '2.0.1', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6+PTX'         }
         - { artname: 'wheel', os: ubuntu-20.04, pyver:  '3.8', cuda: '11.8.0', rocm:    '', torch: '2.3.1', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX' }
         - { artname: 'wheel', os: ubuntu-20.04, pyver:  '3.8', cuda: '12.1.0', rocm:    '', torch: '2.3.1', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX' }
         - { artname: 'wheel', os: ubuntu-20.04, pyver:  '3.9', cuda: '11.7.0', rocm:    '', torch: '2.0.1', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6+PTX'         }
         - { artname: 'wheel', os: ubuntu-20.04, pyver:  '3.9', cuda: '11.8.0', rocm:    '', torch: '2.3.1', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX' }
         - { artname: 'wheel', os: ubuntu-20.04, pyver:  '3.9', cuda: '12.1.0', rocm:    '', torch: '2.3.1', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX' }
         - { artname: 'wheel', os: ubuntu-20.04, pyver: '3.10', cuda: '11.7.0', rocm:    '', torch: '2.0.1', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6+PTX'         }
         - { artname: 'wheel', os: ubuntu-20.04, pyver: '3.10', cuda: '11.8.0', rocm:    '', torch: '2.3.1', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX' }
         - { artname: 'wheel', os: ubuntu-20.04, pyver: '3.10', cuda: '12.1.0', rocm:    '', torch: '2.3.1', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX' }
         - { artname: 'wheel', os: ubuntu-20.04, pyver: '3.11', cuda: '11.7.0', rocm:    '', torch: '2.0.1', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6+PTX'         }
         - { artname: 'wheel', os: ubuntu-20.04, pyver: '3.11', cuda: '11.8.0', rocm:    '', torch: '2.3.1', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX' }
         - { artname: 'wheel', os: ubuntu-20.04, pyver: '3.11', cuda: '12.1.0', rocm:    '', torch: '2.3.1', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX' }
         - { artname: 'wheel', os: ubuntu-20.04, pyver: '3.12', cuda: '11.8.0', rocm:    '', torch: '2.3.1', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX' }
         - { artname: 'wheel', os: ubuntu-20.04, pyver: '3.12', cuda: '12.1.0', rocm:    '', torch: '2.3.1', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX' }
        
        # Windows 2022 CUDA
         - { artname: 'wheel', os: windows-2022, pyver:  '3.8', cuda: '11.7.0', rocm:    '', torch: '2.0.1', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6+PTX'         }
         - { artname: 'wheel', os: windows-2022, pyver:  '3.8', cuda: '11.8.0', rocm:    '', torch: '2.3.1', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX' }
         - { artname: 'wheel', os: windows-2022, pyver:  '3.8', cuda: '12.1.0', rocm:    '', torch: '2.3.1', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX' }
         - { artname: 'wheel', os: windows-2022, pyver:  '3.9', cuda: '11.7.0', rocm:    '', torch: '2.0.1', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6+PTX'         }
         - { artname: 'wheel', os: windows-2022, pyver:  '3.9', cuda: '11.8.0', rocm:    '', torch: '2.3.1', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX' }
         - { artname: 'wheel', os: windows-2022, pyver:  '3.9', cuda: '12.1.0', rocm:    '', torch: '2.3.1', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX' }
         - { artname: 'wheel', os: windows-2022, pyver: '3.10', cuda: '11.7.0', rocm:    '', torch: '2.0.1', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6+PTX'         }
         - { artname: 'wheel', os: windows-2022, pyver: '3.10', cuda: '11.8.0', rocm:    '', torch: '2.3.1', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX' }
         - { artname: 'wheel', os: windows-2022, pyver: '3.10', cuda: '12.1.0', rocm:    '', torch: '2.3.1', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX' }
         - { artname: 'wheel', os: windows-2022, pyver: '3.11', cuda: '11.7.0', rocm:    '', torch: '2.0.1', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6+PTX'         }
         - { artname: 'wheel', os: windows-2022, pyver: '3.11', cuda: '11.8.0', rocm:    '', torch: '2.3.1', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX' }
         - { artname: 'wheel', os: windows-2022, pyver: '3.11', cuda: '12.1.0', rocm:    '', torch: '2.3.1', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX' }
         - { artname: 'wheel', os: windows-2022, pyver: '3.12', cuda: '11.8.0', rocm:    '', torch: '2.3.1', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX' }
         - { artname: 'wheel', os: windows-2022, pyver: '3.12', cuda: '12.1.0', rocm:    '', torch: '2.3.1', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX' }
        
        # Ubuntu 20.04 ROCm
         - { artname: 'wheel', os: ubuntu-20.04, pyver: '3.10', cuda: '',       rocm: '5.6', torch: '2.2.2', cudaarch: ''                                    }
         - { artname: 'wheel', os: ubuntu-20.04, pyver: '3.10', cuda: '',       rocm: '6.0', torch: '2.3.1', cudaarch: ''                                    }
         - { artname: 'wheel', os: ubuntu-20.04, pyver: '3.11', cuda: '',       rocm: '5.6', torch: '2.2.2', cudaarch: ''                                    }
         - { artname: 'wheel', os: ubuntu-20.04, pyver: '3.11', cuda: '',       rocm: '6.0', torch: '2.3.1', cudaarch: ''                                    }
         - { artname: 'wheel', os: ubuntu-20.04, pyver: '3.12', cuda: '',       rocm: '6.0', torch: '2.3.1', cudaarch: ''                                    }

         # sdist
         - { artname: 'sdist', os: ubuntu-20.04, pyver: '3.11', cuda: '',       rocm:    '', torch: '2.3.1', cudaarch: ''                                    }

         # Extra Torch 2.2 wheels for Windows 2022 until PyTorch resolves the shm.dll issue
         - { artname: 'wheel', os: windows-2022, pyver:  '3.8', cuda: '11.8.0', rocm:    '', torch: '2.2.0', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX' }
         - { artname: 'wheel', os: windows-2022, pyver:  '3.8', cuda: '12.1.0', rocm:    '', torch: '2.2.0', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX' }
         - { artname: 'wheel', os: windows-2022, pyver:  '3.9', cuda: '11.8.0', rocm:    '', torch: '2.2.0', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX' }
         - { artname: 'wheel', os: windows-2022, pyver:  '3.9', cuda: '12.1.0', rocm:    '', torch: '2.2.0', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX' }
         - { artname: 'wheel', os: windows-2022, pyver: '3.10', cuda: '11.8.0', rocm:    '', torch: '2.2.0', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX' }
         - { artname: 'wheel', os: windows-2022, pyver: '3.10', cuda: '12.1.0', rocm:    '', torch: '2.2.0', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX' }
         - { artname: 'wheel', os: windows-2022, pyver: '3.11', cuda: '11.8.0', rocm:    '', torch: '2.2.0', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX' }
         - { artname: 'wheel', os: windows-2022, pyver: '3.11', cuda: '12.1.0', rocm:    '', torch: '2.2.0', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX' }
         - { artname: 'wheel', os: windows-2022, pyver: '3.12', cuda: '11.8.0', rocm:    '', torch: '2.2.0', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX' }
         - { artname: 'wheel', os: windows-2022, pyver: '3.12', cuda: '12.1.0', rocm:    '', torch: '2.2.0', cudaarch: '6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX' }

    defaults:
      run:
        shell: pwsh

    steps:
      # Free disk space

      - name: Free Disk Space
        uses: jlumbroso/free-disk-space@v1.3.1
        if: runner.os == 'Linux'
        with:
          tool-cache: true
          android: true
          dotnet: true
          haskell: true
          large-packages: false
          swap-storage: true

      # Setup Python

      - uses: actions/checkout@v3

      - uses: actions/setup-python@v3
        with:
          python-version: ${{ matrix.pyver }}

      # Get version string from package

      - name: Get version string
        id: package_version
        run: | 
          $versionString = Get-Content $(Join-Path 'exllamav2' 'version.py') -raw
          if ($versionString -match '__version__ = "(\d+\.(?:\d+\.?(?:dev\d+)?)*)"') 
          {
            Write-Output $('::notice file=build-wheels-release.yml,line=200,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=203::Could not parse version from exllamav2/version.py! You must upload wheels manually!'
            Write-Output "PACKAGE_VERSION=None" >> "$env:GITHUB_OUTPUT"
          }

      # Pin VS build tools to 17.9 so builds won't fail

      - name: Install VS2022 BuildTools 17.9.7
        run: choco install -y visualstudio2022buildtools --version=117.9.7.0 --params "--add Microsoft.VisualStudio.Component.VC.Tools.x86.x64 --installChannelUri https://aka.ms/vs/17/release/180911598_-255012421/channel"
        if: runner.os == 'Windows'

      # Install ROCm SDK, apparently needs to happen before setting up Python

      - name: Build for ROCm
        if: matrix.rocm != ''
        shell: bash
        run: |
          # --- Install ROCm SDK

          export ROCM_VERSION=${{ matrix.rocm }}
          export TORCH_VERSION=${{ matrix.torch }}

          [ ! -d /etc/apt/keyrings ] && sudo mkdir --parents --mode=0755 /etc/apt/keyrings
          wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null
          echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/$ROCM_VERSION focal main" | sudo tee --append /etc/apt/sources.list.d/rocm.list
          echo -e 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' | sudo tee /etc/apt/preferences.d/rocm-pin-600
          
          sudo apt update
          sudo apt install rocm-hip-sdk -y
          sudo apt clean -y

          echo "/opt/rocm/bin" >> $GITHUB_PATH
          echo "ROCM_PATH=/opt/rocm" >> $GITHUB_ENV
          echo "ROCM_VERSION=$ROCM_VERSION" >> $GITHUB_ENV
          echo "USE_ROCM=1" >> $GITHUB_ENV

          # --- Install dependencies

          python3 -m ensurepip --upgrade
          pip3 install torch==${{ matrix.torch }} --index-url="https://download.pytorch.org/whl/rocm$ROCM_VERSION"
          pip3 install --upgrade setuptools==69.5.1 build wheel safetensors sentencepiece ninja
          pip3 cache purge

          # --- Build wheel 

          python3 -m build -n --wheel -C--build-option=egg_info "-C--build-option=--tag-build=+rocm${{ matrix.rocm }}-torch${{ matrix.torch }}"

      # Build for CUDA

      - name: Setup Mamba
        if: matrix.cuda != ''
        uses: conda-incubator/setup-miniconda@v2.3.0
        with:
          activate-environment: "exllama"
          python-version: ${{ matrix.pyver }}
          miniforge-variant: Mambaforge
          miniforge-version: latest
          use-mamba: true
          add-pip-as-python-dependency: true
          auto-activate-base: false

      - name: Build for CUDA
        if: matrix.cuda != ''
        run: |
          # --- Spawn the VS shell
          if ($IsWindows) {
            Import-Module 'C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\Common7\Tools\Microsoft.VisualStudio.DevShell.dll'
            Enter-VsDevShell -VsInstallPath 'C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools' -DevCmdArguments '-arch=x64 -host_arch=x64'
            $env:DISTUTILS_USE_SDK=1
          }
  
          # --- Install CUDA using Conda
          $cudaVersion = '${{ matrix.cuda }}'
          $cudaVersionPytorch = '${{ matrix.cuda }}'.Remove('${{ matrix.cuda }}'.LastIndexOf('.')).Replace('.','')
          
          mamba install -y -c nvidia/label/cuda-$cudaVersion cuda-toolkit cuda-runtime

          if (!(mamba list cuda)[-1].contains('cuda')) {sleep -s 10; mamba install -y 'cuda' $cudaVersion}
          if (!(mamba list cuda)[-1].contains('cuda')) {throw 'CUDA Toolkit failed to install!'}

          $env:CUDA_PATH = $env:CONDA_PREFIX
          $env:CUDA_HOME = $env:CONDA_PREFIX
          if ($IsLinux) {$env:LD_LIBRARY_PATH = $env:CONDA_PREFIX + '/lib:' + $env:LD_LIBRARY_PATH}
          
          # --- Install dependencies
          
          python -m ensurepip --upgrade
          python -m pip install torch==${{ matrix.torch }} --index-url https://download.pytorch.org/whl/cu$cudaVersionPytorch
          python -m pip install --upgrade setuptools==69.5.1 build wheel safetensors sentencepiece ninja

          # --- Build wheel
                  
          $BUILDTAG = "+cu$cudaVersionPytorch-torch${{ matrix.torch }}"
          $env:TORCH_CUDA_ARCH_LIST = '${{ matrix.cudaarch }}'
          python -m build -n --wheel -C--build-option=egg_info "-C--build-option=--tag-build=$BUILDTAG"

      # Build sdist

      - name: Build sdist
        if: matrix.cuda == '' && matrix.rocm == ''
        run: |
          # --- Spawn the VS shell
          if ($IsWindows) {
            Import-Module 'C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\Common7\Tools\Microsoft.VisualStudio.DevShell.dll'
            Enter-VsDevShell -VsInstallPath 'C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools' -DevCmdArguments '-arch=x64 -host_arch=x64'
            $env:DISTUTILS_USE_SDK=1
          }

          # --- Install dependencies
          
          python -m pip install torch==${{ matrix.torch }} --index-url https://download.pytorch.org/whl/cpu
          python -m pip install build wheel ninja

          # --- Build wheel
           
          $env:EXLLAMA_NOCOMPILE=1
          python -m build -n

      # Upload files

      - uses: actions/upload-artifact@v3
        if: matrix.artname == 'wheel'
        with:
          name: 'wheel'
          path: ./dist/*

      - uses: actions/upload-artifact@v3
        if: matrix.artname == 'sdist'
        with:
          name: 'sdist'
          path: ./dist/*

      - name: Upload files to GitHub release
        if: steps.package_version.outputs.PACKAGE_VERSION != 'None' && inputs.release == '1'
        uses: svenstaro/upload-release-action@2.6.1
        with:
          file: ./dist/*.whl
          tag: ${{ format('v{0}', steps.package_version.outputs.PACKAGE_VERSION) }}
          file_glob: true
          overwrite: true
          release_name: ${{ steps.package_version.outputs.PACKAGE_VERSION }}
