[Refactor] Phaseout LLVM Dependency by Making it Optional (#247)

* remove llvm build * [Refactor] Update kernel compilation and profiling in examples - Replaced `tilelang.lower` with `tilelang.compile` in multiple example scripts to streamline kernel compilation. - Updated profiling calls to utilize the new `get_profiler` method, enhancing performance measurement consistency. - Adjusted assertions and benchmarking methods to align with the new profiling structure across various examples, ensuring correctness and clarity in performance evaluations. * lint fix * License Update * [Refactor] Improve code formatting and documentation in CUDA header and HIP runtime files - Adjusted formatting in `cuda.h` for better readability, including alignment of comments and struct fields. - Cleaned up whitespace and improved comment clarity in `rt_mod_hip.cc` to enhance code maintainability. * [Refactor] Enhance formatting and clarity in CUDA header and HIP runtime files - Improved comment alignment and readability in `cuda.h`. - Cleaned up whitespace and formatting in `rt_mod_hip.cc` to enhance maintainability. * lint fix * lint fix * lint fix * lint fix * fix * License update * [Enhancement] Update JITKernel to use artifact for kernel source - Assigned the generated artifact to `self.artifact` for better management. - Updated kernel source references to use `artifact.kernel_source` for consistency in execution backend handling. * lint fix * Add @tilelang.testing.requires_llvm decorator to vectorization tests * Enhance setup.py and env.py for library management - Added functionality to remove original files after copying in CMakeBuild. - Updated TVM_LIBRARY_PATH in env.py to include the PyPI build library path for better integration. * Refactor TVM_LIBRARY_PATH assignment for improved readability in env.py * Refactor CMakeBuild file handling in setup.py - Added a check to ensure the target library directory exists before copying .so files. - Improved the logic for creating the target directory and copying files to enhance robustness. * bugfix * Rename BuildTLDebug to BuildTileLangCUDAWithoutCompile and update registration. Add @tilelang.testing.requires_llvm decorator to multiple tests for LLVM requirement. * lint fix * Enhance TileLang code generation by adding support for device code generation without compilation. Updated `host_codegen` and `device_codegen` functions to include new transformations and registration for `tilelang_hip_without_compile`. Refactored JIT kernel adapters to accommodate host and device modules, improving overall integration and flexibility. * lint fix * Add support for C target in device code generation - Updated `device_codegen_without_compile` to include handling for the C target by registering the `tilelang_cpp` function. * [Enhancement] Implement auto-clear cache feature based on environment variable * Added TILELANG_CLEAR_CACHE environment variable to control cache clearing. * Updated CI workflow to set TILELANG_CLEAR_CACHE during testing. * Modified cache initialization to clear cache if TILELANG_CLEAR_CACHE is set to true. * [Refactor] Update kernel invocation and import paths in tests and cache * Changed kernel invocation in `test_tilelang_kernel_dequantize_gemm.py` to return the result. * Updated import statements in `test_tilelang_kernel_int4_gemm_mma.py` to use `bitblas` instead of `tilelang`. * Refactored paths for artifact and parameters in `kernel_cache.py` for better maintainability. * [Refactor] Clean up whitespace and improve code formatting in kernel_cache.py * Removed unnecessary blank lines and adjusted spacing for better readability in the KernelCache class. * Enhanced overall code formatting to align with project standards. * [Enhancement] Add bfloat16 test case and improve kernel caching logic * Introduced a new test case for bfloat16 matrix multiplication in `test_tilelang_kernel_gemm_mma_intrinsic.py`. * Updated `KernelCache` to handle multiple kernel source files and improve error handling during saving and loading. * Refactored `JITKernel` to support instantiation from a database, enhancing flexibility in kernel management. * Adjusted `CtypesKernelAdapter` and `CythonKernelAdapter` to utilize the new kernel loading mechanism from the database. * Improved code formatting and readability across several files. * lint fix * Update bfloat16 matrix multiplication test case to use larger dimensions for improved coverage

[Refactor] Phaseout LLVM Dependency by Making it Optional (#247)
* remove llvm build * [Refactor] Update kernel compilation and profiling in examples - Replaced `tilelang.lower` with `tilelang.compile` in multiple example scripts to streamline kernel compilation. - Updated profiling calls to utilize the new `get_profiler` method, enhancing performance measurement consistency. - Adjusted assertions and benchmarking methods to align with the new profiling structure across various examples, ensuring correctness and clarity in performance evaluations. * lint fix * License Update * [Refactor] Improve code formatting and documentation in CUDA header and HIP runtime files - Adjusted formatting in `cuda.h` for better readability, including alignment of comments and struct fields. - Cleaned up whitespace and improved comment clarity in `rt_mod_hip.cc` to enhance code maintainability. * [Refactor] Enhance formatting and clarity in CUDA header and HIP runtime files - Improved comment alignment and readability in `cuda.h`. - Cleaned up whitespace and formatting in `rt_mod_hip.cc` to enhance maintainability. * lint fix * lint fix * lint fix * lint fix * fix * License update * [Enhancement] Update JITKernel to use artifact for kernel source - Assigned the generated artifact to `self.artifact` for better management. - Updated kernel source references to use `artifact.kernel_source` for consistency in execution backend handling. * lint fix * Add @tilelang.testing.requires_llvm decorator to vectorization tests * Enhance setup.py and env.py for library management - Added functionality to remove original files after copying in CMakeBuild. - Updated TVM_LIBRARY_PATH in env.py to include the PyPI build library path for better integration. * Refactor TVM_LIBRARY_PATH assignment for improved readability in env.py * Refactor CMakeBuild file handling in setup.py - Added a check to ensure the target library directory exists before copying .so files. - Improved the logic for creating the target directory and copying files to enhance robustness. * bugfix * Rename BuildTLDebug to BuildTileLangCUDAWithoutCompile and update registration. Add @tilelang.testing.requires_llvm decorator to multiple tests for LLVM requirement. * lint fix * Enhance TileLang code generation by adding support for device code generation without compilation. Updated `host_codegen` and `device_codegen` functions to include new transformations and registration for `tilelang_hip_without_compile`. Refactored JIT kernel adapters to accommodate host and device modules, improving overall integration and flexibility. * lint fix * Add support for C target in device code generation - Updated `device_codegen_without_compile` to include handling for the C target by registering the `tilelang_cpp` function. * [Enhancement] Implement auto-clear cache feature based on environment variable * Added TILELANG_CLEAR_CACHE environment variable to control cache clearing. * Updated CI workflow to set TILELANG_CLEAR_CACHE during testing. * Modified cache initialization to clear cache if TILELANG_CLEAR_CACHE is set to true. * [Refactor] Update kernel invocation and import paths in tests and cache * Changed kernel invocation in `test_tilelang_kernel_dequantize_gemm.py` to return the result. * Updated import statements in `test_tilelang_kernel_int4_gemm_mma.py` to use `bitblas` instead of `tilelang`. * Refactored paths for artifact and parameters in `kernel_cache.py` for better maintainability. * [Refactor] Clean up whitespace and improve code formatting in kernel_cache.py * Removed unnecessary blank lines and adjusted spacing for better readability in the KernelCache class. * Enhanced overall code formatting to align with project standards. * [Enhancement] Add bfloat16 test case and improve kernel caching logic * Introduced a new test case for bfloat16 matrix multiplication in `test_tilelang_kernel_gemm_mma_intrinsic.py`. * Updated `KernelCache` to handle multiple kernel source files and improve error handling during saving and loading. * Refactored `JITKernel` to support instantiation from a database, enhancing flexibility in kernel management. * Adjusted `CtypesKernelAdapter` and `CythonKernelAdapter` to utilize the new kernel loading mechanism from the database. * Improved code formatting and readability across several files. * lint fix * Update bfloat16 matrix multiplication test case to use larger dimensions for improved coverage
f2e99180 · Lei Wang · LeiWang1999 · 43bd9d3e · f2e99180 · f2e99180
Commit f2e99180 authored Mar 20, 2025 by Lei Wang Committed by LeiWang1999 Mar 20, 2025
20 changed files
--- a/install_cuda.sh
+++ b/install_cuda.sh
-#!/bin/bash
-
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT License.
+# Add command line option parsing
+USE_LLVM=false
+while [[ $# -gt 0 ]]; do
+    case $1 in
+        --enable-llvm)
+            USE_LLVM=true
+            shift
+            ;;
+        *)
+            echo "Unknown option: $1"
+            echo "Usage: $0 [--enable-llvm]"
+            exit 1
+            ;;
+    esac
+done

 echo "Starting installation script..."
+echo "LLVM enabled: $USE_LLVM"

 # Step 1: Install Python requirements
 echo "Installing Python requirements from requirements.txt..."
@@ -17,67 +29,69 @@ else
 fi

 # Step 2: Define LLVM version and architecture
-LLVM_VERSION="10.0.1"
-IS_AARCH64=false
-EXTRACT_PATH="3rdparty"
-echo "LLVM version set to ${LLVM_VERSION}."
-echo "Is AARCH64 architecture: $IS_AARCH64"
-
-# Step 3: Determine the correct Ubuntu version based on LLVM version
-UBUNTU_VERSION="16.04"
-if [[ "$LLVM_VERSION" > "17.0.0" ]]; then
-    UBUNTU_VERSION="22.04"
-elif [[ "$LLVM_VERSION" > "16.0.0" ]]; then
-    UBUNTU_VERSION="20.04"
-elif [[ "$LLVM_VERSION" > "13.0.0" ]]; then
-    UBUNTU_VERSION="18.04"
-fi
-echo "Ubuntu version for LLVM set to ${UBUNTU_VERSION}."
-
-# Step 4: Set download URL and file name for LLVM
-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}"
-echo "Download URL for LLVM: ${DOWNLOAD_URL}"
-
-# Step 5: Create extraction directory
-echo "Creating extraction directory at ${EXTRACT_PATH}..."
-mkdir -p "$EXTRACT_PATH"
-if [ $? -ne 0 ]; then
-    echo "Error: Failed to create extraction directory."
-    exit 1
-else
-    echo "Extraction directory created successfully."
-fi
-
-# Step 6: Download LLVM
-echo "Downloading $FILE_NAME from $DOWNLOAD_URL..."
-curl -L -o "${EXTRACT_PATH}/${FILE_NAME}" "$DOWNLOAD_URL"
-if [ $? -ne 0 ]; then
-    echo "Error: Download failed!"
-    exit 1
-else
-    echo "Download completed successfully."
+if $USE_LLVM; then
+    LLVM_VERSION="10.0.1"
+    IS_AARCH64=false
+    EXTRACT_PATH="3rdparty"
+    echo "LLVM version set to ${LLVM_VERSION}."
+    echo "Is AARCH64 architecture: $IS_AARCH64"
+
+    # Step 3: Determine the correct Ubuntu version based on LLVM version
+    UBUNTU_VERSION="16.04"
+    if [[ "$LLVM_VERSION" > "17.0.0" ]]; then
+        UBUNTU_VERSION="22.04"
+    elif [[ "$LLVM_VERSION" > "16.0.0" ]]; then
+        UBUNTU_VERSION="20.04"
+    elif [[ "$LLVM_VERSION" > "13.0.0" ]]; then
+        UBUNTU_VERSION="18.04"
+    fi
+    echo "Ubuntu version for LLVM set to ${UBUNTU_VERSION}."
+
+    # Step 4: Set download URL and file name for LLVM
+    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}"
+    echo "Download URL for LLVM: ${DOWNLOAD_URL}"
+
+    # Step 5: Create extraction directory
+    echo "Creating extraction directory at ${EXTRACT_PATH}..."
+    mkdir -p "$EXTRACT_PATH"
+    if [ $? -ne 0 ]; then
+        echo "Error: Failed to create extraction directory."
+        exit 1
+    else
+        echo "Extraction directory created successfully."
+    fi
+
+    # Step 6: Download LLVM
+    echo "Downloading $FILE_NAME from $DOWNLOAD_URL..."
+    curl -L -o "${EXTRACT_PATH}/${FILE_NAME}" "$DOWNLOAD_URL"
+    if [ $? -ne 0 ]; then
+        echo "Error: Download failed!"
+        exit 1
+    else
+        echo "Download completed successfully."
+    fi
+
+    # Step 7: Extract LLVM
+    echo "Extracting $FILE_NAME to $EXTRACT_PATH..."
+    tar -xJf "${EXTRACT_PATH}/${FILE_NAME}" -C "$EXTRACT_PATH"
+    if [ $? -ne 0 ]; then
+        echo "Error: Extraction failed!"
+        exit 1
+    else
+        echo "Extraction completed successfully."
+    fi
+
+    # Step 8: Determine LLVM config path
+    LLVM_CONFIG_PATH="$(realpath ${EXTRACT_PATH}/$(basename ${FILE_NAME} .tar.xz)/bin/llvm-config)"
+    echo "LLVM config path determined as: $LLVM_CONFIG_PATH"
 fi

