setup.py 3.66 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
import glob
Matthias Fey's avatar
Matthias Fey committed
2
import os
rusty1s's avatar
update  
rusty1s committed
3
import os.path as osp
Amit Aflalo's avatar
Amit Aflalo committed
4
import platform
Matthias Fey's avatar
Matthias Fey committed
5
6
import sys
from itertools import product
rusty1s's avatar
rusty1s committed
7

rusty1s's avatar
cleanup  
rusty1s committed
8
import torch
Matthias Fey's avatar
Matthias Fey committed
9
from setuptools import find_packages, setup
rusty1s's avatar
update  
rusty1s committed
10
from torch.__config__ import parallel_info
Matthias Fey's avatar
Matthias Fey committed
11
12
13
14
15
from torch.utils.cpp_extension import (CUDA_HOME, BuildExtension, CppExtension,
                                       CUDAExtension)

__version__ = '2.0.9'
URL = 'https://github.com/rusty1s/pytorch_scatter'
rusty1s's avatar
rusty1s committed
16

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

rusty1s's avatar
rusty1s committed
26
27
BUILD_DOCS = os.getenv('BUILD_DOCS', '0') == '1'

rusty1s's avatar
rusty1s committed
28

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

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

    for main, suffix in product(main_files, suffices):
36
        define_macros = []
rusty1s's avatar
rusty1s committed
37
38
39
        extra_compile_args = {'cxx': ['-O2']}
        if not os.name == 'nt':  # Not on Windows:
            extra_compile_args['cxx'] += ['-Wno-sign-compare']
rusty1s's avatar
update  
rusty1s committed
40
41
42
        extra_link_args = ['-s']

        info = parallel_info()
rusty1s's avatar
rusty1s committed
43
44
        if ('backend: OpenMP' in info and 'OpenMP not found' not in info
                and sys.platform != 'darwin'):
rusty1s's avatar
update  
rusty1s committed
45
46
47
48
49
50
51
            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...')
rusty1s's avatar
rusty1s committed
52

Amit Aflalo's avatar
Amit Aflalo committed
53
54
55
56
57
        # Compile for mac arm64
        if (sys.platform == 'darwin' and platform.machine() == 'arm64'):
            extra_compile_args['cxx'] += ['-arch', 'arm64']
            extra_link_args += ['-arch', 'arm64']

rusty1s's avatar
update  
rusty1s committed
58
        if suffix == 'cuda':
59
60
61
            define_macros += [('WITH_CUDA', None)]
            nvcc_flags = os.getenv('NVCC_FLAGS', '')
            nvcc_flags = [] if nvcc_flags == '' else nvcc_flags.split(' ')
rusty1s's avatar
rusty1s committed
62
            nvcc_flags += ['--expt-relaxed-constexpr', '-O2']
63
            extra_compile_args['nvcc'] = nvcc_flags
rusty1s's avatar
rusty1s committed
64

rusty1s's avatar
update  
rusty1s committed
65
66
        name = main.split(os.sep)[-1][:-4]
        sources = [main]
rusty1s's avatar
typo  
rusty1s committed
67

rusty1s's avatar
update  
rusty1s committed
68
69
70
        path = osp.join(extensions_dir, 'cpu', f'{name}_cpu.cpp')
        if osp.exists(path):
            sources += [path]
rusty1s's avatar
typo  
rusty1s committed
71

rusty1s's avatar
update  
rusty1s committed
72
73
74
        path = osp.join(extensions_dir, 'cuda', f'{name}_cuda.cu')
        if suffix == 'cuda' and osp.exists(path):
            sources += [path]
rusty1s's avatar
rusty1s committed
75

rusty1s's avatar
update  
rusty1s committed
76
77
78
79
80
81
82
83
84
85
        Extension = CppExtension if suffix == 'cpu' else CUDAExtension
        extension = Extension(
            f'torch_scatter._{name}_{suffix}',
            sources,
            include_dirs=[extensions_dir],
            define_macros=define_macros,
            extra_compile_args=extra_compile_args,
            extra_link_args=extra_link_args,
        )
        extensions += [extension]
rusty1s's avatar
rusty1s committed
86
87

    return extensions
rusty1s's avatar
rusty1s committed
88

rusty1s's avatar
rusty1s committed
89

rusty1s's avatar
rusty1s committed
90
install_requires = []
Matthias Fey's avatar
Matthias Fey committed
91
92
93
94
95

test_require = [
    'pytest',
    'pytest-cov',
]
rusty1s's avatar
rusty1s committed
96

rusty1s's avatar
rusty1s committed
97
setup(
rusty1s's avatar
rusty1s committed
98
    name='torch_scatter',
Matthias Fey's avatar
Matthias Fey committed
99
100
    version=__version__,
    description='PyTorch Extension Library of Optimized Scatter Operations',
rusty1s's avatar
rusty1s committed
101
102
    author='Matthias Fey',
    author_email='matthias.fey@tu-dortmund.de',
Matthias Fey's avatar
Matthias Fey committed
103
104
    url=URL,
    download_url=f'{URL}/archive/{__version__}.tar.gz',
rusty1s's avatar
rusty1s committed
105
    keywords=['pytorch', 'scatter', 'segment', 'gather'],
Matthias Fey's avatar
Matthias Fey committed
106
    python_requires='>=3.7',
rusty1s's avatar
rusty1s committed
107
    install_requires=install_requires,
Matthias Fey's avatar
Matthias Fey committed
108
109
110
    extras_require={
        'test': test_require,
    },
rusty1s's avatar
rusty1s committed
111
    ext_modules=get_extensions() if not BUILD_DOCS else [],
rusty1s's avatar
rusty1s committed
112
    cmdclass={
rusty1s's avatar
rusty1s committed
113
114
        'build_ext':
        BuildExtension.with_options(no_python_abi_suffix=True, use_ninja=False)
rusty1s's avatar
rusty1s committed
115
    },
rusty1s's avatar
rusty1s committed
116
    packages=find_packages(),
Matthias Fey's avatar
Matthias Fey committed
117
    include_package_data=True,
rusty1s's avatar
rusty1s committed
118
)