setup.py 4.3 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
import os
rusty1s's avatar
rusty1s committed
2
import sys
rusty1s's avatar
rusty1s committed
3
import glob
rusty1s's avatar
rusty1s committed
4
5
import os.path as osp
from itertools import product
rusty1s's avatar
rusty1s committed
6
from setuptools import setup, find_packages
rusty1s's avatar
rusty1s committed
7

rusty1s's avatar
rusty1s committed
8
import torch
rusty1s's avatar
rusty1s committed
9
from torch.__config__ import parallel_info
rusty1s's avatar
rusty1s committed
10
from torch.utils.cpp_extension import BuildExtension
rusty1s's avatar
rusty1s committed
11
from torch.utils.cpp_extension import CppExtension, CUDAExtension, CUDA_HOME
rusty1s's avatar
rusty1s committed
12

rusty1s's avatar
rusty1s committed
13
WITH_CUDA = torch.cuda.is_available() and CUDA_HOME is not None
rusty1s's avatar
rusty1s committed
14
suffices = ['cpu', 'cuda'] if WITH_CUDA else ['cpu']
rusty1s's avatar
rusty1s committed
15
if os.getenv('FORCE_CUDA', '0') == '1':
rusty1s's avatar
rusty1s committed
16
17
18
19
20
    suffices = ['cuda', 'cpu']
if os.getenv('FORCE_ONLY_CUDA', '0') == '1':
    suffices = ['cuda']
if os.getenv('FORCE_ONLY_CPU', '0') == '1':
    suffices = ['cpu']
rusty1s's avatar
rusty1s committed
21
22
23

BUILD_DOCS = os.getenv('BUILD_DOCS', '0') == '1'

rusty1s's avatar
rusty1s committed
24
25
WITH_METIS = True if os.getenv('WITH_METIS', '0') == '1' else False
WITH_MTMETIS = True if os.getenv('WITH_MTMETIS', '0') == '1' else False
rusty1s's avatar
rusty1s committed
26

rusty1s's avatar
rusty1s committed
27
28

def get_extensions():
rusty1s's avatar
rusty1s committed
29
    extensions = []
rusty1s's avatar
rusty1s committed
30

rusty1s's avatar
rusty1s committed
31
    extensions_dir = osp.join('csrc')
rusty1s's avatar
rusty1s committed
32
33
    main_files = glob.glob(osp.join(extensions_dir, '*.cpp'))

rusty1s's avatar
rusty1s committed
34
35
36
37
38
39
40
41
42
43
44
45
46
47
    for main, suffix in product(main_files, suffices):
        define_macros = []
        libraries = []
        if WITH_METIS:
            define_macros += [('WITH_METIS', None)]
            libraries += ['metis']
        if WITH_MTMETIS:
            define_macros += [('WITH_MTMETIS', None)]
            define_macros += [('MTMETIS_64BIT_VERTICES', None)]
            define_macros += [('MTMETIS_64BIT_EDGES', None)]
            define_macros += [('MTMETIS_64BIT_WEIGHTS', None)]
            define_macros += [('MTMETIS_64BIT_PARTITIONS', None)]
            libraries += ['mtmetis', 'wildriver']
        extra_compile_args = {'cxx': ['-O2']}
rusty1s's avatar
rusty1s committed
48
49
        if not os.name == 'nt':  # Not on Windows:
            extra_compile_args['cxx'] += ['-Wno-sign-compare']
rusty1s's avatar
rusty1s committed
50
51
52
        extra_link_args = ['-s']

        info = parallel_info()
rusty1s's avatar
rusty1s committed
53
54
        if ('backend: OpenMP' in info and 'OpenMP not found' not in info
                and sys.platform != 'darwin'):
