setup.py 8.47 KB
Newer Older
Christian Sarofeen's avatar
Christian Sarofeen committed
1
2
import re

3
import sys
Christian Sarofeen's avatar
Christian Sarofeen committed
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
import os
import shutil
import inspect

import distutils
import distutils.spawn
from distutils.command.clean import clean

from setuptools import setup, Extension, find_packages
from setuptools.command.install import install

import subprocess
import ctypes.util

import torch

def find(path, regex_func, collect=False):
21
22
23
24
25
26
27
28
29
    """
    Recursively searches through a directory with regex_func and
    either collects all instances or returns the first instance.
    
    Args:
        path: Directory to search through
        regex_function: A function to run on each file to decide if it should be returned/collected
        collect (False) : If True will collect all instances of matching, else will return first instance only 
    """
Christian Sarofeen's avatar
Christian Sarofeen committed
30
31
32
33
34
35
36
37
38
39
    collection = [] if collect else None
    for root, dirs, files in os.walk(path):
        for file in files:
            if regex_func(file):
                if collect:
                    collection.append(os.path.join(root, file))
                else:
                    return os.path.join(root, file)
    return list(set(collection))

40

Christian Sarofeen's avatar
Christian Sarofeen committed
41
def findcuda():
42
43
44
45
46
47
48
49
    """
    Based on PyTorch build process. Will look for nvcc for compilation.
    Either will set cuda home by enviornment variable CUDA_HOME or will search
    for nvcc. Returns NVCC executable, cuda major version and cuda home directory.
    """
    cuda_path = None
    CUDA_HOME = None

Christian Sarofeen's avatar
Christian Sarofeen committed
50
51
52
    CUDA_HOME = os.getenv('CUDA_HOME', '/usr/local/cuda')
    if not os.path.exists(CUDA_HOME):
        # We use nvcc path on Linux and cudart path on macOS
53
54
55
        cudart_path = ctypes.util.find_library('cudart')
        if cudart_path is not None:
            cuda_path = os.path.dirname(cudart_path)
Christian Sarofeen's avatar
Christian Sarofeen committed
56
57
        if cuda_path is not None:
            CUDA_HOME = os.path.dirname(cuda_path)
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
            
    if not cuda_path and not CUDA_HOME:
        nvcc_path = find('/usr/local/', re.compile("nvcc").search, False)
        if nvcc_path:
            CUDA_HOME = os.path.dirname(nvcc_path)
            if CUDA_HOME:
                os.path.dirname(CUDA_HOME)

        if (not os.path.exists(CUDA_HOME+os.sep+"lib64")
            or not os.path.exists(CUDA_HOME+os.sep+"include") ):
            raise RuntimeError("Error: found NVCC at ", nvcc_path ," but could not locate CUDA libraries"+
                               " or include directories.")
        
        raise RuntimeError("Error: Could not find cuda on this system."+
                           " Please set your CUDA_HOME enviornment variable to the CUDA base directory.")
    
74
75
76
77
78
79
80
81
82
    NVCC = find(CUDA_HOME+os.sep+"bin", 
                re.compile('nvcc$').search)
    print("Found NVCC = ", NVCC)

    # Parse output of nvcc to get cuda major version
    nvcc_output = subprocess.check_output([NVCC, '--version']).decode("utf-8")
    CUDA_LIB = re.compile(', V[0-9]+\.[0-9]+\.[0-9]+').search(nvcc_output).group(0).split('V')[1]
    print("Found CUDA_LIB = ", CUDA_LIB)

83
84
    if CUDA_LIB:
        try:
85
            CUDA_VERSION = int(CUDA_LIB.split('.')[0])
86
87
88
89
90
91
92
93
94
95
        except (ValueError, TypeError):
            CUDA_VERSION = 9
    else:
         CUDA_VERSION = 9   

    if CUDA_VERSION < 8:
        raise RuntimeError("Error: APEx requires CUDA 8 or newer")

    
    return NVCC, CUDA_VERSION, CUDA_HOME
Christian Sarofeen's avatar
Christian Sarofeen committed
96
97
98
99
100
101
102
103
104
105
106
107
108

#Get some important paths
curdir = os.path.dirname(os.path.abspath(inspect.stack()[0][1]))
buildir = curdir+os.sep+"build"
if not os.path.exists(buildir):
    os.makedirs(buildir)

torch_dir = os.path.split(torch.__file__)[0] + os.sep + "lib"

cuda_files = find(curdir, lambda file: file.endswith(".cu"), True)
cuda_headers = find(curdir, lambda file: file.endswith(".cuh"), True)
headers = find(curdir, lambda file: file.endswith(".h"), True)

109
110
111
112
113
114
libaten = list(set(find(torch_dir, re.compile("libaten", re.IGNORECASE).search, True)))
libaten_names = [os.path.splitext(os.path.basename(entry))[0] for entry in libaten]
for i, entry in enumerate(libaten_names):
    if entry[:3]=='lib':
        libaten_names[i] = entry[3:]

Christian Sarofeen's avatar
Christian Sarofeen committed
115
116
aten_h = find(torch_dir, re.compile("aten.h", re.IGNORECASE).search, False)

117
118
torch_inc = os.path.dirname(os.path.dirname(aten_h))
include_dirs = [torch_inc]
Christian Sarofeen's avatar
Christian Sarofeen committed
119
120
121
122
123
124
library_dirs = []
for file in cuda_headers+headers:
    dir = os.path.dirname(file)
    if dir not in include_dirs:
        include_dirs.append(dir)

125
126
127
128
# Object files that use the PyTorch cffi-extension interface
# They need special handling during compilation
cffi_objects = ['scale_kernel.o']

