torch-infini/setup.py at master · InfiniTensor/torch-infini · 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
from pathlib import Path
import os

from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDA_HOME


PACKAGE_ROOT = Path(__file__).resolve().parent


def _split_paths(value):
    return [path for path in value.split(os.pathsep) if path]


def _cuda_include_dirs():
    # Some InfiniRT installations expose CUDA headers from their public API.
    # Keep this transitive dependency isolated until InfiniRT exports it.
    include_dirs = []
    include_dirs.extend(_split_paths(os.environ.get("CUDA_INCLUDE_DIRS", "")))

    for root in (CUDA_HOME, os.environ.get("CUDA_HOME"), os.environ.get("CUDA_PATH")):
        if not root:
            continue
        root_path = Path(root)
        include_dirs.append(str(root_path / "include"))
        include_dirs.append(str(root_path / "targets" / "x86_64-linux" / "include"))

    for path in (
        Path("/usr/local/cuda/include"),
        Path("/usr/local/cuda/targets/x86_64-linux/include"),
    ):
        include_dirs.append(str(path))

    existing_dirs = []
    seen = set()
    for path in include_dirs:
        if path in seen or not Path(path).exists():
            continue
        seen.add(path)
        existing_dirs.append(path)
    return existing_dirs


def _infini_paths():
    include_dirs = []
    library_dirs = []
    runtime_library_dirs = []

    prefix = os.environ.get("INFINI_RT_PREFIX")
    if prefix:
        prefix_path = Path(prefix)
        include_dirs.append(str(prefix_path / "include"))
        library_dirs.extend(
            str(path)
            for path in (prefix_path / "lib", prefix_path / "lib64")
            if path.exists()
        )
        runtime_library_dirs.extend(library_dirs)

    include_dirs.extend(_split_paths(os.environ.get("INFINI_RT_INCLUDE_DIRS", "")))
    library_dirs.extend(_split_paths(os.environ.get("INFINI_RT_LIBRARY_DIRS", "")))
    runtime_library_dirs.extend(
        _split_paths(os.environ.get("INFINI_RT_RUNTIME_LIBRARY_DIRS", ""))
    )

    if not include_dirs:
        raise RuntimeError(
            "InfiniRT headers were not found. Set INFINI_RT_PREFIX to an "
            "installed InfiniRT prefix, or set INFINI_RT_INCLUDE_DIRS."
        )

    return include_dirs, library_dirs, runtime_library_dirs


include_dirs, library_dirs, runtime_library_dirs = _infini_paths()
include_dirs.extend(_cuda_include_dirs())

setup(
    packages=["torch_infini"],
    ext_modules=[
        CppExtension(
            name="torch_infini._C",
            sources=[
                "csrc/allocator.cpp",
                "csrc/copy.cpp",
                "csrc/device_guard.cpp",
                "csrc/empty.cpp",
                "csrc/init.cpp",
                "csrc/runtime.cpp",
            ],
            include_dirs=[*include_dirs, str(PACKAGE_ROOT / "csrc")],
            libraries=["infinirt"],
            library_dirs=library_dirs,
            runtime_library_dirs=runtime_library_dirs,
            extra_compile_args={"cxx": ["-std=c++17"]},
        )
    ],
    cmdclass={"build_ext": BuildExtension},
)