[CI] Auto Format Checking and test checking.

.github/workflows/ci.yml

@@ -0,0 +1,67 @@
+name: CI
+
+on: [push, pull_request]
+
+jobs:
+  format-check:
+    runs-on: self-hosted
+
+    steps:
+    - name: Checkout repository
+      uses: actions/checkout@v2
+      with:
+        fetch-depth: 0
+
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: '3.9'
+
+    - name: Create virtual environment
+      run: python -m venv bitblas_ci
+
+    - name: Activate virtual environment and install dependencies
+      run: |
+        source bitblas_ci/bin/activate
+        python -m pip install --upgrade pip
+        if [ -f requirements-dev.txt ]; then python -m pip install -r requirements-dev.txt; fi
+
+    - name: Run format check
+      run: |
+        source bitblas_ci/bin/activate
+        ./format.sh
+
+  build-test:
+    runs-on: self-hosted
+    needs: format-check
+
+    steps:
+    - name: Checkout repository
+      uses: actions/checkout@v2
+      with:
+        fetch-depth: 0
+
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: '3.9'
+
+    - name: Create virtual environment
+      run: python -m venv bitblas_ci
+
+    - name: Activate virtual environment and install dependencies
+      run: |
+        source bitblas_ci/bin/activate
+        python -m pip install --upgrade pip
+        if [ -f requirements-dev.txt ]; then python -m pip install -r requirements-dev.txt; fi
+
+    - name: Install project in wheel mode
+      run: |
+        source bitblas_ci/bin/activate
+        python -m pip install .
+
+    - name: Run tests
+      run: |
+        source bitblas_ci/bin/activate
+        cd testing/python
+        python -m pytest

bitblas/ops/impl/ladder_permutate_impl.py

