CuMesh/setup.py at main · visualbruno/CuMesh · 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
from setuptools import setup
from torch.utils.cpp_extension import CUDAExtension, BuildExtension, IS_HIP_EXTENSION
import os

ROOT = os.path.dirname(os.path.abspath(__file__))
BUILD_TARGET = os.environ.get("BUILD_TARGET", "auto")

if BUILD_TARGET == "auto":
    if IS_HIP_EXTENSION:
        IS_HIP = True
    else:
        IS_HIP = False
else:
    if BUILD_TARGET == "cuda":
        IS_HIP = False
    elif BUILD_TARGET == "rocm":
        IS_HIP = True

if not IS_HIP:
    cc_flag = [f"-allow-unsupported-compiler"]
else:
    archs = os.getenv("GPU_ARCHS", "native").split(";")
    cc_flag = [f"--offload-arch={arch}" for arch in archs]

setup(
    name="cumesh",
    packages=[
        'cumesh',
    ],
    ext_modules=[
        CUDAExtension(
            name="cumesh._C",
            sources=[
                # Hashmap functions
                "src/hash/hash.cu",

                # CuMesh
                "src/atlas.cu",
                "src/clean_up.cu",
                "src/cumesh.cu",
                "src/connectivity.cu",
                "src/geometry.cu",
                "src/io.cu",
                "src/simplify.cu",
                "src/shared.cu",

                # Remeshing
                "src/remesh/simple_dual_contour.cu",
                "src/remesh/svox2vert.cu",

                # main
                "src/ext.cpp",
            ],
            extra_compile_args={
                "cxx": ["-O3", "-std=c++20"],
                "nvcc": ["-O3","-std=c++20"] + cc_flag,
            }
        ),
        CUDAExtension(
            name='cumesh._cubvh',
            sources=[
                'third_party/cubvh/src/bvh.cu',
                'third_party/cubvh/src/api_gpu.cu',
                'third_party/cubvh/src/bindings.cpp',
            ],
            include_dirs=[
                os.path.join(ROOT, "third_party/cubvh/include"),
                os.path.join(ROOT, "third_party/cubvh/third_party/eigen"),
            ],
            extra_compile_args={
                "cxx": ["-O3", "-std=c++20"],
                "nvcc": ["-O3","-std=c++20"] + cc_flag + [
                    "--extended-lambda",
                    "--expt-relaxed-constexpr",
                    # The following definitions must be undefined
                    # since we need half-precision operation.
                    "-U__CUDA_NO_HALF_OPERATORS__",
                    "-U__CUDA_NO_HALF_CONVERSIONS__",
                    "-U__CUDA_NO_HALF2_OPERATORS__",
                ]
            }
        ),
        CUDAExtension(
            name='cumesh._xatlas',
            sources=[
                'third_party/xatlas/xatlas_mod.cpp',
                'third_party/xatlas/binding.cpp',
            ],
            extra_compile_args={
                "cxx": ["-O3", "-std=c++20"],
            }
        ),
    ],
    cmdclass={
        'build_ext': BuildExtension
    },
)