forked from NVlabs/curobo
-
Notifications
You must be signed in to change notification settings - Fork 0
/
setup.py
97 lines (89 loc) · 2.95 KB
/
setup.py
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
#
# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
#
"""curobo package setuptools."""
# NOTE: This file is still needed to allow the package to be
# installed in editable mode.
#
# References:
# * https://setuptools.pypa.io/en/latest/setuptools.html#setup-cfg-only-projects
# Standard Library
import sys
# Third Party
import setuptools
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
print("*********************************")
print("CuRobo: Compiling CUDA kernels...")
print(
"This will take 20+ minutes, to speedup compilation set TORCH_CUDA_ARCH_LIST={X}+PTX to "
+ " compile only for a specific architecture. e.g., export TORCH_CUDA_ARCH_LIST=7.0+PTX"
)
extra_cuda_args = {
"nvcc": [
"--threads=6",
"-O3",
"--ftz=true",
"--fmad=true",
"--prec-div=false",
"--prec-sqrt=false",
]
}
if sys.platform == "win32":
extra_cuda_args["nvcc"].append("--allow-unsupported-compiler")
# create a list of modules to be compiled:
ext_modules = [
CUDAExtension(
"curobo.curobolib.lbfgs_step_cu",
[
"src/curobo/curobolib/cpp/lbfgs_step_cuda.cpp",
"src/curobo/curobolib/cpp/lbfgs_step_kernel.cu",
],
extra_compile_args=extra_cuda_args,
),
CUDAExtension(
"curobo.curobolib.kinematics_fused_cu",
[
"src/curobo/curobolib/cpp/kinematics_fused_cuda.cpp",
"src/curobo/curobolib/cpp/kinematics_fused_kernel.cu",
],
extra_compile_args=extra_cuda_args,
),
CUDAExtension(
"curobo.curobolib.geom_cu",
[
"src/curobo/curobolib/cpp/geom_cuda.cpp",
"src/curobo/curobolib/cpp/sphere_obb_kernel.cu",
"src/curobo/curobolib/cpp/pose_distance_kernel.cu",
"src/curobo/curobolib/cpp/self_collision_kernel.cu",
],
extra_compile_args=extra_cuda_args,
),
CUDAExtension(
"curobo.curobolib.line_search_cu",
[
"src/curobo/curobolib/cpp/line_search_cuda.cpp",
"src/curobo/curobolib/cpp/line_search_kernel.cu",
"src/curobo/curobolib/cpp/update_best_kernel.cu",
],
extra_compile_args=extra_cuda_args,
),
CUDAExtension(
"curobo.curobolib.tensor_step_cu",
[
"src/curobo/curobolib/cpp/tensor_step_cuda.cpp",
"src/curobo/curobolib/cpp/tensor_step_kernel.cu",
],
extra_compile_args=extra_cuda_args,
),
]
setuptools.setup(
ext_modules=ext_modules,
cmdclass={"build_ext": BuildExtension},
)