-# Step 7: Extract LLVM
-echo "Extracting $FILE_NAME to $EXTRACT_PATH..."
-tar -xJf "${EXTRACT_PATH}/${FILE_NAME}" -C "$EXTRACT_PATH"
-if [ $? -ne 0 ]; then
-    echo "Error: Extraction failed!"
-    exit 1
-else
-    echo "Extraction completed successfully."
-fi
-
-# Step 8: Determine LLVM config path
-LLVM_CONFIG_PATH="$(realpath ${EXTRACT_PATH}/$(basename ${FILE_NAME} .tar.xz)/bin/llvm-config)"
-echo "LLVM config path determined as: $LLVM_CONFIG_PATH"
-
 # Step 9: Clone and build TVM
 echo "Cloning TVM repository and initializing submodules..."
 # clone and build tvm
@@ -91,10 +105,11 @@ mkdir build
 cp 3rdparty/tvm/cmake/config.cmake build
 cd build

-
-echo "Configuring TVM build with LLVM and CUDA paths..."
-echo "set(USE_LLVM \"$LLVM_CONFIG_PATH\")" >> config.cmake && \
-    CUDA_HOME=$(python -c "import sys; sys.path.append('../tilelang'); from env import CUDA_HOME; print(CUDA_HOME)") || \
+echo "Configuring TVM build with CUDA paths..."
+if $USE_LLVM; then
+    echo "set(USE_LLVM \"$LLVM_CONFIG_PATH\")" >> config.cmake
+fi
+CUDA_HOME=$(python -c "import sys; sys.path.append('../tilelang'); from env import CUDA_HOME; print(CUDA_HOME)") || \
    { echo "ERROR: Failed to retrieve CUDA_HOME via Python script." >&2; exit 1; } && \
    { [ -n "$CUDA_HOME" ] || { echo "ERROR: CUDA_HOME is empty, check CUDA installation or _find_cuda_home() in setup.py" >&2; exit 1; }; } && \
    echo "set(USE_CUDA \"$CUDA_HOME\")" >> config.cmake

