InvariantBitPacking/setup.py at devel · AKKamath/InvariantBitPacking · 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
import os
from pathlib import Path
from packaging.version import parse, Version

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

# ninja build does not work unless include_dirs are abs path
this_dir = os.path.dirname(os.path.abspath(__file__))

PACKAGE_NAME = "ibp"

cmdclass = {}
ext_modules = []

cc_flag = []
cc_flag.append("-gencode")
cc_flag.append("arch=compute_70,code=sm_70")
cc_flag.append("-gencode")
cc_flag.append("arch=compute_80,code=sm_80")
#cc_flag.append("-gencode")
#cc_flag.append("arch=compute_90,code=sm_90")

repo_dir = Path(this_dir).parent
sources = [
    "src/ibp_api.cu",
    "src/preprocess.cu",
    "src/compress.cu",
    "src/decompress.cu",
]
include_dirs = [
    "include/",
    "include/preproc/",
    "include/misc/",
    "include/compress/",
    os.path.join(os.environ["CONDA_PREFIX"], "include"),
]
include_files = [
    "include/ibp_helpers.cuh",
    "include/decompress/ibp_decompress_kernel.cuh",
    "include/decompress/ibp_decompress_host.cuh",
    "include/decompress/ibp_decompress_dev.cuh",
    "include/preproc/ibp_preproc_kmeans.cuh",
    "include/preproc/ibp_preproc_host.cuh",
    "include/preproc/ibp_preproc_kernels.cuh",
    "include/ibp_dev_func.cuh",
    "include/misc/ibp_misc_dev.cuh",
    "include/misc/ibp_misc_kernels.cuh",
    "include/compress/ibp_compress_dev.cuh",
    "include/compress/ibp_compress_kernel.cuh",
    "include/compress/ibp_compress_host.cuh",
]
nvcc_flags = [
    "-O3",
    "-std=c++17",
    "--use_fast_math",
    "-lineinfo",
]

ext_modules.append(
    CUDAExtension(
        name="ibp_cuda",
        sources=sources,
        include_dirs=[Path(this_dir) / i for i in include_dirs],
        depends=sources + include_files,
        extra_compile_args={
            "cxx": ["-O3", "-std=c++17"],
            "nvcc": nvcc_flags + cc_flag,
        },
    )
)

setup(
    name=PACKAGE_NAME,
    packages=find_packages(
        exclude=(
            "build",
            "include",
            "ibp.egg-info",
            "ibp_cuda_test",
        )
    ),
    description="Invariant Bit Packing",
    ext_modules=ext_modules,
    cmdclass={"build_ext": BuildExtension},
    python_requires=">=3.8",
    depends=sources + include_files,
    install_requires=[
        "packaging",
        "ninja",
        "torch",
    ],
)