setup.py 5.36 KB
Newer Older
Boris Bonev's avatar
Boris Bonev committed
1
2
# coding=utf-8

Boris Bonev's avatar
Boris Bonev committed
3
# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.
Boris Bonev's avatar
Boris Bonev committed
4
# SPDX-License-Identifier: BSD-3-Clause
Boris Bonev's avatar
Boris Bonev committed
5
#
Boris Bonev's avatar
Boris Bonev committed
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
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
Boris Bonev's avatar
Boris Bonev committed
31

32
33
import os, sys
import warnings
Boris Bonev's avatar
Boris Bonev committed
34

35
36
from setuptools import setup, find_packages
from setuptools.command.install import install
Boris Bonev's avatar
Boris Bonev committed
37

38
39
40
# some code to handle the building of custom modules
FORCE_CUDA_EXTENSION = os.getenv("FORCE_CUDA_EXTENSION", "0") == "1"
BUILD_CPP = BUILD_CUDA = False
Boris Bonev's avatar
Boris Bonev committed
41

42
43
44
# try to import torch
try:
    import torch
Boris Bonev's avatar
Boris Bonev committed
45

46
47
    print(f"setup.py with torch {torch.__version__}")
    from torch.utils.cpp_extension import BuildExtension, CppExtension
Boris Bonev's avatar
Boris Bonev committed
48

49
50
    BUILD_CPP = True
    from torch.utils.cpp_extension import CUDA_HOME, CUDAExtension
Boris Bonev's avatar
Boris Bonev committed
51

52
53
54
    BUILD_CUDA = FORCE_CUDA_EXTENSION or (torch.cuda.is_available() and (CUDA_HOME is not None))
except (ImportError, TypeError, AssertionError, AttributeError) as e:
    warnings.warn(f"building custom extensions skipped: {e}")
Boris Bonev's avatar
Boris Bonev committed
55

Boris Bonev's avatar
Boris Bonev committed
56
def get_compile_args(module_name):
apaaris's avatar
apaaris committed
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
    """
    Get compilation arguments based on environment variables.
    
    If user runs build with TORCH_HARMONICS_DEBUG=1 set, it will use debugging flags to build.
    If TORCH_HARMONICS_PROFILE=1 is set, it will include profiling flags.
    
    Parameters
    -----------
    module_name : str
        Name of the module being compiled
        
    Returns
    -------
    dict
        Dictionary containing compilation flags for 'cxx' and 'nvcc' compilers
    """
Boris Bonev's avatar
Boris Bonev committed
73
    debug_mode = os.environ.get('TORCH_HARMONICS_DEBUG', '0') == '1'
74
75
76
77
78
    profile_mode = os.environ.get('TORCH_HARMONICS_PROFILE', '0') == '1'

    nvcc_extra_flags = []
    if profile_mode:
        nvcc_extra_flags.append("-lineinfo")
79
        nvcc_extra_flags.append("-Xptxas=-v")
80
        
Boris Bonev's avatar
Boris Bonev committed
81
82
83
84
    if debug_mode:
        print(f"WARNING: Compiling {module_name} with debugging flags")
        return {
            'cxx': ['-g', '-O0', '-Wall'],
85
            'nvcc': ['-g', '-G', '-O0'] + nvcc_extra_flags
Boris Bonev's avatar
Boris Bonev committed
86
87
88
89
90
        }
    else:
        print(f"NOTE: Compiling {module_name} with release flags")
        return {
            'cxx': ['-O3', "-DNDEBUG"],
91
            'nvcc': ['-O3', "-DNDEBUG"] + nvcc_extra_flags
Boris Bonev's avatar
Boris Bonev committed
92
93
        }

94
def get_ext_modules():
apaaris's avatar
apaaris committed
95
96
97
98
99
100
101
102
103
    """
    Get list of extension modules to compile.
    
    Returns
    -------
    tuple
        (ext_modules, cmdclass) where ext_modules is a list of extension modules
        and cmdclass is a dictionary of build commands
    """
104
105
    ext_modules = []
    cmdclass = {}
Boris Bonev's avatar
Boris Bonev committed
106

107
108
109
110
    if BUILD_CPP:
        print(f"Compiling helper routines for torch-harmonics.")
        ext_modules.append(CppExtension("disco_helpers", ["torch_harmonics/csrc/disco/disco_helpers.cpp"]))
        cmdclass["build_ext"] = BuildExtension
Boris Bonev's avatar
Boris Bonev committed
111

112
113
    if BUILD_CUDA:
        print(f"Compiling custom CUDA kernels for torch-harmonics.")
Boris Bonev's avatar
Boris Bonev committed
114
        ext_modules.append(
115
            CUDAExtension(
Boris Bonev's avatar
Boris Bonev committed
116
117
118
119
120
121
                "disco_cuda_extension",
                [
                    "torch_harmonics/csrc/disco/disco_interface.cu",
                    "torch_harmonics/csrc/disco/disco_cuda_fwd.cu",
                    "torch_harmonics/csrc/disco/disco_cuda_bwd.cu",
                ],
Boris Bonev's avatar
Boris Bonev committed
122
123
124
125
126
127
128
                extra_compile_args=get_compile_args("disco")
            )
        )
        ext_modules.append(
            CUDAExtension(
                name="attention_cuda_extension",
                sources=[
129
                    "torch_harmonics/csrc/attention/attention_utils.cu",
Boris Bonev's avatar
Boris Bonev committed
130
131
132
133
134
                    "torch_harmonics/csrc/attention/attention_fwd_cuda.cu",
                    "torch_harmonics/csrc/attention/attention_bwd_cuda.cu",
                    "torch_harmonics/csrc/attention/attention_interface.cu",
                ],
                extra_compile_args=get_compile_args("neighborhood_attention")
Boris Bonev's avatar
Boris Bonev committed
135
136
            )
        )
137
138
139
140
141
142
143
        cmdclass["build_ext"] = BuildExtension

    return ext_modules, cmdclass

if __name__ == "__main__":

    ext_modules, cmdclass = get_ext_modules()
Boris Bonev's avatar
Boris Bonev committed
144

145
146
147
148
    setup(
        packages=find_packages(),
        ext_modules=ext_modules,
        cmdclass=cmdclass,
Boris Bonev's avatar
Boris Bonev committed
149
    )