--- a/install_rocm.sh
+++ b/install_rocm.sh
-#!/bin/bash
-
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT License.
-
 echo "Starting installation script..."

 # install requirements

--- a/maint/scripts/apply_mit_license.sh
+++ b/maint/scripts/apply_mit_license.sh
-#!/bin/bash
-
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT License.
-
 echo "Add MIT license boilerplate..."
 PWD="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
 # TO source code root
@@ -13,7 +8,7 @@ EXITCODE=0
 for SRC_FILE in $(find . -path './3rdparty' -prune -false -o -path './build' -prune -false -o -type f -not -name \
    '*apply_mit_liscense.sh' -not -name '*check_mit_liscense.sh' -and \( -name '*.cpp' -or -name '*.h*' -or -name '*.cu' -or -name '*.in' \) ); do
    sed -i '/\/\/\s*Microsoft\s*(c)/Id' ${SRC_FILE}
-    if !(grep -q "Copyright (c) Microsoft Corporation." "${SRC_FILE}"); then
+    if !(grep -q "Copyright (c) Tile-AI Corporation." "${SRC_FILE}"); then
        cat maint/scripts/mit_liscense1.txt ${SRC_FILE} > ${SRC_FILE}.new
        mv ${SRC_FILE}.new ${SRC_FILE}
    fi
