setup.py 2.75 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
import os
rusty1s's avatar
rusty1s committed
2
import os.path as osp
rusty1s's avatar
rusty1s committed
3
import glob
rusty1s's avatar
rusty1s committed
4
from setuptools import setup, find_packages
rusty1s's avatar
rusty1s committed
5

rusty1s's avatar
cleanup  
rusty1s committed
6
import torch
rusty1s's avatar
rusty1s committed
7
from torch.utils.cpp_extension import BuildExtension
8
from torch.utils.cpp_extension import CppExtension, CUDAExtension, CUDA_HOME
rusty1s's avatar
rusty1s committed
9

10
WITH_CUDA = CUDA_HOME is not None
rusty1s's avatar
rusty1s committed
11
12
if os.getenv('FORCE_CUDA', '0') == '1':
    WITH_CUDA = True
rusty1s's avatar
path  
rusty1s committed
13
if os.getenv('FORCE_CPU', '0') == '1':
rusty1s's avatar
rusty1s committed
14
    WITH_CUDA = False
15

rusty1s's avatar
rusty1s committed
16
17
BUILD_DOCS = os.getenv('BUILD_DOCS', '0') == '1'

rusty1s's avatar
rusty1s committed
18

rusty1s's avatar
rusty1s committed
19
def get_extensions():
20
21
22
23
24
25
26
27
28
29
    extensions = []
    for with_cuda, supername in [
        (False, "cpu"),
        (True, "gpu"),
    ]:
        if with_cuda and not WITH_CUDA:
            continue
        Extension = CppExtension
        define_macros = []
        extra_compile_args = {'cxx': []}
rusty1s's avatar
rusty1s committed
30

31
32
33
34
35
36
37
        if with_cuda:
            Extension = CUDAExtension
            define_macros += [('WITH_CUDA', None)]
            nvcc_flags = os.getenv('NVCC_FLAGS', '')
            nvcc_flags = [] if nvcc_flags == '' else nvcc_flags.split(' ')
            nvcc_flags += ['-arch=sm_35', '--expt-relaxed-constexpr']
            extra_compile_args['nvcc'] = nvcc_flags
rusty1s's avatar
rusty1s committed
38

39
40
41
42
        extensions_dir = osp.join(osp.dirname(osp.abspath(__file__)), 'csrc')
        main_files = glob.glob(osp.join(extensions_dir, '*.cpp'))
        for main in main_files:
            name = main.split(os.sep)[-1][:-4]
rusty1s's avatar
rusty1s committed
43

44
            sources = [main]
rusty1s's avatar
typo  
rusty1s committed
45

46
47
48
            path = osp.join(extensions_dir, 'cpu', f'{name}_cpu.cpp')
            if osp.exists(path):
                sources += [path]
rusty1s's avatar
typo  
rusty1s committed
49

50
51
52
            path = osp.join(extensions_dir, 'cuda', f'{name}_cuda.cu')
            if with_cuda and osp.exists(path):
                sources += [path]
rusty1s's avatar
rusty1s committed
53

54
55
56
57
58
59
60
61
            extension = Extension(
                'torch_scatter._%s_%s' % (name, supername),
                sources,
                include_dirs=[extensions_dir],
                define_macros=define_macros,
                extra_compile_args=extra_compile_args,
            )
            extensions += [extension]
rusty1s's avatar
rusty1s committed
62
63

    return extensions
rusty1s's avatar
rusty1s committed
64

rusty1s's avatar
rusty1s committed
65

rusty1s's avatar
rusty1s committed
66
install_requires = []
rusty1s's avatar
rusty1s committed
67
setup_requires = ['pytest-runner']
rusty1s's avatar
rusty1s committed
68
tests_require = ['pytest', 'pytest-cov']
rusty1s's avatar
rusty1s committed
69

rusty1s's avatar
rusty1s committed
70
setup(
rusty1s's avatar
rusty1s committed
71
    name='torch_scatter',
rusty1s's avatar
rusty1s committed
72
    version='2.0.5',
rusty1s's avatar
rusty1s committed
73
74
    author='Matthias Fey',
    author_email='matthias.fey@tu-dortmund.de',
rusty1s's avatar
rusty1s committed
75
76
77
78
    url='https://github.com/rusty1s/pytorch_scatter',
    description='PyTorch Extension Library of Optimized Scatter Operations',
    keywords=['pytorch', 'scatter', 'segment', 'gather'],
    license='MIT',
rusty1s's avatar
rusty1s committed
79
    python_requires='>=3.6',
rusty1s's avatar
rusty1s committed
80
81
    install_requires=install_requires,
    setup_requires=setup_requires,
rusty1s's avatar
rusty1s committed
82
    tests_require=tests_require,
rusty1s's avatar
rusty1s committed
83
    extras_require={'test': tests_require},
rusty1s's avatar
rusty1s committed
84
    ext_modules=get_extensions() if not BUILD_DOCS else [],
rusty1s's avatar
rusty1s committed
85
    cmdclass={
rusty1s's avatar
rusty1s committed
86
87
        'build_ext':
        BuildExtension.with_options(no_python_abi_suffix=True, use_ninja=False)
rusty1s's avatar
rusty1s committed
88
    },
rusty1s's avatar
rusty1s committed
89
90
    packages=find_packages(),
)