Christian Sarofeen's avatar
Christian Sarofeen committed
129
130
131
assert libaten, "Could not find PyTorch's libATen."
assert aten_h, "Could not find PyTorch's ATen header."

132
library_dirs.append(os.path.dirname(libaten[0]))
Christian Sarofeen's avatar
Christian Sarofeen committed
133
134
135
136
137

#create some places to collect important things
object_files = []
extra_link_args=[]
main_libraries = []
138
main_libraries += ['cudart',]+libaten_names
Christian Sarofeen's avatar
Christian Sarofeen committed
139
140
141
142
extra_compile_args = ["--std=c++11",]

#findcuda returns root dir of CUDA
#include cuda/include and cuda/lib64 for python module build.
143
NVCC, CUDA_VERSION, CUDA_HOME=findcuda()
Christian Sarofeen's avatar
Christian Sarofeen committed
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
library_dirs.append(os.path.join(CUDA_HOME, "lib64"))
include_dirs.append(os.path.join(CUDA_HOME, 'include'))

class RMBuild(clean):
    def run(self):
        #BE VERY CAUTIOUS WHEN USING RMTREE!!!
        #These are some carefully written/crafted directories
        if os.path.exists(buildir):
            shutil.rmtree(buildir)
            
        distdir = curdir+os.sep+"dist"
        if os.path.exists(distdir):
            shutil.rmtree(distdir)

        eggdir = curdir+os.sep+"apex.egg-info"
        if os.path.exists(eggdir):
            shutil.rmtree(eggdir)
        clean.run(self)

163
def CompileCudaFiles(NVCC, CUDA_VERSION):
Christian Sarofeen's avatar
Christian Sarofeen committed
164
165
166

        print()
        print("Compiling cuda modules with nvcc:")
167
        gencodes =  ['-gencode', 'arch=compute_52,code=sm_52',
Christian Sarofeen's avatar
Christian Sarofeen committed
168
                    '-gencode', 'arch=compute_60,code=sm_60',
169
170
171
172
173
174
175
176
177
178
179
180
181
                    '-gencode', 'arch=compute_61,code=sm_61',]
        
        if CUDA_VERSION > 8:
            gencodes += ['-gencode', 'arch=compute_70,code=sm_70',
                         '-gencode', 'arch=compute_70,code=compute_70',]
        #Need arches to compile for. Compiles for 70 which requires CUDA9
        nvcc_cmd = [NVCC, 
                       '-Xcompiler', 
                       '-fPIC'
                   ] + gencodes + [
                       '--std=c++11',
                       '-O3',
                   ]
Christian Sarofeen's avatar
Christian Sarofeen committed
182
183
184
185
        
        for dir in include_dirs:
            nvcc_cmd.append("-I"+dir)

186
187
188
189
190
        # Hack: compiling the cffi kernel code needs the TH{C}
        #       subdirs of include on path as well
        for suffix in ['TH', 'THC']:
            nvcc_cmd.append('-I{}/{}'.format(torch_inc, suffix))

Christian Sarofeen's avatar
Christian Sarofeen committed
191
192
193
194
        for file in cuda_files:
            object_name = os.path.basename(
                os.path.splitext(file)[0]+".o"
            )
195

Christian Sarofeen's avatar
Christian Sarofeen committed
196
197
198
199
            object_file = os.path.join(buildir, object_name)
            object_files.append(object_file)
    
            file_opts = ['-c', file, '-o', object_file]
200
201
202
203
204
205
206
207
208

            extra_args = []
            if object_name in cffi_objects:
                for module in ['TH', 'THC']:
                    extra_args.append('-I{}/{}'.format(torch_inc, module))

            build_args = nvcc_cmd + extra_args + file_opts
            print(' '.join(build_args))
            subprocess.check_call(build_args)
Christian Sarofeen's avatar
Christian Sarofeen committed
209
210
211
            
        for object_file in object_files:
            extra_link_args.append(object_file)
212
213
214
215
216
217
218
219
220
221
if 'clean' not in sys.argv:

    print()
    print("Arguments used to build CUDA extension:")
    print("extra_compile_args :", extra_compile_args)
    print("include_dirs: ", include_dirs)
    print("extra_link_args: ", extra_link_args)
    print("library_dirs: ", library_dirs)
    print("libraries: ", main_libraries)
    print()
222
    CompileCudaFiles(NVCC, CUDA_VERSION)
223
224
225

    print("Building CUDA extension.")
    
Christian Sarofeen's avatar
Christian Sarofeen committed
226
227
228
229
230
231
232
233
cuda_ext = Extension('apex._C',
                 [os.path.join('csrc', 'Module.cpp')],
                 extra_compile_args = extra_compile_args,
                 include_dirs=include_dirs,
                 extra_link_args=extra_link_args,
                 library_dirs=library_dirs,
                 runtime_library_dirs = library_dirs,
                 libraries=main_libraries
234
    )
Christian Sarofeen's avatar
Christian Sarofeen committed
235

236
237
238
if 'clean' not in sys.argv:
    print("Building module.")
    
Christian Sarofeen's avatar
Christian Sarofeen committed
239
240
241
242
243
244
245
246
setup(
    name='apex', version='0.1',
    cmdclass={
        'clean' : RMBuild,
    },  
    ext_modules=[cuda_ext,],
    description='PyTorch Extensions written by NVIDIA',
    packages=find_packages(exclude=("build", "csrc", "include", "tests")),
247
248
249
250
251
252

    # Require cffi
    install_requires=["cffi>=1.0.0"],
    setup_requires=["cffi>=1.0.0"],
    cffi_modules=[os.path.join(os.path.dirname(__file__),
                               'build_cffi.py:extension')],
Christian Sarofeen's avatar
Christian Sarofeen committed
253
)