@@ -23,7 +18,7 @@ for SRC_FILE in $(find . -path './3rdparty' -prune -false -o -path './build' -pr
    '*apply_mit_liscense.sh' -not -name '*check_mit_liscense.sh' -and \( -name 'CMakeLists.txt' -or -name '*.cmake' \
    -or -name '*.py' -or -name '*.dockerfile' -or -name '*.yaml' \) ); do
    sed -i '/\#\s*Microsoft\s*(c)/Id' ${SRC_FILE} 
-    if !(grep -q "Copyright (c) Microsoft Corporation" "${SRC_FILE}"); then       
+    if !(grep -q "Copyright (c) Tile-AI Corporation." "${SRC_FILE}"); then       
        cat maint/scripts/mit_liscense2.txt ${SRC_FILE} > ${SRC_FILE}.new
        mv ${SRC_FILE}.new ${SRC_FILE}
    fi
@@ -32,7 +27,7 @@ done
 for SRC_FILE in $(find . -path './3rdparty' -prune -false -o -path './build' -prune -false -o -type f -not -name \
    '*apply_mit_liscense.sh' -not -name '*check_mit_liscense.sh' -name '*.sh' ); do
    sed -i '/\#\s*Microsoft\s*(c)/Id' ${SRC_FILE} 
-    if !(grep -q "Copyright (c) Microsoft Corporation" "${SRC_FILE}"); then
+    if !(grep -q "Copyright (c) Tile-AI Corporation." "${SRC_FILE}"); then
        line=$(head -n 1 ${SRC_FILE})
        if [[ $line == "#!/bin/bash"* ]]; then
            (echo ${line}; echo ''; cat maint/scripts/mit_liscense2.txt; echo "$(tail -n +2 "${SRC_FILE}")" ) > ${SRC_FILE}.new

