ExpertRT/setup.py at main · Sunnylincc/ExpertRT · 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
from pathlib import Path
from setuptools import setup, find_packages

import torch
from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtension, CUDA_HOME

ROOT = Path(__file__).parent

sources = [
    "src/bindings.cpp",
    "src/routing.cpp",
    "src/dispatch.cpp",
    "src/combine.cpp",
]

extra_compile_args = {"cxx": ["-O3", "-fvisibility=hidden", "-std=c++17"]}

has_cuda_toolchain = CUDA_HOME is not None and Path(CUDA_HOME).exists()
if has_cuda_toolchain:
    ext = CUDAExtension(
        "expertrt._C",
        sources=sources + [
            "cuda/routing_kernels.cu",
            "cuda/dispatch_kernels.cu",
            "cuda/combine_kernels.cu",
        ],
        include_dirs=["include", "cuda"],
        extra_compile_args={
            "cxx": extra_compile_args["cxx"],
            "nvcc": ["-O3", "--use_fast_math", "-lineinfo"],
        },
    )
else:
    ext = CppExtension(
        "expertrt._C",
        sources=sources,
        include_dirs=["include"],
        extra_compile_args=extra_compile_args["cxx"],
    )

setup(
    name="expertrt",
    version="0.1.0",
    package_dir={"": "python"},
    packages=find_packages(where="python"),
    ext_modules=[ext],
    cmdclass={"build_ext": BuildExtension.with_options(no_python_abi_suffix=True)},
)