From 7108f9331444ded955c994a600e9f88955ba9092 Mon Sep 17 00:00:00 2001
From: Casper <casperbh.96@gmail.com>
Date: Thu, 15 Feb 2024 20:12:36 +0100
Subject: [PATCH] Add BuildExtension from torch

---
 .github/workflows/build.yaml | 96 +++++++++++++++++++++++-------------
 setup.py                     |  2 +
 2 files changed, 63 insertions(+), 35 deletions(-)

diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml
index eae16a25..e5d1e769 100644
--- a/.github/workflows/build.yaml
+++ b/.github/workflows/build.yaml
@@ -2,40 +2,42 @@ name: Build AutoAWQ Wheels with CUDA
 
 on:
   push:
-    tags:
-      - "v*"
+    branches:
+      - "fix_workflow"
+    # 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
+  # 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: 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)
+  #     - 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
+    # needs: release
     
     strategy:
       matrix:
@@ -114,16 +116,28 @@ jobs:
 
           python setup.py sdist bdist_wheel
 
-      - name: Upload Assets
-        uses: shogo82148/actions-upload-release-asset@v1
+      - uses: actions/upload-artifact@v3
+        if: runner.os == 'Linux'
         with:
-          upload_url: ${{ needs.release.outputs.upload_url }}
-          asset_path: ./dist/*.whl
+          name: 'linux-cuda-wheels'
+          path: ./dist/*.whl
+
+      - uses: actions/upload-artifact@v3
+        if: runner.os == 'Windows'
+        with:
+          name: 'windows-cuda-wheels'
+          path: ./dist/*.whl
+  
+      # - 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
+    # needs: release
 
     strategy:
       matrix:
@@ -225,8 +239,20 @@ jobs:
 
           ROCM_VERSION=${{ matrix.rocm }} python setup.py sdist bdist_wheel
 
-      - name: Upload Assets
-        uses: shogo82148/actions-upload-release-asset@v1
+      - uses: actions/upload-artifact@v3
+        if: runner.os == 'Linux'
         with:
-          upload_url: ${{ needs.release.outputs.upload_url }}
-          asset_path: ./dist/*.whl
\ No newline at end of file
+          name: 'linux-cuda-wheels'
+          path: ./dist/*.whl
+
+      - uses: actions/upload-artifact@v3
+        if: runner.os == 'Windows'
+        with:
+          name: 'windows-cuda-wheels'
+          path: ./dist/*.whl
+
+      # - name: Upload Assets
+      #   uses: shogo82148/actions-upload-release-asset@v1
+      #   with:
+      #     upload_url: ${{ needs.release.outputs.upload_url }}
+      #     asset_path: ./dist/*.whl
\ No newline at end of file
diff --git a/setup.py b/setup.py
index b0bcfd33..68969e3c 100644
--- a/setup.py
+++ b/setup.py
@@ -4,6 +4,7 @@
 import requests
 from pathlib import Path
 from setuptools import setup, find_packages
+from torch.utils.cpp_extension import BuildExtension
 
 
 def get_latest_kernels_version(repo):
@@ -136,5 +137,6 @@ def get_kernels_whl_url(
         "eval": ["lm_eval>=0.4.0", "tabulate", "protobuf", "evaluate", "scipy"],
         "dev": ["black", "mkdocstrings-python", "mkdocs-material", "griffe-typingdoc"]
     },
+    cmdclass={"build_ext": BuildExtension},
     **common_setup_kwargs,
 )