--- a/maint/scripts/build_docs.sh
+++ b/maint/scripts/build_docs.sh
-#!/bin/bash
-
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT License.
-
 cd docs

 pip install -r requirements.txt

--- a/maint/scripts/check_mit_license.sh
+++ b/maint/scripts/check_mit_license.sh
-#!/bin/bash
-
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT License.
-
 echo "Check MIT License boilerplate..."
 PWD="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
 # To source code root
@@ -19,7 +14,7 @@ for SRC_FILE in $(find . -path './3rdparty' -prune -false -o -path './build' -pr
        continue
    fi

-    if !(grep -q "Copyright (c) Microsoft Corporation." "${SRC_FILE}") || !(grep -q "Licensed under the MIT License." "${SRC_FILE}") \
+    if !(grep -q "Copyright (c) Tile-AI Corporation." "${SRC_FILE}") || !(grep -q "Licensed under the MIT License." "${SRC_FILE}") \
    || (grep -q -i -P "Microsoft( |)\(c\)" "${SRC_FILE}"); then
        echo "[ERROR] Require: MIT License boilerplate" "${SRC_FILE}"
        EXITCODE=1

--- a/maint/scripts/docker_local_distribute.sh
+++ b/maint/scripts/docker_local_distribute.sh
-#!/bin/bash
-
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT License.
-
 # Get the CUDA version from the command line
 IMAGE=nvidia/cuda:12.1.0-devel-ubuntu18.04


--- a/maint/scripts/docker_pypi_distribute.sh
+++ b/maint/scripts/docker_pypi_distribute.sh
-#!/bin/bash
-
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT License.
-
 # Get the CUDA version from the command line
 IMAGE=nvidia/cuda:12.1.0-devel-ubuntu18.04


--- a/maint/scripts/local_distribution.sh
+++ b/maint/scripts/local_distribution.sh
-#!/bin/bash
-
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT License.
-
 # if dist and build directories exist, remove them
 if [ -d dist ]; then
    rm -r dist

--- a/maint/scripts/local_distribution_tox.sh
+++ b/maint/scripts/local_distribution_tox.sh
-#!/bin/bash
-
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT License.
-
 multi_python_version=("3.8" "3.9" "3.10" "3.11" "3.12")
 for python_version in "${multi_python_version[@]}"; do
    echo "Installing Python ${python_version}..."

--- a/maint/scripts/mit_liscense1.txt
+++ b/maint/scripts/mit_liscense1.txt
-// Copyright (c) Microsoft Corporation.
-// Licensed under the MIT License.
--- a/maint/scripts/mit_liscense2.txt
+++ b/maint/scripts/mit_liscense2.txt
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT License.
--- a/maint/scripts/pypi_distribution.sh
+++ b/maint/scripts/pypi_distribution.sh
-#!/bin/bash
-
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT License.
-
 # if dist and build directories exist, remove them
 if [ -d dist ]; then
    rm -r dist

--- a/maint/scripts/pypi_distribution_tox.sh
+++ b/maint/scripts/pypi_distribution_tox.sh
-#!/bin/bash
-
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT License.
-
 multi_python_version=("3.8" "3.9" "3.10" "3.11" "3.12")
 for python_version in "${multi_python_version[@]}"; do
    echo "Installing Python ${python_version}..."

--- a/setup.py
+++ b/setup.py
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT License.
-
 import io
 import subprocess
 import shutil
@@ -27,6 +24,9 @@ PYPI_BUILD = os.environ.get("PYPI_BUILD", "False").lower() == "true"
 PACKAGE_NAME = "tilelang"
 ROOT_DIR = os.path.dirname(__file__)

+# Add LLVM control environment variable
+USE_LLVM = os.environ.get("USE_LLVM", "False").lower() == "true"
+

 def load_module_from_path(module_name, path):
    spec = importlib.util.spec_from_file_location(module_name, path)
@@ -212,7 +212,7 @@ def build_csrc(llvm_config_path):
    # Run CMake and make
    try:
        subprocess.check_call(["cmake", ".."])
-        num_jobs = max(1, int(multiprocessing.cpu_count() * 0.9))
+        num_jobs = max(1, int(multiprocessing.cpu_count() * 0.75))
        subprocess.check_call(["make", f"-j{num_jobs}"])
    except subprocess.CalledProcessError as error:
        raise RuntimeError("Failed to build TileLang C Source") from error
@@ -423,6 +423,7 @@ class TileLangDevelopCommand(develop):
    """

    def run(self):
+        print("Running TileLangDevelopCommand")
        # 1. Build the C/C++ extension modules
        self.run_command("build_ext")

@@ -500,15 +501,30 @@ class CMakeBuild(build_ext):
        for ext in self.extensions:
            self.build_cmake(ext)

+        # To make it works with editable install,
+        # we need to copy the lib*.so files to the tilelang/lib directory
+        import glob
+        files = glob.glob("*.so")
+        if os.path.exists(PACKAGE_NAME):
+            target_lib_dir = os.path.join(PACKAGE_NAME, "lib")
+            for file in files:
+                if not os.path.exists(target_lib_dir):
+                    os.makedirs(target_lib_dir)
+                shutil.copy(file, target_lib_dir)
+                # remove the original file
+                os.remove(file)
+
    def build_cmake(self, ext):
        """
        Build a single CMake-based extension.

        :param ext: The extension (an instance of CMakeExtension).
        """
-        # Setup LLVM for TVM and retrieve the path to llvm-config.
-        # We assume the function returns (_, llvm_config_path).
-        _, llvm_config_path = setup_llvm_for_tvm()
+        # Only setup LLVM if it's enabled
+        llvm_config_path = "OFF"
+        if USE_LLVM:
+            # Setup LLVM for TVM and retrieve the path to llvm-config
+            _, llvm_config_path = setup_llvm_for_tvm()

        # Determine the directory where the final .so or .pyd library should go.
        extdir = os.path.abspath(os.path.dirname(self.get_ext_fullpath(ext.name)))
@@ -530,8 +546,7 @@ class CMakeBuild(build_ext):
        dst_config_cmake = os.path.join(build_temp, "config.cmake")
        shutil.copy(src_config_cmake, dst_config_cmake)

-        # Append some configuration variables to 'config.cmake'.
-        # Here, we set USE_LLVM and USE_CUDA, for example.
+        # Append some configuration variables to 'config.cmake'
        with open(dst_config_cmake, "a") as config_file:
            config_file.write(f"set(USE_LLVM {llvm_config_path})\n")
            config_file.write(f"set(USE_CUDA {CUDA_HOME})\n")

--- a/src/ir.cc
+++ b/src/ir.cc
-// Copyright (c) Microsoft Corporation.
-// Licensed under the MIT License.
-
 /*!
 * \file tl/ir.cc
 * \brief Extension for the tvm script frontend.

--- a/src/layout/gemm_layouts.cc
+++ b/src/layout/gemm_layouts.cc
-// Copyright (c) Microsoft Corporation.
-// Licensed under the MIT License.
-
 /*!
 * \file layout/gemm_layouts.cc
 * \brief Define Layout used in MMA and other operations.

--- a/src/layout/layout.cc
+++ b/src/layout/layout.cc
-// Copyright (c) Microsoft Corporation.
-// Licensed under the MIT License.
-
 /*!
 * \file layout/layout.cc
 *

--- a/src/layout/layout.h
+++ b/src/layout/layout.h
-// Copyright (c) Microsoft Corporation.
-// Licensed under the MIT License.
-
 /*!
 * \file Layout.h
 *

--- a/src/layout/swizzle.cc
+++ b/src/layout/swizzle.cc
-// Copyright (c) Microsoft Corporation.
-// Licensed under the MIT License.
-
 /*!
 * \file layout/swizzle.cc
 * \brief Define swizzled layout

--- a/src/layout/swizzle.h
+++ b/src/layout/swizzle.h
-// Copyright (c) Microsoft Corporation.
-// Licensed under the MIT License.
-
 /*!
 * \file swizzle.h
 * \brief Define swizzled layout