-
Notifications
You must be signed in to change notification settings - Fork 61
/
setup.py
162 lines (139 loc) · 5.04 KB
/
setup.py
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
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
#!/usr/bin/env python
import os
import glob
import re
import numpy
from setuptools import setup, find_packages, Feature, Command
from Cython.Build import cythonize
from Cython.Distutils import build_ext
from Cython.Distutils.extension import Extension
def read(fname):
return open(os.path.join(os.path.dirname(__file__), fname)).read()
def cuda_extensions():
cuda_dir = os.getenv('CUDA_PREFIX', '/usr/local/cuda')
cuda_include_dir = os.path.join(cuda_dir, 'include')
cuda_library_dir = os.path.join(cuda_dir, 'lib64')
if not os.path.exists(cuda_library_dir):
# Use lib if lib64 does not exist
cuda_library_dir = os.path.join(cuda_dir, 'lib')
library_dirs = [cuda_library_dir]
prefix = os.getenv('INSTALL_PREFIX')
if prefix is not None:
library_dirs.append(os.path.join(prefix, 'lib'))
cudarray_dir = './cudarray'
cudarray_include_dir = './include'
include_dirs = [cuda_include_dir, cudarray_include_dir,
numpy.get_include()]
cython_include_dirs = ['./cudarray/wrap']
extra_compile_args = ['-O3', '-fPIC', '-Wall', '-Wfatal-errors']
libraries = ['cudart', 'cudarray']
extra_link_args = ['-fPIC']
language = 'c++'
def make_extension(name):
return Extension(
name='cudarray.wrap.' + name,
sources=[os.path.join(cudarray_dir, 'wrap', name + '.pyx')],
language=language,
include_dirs=include_dirs,
cython_include_dirs=cython_include_dirs,
extra_compile_args=extra_compile_args,
library_dirs=library_dirs,
libraries=libraries,
extra_link_args=extra_link_args,
)
ext_names = ['cudart', 'array_data', 'array_ops', 'elementwise',
'reduction', 'blas', 'random', 'nnet', 'image']
exts = list(map(make_extension, ext_names))
if os.getenv('CUDNN_ENABLED') == '1':
cudnn_ext = Extension(
name='cudarray.wrap.cudnn',
sources=[os.path.join(cudarray_dir, 'wrap', 'cudnn.pyx')],
language=language,
include_dirs=include_dirs,
cython_include_dirs=cython_include_dirs,
extra_compile_args=['-DCUDNN_ENABLED'] + extra_compile_args,
library_dirs=library_dirs,
libraries=libraries+['cudnn'],
extra_link_args=extra_link_args,
)
exts.append(cudnn_ext)
return exts
def numpy_extensions():
cython_srcs = [
'cudarray/numpy_backend/nnet/conv_bc01.pyx',
'cudarray/numpy_backend/nnet/pool_bc01.pyx',
'cudarray/numpy_backend/nnet/lrnorm_bc01.pyx',
]
return cythonize(cython_srcs, include_path=[numpy.get_include()])
class Clean(Command):
description = 'Remove Cython generated files.'
user_options = []
def initialize_options(self):
self.report = 'clean'
def finalize_options(self):
pass
def run(self):
for ext in (cuda_extensions() + numpy_extensions()):
for pyx in ext.sources:
if pyx.endswith('.pyx'):
prefix = pyx[:-4]
for suffix in ['.c', '.cpp', '.so', '.dll']:
files = glob.glob(prefix + '*' + suffix)
for f in files:
if os.path.exists(f):
os.unlink(f)
with open('requirements.txt') as f:
install_requires = [l.strip() for l in f]
setup_requires = [r for r in install_requires if r.startswith('cython')]
version = None
regex = re.compile(r'''^__version__ = ['"]([^'"]*)['"]''')
with open(os.path.join('cudarray', '__init__.py')) as f:
for line in f:
mo = regex.search(line)
if mo is not None:
version = mo.group(1)
break
if version is None:
raise RuntimeError('Could not find version number')
setup(
name='cudarray',
version=version,
author='Anders Boesen Lindbo Larsen',
author_email='[email protected]',
description='CUDA-based Numpy array and operations',
license='MIT',
url='http://compute.dtu.dk/~abll',
packages=find_packages(),
setup_requires=setup_requires,
install_requires=install_requires,
long_description=read('README.md'),
classifiers=[
'Development Status :: 4 - Beta',
'Intended Audience :: Developers',
'Intended Audience :: Science/Research',
'License :: OSI Approved :: MIT License',
'Operating System :: OS Independent',
'Programming Language :: Python',
'Topic :: Scientific/Engineering',
],
include_dirs=[numpy.get_include()],
features={
'cuda': Feature(
description='CUDA back-end',
standard=True,
remove=['cudarray.wrap'],
ext_modules=cuda_extensions(),
),
'numpy': Feature(
description='Numpy back-end',
standard=True,
remove=['cudarray.numpy_backend'],
ext_modules=numpy_extensions(),
),
},
cmdclass={
'build_ext': build_ext,
'clean': Clean,
},
zip_safe=False,
)