name: Build AutoAWQ Wheels with CUDA

on:
  push:
    tags:
      - "v*"

jobs:
  release:
    # Retrieve tag and create release
    name: Create Release
    runs-on: ubuntu-latest
    outputs:
      upload_url: ${{ steps.create_release.outputs.upload_url }}
    steps:
      - name: Checkout
        uses: actions/checkout@v3

      - name: Extract branch info
        shell: bash
        run: |
          echo "release_tag=${GITHUB_REF#refs/*/}" >> $GITHUB_ENV

      - name: Create Release
        id: create_release
        uses: "actions/github-script@v6"
        env:
          RELEASE_TAG: ${{ env.release_tag }}
        with:
          github-token: "${{ secrets.GITHUB_TOKEN }}"
          script: |
            const script = require('.github/workflows/scripts/github_create_release.js')
            await script(github, context, core)

  build_cuda_wheels:
    name: Build AWQ with CUDA
    runs-on: ${{ matrix.os }}
    needs: release

    strategy:
      matrix:
        os: [ubuntu-20.04, windows-latest]
        pyver: ["3.8", "3.9", "3.10", "3.11", "3.12"]
        cuda: ["11.8.0", "12.1.1"]
    defaults:
      run:
        shell: pwsh
    env:
      PYPI_CUDA_VERSION: "12.1.1"
      CUDA_VERSION: ${{ matrix.cuda }}

    steps:
      - name: Free Disk Space
        uses: jlumbroso/free-disk-space@v1.3.0
        if: runner.os == 'Linux'
        with:
          tool-cache: false
          android: true
          dotnet: true
          haskell: true
          large-packages: false
          docker-images: true
          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: |
          # Install CUDA toolkit
          mamba install -y 'cuda' -c "nvidia/label/cuda-${env:CUDA_VERSION}"

          # Env variables
          $env:CUDA_PATH = $env:CONDA_PREFIX
          $env:CUDA_HOME = $env:CONDA_PREFIX

          # Install torch
          $cudaVersion = $env:CUDA_VERSION.Replace('.', '')
          $cudaVersionPytorch = $cudaVersion.Substring(0, $cudaVersion.Length - 1)
          $pytorchVersion = "torch==2.3.1"
          python -m pip install --upgrade --no-cache-dir $pytorchVersion+cu$cudaVersionPytorch --index-url https://download.pytorch.org/whl/cu$cudaVersionPytorch
          python -m pip install build setuptools wheel ninja

          # Print version information
          python --version
          python -c "import torch; print('PyTorch:', torch.__version__)"
          python -c "import torch; print('CUDA:', torch.version.cuda)"
          python -c "import os; print('CUDA_HOME:', os.getenv('CUDA_HOME', None))"
          python -c "from torch.utils import cpp_extension; print (cpp_extension.CUDA_HOME)"

      - name: Build Wheel
        run: |
          $env:CUDA_PATH = $env:CONDA_PREFIX
          $env:CUDA_HOME = $env:CONDA_PREFIX

          # Only add +cu118 to wheel if not releasing on PyPi
          if ( $env:CUDA_VERSION -eq $env:PYPI_CUDA_VERSION ){
            $env:PYPI_BUILD = 1
          }

          python setup.py sdist bdist_wheel

      - name: Upload Assets
        uses: shogo82148/actions-upload-release-asset@v1
        with:
          upload_url: ${{ needs.release.outputs.upload_url }}
          asset_path: ./dist/*.whl

  build_rocm_wheels:
    name: Build AWQ with ROCm
    runs-on: ${{ matrix.os }}
    needs: release

    strategy:
      matrix:
        os: [ubuntu-20.04]
        python: ["3.8", "3.9", "3.10", "3.11"]
        rocm: ["5.6.1", "5.7.1"] # we build only for rocm5.6 & 5.7 to match PyTorch 2.1.0 and PyTorch 2.2 nightly
    defaults:
      run:
        shell: bash
    env:
      ROCM_VERSION: ${{ matrix.rocm }}

    steps:
      - uses: actions/checkout@v3

      - name: Free Disk Space
        run: |
          df -h
          echo "Removing large packages"
          sudo apt-get remove -y '^dotnet-.*'
          sudo apt-get remove -y 'php.*'
          sudo apt-get remove -y azure-cli google-chrome-stable firefox powershell mono-devel
          df -h
          sudo apt-get autoremove -y >/dev/null 2>&1
          sudo apt-get clean
          sudo apt-get autoremove -y >/dev/null 2>&1
          sudo apt-get autoclean -y >/dev/null 2>&1
          df -h
          echo "https://github.com/actions/virtual-environments/issues/709"
          sudo rm -rf "$AGENT_TOOLSDIRECTORY"
          df -h
          echo "remove big /usr/local"
          sudo rm -rf "/usr/local/share/boost"
          sudo rm -rf /usr/local/lib/android >/dev/null 2>&1
          df -h
          sudo rm -rf /usr/share/dotnet/sdk > /dev/null 2>&1
          sudo rm -rf /usr/share/dotnet/shared > /dev/null 2>&1
          sudo rm -rf /usr/share/swift > /dev/null 2>&1
          df -h

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

      - name: Setup Mamba
        uses: conda-incubator/setup-miniconda@v2.2.0
        with:
          activate-environment: "build"
          python-version: ${{ matrix.python }}
          mamba-version: "*"
          use-mamba: false
          channels: conda-forge,defaults
          channel-priority: true
          add-pip-as-python-dependency: true
          auto-activate-base: false

      - name: Set up ROCm
        run: |
          echo "Using python:"
          python --version
          which python

          if [[ "${{ matrix.rocm }}" == "5.4.2" ]]; then
            export ROCM_DL_FILE=amdgpu-install_5.4.50402-1_all.deb
          elif [[ "${{ matrix.rocm }}" == "5.6.1" ]]; then
            export ROCM_DL_FILE=amdgpu-install_5.6.50601-1_all.deb
          elif [[ "${{ matrix.rocm }}" == "5.7.1" ]]; then
            export ROCM_DL_FILE=amdgpu-install_5.7.50701-1_all.deb
          else
            echo Unknown rocm version
            exit 1
          fi

          curl -O https://repo.radeon.com/amdgpu-install/${{ matrix.rocm }}/ubuntu/focal/$ROCM_DL_FILE
          sudo dpkg -i $ROCM_DL_FILE
          sudo DEBIAN_FRONTEND=noninteractive amdgpu-install --usecase=rocm --no-dkms --no-32 -y

      - name: Install Dependencies
        run: |
          sudo apt-get update
          sudo apt-get install -y --no-install-recommends rocsparse-dev rocthrust-dev rocblas-dev hipblas-dev hipsparse-dev

          python -m pip install --upgrade build setuptools wheel

          if [[ "${{ matrix.rocm }}" == "5.7.1" ]]; then
            python -m pip install torch==2.3.1 --index-url https://download.pytorch.org/whl/rocm5.7
          elif [[ "${{ matrix.rocm }}" == "5.6.1" ]]; then
            python -m pip install torch==2.3.1 --index-url https://download.pytorch.org/whl/rocm5.6
          else
            echo Unknown rocm version for python install
            exit 1
          fi

      - name: Build Wheel
        run: |
          echo "Using python for build:"
          python --version
          which python

          ROCM_VERSION=${{ matrix.rocm }} python setup.py sdist bdist_wheel

      - name: Upload Assets
        uses: shogo82148/actions-upload-release-asset@v1
        with:
          upload_url: ${{ needs.release.outputs.upload_url }}
          asset_path: ./dist/*.whl