setup.py 7.15 KB
Newer Older
1
import torch
2
from setuptools import setup, find_packages
mcarilli's avatar
mcarilli committed
3
import subprocess
4

jjsjann123's avatar
jjsjann123 committed
5
6
import sys

7
if not torch.cuda.is_available():
mcarilli's avatar
mcarilli committed
8
9
    print("\nWarning: Torch did not find available GPUs on this system.\n",
          "If your intention is to cross-compile, this is not an error.\n")
10

11
12
13
14
15
print("torch.__version__  = ", torch.__version__)
TORCH_MAJOR = int(torch.__version__.split('.')[0])
TORCH_MINOR = int(torch.__version__.split('.')[1])

if TORCH_MAJOR == 0 and TORCH_MINOR < 4:
Michael Carilli's avatar
Michael Carilli committed
16
      raise RuntimeError("Apex requires Pytorch 0.4 or newer.\n" +
17
18
                         "The latest stable release can be obtained from https://pytorch.org/")

jjsjann123's avatar
jjsjann123 committed
19
20
21
cmdclass = {}
ext_modules = []

22
if "--cpp_ext" in sys.argv or "--cuda_ext" in sys.argv:
Michael Carilli's avatar
Michael Carilli committed
23
24
    if TORCH_MAJOR == 0:
        raise RuntimeError("--cpp_ext requires Pytorch 1.0 or later, "
25
                           "found torch.__version__ = {}".format(torch.__version__))
26
27
28
29
30
31
32
33
34
35
    from torch.utils.cpp_extension import BuildExtension
    cmdclass['build_ext'] = BuildExtension

if "--cpp_ext" in sys.argv:
    from torch.utils.cpp_extension import CppExtension
    sys.argv.remove("--cpp_ext")
    ext_modules.append(
        CppExtension('apex_C',
                     ['csrc/flatten_unflatten.cpp',]))

mcarilli's avatar
mcarilli committed
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
def check_cuda_torch_binary_vs_bare_metal(cuda_dir):
    raw_output = subprocess.check_output([cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True)
    output = raw_output.split()
    release_idx = output.index("release") + 1
    release = output[release_idx].split(".")
    bare_metal_major = release[0]
    bare_metal_minor = release[1][0]
    torch_binary_major = torch.version.cuda.split(".")[0]
    torch_binary_minor = torch.version.cuda.split(".")[1]

    print("\nCompiling cuda extensions with")
    print(raw_output + "from " + cuda_dir + "/bin\n")

    if (bare_metal_major != torch_binary_major) or (bare_metal_minor != torch_binary_minor):
        # TODO:  make this a hard error?
        print("\nWarning:  Cuda extensions are being compiled with a version of Cuda that does "
              "not match the version used to compile Pytorch binaries.\n")
    print("Pytorch binaries were compiled with Cuda {}\n".format(torch.version.cuda))

jjsjann123's avatar
jjsjann123 committed
55
if "--cuda_ext" in sys.argv:
56
    from torch.utils.cpp_extension import CUDAExtension
jjsjann123's avatar
jjsjann123 committed
57
    sys.argv.remove("--cuda_ext")
58
59

    if torch.utils.cpp_extension.CUDA_HOME is None:
Michael Carilli's avatar
Michael Carilli committed
60
        raise RuntimeError("--cuda_ext was requested, but nvcc was not found.  Are you sure your environment has nvcc available?  If you're installing within a container from https://hub.docker.com/r/pytorch/pytorch, only images whose names contain 'devel' will provide nvcc.")
61
    else:
mcarilli's avatar
mcarilli committed
62
63
        check_cuda_torch_binary_vs_bare_metal(torch.utils.cpp_extension.CUDA_HOME)

Michael Carilli's avatar
Michael Carilli committed
64
65
66
67
68
69
        # Set up macros for forward/backward compatibility hack around
        # https://github.com/pytorch/pytorch/commit/4404762d7dd955383acee92e6f06b48144a0742e
        version_ge_1_1 = []
        if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 0):
            version_ge_1_1 = ['-DVERSION_GE_1_1']

