setup.py 2.32 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
import platform
rusty1s's avatar
rusty1s committed
2
from setuptools import setup, find_packages
Dmitry Leonov's avatar
Dmitry Leonov committed
3
from sys import argv
rusty1s's avatar
rusty1s committed
4
5
import torch
from torch.utils.cpp_extension import CppExtension, CUDAExtension, CUDA_HOME
rusty1s's avatar
rusty1s committed
6

rusty1s's avatar
rusty1s committed
7
8
9
10
11
12
13
14
15
16
17
TORCH_MAJOR = int(torch.__version__.split('.')[0])
TORCH_MINOR = int(torch.__version__.split('.')[1])

extra_compile_args = []
if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 2):
    extra_compile_args += ['-DVERSION_GE_1_3']

ext_modules = [
    CppExtension('torch_sparse.spspmm_cpu', ['cpu/spspmm.cpp'],
                 extra_compile_args=extra_compile_args)
]
rusty1s's avatar
rusty1s committed
18
cmdclass = {'build_ext': torch.utils.cpp_extension.BuildExtension}
rusty1s's avatar
rusty1s committed
19

Dmitry Leonov's avatar
Dmitry Leonov committed
20
21
22
23
24
25
26
GPU = True
for arg in argv:
    if arg == '--cpu':
        GPU = False
        argv.remove(arg)

if CUDA_HOME is not None and GPU:
rusty1s's avatar
rusty1s committed
27
    if platform.system() == 'Windows':
28
        extra_link_args = ['cusparse.lib']
rusty1s's avatar
rusty1s committed
29
    else:
rusty1s's avatar
rusty1s committed
30
        extra_link_args = ['-lcusparse', '-l', 'cusparse']
rusty1s's avatar
rusty1s committed
31

rusty1s's avatar
to csr  
rusty1s committed
32
    ext_modules += [
rusty1s's avatar
rusty1s committed
33
34
35
36
        CUDAExtension('torch_sparse.spmm_cuda',
                      ['cuda/spmm.cpp', 'cuda/spmm_kernel.cu'],
                      extra_link_args=extra_link_args,
                      extra_compile_args=extra_compile_args),
rusty1s's avatar
rusty1s committed
37
38
39
40
        CUDAExtension('torch_sparse.spspmm_cuda',
                      ['cuda/spspmm.cpp', 'cuda/spspmm_kernel.cu'],
                      extra_link_args=extra_link_args,
                      extra_compile_args=extra_compile_args),
41
        CUDAExtension('torch_sparse.unique_cuda',
rusty1s's avatar
rusty1s committed
42
43
                      ['cuda/unique.cpp', 'cuda/unique_kernel.cu'],
                      extra_compile_args=extra_compile_args),
rusty1s's avatar
to csr  
rusty1s committed
44
    ]
rusty1s's avatar
rusty1s committed
45

rusty1s's avatar
rusty1s committed
46
__version__ = '0.4.3'
rusty1s's avatar
rusty1s committed
47
48
49
50
51
url = 'https://github.com/rusty1s/pytorch_sparse'

install_requires = ['scipy']
setup_requires = ['pytest-runner']
tests_require = ['pytest', 'pytest-cov']
rusty1s's avatar
rusty1s committed
52
53
54
55

setup(
    name='torch_sparse',
    version=__version__,
56
57
    description=('PyTorch Extension Library of Optimized Autograd Sparse '
                 'Matrix Operations'),
rusty1s's avatar
rusty1s committed
58
59
60
61
    author='Matthias Fey',
    author_email='matthias.fey@tu-dortmund.de',
    url=url,
    download_url='{}/archive/{}.tar.gz'.format(url, __version__),
rusty1s's avatar
rusty1s committed
62
63
64
65
66
67
    keywords=[
        'pytorch',
        'sparse',
        'sparse-matrices',
        'autograd',
    ],
rusty1s's avatar
rusty1s committed
68
69
70
71
72
    install_requires=install_requires,
    setup_requires=setup_requires,
    tests_require=tests_require,
    ext_modules=ext_modules,
    cmdclass=cmdclass,
73
74
    packages=find_packages(),
)