flashinfer/custom_backend.py at main · xiaopoc/flashinfer · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
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
111
112
113
import shutil
from pathlib import Path

from setuptools import build_meta as orig

_root = Path(__file__).parent.resolve()
_data_dir = _root / "flashinfer" / "data"
_aot_ops_dir = _root / "aot-ops"
_aot_ops_package_dir = _root / "build" / "aot-ops-package-dir"

_requires_for_aot = ["torch", "ninja", "numpy"]


def _rm_aot_ops_package_dir():
    if _aot_ops_package_dir.is_symlink():
        _aot_ops_package_dir.unlink()
    elif _aot_ops_package_dir.exists():
        shutil.rmtree(_aot_ops_package_dir)


def _create_data_dir():
    _data_dir.mkdir(parents=True, exist_ok=True)

    def ln(source: str, target: str) -> None:
        src = _root / source
        dst = _data_dir / target
        if dst.exists():
            if dst.is_symlink():
                dst.unlink()
            elif dst.is_dir():
                dst.rmdir()
        dst.symlink_to(src, target_is_directory=True)

    ln("3rdparty/cutlass", "cutlass")
    ln("3rdparty/spdlog", "spdlog")
    ln("csrc", "csrc")
    ln("include", "include")
    ln("tvm_binding", "tvm_binding")


def _prepare_for_wheel():
    # Remove data directory
    if _data_dir.exists():
        shutil.rmtree(_data_dir)

    # Link AOT ops directory to "aot-ops"
    _rm_aot_ops_package_dir()
    if not _aot_ops_dir.exists():
        _aot_ops_dir.mkdir()
    num_ops = len(list(_aot_ops_dir.glob("*/*.so")))
    print(f"{num_ops} AOT ops found in {_aot_ops_dir}")
    _aot_ops_package_dir.parent.mkdir(parents=True, exist_ok=True)
    _aot_ops_package_dir.symlink_to(_aot_ops_dir)


def _prepare_for_editable():
    _create_data_dir()

    _rm_aot_ops_package_dir()
    _aot_ops_dir.mkdir(parents=True, exist_ok=True)
    _aot_ops_package_dir.parent.mkdir(parents=True, exist_ok=True)
    _aot_ops_package_dir.symlink_to(_aot_ops_dir)


def _prepare_for_sdist():
    # Remove data directory
    if _data_dir.exists():
        shutil.rmtree(_data_dir)

    # Create an empty directory for AOT ops
    _rm_aot_ops_package_dir()
    _aot_ops_package_dir.parent.mkdir(parents=True, exist_ok=True)
    _aot_ops_package_dir.mkdir(parents=True)


def get_requires_for_build_wheel(config_settings=None):
    _prepare_for_wheel()
    return _requires_for_aot


def get_requires_for_build_sdist(config_settings=None):
    _prepare_for_sdist()
    return []


def get_requires_for_build_editable(config_settings=None):
    _prepare_for_editable()
    return _requires_for_aot


def prepare_metadata_for_build_wheel(metadata_directory, config_settings=None):
    _prepare_for_wheel()
    return orig.prepare_metadata_for_build_wheel(metadata_directory, config_settings)


def prepare_metadata_for_build_editable(metadata_directory, config_settings=None):
    _prepare_for_editable()
    return orig.prepare_metadata_for_build_editable(metadata_directory, config_settings)


def build_editable(wheel_directory, config_settings=None, metadata_directory=None):
    _prepare_for_editable()
    return orig.build_editable(wheel_directory, config_settings, metadata_directory)


def build_sdist(sdist_directory, config_settings=None):
    _prepare_for_sdist()
    return orig.build_sdist(sdist_directory, config_settings)


def build_wheel(wheel_directory, config_settings=None, metadata_directory=None):
    _prepare_for_wheel()
    return orig.build_wheel(wheel_directory, config_settings, metadata_directory)