"src/git@developer.sourcefind.cn:gaoqiong/migraphx.git" did not exist on "3981467bf2b1fea77b8faa2f485e708928bee775"
setup.py 26.1 KB
Newer Older
1
2
3
4
5
6
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
7
from setuptools.command.develop import develop
8
9
10
11
12
13
14
15
16
17
18
19
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
20
import importlib
21
22
23
24
25
26
27
28
29
import logging

# Configure logging with basic settings
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S')

logger = logging.getLogger(__name__)
30
31
32
33
34
35

# Environment variables False/True
PYPI_BUILD = os.environ.get("PYPI_BUILD", "False").lower() == "true"
PACKAGE_NAME = "tilelang"
ROOT_DIR = os.path.dirname(__file__)

36
37
# Add LLVM control environment variable
USE_LLVM = os.environ.get("USE_LLVM", "False").lower() == "true"
38
39
# Add ROCM control environment variable
USE_ROCM = os.environ.get("USE_ROCM", "False").lower() == "true"
40

41
42
43
44
45
46
47
48
49
50

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'))
51

52
CUDA_HOME = envs.CUDA_HOME
53
54
55
56
ROCM_HOME = envs.ROCM_HOME

# Check if both CUDA and ROCM are enabled
if USE_ROCM and not ROCM_HOME:
57
58
    raise ValueError(
        "ROCM support is enabled (USE_ROCM=True) but ROCM_HOME is not set or detected.")
59
60

if not USE_ROCM and not CUDA_HOME:
61
62
    raise ValueError(
        "CUDA support is enabled by default (USE_ROCM=False) but CUDA_HOME is not set or detected.")
63

64
65
# Ensure one of CUDA or ROCM is available
if not (CUDA_HOME or ROCM_HOME):
66
67
68
    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)."
    )
69

70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
# 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


111
112
113
114
115
116
117
118
119
120
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")
121
122
        rocm_version_file = os.path.join(rocm_path, "lib", "cmake", "rocm",
                                         "rocm-config-version.cmake")
123
124
125
126
127
128
129
130
131
132
        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")


133
134
135
136
137
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("-", "."))
138

139
    if with_cuda:
140
141
142
143
144
145
146
147
148
149
150
        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}")

151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
    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
204
    logger.info(f"Downloading {file_name} from {download_url}")
205
206
207
208
209
210
211
212
213
214
215
216
    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
217
    logger.info(f"Extracting {file_name} to {extract_path}")
218
219
220
    with tarfile.open(fileobj=BytesIO(file_content), mode="r:xz") as tar:
        tar.extractall(path=extract_path)

221
    logger.info("Download and extraction completed successfully.")
222
223
224
225
    return os.path.abspath(os.path.join(extract_path, file_name.replace(".tar.xz", "")))


package_data = {
226
    "tilelang": ["py.typed", "*pyx"],
227
228
229
230
231
232
233
234
}

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():
235
236
237
238
239
240
241
242
243
244
245
246
    """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():
247
        logger.info("Info: Not a git repository, skipping submodule update.")
248
249
        return

250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
    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")
265
    # Set LLVM path and enable CUDA or ROCM in config.cmake
266
267
    with open("config.cmake", "a") as config_file:
        config_file.write(f"set(USE_LLVM {llvm_config_path})\n")
268
269
270
271
272
273
        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")
274
275
276
    # Run CMake and make
    try:
        subprocess.check_call(["cmake", ".."])
277
        num_jobs = max(1, int(multiprocessing.cpu_count() * 0.75))
278
279
280
281
282
283
284
285
286
287
288
289
290
        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


291
292
293
294
295
296
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.
    """
297
298
299
300
301
302
303
304
305
    # check if patchelf is installed
    # find patchelf in the system
    patchelf_path = shutil.which("patchelf")
    if not patchelf_path:
        logger.warning(
            "patchelf is not installed, which is required for auditwheel to work for compatible wheels."
        )
        return
    subprocess.run([patchelf_path, '--set-rpath', '$ORIGIN', libpath])
306
307


308
309
310
311
312
313
314
315
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
316
        ext_modules = build_ext_cmd.extensions
