curobo/setup.py at main · NVlabs/curobo · 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
100
101
102
103
104
105
106
107
108
109
# SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#

"""curobo package setuptools.

This file handles optional CUDA extension compilation via pybind11.
All package metadata is now in pyproject.toml following PEP 621.
"""

# Standard Library
import os
import sys
from pathlib import Path

# Third Party
import setuptools

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

# Check if user wants to use pybind compiled extensions. This is deprecated.
# Default: False (use cuda.core runtime compilation)
# Enable with: CUROBO_USE_PYBIND=1 pip install -e . --no-build-isolation
USE_PYBIND = os.environ.get("CUROBO_USE_PYBIND", "0") == "1"
ext_modules = []
cmdclass = {}

if USE_PYBIND:
    # Only import torch if compiling
    try:
        from torch.utils.cpp_extension import BuildExtension, CUDAExtension

        extra_cuda_args = {
            "nvcc": [
                "--threads=8",
                "-O3",
                "--ftz=true",
                "--fmad=true",
                "--prec-div=false",
                "--prec-sqrt=false",
                "--generate-line-info",
            ]
        }

        if sys.platform == "win32":
            extra_cuda_args["nvcc"].append("--allow-unsupported-compiler")

        # Kernel headers live in nested subdirs (e.g. optimization/lbfgs/) but the .cu
        # files #include them by bare filename, so every subdir must be on the include path.
        _kernels_root = source_dir / "curobo/_src/curobolib/kernels"
        kernel_include_dirs = [str(_kernels_root)] + [
            str(p) for p in _kernels_root.rglob("*") if p.is_dir()
        ]

        # Create a list of modules to be compiled
        ext_modules = [
            CUDAExtension(
                "curobo._src.curobolib.backends.pybind.kinematics",
                [
                    "curobo/_src/curobolib/backends/pybind/kinematics_bindings.cpp",
                    "curobo/_src/curobolib/backends/pybind/kinematics_forward_kernel_launch.cu",
                    "curobo/_src/curobolib/backends/pybind/kinematics_backward_kernel_launch.cu",
                ],
                include_dirs=kernel_include_dirs,
                extra_compile_args=extra_cuda_args,
            ),
            CUDAExtension(
                "curobo._src.curobolib.backends.pybind.optimization",
                [
                    "curobo/_src/curobolib/backends/pybind/optimization_bindings.cpp",
                    "curobo/_src/curobolib/backends/pybind/line_search_kernel_launch.cu",
                    "curobo/_src/curobolib/backends/pybind/lbfgs_step_kernel_launch.cu",
                ],
                include_dirs=kernel_include_dirs,
                extra_compile_args=extra_cuda_args,
            ),
            CUDAExtension(
                "curobo._src.curobolib.backends.pybind.trajectory",
                [
                    "curobo/_src/curobolib/backends/pybind/trajectory_bindings.cpp",
                    "curobo/_src/curobolib/backends/pybind/trajectory_kernel_launch.cu",
                ],
                include_dirs=kernel_include_dirs,
                extra_compile_args=extra_cuda_args,
            ),
            CUDAExtension(
                "curobo._src.curobolib.backends.pybind.geometry",
                [
                    "curobo/_src/curobolib/backends/pybind/geometry_bindings.cpp",
                    "curobo/_src/curobolib/backends/pybind/self_collision_kernel_launch.cu",
                ],
                include_dirs=kernel_include_dirs,
                extra_compile_args=extra_cuda_args,
            ),
        ]

        cmdclass = {"build_ext": BuildExtension}

    except ImportError as e:
        print(f"Warning: Could not import torch for compilation: {e}")
        print("Installing without compiled extensions - will use cuda.core backend")


# All package metadata is now in pyproject.toml (PEP 621)
# This setup.py only handles optional CUDA extension compilation
setuptools.setup(
    ext_modules=ext_modules,
    cmdclass=cmdclass,
)