import io import subprocess import shutil from setuptools import setup, find_packages, Extension from setuptools.command.build_py import build_py from setuptools.command.sdist import sdist from setuptools.command.develop import develop import distutils.dir_util from typing import List import re import tarfile from io import BytesIO import os import sys import urllib.request from distutils.version import LooseVersion import platform import multiprocessing from setuptools.command.build_ext import build_ext import importlib # Environment variables False/True 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" # Add ROCM control environment variable USE_ROCM = os.environ.get("USE_ROCM", "False").lower() == "true" def load_module_from_path(module_name, path): spec = importlib.util.spec_from_file_location(module_name, path) module = importlib.util.module_from_spec(spec) sys.modules[module_name] = module spec.loader.exec_module(module) return module envs = load_module_from_path('env', os.path.join(ROOT_DIR, PACKAGE_NAME, 'env.py')) CUDA_HOME = envs.CUDA_HOME ROCM_HOME = envs.ROCM_HOME # Check if both CUDA and ROCM are enabled if USE_ROCM and not ROCM_HOME: raise ValueError( "ROCM support is enabled (USE_ROCM=True) but ROCM_HOME is not set or detected.") if not USE_ROCM and not CUDA_HOME: raise ValueError( "CUDA support is enabled by default (USE_ROCM=False) but CUDA_HOME is not set or detected.") # Ensure one of CUDA or ROCM is available if not (CUDA_HOME or ROCM_HOME): raise ValueError( "Failed to automatically detect CUDA or ROCM installation. Please set the CUDA_HOME or ROCM_HOME environment variable manually (e.g., export CUDA_HOME=/usr/local/cuda or export ROCM_HOME=/opt/rocm)." ) # TileLang only supports Linux platform assert sys.platform.startswith("linux"), "TileLang only supports Linux platform (including WSL)." def get_path(*filepath) -> str: return os.path.join(ROOT_DIR, *filepath) def get_requirements(file_path: str = "requirements.txt") -> List[str]: """Get Python package dependencies from requirements.txt.""" with open(get_path(file_path)) as f: requirements = f.read().strip().split("\n") return requirements def find_version(version_file_path: str) -> str: """Extract version information from the given filepath. Adapted from https://github.com/ray-project/ray/blob/0b190ee1160eeca9796bc091e07eaebf4c85b511/python/setup.py """ # Read and store the version information from the VERSION file # Use 'strip()' to remove any leading/trailing whitespace or newline characters if not os.path.exists(version_file_path): raise FileNotFoundError(f"Version file not found at {version_file_path}") with open(version_file_path, "r") as version_file: version = version_file.read().strip() return version def get_nvcc_cuda_version(): """Get the CUDA version from nvcc. Adapted from https://github.com/NVIDIA/apex/blob/8b7a1ff183741dd8f9b87e7bafd04cfde99cea28/setup.py """ nvcc_output = subprocess.check_output(["nvcc", "-V"], universal_newlines=True) output = nvcc_output.split() release_idx = output.index("release") + 1 nvcc_cuda_version = LooseVersion(output[release_idx].split(",")[0]) return nvcc_cuda_version def get_rocm_version(): """Get the ROCM version from rocminfo.""" rocm_output = subprocess.check_output(["rocminfo"], universal_newlines=True) # Parse ROCM version from output # Example output: ROCM version: x.y.z-... match = re.search(r'ROCm Version: (\d+\.\d+\.\d+)', rocm_output) if match: return LooseVersion(match.group(1)) else: rocm_path = os.environ.get("ROCM_PATH", "/opt/rocm") rocm_version_file = os.path.join(rocm_path, "lib", "cmake", "rocm", "rocm-config-version.cmake") if os.path.exists(rocm_version_file): with open(rocm_version_file, "r") as f: content = f.read() match = re.search(r'set\(PACKAGE_VERSION "(\d+\.\d+\.\d+)"', content) if match: return LooseVersion(match.group(1)) # return a default return LooseVersion("5.0.0") def get_tilelang_version(with_cuda=True, with_system_info=True) -> str: version = find_version(get_path(".", "VERSION")) local_version_parts = [] if with_system_info: local_version_parts.append(get_system_info().replace("-", ".")) if with_cuda: if USE_ROCM: if ROCM_HOME: rocm_version = str(get_rocm_version()) rocm_version_str = rocm_version.replace(".", "")[:3] local_version_parts.append(f"rocm{rocm_version_str}") else: if CUDA_HOME: cuda_version = str(get_nvcc_cuda_version()) cuda_version_str = cuda_version.replace(".", "")[:3] local_version_parts.append(f"cu{cuda_version_str}") if local_version_parts: version += f"+{'.'.join(local_version_parts)}" return version def get_system_info(): system = platform.system().lower() if system == "linux": try: with open("/etc/os-release") as f: os_release = f.read() version_id_match = re.search(r'VERSION_ID="(\d+\.\d+)"', os_release) if version_id_match: version_id = version_id_match.group(1) distro = "ubuntu" return f"{distro}-{version_id}" except FileNotFoundError: pass return system def read_readme() -> str: """Read the README file if present.""" p = get_path("README.md") if os.path.isfile(p): return io.open(get_path("README.md"), "r", encoding="utf-8").read() else: return "" def download_and_extract_llvm(version, is_aarch64=False, extract_path="3rdparty"): """ Downloads and extracts the specified version of LLVM for the given platform. Args: version (str): The version of LLVM to download. is_aarch64 (bool): True if the target platform is aarch64, False otherwise. extract_path (str): The directory path where the archive will be extracted. Returns: str: The path where the LLVM archive was extracted. """ ubuntu_version = "16.04" if version >= "16.0.0": ubuntu_version = "20.04" elif version >= "13.0.0": ubuntu_version = "18.04" base_url = (f"https://github.com/llvm/llvm-project/releases/download/llvmorg-{version}") file_name = f"clang+llvm-{version}-{'aarch64-linux-gnu' if is_aarch64 else f'x86_64-linux-gnu-ubuntu-{ubuntu_version}'}.tar.xz" download_url = f"{base_url}/{file_name}" # Download the file print(f"Downloading {file_name} from {download_url}") with urllib.request.urlopen(download_url) as response: if response.status != 200: raise Exception(f"Download failed with status code {response.status}") file_content = response.read() # Ensure the extract path exists os.makedirs(extract_path, exist_ok=True) # if the file already exists, remove it if os.path.exists(os.path.join(extract_path, file_name)): os.remove(os.path.join(extract_path, file_name)) # Extract the file print(f"Extracting {file_name} to {extract_path}") with tarfile.open(fileobj=BytesIO(file_content), mode="r:xz") as tar: tar.extractall(path=extract_path) print("Download and extraction completed successfully.") return os.path.abspath(os.path.join(extract_path, file_name.replace(".tar.xz", ""))) package_data = { "tilelang": ["py.typed", "*pyx"], } LLVM_VERSION = "10.0.1" IS_AARCH64 = False # Set to True if on an aarch64 platform EXTRACT_PATH = "3rdparty" # Default extraction path def update_submodules(): """Updates git submodules if in a git repository.""" def is_git_repo(): try: # Check if current directory is a git repository subprocess.check_output(["git", "rev-parse", "--is-inside-work-tree"], stderr=subprocess.STDOUT) return True except (subprocess.CalledProcessError, FileNotFoundError): return False if not is_git_repo(): print("Info: Not a git repository, skipping submodule update.") return try: subprocess.check_call(["git", "submodule", "update", "--init", "--recursive"]) except subprocess.CalledProcessError as error: raise RuntimeError("Failed to update submodules") from error def build_csrc(llvm_config_path): """Configures and builds TVM.""" if not os.path.exists("build"): os.makedirs("build") os.chdir("build") # Copy the config.cmake as a baseline if not os.path.exists("config.cmake"): shutil.copy("../3rdparty/tvm/cmake/config.cmake", "config.cmake") # Set LLVM path and enable CUDA or ROCM in config.cmake with open("config.cmake", "a") as config_file: config_file.write(f"set(USE_LLVM {llvm_config_path})\n") if USE_ROCM: config_file.write(f"set(USE_ROCM {ROCM_HOME})\n") config_file.write("set(USE_CUDA OFF)\n") else: config_file.write(f"set(USE_CUDA {CUDA_HOME})\n") config_file.write("set(USE_ROCM OFF)\n") # Run CMake and make try: subprocess.check_call(["cmake", ".."]) 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 def setup_llvm_for_tvm(): """Downloads and extracts LLVM, then configures TVM to use it.""" # Assume the download_and_extract_llvm function and its dependencies are defined elsewhere in this script extract_path = download_and_extract_llvm(LLVM_VERSION, IS_AARCH64, EXTRACT_PATH) llvm_config_path = os.path.join(extract_path, "bin", "llvm-config") return extract_path, llvm_config_path def patch_libs(libpath): """ tvm and tilelang libs are copied from elsewhere into wheels and have a hard-coded rpath. Set rpath to the directory of libs so auditwheel works well. """ subprocess.run(['patchelf', '--set-rpath', '$ORIGIN', libpath]) class TileLangBuilPydCommand(build_py): """Customized setuptools install command - builds TVM after setting up LLVM.""" def run(self): build_py.run(self) self.run_command("build_ext") build_ext_cmd = self.get_finalized_command("build_ext") build_temp_dir = build_ext_cmd.build_temp ext_modules = build_ext_cmd.extensions for ext in ext_modules: extdir = build_ext_cmd.get_ext_fullpath(ext.name) print(f"Extension {ext.name} output directory: {extdir}") ext_output_dir = os.path.dirname(extdir) print(f"Extension output directory (parent): {ext_output_dir}") print(f"Build temp directory: {build_temp_dir}") # copy cython files CYTHON_SRC = [ "tilelang/jit/adapter/cython/cython_wrapper.pyx", ] for item in CYTHON_SRC: source_dir = os.path.join(ROOT_DIR, item) target_dir = os.path.join(self.build_lib, item) if os.path.isdir(source_dir): self.mkpath(target_dir) distutils.dir_util.copy_tree(source_dir, target_dir) else: target_dir = os.path.dirname(target_dir) if not os.path.exists(target_dir): os.makedirs(target_dir) if not os.path.exists(os.path.join(target_dir, os.path.basename(source_dir))): # if not exists, copy the file # as tox will copy the file to the build # directory based on manifest file shutil.copy2(source_dir, target_dir) # copy the tl_templates TILELANG_SRC = [ "src/tl_templates", ] for item in TILELANG_SRC: source_dir = os.path.join(ROOT_DIR, item) target_dir = os.path.join(self.build_lib, PACKAGE_NAME, item) if os.path.isdir(source_dir): self.mkpath(target_dir) distutils.dir_util.copy_tree(source_dir, target_dir) else: target_dir = os.path.dirname(target_dir) if not os.path.exists(target_dir): os.makedirs(target_dir) shutil.copy2(source_dir, target_dir) TVM_PREBUILD_ITEMS = [ "libtvm_runtime.so", "libtvm.so", "libtilelang.so", "libtilelang_module.so", ] potential_dirs = [ ext_output_dir, self.build_lib, build_temp_dir, os.path.join(ROOT_DIR, "build"), ] for item in TVM_PREBUILD_ITEMS: source_lib_file = None for dir in potential_dirs: candidate = os.path.join(dir, item) if os.path.exists(candidate): source_lib_file = candidate break if source_lib_file: patch_libs(source_lib_file) target_dir_release = os.path.join(self.build_lib, PACKAGE_NAME, "lib") target_dir_develop = os.path.join(PACKAGE_NAME, "lib") os.makedirs(target_dir_release, exist_ok=True) os.makedirs(target_dir_develop, exist_ok=True) shutil.copy2(source_lib_file, target_dir_release) print(f"Copied {source_lib_file} to {target_dir_release}") shutil.copy2(source_lib_file, target_dir_develop) print(f"Copied {source_lib_file} to {target_dir_develop}") os.remove(source_lib_file) else: print(f"WARNING: {item} not found in any expected directories!") TVM_CONFIG_ITEMS = [ f"{build_temp_dir}/config.cmake", ] for item in TVM_CONFIG_ITEMS: source_dir = os.path.join(ROOT_DIR, item) # only copy the file file_name = os.path.basename(item) target_dir = os.path.join(self.build_lib, PACKAGE_NAME, file_name) target_dir = os.path.dirname(target_dir) if not os.path.exists(target_dir): os.makedirs(target_dir) if os.path.exists(source_dir): shutil.copy2(source_dir, target_dir) else: print(f"INFO: {source_dir} does not exist.") TVM_PACAKGE_ITEMS = [ "3rdparty/tvm/src", "3rdparty/tvm/python", "3rdparty/tvm/licenses", "3rdparty/tvm/conftest.py", "3rdparty/tvm/CONTRIBUTORS.md", "3rdparty/tvm/KEYS", "3rdparty/tvm/LICENSE", "3rdparty/tvm/README.md", "3rdparty/tvm/mypy.ini", "3rdparty/tvm/pyproject.toml", "3rdparty/tvm/version.py", ] for item in TVM_PACAKGE_ITEMS: source_dir = os.path.join(ROOT_DIR, item) target_dir = os.path.join(self.build_lib, PACKAGE_NAME, item) if os.path.isdir(source_dir): self.mkpath(target_dir) distutils.dir_util.copy_tree(source_dir, target_dir) else: target_dir = os.path.dirname(target_dir) if not os.path.exists(target_dir): os.makedirs(target_dir) shutil.copy2(source_dir, target_dir) # Copy CUTLASS to the package directory CUTLASS_PREBUILD_ITEMS = [ "3rdparty/cutlass/include", "3rdparty/cutlass/tools", ] for item in CUTLASS_PREBUILD_ITEMS: source_dir = os.path.join(ROOT_DIR, item) target_dir = os.path.join(self.build_lib, PACKAGE_NAME, item) if os.path.isdir(source_dir): self.mkpath(target_dir) distutils.dir_util.copy_tree(source_dir, target_dir) else: target_dir = os.path.dirname(target_dir) if not os.path.exists(target_dir): os.makedirs(target_dir) shutil.copy2(source_dir, target_dir) # copy compoable kernel to the package directory CK_PREBUILD_ITEMS = [ "3rdparty/composable_kernel/include", "3rdparty/composable_kernel/library", ] for item in CK_PREBUILD_ITEMS: source_dir = os.path.join(ROOT_DIR, item) target_dir = os.path.join(self.build_lib, PACKAGE_NAME, item) if os.path.isdir(source_dir): self.mkpath(target_dir) distutils.dir_util.copy_tree(source_dir, target_dir) else: target_dir = os.path.dirname(target_dir) if not os.path.exists(target_dir): os.makedirs(target_dir) shutil.copy2(source_dir, target_dir) # copy compoable kernel to the package directory TL_CONFIG_ITEMS = ["CMakeLists.txt", "VERSION", "README.md", "LICENSE"] for item in TL_CONFIG_ITEMS: source_dir = os.path.join(ROOT_DIR, item) target_dir = os.path.join(self.build_lib, PACKAGE_NAME, item) if os.path.isdir(source_dir): self.mkpath(target_dir) distutils.dir_util.copy_tree(source_dir, target_dir) else: target_dir = os.path.dirname(target_dir) if not os.path.exists(target_dir): os.makedirs(target_dir) shutil.copy2(source_dir, target_dir) class TileLangSdistCommand(sdist): """Customized setuptools sdist command - includes the pyproject.toml file.""" def make_distribution(self): self.distribution.metadata.name = PACKAGE_NAME self.distribution.metadata.version = get_tilelang_version( with_cuda=False, with_system_info=False) super().make_distribution() # ------------------------------------------------------------------------ # NEW: Add a custom 'develop' command so that `pip install -e .` works. # ------------------------------------------------------------------------ class TileLangDevelopCommand(develop): """ Customized setuptools 'develop' command for an editable install. Ensures the extension is built and all necessary assets are copied. """ def run(self): print("Running TileLangDevelopCommand") # 1. Build the C/C++ extension modules self.run_command("build_ext") build_ext_cmd = self.get_finalized_command("build_ext") ext_modules = build_ext_cmd.extensions for ext in ext_modules: extdir = build_ext_cmd.get_ext_fullpath(ext.name) print(f"Extension {ext.name} output directory: {extdir}") ext_output_dir = os.path.dirname(extdir) print(f"Extension output directory (parent): {ext_output_dir}") # Copy the built TVM to the package directory TVM_PREBUILD_ITEMS = [ f"{ext_output_dir}/libtvm_runtime.so", f"{ext_output_dir}/libtvm.so", f"{ext_output_dir}/libtilelang.so", f"{ext_output_dir}/libtilelang_module.so", ] for item in TVM_PREBUILD_ITEMS: source_lib_file = os.path.join(ROOT_DIR, item) # only copy the file file_name = os.path.basename(item) target_dir = os.path.join(PACKAGE_NAME, file_name) target_dir = os.path.dirname(target_dir) target_dir = os.path.join(target_dir, "lib") if not os.path.exists(target_dir): os.makedirs(target_dir) if os.path.exists(source_lib_file): patch_libs(source_lib_file) shutil.copy2(source_lib_file, target_dir) # remove the original file os.remove(source_lib_file) else: print(f"INFO: {source_lib_file} does not exist.") class CMakeExtension(Extension): """ A specialized setuptools Extension class for building a CMake project. :param name: Name of the extension module. :param sourcedir: Directory containing the top-level CMakeLists.txt. """ def __init__(self, name, sourcedir=""): # We pass an empty 'sources' list because # the actual build is handled by CMake, not setuptools. super().__init__(name=name, sources=[]) # Convert the source directory to an absolute path # so that CMake can correctly locate the CMakeLists.txt. self.sourcedir = os.path.abspath(sourcedir) class CMakeBuild(build_ext): """ Custom build_ext command for CMake-based projects. This class overrides the 'run' method to ensure that CMake is available, and then iterates over all extensions defined as CMakeExtension, delegating the actual build logic to 'build_cmake'. """ def run(self): # Check if CMake is installed and accessible by attempting to run 'cmake --version'. try: subprocess.check_output(["cmake", "--version"]) except OSError as e: # If CMake is not found, raise an error. raise RuntimeError("CMake must be installed to build the following extensions") from e update_submodules() # Build each extension (of type CMakeExtension) using our custom method. 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). """ # 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))) # Prepare arguments for the CMake configuration step. # -DCMAKE_LIBRARY_OUTPUT_DIRECTORY sets where built libraries go # -DPYTHON_EXECUTABLE ensures that the correct Python is used cmake_args = [ f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={extdir}", f"-DPYTHON_EXECUTABLE={sys.executable}", ] # Create the temporary build directory (if it doesn't exist). build_temp = os.path.abspath(self.build_temp) os.makedirs(build_temp, exist_ok=True) # Copy the default 'config.cmake' from the source tree into our build directory. src_config_cmake = os.path.join(ext.sourcedir, "3rdparty", "tvm", "cmake", "config.cmake") 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' with open(dst_config_cmake, "a") as config_file: config_file.write(f"set(USE_LLVM {llvm_config_path})\n") if USE_ROCM: config_file.write(f"set(USE_ROCM {ROCM_HOME})\n") config_file.write("set(USE_CUDA OFF)\n") else: config_file.write(f"set(USE_CUDA {CUDA_HOME})\n") config_file.write("set(USE_ROCM OFF)\n") # Run CMake to configure the project with the given arguments. subprocess.check_call(["cmake", ext.sourcedir] + cmake_args, cwd=build_temp) # Build the project in "Release" mode with all available CPU cores ("-j"). subprocess.check_call(["cmake", "--build", ".", "--config", "Release", "-j"], cwd=build_temp) setup( name=PACKAGE_NAME, version=(get_tilelang_version(with_cuda=False, with_system_info=False) if PYPI_BUILD else get_tilelang_version()), packages=find_packages(where="."), package_dir={"": "."}, author="Tile-AI", description="A tile level programming language to generate high performance code.", long_description=read_readme(), long_description_content_type="text/markdown", platforms=[ "Environment :: GPU :: NVIDIA CUDA" if not USE_ROCM else "Environment :: GPU :: AMD ROCm", "Operating System :: POSIX :: Linux", ], license="MIT", keywords="BLAS, CUDA, HIP, Code Generation, TVM", url="https://github.com/tile-ai/tilelang", classifiers=[ "Programming Language :: Python :: 3.8", "License :: OSI Approved :: MIT License", "Operating System :: OS Independent", "Intended Audience :: Developers", "Intended Audience :: Science/Research", ], python_requires=">=3.8", install_requires=get_requirements(), package_data=package_data, include_package_data=False, ext_modules=[CMakeExtension("TileLangCXX", sourcedir=".")], cmdclass={ "build_py": TileLangBuilPydCommand, "sdist": TileLangSdistCommand, "build_ext": CMakeBuild, "develop": TileLangDevelopCommand, }, )