317
        for ext in ext_modules:
318
            extdir = build_ext_cmd.get_ext_fullpath(ext.name)
319
            logger.info(f"Extension {ext.name} output directory: {extdir}")
320
321

        ext_output_dir = os.path.dirname(extdir)
322
323
        logger.info(f"Extension output directory (parent): {ext_output_dir}")
        logger.info(f"Build temp directory: {build_temp_dir}")
324

325
326
327
328
329
330
331
332
333
334
335
336
337
338
        # 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)
339
340
341
342
343
344
                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)

345
        # copy the tl_templates
346
347
348
349
350
351
352
353
354
355
356
357
358
359
        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)
360

361
        TVM_PREBUILD_ITEMS = [
362
363
364
365
            "libtvm_runtime.so",
            "libtvm.so",
            "libtilelang.so",
            "libtilelang_module.so",
366
        ]
367

368
369
370
371
372
373
374
        potential_dirs = [
            ext_output_dir,
            self.build_lib,
            build_temp_dir,
            os.path.join(ROOT_DIR, "build"),
        ]

375
        for item in TVM_PREBUILD_ITEMS:
376
377
378
379
380
381
382
383
            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:
384
                patch_libs(source_lib_file)
385
386
387
388
389
                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)
390
                logger.info(f"Copied {source_lib_file} to {target_dir_release}")
391
                shutil.copy2(source_lib_file, target_dir_develop)
392
                logger.info(f"Copied {source_lib_file} to {target_dir_develop}")
393
394
                os.remove(source_lib_file)
            else:
395
                logger.info(f"WARNING: {item} not found in any expected directories!")
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410

        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:
411
                logger.info(f"INFO: {source_dir} does not exist.")
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439

        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 = [
Lei Wang's avatar
Lei Wang committed
440
441
            "3rdparty/cutlass/include",
            "3rdparty/cutlass/tools",
442
443
444
445
446
447
448
449
450
451
452
453
454
455
        ]
        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 = [
Lei Wang's avatar
Lei Wang committed
456
457
            "3rdparty/composable_kernel/include",
            "3rdparty/composable_kernel/library",
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
        ]
        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()


496
497
498
499
500
501
502
503
504
505
# ------------------------------------------------------------------------
# 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):
506
        logger.info("Running TileLangDevelopCommand")
507
508
509
510
511
512
513
        # 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)
514
            logger.info(f"Extension {ext.name} output directory: {extdir}")
515
516

        ext_output_dir = os.path.dirname(extdir)
517
        logger.info(f"Extension output directory (parent): {ext_output_dir}")
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535

        # 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):
536
                patch_libs(source_lib_file)
537
538
539
540
                shutil.copy2(source_lib_file, target_dir)
                # remove the original file
                os.remove(source_lib_file)
            else:
541
                logger.info(f"INFO: {source_lib_file} does not exist.")
542
543


544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
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)

585
586
587
588
589
590
591
592
593
594
595
596
597
        # 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)

598
599
600
601
602
603
    def build_cmake(self, ext):
        """
        Build a single CMake-based extension.

        :param ext: The extension (an instance of CMakeExtension).
        """
604
605
606
607
608
        # 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()
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629

        # 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)

630
        # Append some configuration variables to 'config.cmake'
631
632
        with open(dst_config_cmake, "a") as config_file:
            config_file.write(f"set(USE_LLVM {llvm_config_path})\n")
633
634
635
636
637
638
            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")
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653

        # 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={"": "."},
654
    author="Tile-AI",
655
656
657
658
    description="A tile level programming language to generate high performance code.",
    long_description=read_readme(),
    long_description_content_type="text/markdown",
    platforms=[
659
        "Environment :: GPU :: NVIDIA CUDA" if not USE_ROCM else "Environment :: GPU :: AMD ROCm",
660
661
662
663
        "Operating System :: POSIX :: Linux",
    ],
    license="MIT",
    keywords="BLAS, CUDA, HIP, Code Generation, TVM",
664
    url="https://github.com/tile-ai/tilelang",
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
    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,
681
        "develop": TileLangDevelopCommand,
682
683
    },
)