@@ -49,6 +49,7 @@ def select_implementation(
     inp = te.placeholder((M, N // scaling_factor), name="inp", dtype=storage_dtype)
     args = [inp]
 
+    assert transform_kind != 0, "Permute only apply when transform_kind >= 1"
     if transform_kind >= 1:
         arg = args[-1]
 

bitblas/ops/operator.py

@@ -106,8 +106,7 @@ def tvm_callback_cuda_postproc(code, _):
                         **self.pass_context
                 }):
                     rt_mod = tvm.build(self.optimized_func, target=target, name=self.name)
-            except Exception as e:
-                rt_build_error = e  # noqa
+            except Exception: # noqa: F841
                 logger.debug(
                     "Failed to build optimized function for CUDA target with default schedule, Please consider enable hardware aware tuning!"
                 )

format.sh

@@ -64,7 +64,13 @@ format_changed() {
     #
     # `diff-filter=ACM` and $MERGEBASE is to ensure we only format files that
     # exist on both branches.
-    MERGEBASE="$(git merge-base origin/main HEAD)"
+    if git show-ref --verify --quiet refs/remotes/origin/main; then
+        BASE_BRANCH="origin/main"
+    else
+        BASE_BRANCH="main"
+    fi
+
+    MERGEBASE="$(git merge-base $BASE_BRANCH HEAD)"
 
     if ! git diff --diff-filter=ACM --quiet --exit-code "$MERGEBASE" -- '*.py' '*.pyi' &>/dev/null; then
         git diff --name-only --diff-filter=ACM "$MERGEBASE" -- '*.py' '*.pyi' | xargs -P 5 \
@@ -110,7 +116,13 @@ spell_check_changed() {
     #
     # `diff-filter=ACM` and $MERGEBASE is to ensure we only lint files that
     # exist on both branches.
-    MERGEBASE="$(git merge-base origin/main HEAD)"
+    if git show-ref --verify --quiet refs/remotes/origin/main; then
+        BASE_BRANCH="origin/main"
+    else
+        BASE_BRANCH="main"
+    fi
+
+    MERGEBASE="$(git merge-base $BASE_BRANCH HEAD)"
 
     if ! git diff --diff-filter=ACM --quiet --exit-code "$MERGEBASE" -- '*.py' '*.pyi' &>/dev/null; then
         git diff --name-only --diff-filter=ACM "$MERGEBASE" -- '*.py' '*.pyi' | xargs \
@@ -148,7 +160,13 @@ lint_changed() {
     #
     # `diff-filter=ACM` and $MERGEBASE is to ensure we only lint files that
     # exist on both branches.
-    MERGEBASE="$(git merge-base origin/main HEAD)"
+    if git show-ref --verify --quiet refs/remotes/origin/main; then
+        BASE_BRANCH="origin/main"
+    else
+        BASE_BRANCH="main"
+    fi
+
+    MERGEBASE="$(git merge-base $BASE_BRANCH HEAD)"
 
     if ! git diff --diff-filter=ACM --quiet --exit-code "$MERGEBASE" -- '*.py' '*.pyi' &>/dev/null; then
         git diff --name-only --diff-filter=ACM "$MERGEBASE" -- '*.py' '*.pyi' | xargs \

install.sh

@@ -7,7 +7,47 @@
 pip install -r requirements.txt
 
 # install llvm
-apt-get install llvm-10
+LLVM_VERSION="10.0.1"
+IS_AARCH64=false
+EXTRACT_PATH="3rdparty"
+
+UBUNTU_VERSION="16.04"
+if [[ "$LLVM_VERSION" > "16.0.0" ]]; then
+    UBUNTU_VERSION="20.04"
+elif [[ "$LLVM_VERSION" > "13.0.0" ]]; then
+    UBUNTU_VERSION="18.04"
+fi
+
+BASE_URL="https://github.com/llvm/llvm-project/releases/download/llvmorg-${LLVM_VERSION}"
+if $IS_AARCH64; then
+    FILE_NAME="clang+llvm-${LLVM_VERSION}-aarch64-linux-gnu.tar.xz"
+else
+    FILE_NAME="clang+llvm-${LLVM_VERSION}-x86_64-linux-gnu-ubuntu-${UBUNTU_VERSION}.tar.xz"
+fi
+DOWNLOAD_URL="${BASE_URL}/${FILE_NAME}"
+
+mkdir -p "$EXTRACT_PATH"
+
+echo "Downloading $FILE_NAME from $DOWNLOAD_URL"
+curl -L -o "${EXTRACT_PATH}/${FILE_NAME}" "$DOWNLOAD_URL"
+
+if [ $? -ne 0 ]; then
+    echo "Download failed!"
+    exit 1
+fi
+
+echo "Extracting $FILE_NAME to $EXTRACT_PATH"
+tar -xJf "${EXTRACT_PATH}/${FILE_NAME}" -C "$EXTRACT_PATH"
+
+if [ $? -ne 0 ]; then
+    echo "Extraction failed!"
+    exit 1
+fi
+
+echo "Download and extraction completed successfully."
+
+LLVM_CONFIG_PATH="${EXTRACT_PATH}/$(basename ${FILE_NAME} .tar.xz)/bin/llvm-config"
+echo "LLVM config path: $LLVM_CONFIG_PATH"
 
 # clone and build tvm
 git submodule update --init --recursive
@@ -16,7 +56,7 @@ cd 3rdparty/tvm
 mkdir build
 cp cmake/config.cmake build
 cd build
-echo "set(USE_LLVM llvm-config-10)" >> config.cmake && echo "set(USE_CUDA ON)" >> config.cmake
+echo "set(USE_LLVM $LLVM_CONFIG_PATH)" >> config.cmake && echo "set(USE_CUDA ON)" >> config.cmake
 
 cmake .. && make -j && cd ../../..
 

integration/BitNet/modeling_bitnet.py

@@ -54,7 +54,7 @@
 
 if is_flash_attn_2_available():
     from flash_attn import flash_attn_func, flash_attn_varlen_func
-    from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
+    from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa: F401
 
 
 logger = logging.get_logger(__name__)

maint/scripts/installation.sh

@@ -7,7 +7,47 @@
 pip install -r requirements.txt
 
 # install llvm
-apt-get install llvm-10
+LLVM_VERSION="10.0.1"
+IS_AARCH64=false
+EXTRACT_PATH="3rdparty"
+
+UBUNTU_VERSION="16.04"
+if [[ "$LLVM_VERSION" > "16.0.0" ]]; then
+    UBUNTU_VERSION="20.04"
+elif [[ "$LLVM_VERSION" > "13.0.0" ]]; then
+    UBUNTU_VERSION="18.04"
+fi
+
+BASE_URL="https://github.com/llvm/llvm-project/releases/download/llvmorg-${LLVM_VERSION}"
+if $IS_AARCH64; then
+    FILE_NAME="clang+llvm-${LLVM_VERSION}-aarch64-linux-gnu.tar.xz"
+else
+    FILE_NAME="clang+llvm-${LLVM_VERSION}-x86_64-linux-gnu-ubuntu-${UBUNTU_VERSION}.tar.xz"
+fi
+DOWNLOAD_URL="${BASE_URL}/${FILE_NAME}"
+
+mkdir -p "$EXTRACT_PATH"
+
+echo "Downloading $FILE_NAME from $DOWNLOAD_URL"
+curl -L -o "${EXTRACT_PATH}/${FILE_NAME}" "$DOWNLOAD_URL"
+
+if [ $? -ne 0 ]; then
+    echo "Download failed!"
+    exit 1
+fi
+
+echo "Extracting $FILE_NAME to $EXTRACT_PATH"
+tar -xJf "${EXTRACT_PATH}/${FILE_NAME}" -C "$EXTRACT_PATH"
+
+if [ $? -ne 0 ]; then
+    echo "Extraction failed!"
+    exit 1
+fi
+
+echo "Download and extraction completed successfully."
+
+LLVM_CONFIG_PATH="${EXTRACT_PATH}/$(basename ${FILE_NAME} .tar.xz)/bin/llvm-config"
+echo "LLVM config path: $LLVM_CONFIG_PATH"
 
 # clone and build tvm
 git submodule update --init --recursive
@@ -16,7 +56,7 @@ cd 3rdparty/tvm
 mkdir build
 cp cmake/config.cmake build
 cd build
-echo "set(USE_LLVM llvm-config-10)" >> config.cmake && echo "set(USE_CUDA ON)" >> config.cmake
+echo "set(USE_LLVM $LLVM_CONFIG_PATH)" >> config.cmake && echo "set(USE_CUDA ON)" >> config.cmake
 
 cmake .. && make -j && cd ../../..
 

requirements-dev.txt

@@ -28,3 +28,5 @@ tornado
 torch
 thefuzz
 tabulate
+wheel
+setuptools

setup.py

@@ -8,7 +8,6 @@
 from setuptools.command.install import install
 from setuptools.command.build_py import build_py
 from setuptools.command.sdist import sdist
-from wheel.bdist_wheel import bdist_wheel
 import distutils.dir_util
 from typing import List
 import re
@@ -66,7 +65,7 @@ def get_nvcc_cuda_version():
 
 
 def get_bitblas_version(with_cuda=True, with_system_info=True) -> str:
-    version = find_version(get_path("python/bitblas", "__init__.py"))
+    version = find_version(get_path("bitblas", "__init__.py"))
     local_version_parts = []
     if with_system_info:
         local_version_parts.append(get_system_info().replace("-", "."))
@@ -209,8 +208,6 @@ def run(self):
         build_tvm(llvm_path)
         # Continue with the standard installation process
         install.run(self)
-        # Create softlink for bitblas
-        create_softlink(tvm_path="../3rdparty/tvm/python/tvm", bitblas_path="bitblas/tvm")
 
 
 class BitBLASBuilPydCommand(build_py):
@@ -224,8 +221,6 @@ def run(self):
         _, llvm_path = setup_llvm_for_tvm()
         # Build TVM
         build_tvm(llvm_path)
-        # Create softlink for bitblas
-        create_softlink(tvm_path="../3rdparty/tvm/python/tvm", bitblas_path="bitblas/tvm")
 
         # Copy the built TVM to the package directory
         TVM_PREBUILD_ITEMS = [
@@ -268,15 +263,15 @@ def make_distribution(self):
 
 setup(
     name=PACKAGE_NAME,
-    version=get_bitblas_version(with_cuda=False, with_system_info=False) if PYPI_BUILD else get_bitblas_version(),
-    packages=find_packages(where="python"),
-    package_dir={"": "python"},
+    version=get_bitblas_version(with_cuda=False, with_system_info=False)
+    if PYPI_BUILD else get_bitblas_version(),
+    packages=find_packages(where="."),
+    package_dir={"": "."},
     author="Microsoft Research",
     description="A light weight framework to generate high performance CUDA/HIP code for BLAS operators.",
     long_description=read_readme(),
     long_description_content_type='text/markdown',
-    platforms=["Environment :: GPU :: NVIDIA CUDA", 
-               "Operating System :: POSIX :: Linux"],
+    platforms=["Environment :: GPU :: NVIDIA CUDA", "Operating System :: POSIX :: Linux"],
     license="MIT",
     keywords="BLAS, CUDA, HIP, Code Generation, TVM",
     url="https://github.com/microsoft/BitBLAS",