bae/setup.py at release · cre185/bae · 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
import os
from setuptools import setup, find_packages
from torch.utils.cpp_extension import CppExtension, CUDAExtension, BuildExtension

VERSION = "0.1"

def readme():
    """Read the README.md file for long description"""
    try:
        with open("README.md", "r", encoding="utf-8") as f:
            return f.read()
    except:
        return "PyTorch implementation of BA"

# Check if CUDSS should be enabled
USE_CUDSS = os.environ.get("USE_CUDSS", "1").lower() in ("1", "true", "yes", "y")
# Optionally get CUDSS directory from environment variable
CUDSS_DIR = os.environ.get("CUDSS_DIR", "")

if __name__ == '__main__':
    # Common extensions
    ext_modules = [
        CppExtension(
            'bae.sparse.bsr',
            [os.path.join('bae', 'sparse', 'sparse_op_cpp.cpp')]
        ),
        CUDAExtension(
            'bae.sparse.bsr_cuda',
            [
                os.path.join('bae', 'sparse', 'sparse_op_cuda.cpp'),
                os.path.join('bae', 'sparse', 'sparse_op_cuda_kernel.cu')
            ]
        ),
        CUDAExtension(
            'bae.sparse.spgemm',
            [os.path.join('bae', 'sparse', 'cusparse_wrapper.cpp')]
        ),
        CUDAExtension(
            'bae.sparse.conversion',
            [os.path.join('bae', 'sparse', 'sparse_conversion.cu')]
        ),
    ]

    # Add CUDSS-dependent extension conditionally
    if USE_CUDSS:
        libraries = ['cusolver', 'cusparse', 'cudss']
        extra_compile_args = {
            'nvcc': [
                '-lcusolver',
                '-lcusparse',
                '-lcudss',
            ]
        }

        # Add CUDSS directory paths if specified
        if CUDSS_DIR:
            extra_compile_args['nvcc'].extend([
                f'-I{CUDSS_DIR}/include',
                f'-L{CUDSS_DIR}/lib',
                f'-Xlinker={CUDSS_DIR}/lib/libcudss_static.a',
            ])

        ext_modules.append(
            CUDAExtension(
                'bae.sparse.solve',
                [os.path.join('bae', 'sparse', 'sparse_cusolve.cu')],
                libraries=libraries,
                extra_compile_args=extra_compile_args,
                include_dirs=['/usr/include/libcudss/12/'],
                library_dirs=['/usr/lib/x86_64-linux-gnu/libcudss/12/']
            )
        )

    setup(
        name = 'bae',
        version = VERSION,
        description = 'PyTorch implementation of BA',
        long_description = readme(),
        long_description_content_type = "text/markdown",
        python_requires = ">=3.8",
        install_requires=[
            'torch',
            'torchvision',
            'warp-lang',
        ],
        packages=find_packages(exclude=['./ba_example.py',
                                        './setup.py',
                                        './README.md',
                                        './data',
                                        './1dsfm_bal',
                                        './bal_data',
                                        './colmap_helpers',
                                        './datapipes',
                                        './examples',
                                        './tests',]),
        ext_modules=ext_modules,
        cmdclass={'build_ext': BuildExtension}
    )