involution/setup.py at 7734b83ae7ff7eeefbb6f8c7b545a4fd03fa93cf · d-li14/involution · 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
import os
from os.path import abspath, dirname, join
from setuptools import setup, find_packages
from torch.utils.cpp_extension import CppExtension, CUDAExtension, BuildExtension

INCLUDE_DIR = join(dirname(abspath(__file__)), 'include')
EXTRA_COMPILE_ARGS = ['-O3']

EXTENSION = []

CC = ['52', '53', '60', '61', '62', '70', '72', '75', '80']

if os.getenv('USE_OPENMP', '1') == '1':
    EXTRA_COMPILE_ARGS.append('-fopenmp')

if os.getenv('USE_CUDA', '1') == '1':
    EXTRA_COMPILE_ARGS.append('-DUSE_CUDA')

    GENERATE_CODES = []

    for cc in CC:
        GENERATE_CODES.append('--generate-code')
        GENERATE_CODES.append(f'arch=compute_{cc},code=compute_{cc}')

    EXTENSION.append(
        CUDAExtension(
            name='involution',
            sources=[
                'src/involution2d_cpu.cpp',
                'src/involution2d_cuda.cu',
                'src/pytorch_wrapper.cpp',
            ],
            include_dirs=[
                INCLUDE_DIR
            ],
            extra_compile_args={
                'cxx': EXTRA_COMPILE_ARGS,
                'nvcc': ['-O3'] + GENERATE_CODES,
            }
        )
    )
else:
    EXTENSION.append(
        CppExtension(
            name='involution',
            sources=[
                'src/involution2d_cpu.cpp',
                'src/pytorch_wrapper.cpp',
            ],
            include_dirs=[
                INCLUDE_DIR
            ],
            extra_compile_args=EXTRA_COMPILE_ARGS
        )
    )

setup(
    name='involution-pytorch',
    version="0.1.0",
    url="https://github.com/shikishima-TasakiLab/Involution-PyTorch",
    license="MIT License",
    author="Junya Shikishima",
    author_email="160442065@ccalumni.meijo-u.ac.jp",
    description="PyTorch Involution",
    packages=find_packages(),
    ext_modules=EXTENSION,
    cmdclass={
        'build_ext': BuildExtension,
    }
)