70
71
        ext_modules.append(
            CUDAExtension(name='amp_C',
72
                          sources=['csrc/amp_C_frontend.cpp',
Michael Carilli's avatar
Michael Carilli committed
73
                                   'csrc/multi_tensor_sgd_kernel.cu',
Simon Layton's avatar
Simon Layton committed
74
                                   'csrc/multi_tensor_scale_kernel.cu',
75
76
                                   'csrc/multi_tensor_axpby_kernel.cu',
                                   'csrc/multi_tensor_l2norm_kernel.cu'],
77
                          extra_compile_args={'cxx': ['-O3'],
Michael Carilli's avatar
Michael Carilli committed
78
                                              'nvcc':['-lineinfo',
79
                                                      '-O3',
80
                                                      # '--resource-usage',
Michael Carilli's avatar
Michael Carilli committed
81
                                                      '--use_fast_math']}))
82
83
        ext_modules.append(
            CUDAExtension(name='fused_adam_cuda',
84
85
                          sources=['csrc/fused_adam_cuda.cpp',
                                   'csrc/fused_adam_cuda_kernel.cu'],
86
                          extra_compile_args={'cxx': ['-O3',],
mcarilli's avatar
mcarilli committed
87
                                              'nvcc':['-O3',
88
89
90
91
92
93
94
                                                      '--use_fast_math']}))
        ext_modules.append(
            CUDAExtension(name='syncbn',
                          sources=['csrc/syncbn.cpp',
                                   'csrc/welford.cu']))
        ext_modules.append(
            CUDAExtension(name='fused_layer_norm_cuda',
95
96
                          sources=['csrc/layer_norm_cuda.cpp',
                                   'csrc/layer_norm_cuda_kernel.cu'],
Michael Carilli's avatar
Michael Carilli committed
97
                          extra_compile_args={'cxx': ['-O3'] + version_ge_1_1,
98
                                              'nvcc':['-maxrregcount=50',
mcarilli's avatar
mcarilli committed
99
                                                      '-O3',
Michael Carilli's avatar
Michael Carilli committed
100
                                                      '--use_fast_math'] + version_ge_1_1}))
jjsjann123's avatar
jjsjann123 committed
101

jjsjann123's avatar
jjsjann123 committed
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
if "--bnp" in sys.argv:
    from torch.utils.cpp_extension import CUDAExtension
    sys.argv.remove("--bnp")

    from torch.utils.cpp_extension import BuildExtension
    cmdclass['build_ext'] = BuildExtension

    if torch.utils.cpp_extension.CUDA_HOME is None:
        raise RuntimeError("--cuda_ext was requested, but nvcc was not found.  Are you sure your environment has nvcc available?  If you're installing within a container from https://hub.docker.com/r/pytorch/pytorch, only images whose names contain 'devel' will provide nvcc.")
    else:
        # Set up macros for forward/backward compatibility hack around
        # https://github.com/pytorch/pytorch/commit/4404762d7dd955383acee92e6f06b48144a0742e
        version_ge_1_1 = []
        if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 0):
            version_ge_1_1 = ['-DVERSION_GE_1_1']
        ext_modules.append(
            CUDAExtension(name='bnp',
                          sources=['apex/contrib/csrc/groupbn/batch_norm.cu',
                                   'apex/contrib/csrc/groupbn/ipc.cu',
                                   'apex/contrib/csrc/groupbn/interface.cpp',
                                   'apex/contrib/csrc/groupbn/batch_norm_add_relu.cu'],
                          extra_compile_args={'cxx': [] + version_ge_1_1,
                                              'nvcc':['-DCUDA_HAS_FP16=1',
                                                      '-D__CUDA_NO_HALF_OPERATORS__',
                                                      '-D__CUDA_NO_HALF_CONVERSIONS__',
                                                      '-D__CUDA_NO_HALF2_OPERATORS__',
                                                      '-gencode',
                                                      'arch=compute_70,code=sm_70'] + version_ge_1_1}))

Christian Sarofeen's avatar
Christian Sarofeen committed
131
setup(
132
133
    name='apex',
    version='0.1',
134
135
136
137
    packages=find_packages(exclude=('build',
                                    'csrc',
                                    'include',
                                    'tests',
138
139
140
141
142
                                    'dist',
                                    'docs',
                                    'tests',
                                    'examples',
                                    'apex.egg-info',)),
Christian Sarofeen's avatar
Christian Sarofeen committed
143
    description='PyTorch Extensions written by NVIDIA',
jjsjann123's avatar
jjsjann123 committed
144
145
    ext_modules=ext_modules,
    cmdclass=cmdclass,
Christian Sarofeen's avatar
Christian Sarofeen committed
146
)