rusty1s's avatar
rusty1s committed
55
56
57
58
59
60
61
62
63
64
65
66
            extra_compile_args['cxx'] += ['-DAT_PARALLEL_OPENMP']
            if sys.platform == 'win32':
                extra_compile_args['cxx'] += ['/openmp']
            else:
                extra_compile_args['cxx'] += ['-fopenmp']
        else:
            print('Compiling without OpenMP...')

        if suffix == 'cuda':
            define_macros += [('WITH_CUDA', None)]
            nvcc_flags = os.getenv('NVCC_FLAGS', '')
            nvcc_flags = [] if nvcc_flags == '' else nvcc_flags.split(' ')
67
            nvcc_flags += ['--expt-relaxed-constexpr', '-O2']
rusty1s's avatar
rusty1s committed
68
69
70
71
72
73
74
75
            extra_compile_args['nvcc'] = nvcc_flags

            if sys.platform == 'win32':
                extra_link_args += ['cusparse.lib']
            else:
                extra_link_args += ['-lcusparse', '-l', 'cusparse']

        name = main.split(os.sep)[-1][:-4]
rusty1s's avatar
rusty1s committed
76
77
        sources = [main]

rusty1s's avatar
rusty1s committed
78
        path = osp.join(extensions_dir, 'cpu', f'{name}_cpu.cpp')
rusty1s's avatar
rusty1s committed
79
80
81
        if osp.exists(path):
            sources += [path]

rusty1s's avatar
rusty1s committed
82
        path = osp.join(extensions_dir, 'cuda', f'{name}_cuda.cu')
rusty1s's avatar
rusty1s committed
83
        if suffix == 'cuda' and osp.exists(path):
rusty1s's avatar
rusty1s committed
84
            sources += [path]
rusty1s's avatar
rusty1s committed
85

rusty1s's avatar
rusty1s committed
86
        Extension = CppExtension if suffix == 'cpu' else CUDAExtension
rusty1s's avatar
rusty1s committed
87
        extension = Extension(
rusty1s's avatar
rusty1s committed
88
            f'torch_sparse._{name}_{suffix}',
rusty1s's avatar
rusty1s committed
89
90
91
92
            sources,
            include_dirs=[extensions_dir],
            define_macros=define_macros,
            extra_compile_args=extra_compile_args,
93
            extra_link_args=extra_link_args,
rusty1s's avatar
rusty1s committed
94
            libraries=libraries,
rusty1s's avatar
rusty1s committed
95
96
97
98
99
        )
        extensions += [extension]

    return extensions

rusty1s's avatar
rusty1s committed
100
101

install_requires = ['scipy']
rusty1s's avatar
rusty1s committed
102
103
setup_requires = []
tests_require = ['pytest', 'pytest-runner', 'pytest-cov']
rusty1s's avatar
rusty1s committed
104
105
106

setup(
    name='torch_sparse',
rusty1s's avatar
rusty1s committed
107
    version='0.6.12',
rusty1s's avatar
rusty1s committed
108
109
    author='Matthias Fey',
    author_email='matthias.fey@tu-dortmund.de',
rusty1s's avatar
rusty1s committed
110
111
112
    url='https://github.com/rusty1s/pytorch_sparse',
    description=('PyTorch Extension Library of Optimized Autograd Sparse '
                 'Matrix Operations'),
rusty1s's avatar
rusty1s committed
113
    keywords=['pytorch', 'sparse', 'sparse-matrices', 'autograd'],
rusty1s's avatar
rusty1s committed
114
    license='MIT',
rusty1s's avatar
rusty1s committed
115
    python_requires='>=3.6',
rusty1s's avatar
rusty1s committed
116
117
118
    install_requires=install_requires,
    setup_requires=setup_requires,
    tests_require=tests_require,
rusty1s's avatar
rusty1s committed
119
    extras_require={'test': tests_require},
rusty1s's avatar
rusty1s committed
120
121
    ext_modules=get_extensions() if not BUILD_DOCS else [],
    cmdclass={
rusty1s's avatar
rusty1s committed
122
123
        'build_ext':
        BuildExtension.with_options(no_python_abi_suffix=True, use_ninja=False)
rusty1s's avatar
rusty1s committed
124
    },
125
126
    packages=find_packages(),
)