setup.py 15.9 KB
Newer Older
1
2
import glob
import os
3
import platform
Rinat Shigapov's avatar
Rinat Shigapov committed
4
import re
5
from pkg_resources import DistributionNotFound, get_distribution, parse_version
6
from setuptools import find_packages, setup
xiabo's avatar
xiabo committed
7
8
9
import subprocess
from typing import Optional, Union
from pathlib import Path
Yuanhao Zhu's avatar
Yuanhao Zhu committed
10
11
12
13
14
15
16

EXT_TYPE = ''
try:
    import torch
    if torch.__version__ == 'parrots':
        from parrots.utils.build_extension import BuildExtension
        EXT_TYPE = 'parrots'
17
18
19
20
    elif (hasattr(torch, 'is_mlu_available') and torch.is_mlu_available()) or \
            os.getenv('FORCE_MLU', '0') == '1':
        from torch_mlu.utils.cpp_extension import BuildExtension
        EXT_TYPE = 'pytorch'
Yuanhao Zhu's avatar
Yuanhao Zhu committed
21
22
23
    else:
        from torch.utils.cpp_extension import BuildExtension
        EXT_TYPE = 'pytorch'
24
    cmd_class = {'build_ext': BuildExtension}
Yuanhao Zhu's avatar
Yuanhao Zhu committed
25
except ModuleNotFoundError:
26
    cmd_class = {}
Yuanhao Zhu's avatar
Yuanhao Zhu committed
27
    print('Skip building ext ops due to the absence of torch.')
28

Rinat Shigapov's avatar
Rinat Shigapov committed
29
30

def choose_requirement(primary, secondary):
Kai Chen's avatar
Kai Chen committed
31
32
    """If some version of primary requirement installed, return primary, else
    return secondary."""
Rinat Shigapov's avatar
Rinat Shigapov committed
33
34
35
36
37
38
39
40
    try:
        name = re.split(r'[!<>=]', primary)[0]
        get_distribution(name)
    except DistributionNotFound:
        return secondary

    return str(primary)

xiabo's avatar
xiabo committed
41
42
43
44
45
46
47
def get_sha(pytorch_root: Union[str, Path]) -> str:
    try:
        return subprocess.check_output(['git', 'rev-parse', 'HEAD'], cwd=pytorch_root).decode('ascii').strip()
    except Exception:
        return 'Unknown'


xiabo's avatar
xiabo committed
48
49
50
51
52
53
54
55
56
57
58
def get_abi():
    try:
        command = "echo '#include <string>' | gcc -x c++ -E -dM - | fgrep _GLIBCXX_USE_CXX11_ABI" 
        result = subprocess.run(command, shell=True, capture_output=True, text=True) 
        output = result.stdout.strip() 
        abi = "abi" + output.split(" ")[-1]
        return abi
    except Exception:
        return 'abiUnknown'


xiabo's avatar
xiabo committed
59
def get_version_add(sha: Optional[str] = None) -> str:
xiabo's avatar
xiabo committed
60
    version=''
xiabo's avatar
xiabo committed
61
62
63
64
65
66
67
    mmcv_root = os.path.dirname(os.path.abspath(__file__))
    add_version_path = os.path.join(os.path.join(mmcv_root, "mmcv"), "version.py")
    if sha != 'Unknown':
        if sha is None:
            sha = get_sha(mmcv_root)
        version = 'git' + sha[:7]

xiabo's avatar
xiabo committed
68
69
70
71
72
73
74
75
76
77
78
79
80
81
    # abi
    version += "." + get_abi()

    # dtk version
    if os.getenv("ROCM_PATH"):
        rocm_path = os.getenv('ROCM_PATH', "")
        rocm_version_path = os.path.join(rocm_path, '.info', "rocm_version")
        with open(rocm_version_path, 'r',encoding='utf-8') as file:
            lines = file.readlines()
        rocm_version=lines[0][:-2].replace(".", "")
        version += ".dtk" + rocm_version
    
    # torch version
    version += ".torch" + torch.__version__[:4]
xiabo's avatar
xiabo committed
82

xiabo's avatar
xiabo committed
83
84
85
    lines=[]
    with open(add_version_path, 'r',encoding='utf-8') as file:
        lines = file.readlines()
xiabo's avatar
xiabo committed
86
    lines[1] = "__version__='2.0.0+{}'\n".format(version)
xiabo's avatar
xiabo committed
87
88
    with open(add_version_path, encoding="utf-8",mode="w") as file:
        file.writelines(lines)
xiabo's avatar
xiabo committed
89
90
    file.close()

Rinat Shigapov's avatar
Rinat Shigapov committed
91

92
def get_version():
xiabo's avatar
xiabo committed
93
    get_version_add()
94
    version_file = 'mmcv/version.py'
95
    with open(version_file, encoding='utf-8') as f:
96
97
98
99
        exec(compile(f.read(), version_file, 'exec'))
    return locals()['__version__']


100
def parse_requirements(fname='requirements/runtime.txt', with_version=True):
Kai Chen's avatar
Kai Chen committed
101
    """Parse the package dependencies listed in a requirements file but strips
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
    specific versioning information.

    Args:
        fname (str): path to requirements file
        with_version (bool, default=False): if True include version specs

    Returns:
        List[str]: list of requirements items

    CommandLine:
        python -c "import setup; print(setup.parse_requirements())"
    """
    import sys
    from os.path import exists
    require_fpath = fname

    def parse_line(line):
Kai Chen's avatar
Kai Chen committed
119
        """Parse information from a line in a requirements text file."""
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
        if line.startswith('-r '):
            # Allow specifying requirements in other files
            target = line.split(' ')[1]
            for info in parse_require_file(target):
                yield info
        else:
            info = {'line': line}
            if line.startswith('-e '):
                info['package'] = line.split('#egg=')[1]
            else:
                # Remove versioning from the package
                pat = '(' + '|'.join(['>=', '==', '>']) + ')'
                parts = re.split(pat, line, maxsplit=1)
                parts = [p.strip() for p in parts]

                info['package'] = parts[0]
                if len(parts) > 1:
                    op, rest = parts[1:]
                    if ';' in rest:
                        # Handle platform specific dependencies
                        # http://setuptools.readthedocs.io/en/latest/setuptools.html#declaring-platform-specific-dependencies
                        version, platform_deps = map(str.strip,
                                                     rest.split(';'))
                        info['platform_deps'] = platform_deps
                    else:
                        version = rest  # NOQA
                    info['version'] = (op, version)
            yield info

    def parse_require_file(fpath):
150
        with open(fpath) as f:
151
152
153
            for line in f.readlines():
                line = line.strip()
                if line and not line.startswith('#'):
154
                    yield from parse_line(line)
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174

    def gen_packages_items():
        if exists(require_fpath):
            for info in parse_require_file(require_fpath):
                parts = [info['package']]
                if with_version and 'version' in info:
                    parts.extend(info['version'])
                if not sys.version.startswith('3.4'):
                    # apparently package_deps are broken in 3.4
                    platform_deps = info.get('platform_deps')
                    if platform_deps is not None:
                        parts.append(';' + platform_deps)
                item = ''.join(parts)
                yield item

    packages = list(gen_packages_items())
    return packages


install_requires = parse_requirements()
175
176
177
178
179
180
181
182
183
184
185
186
187
188

try:
    # OpenCV installed via conda.
    import cv2  # NOQA: F401
    major, minor, *rest = cv2.__version__.split('.')
    if int(major) < 3:
        raise RuntimeError(
            f'OpenCV >=3 is required but {cv2.__version__} is installed')
except ImportError:
    # If first not installed install second package
    CHOOSE_INSTALL_REQUIRES = [('opencv-python-headless>=3',
                                'opencv-python>=3')]
    for main, secondary in CHOOSE_INSTALL_REQUIRES:
        install_requires.append(choose_requirement(main, secondary))
189

Kai Chen's avatar
Kai Chen committed
190

191
192
193
def get_extensions():
    extensions = []

Zaida Zhou's avatar
Zaida Zhou committed
194
    if os.getenv('MMCV_WITH_OPS', '1') == '0':
195
196
        return extensions

Yuanhao Zhu's avatar
Yuanhao Zhu committed
197
198
199
    if EXT_TYPE == 'parrots':
        ext_name = 'mmcv._ext'
        from parrots.utils.build_extension import Extension
200

201
202
203
        # new parrots op impl do not use MMCV_USE_PARROTS
        # define_macros = [('MMCV_USE_PARROTS', None)]
        define_macros = []
204
205
        include_dirs = []
        op_files = glob.glob('./mmcv/ops/csrc/pytorch/cuda/*.cu') +\
206
            glob.glob('./mmcv/ops/csrc/pytorch/cpu/*.cpp') +\
207
            glob.glob('./mmcv/ops/csrc/parrots/*.cpp')
208
209
        include_dirs.append(os.path.abspath('./mmcv/ops/csrc/common'))
        include_dirs.append(os.path.abspath('./mmcv/ops/csrc/common/cuda'))
Yuanhao Zhu's avatar
Yuanhao Zhu committed
210
        cuda_args = os.getenv('MMCV_CUDA_ARGS')
211
        extra_compile_args = {
212
213
            'nvcc': [cuda_args, '-std=c++14'] if cuda_args else ['-std=c++14'],
            'cxx': ['-std=c++14'],
214
215
216
217
218
219
220
221
        }
        if torch.cuda.is_available() or os.getenv('FORCE_CUDA', '0') == '1':
            define_macros += [('MMCV_WITH_CUDA', None)]
            extra_compile_args['nvcc'] += [
                '-D__CUDA_NO_HALF_OPERATORS__',
                '-D__CUDA_NO_HALF_CONVERSIONS__',
                '-D__CUDA_NO_HALF2_OPERATORS__',
            ]
Yuanhao Zhu's avatar
Yuanhao Zhu committed
222
223
224
        ext_ops = Extension(
            name=ext_name,
            sources=op_files,
225
            include_dirs=include_dirs,
Yuanhao Zhu's avatar
Yuanhao Zhu committed
226
            define_macros=define_macros,
227
228
229
            extra_compile_args=extra_compile_args,
            cuda=True,
            pytorch=True)
Yuanhao Zhu's avatar
Yuanhao Zhu committed
230
231
        extensions.append(ext_ops)
    elif EXT_TYPE == 'pytorch':
232
        ext_name = 'mmcv._ext'
233
234
        from torch.utils.cpp_extension import CppExtension, CUDAExtension

Yuanhao Zhu's avatar
Yuanhao Zhu committed
235
        # prevent ninja from using too many resources
lizz's avatar
lizz committed
236
237
238
239
240
241
242
243
        try:
            import psutil
            num_cpu = len(psutil.Process().cpu_affinity())
            cpu_use = max(4, num_cpu - 1)
        except (ModuleNotFoundError, AttributeError):
            cpu_use = 4

        os.environ.setdefault('MAX_JOBS', str(cpu_use))
Yuanhao Zhu's avatar
Yuanhao Zhu committed
244
        define_macros = []
245
246
247
248
249
250
251
252
253

        # Before PyTorch1.8.0, when compiling CUDA code, `cxx` is a
        # required key passed to PyTorch. Even if there is no flag passed
        # to cxx, users also need to pass an empty list to PyTorch.
        # Since PyTorch1.8.0, it has a default value so users do not need
        # to pass an empty list anymore.
        # More details at https://github.com/pytorch/pytorch/pull/45956
        extra_compile_args = {'cxx': []}

254
255
        if platform.system() != 'Windows':
            extra_compile_args['cxx'] = ['-std=c++14']
256
257
258
259
260
261
        else:
            # TODO: In Windows, C++17 is chosen to compile extensions in
            # PyTorch2.0 , but a compile error will be reported.
            # As a temporary solution, force the use of C++14.
            if parse_version(torch.__version__) >= parse_version('2.0.0'):
                extra_compile_args['cxx'] = ['/std:c++14']
262

263
        include_dirs = []
Yuanhao Zhu's avatar
Yuanhao Zhu committed
264

265
        is_rocm_pytorch = False
266
        try:
267
268
269
            from torch.utils.cpp_extension import ROCM_HOME
            is_rocm_pytorch = True if ((torch.version.hip is not None) and
                                       (ROCM_HOME is not None)) else False
270
271
        except ImportError:
            pass
272

q.yao's avatar
q.yao committed
273
274
275
        if is_rocm_pytorch or torch.cuda.is_available() or os.getenv(
                'FORCE_CUDA', '0') == '1':
            if is_rocm_pytorch:
Zachary Streeter's avatar
Zachary Streeter committed
276
                define_macros += [('MMCV_WITH_HIP', None)]
Yuanhao Zhu's avatar
Yuanhao Zhu committed
277
            define_macros += [('MMCV_WITH_CUDA', None)]
zhuyuanhao's avatar
zhuyuanhao committed
278
            cuda_args = os.getenv('MMCV_CUDA_ARGS')
Yuanhao Zhu's avatar
Yuanhao Zhu committed
279
            extra_compile_args['nvcc'] = [cuda_args] if cuda_args else []
xiabo's avatar
xiabo committed
280
281
282
            if is_rocm_pytorch and platform.system() != 'Windows':
                extra_compile_args['nvcc'] += \
                    ['--gpu-max-threads-per-block=1024']
283
            op_files = glob.glob('./mmcv/ops/csrc/pytorch/*.cpp') + \
284
285
286
                glob.glob('./mmcv/ops/csrc/pytorch/cpu/*.cpp') + \
                glob.glob('./mmcv/ops/csrc/pytorch/cuda/*.cu') + \
                glob.glob('./mmcv/ops/csrc/pytorch/cuda/*.cpp')
Yuanhao Zhu's avatar
Yuanhao Zhu committed
287
            extension = CUDAExtension
Zachary Streeter's avatar
Zachary Streeter committed
288
            include_dirs.append(os.path.abspath('./mmcv/ops/csrc/pytorch'))
289
290
            include_dirs.append(os.path.abspath('./mmcv/ops/csrc/common'))
            include_dirs.append(os.path.abspath('./mmcv/ops/csrc/common/cuda'))
291
292
293
294
295
296
297
298
299
        elif (hasattr(torch, 'is_mlu_available') and
                torch.is_mlu_available()) or \
                os.getenv('FORCE_MLU', '0') == '1':
            from torch_mlu.utils.cpp_extension import MLUExtension
            define_macros += [('MMCV_WITH_MLU', None)]
            mlu_args = os.getenv('MMCV_MLU_ARGS')
            extra_compile_args['cncc'] = [mlu_args] if mlu_args else []
            op_files = glob.glob('./mmcv/ops/csrc/pytorch/*.cpp') + \
                glob.glob('./mmcv/ops/csrc/pytorch/cpu/*.cpp') + \
300
301
                glob.glob('./mmcv/ops/csrc/pytorch/mlu/*.cpp') + \
                glob.glob('./mmcv/ops/csrc/common/mlu/*.mlu')
302
303
304
            extension = MLUExtension
            include_dirs.append(os.path.abspath('./mmcv/ops/csrc/common'))
            include_dirs.append(os.path.abspath('./mmcv/ops/csrc/common/mlu'))
Zaida Zhou's avatar
Zaida Zhou committed
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
        elif (hasattr(torch.backends, 'mps')
              and torch.backends.mps.is_available()) or os.getenv(
                  'FORCE_MPS', '0') == '1':
            # objc compiler support
            from distutils.unixccompiler import UnixCCompiler
            if '.mm' not in UnixCCompiler.src_extensions:
                UnixCCompiler.src_extensions.append('.mm')
                UnixCCompiler.language_map['.mm'] = 'objc'

            define_macros += [('MMCV_WITH_MPS', None)]
            extra_compile_args = {}
            extra_compile_args['cxx'] = ['-Wall', '-std=c++17']
            extra_compile_args['cxx'] += [
                '-framework', 'Metal', '-framework', 'Foundation'
            ]
            extra_compile_args['cxx'] += ['-ObjC++']
            # src
            op_files = glob.glob('./mmcv/ops/csrc/pytorch/*.cpp') + \
                glob.glob('./mmcv/ops/csrc/pytorch/cpu/*.cpp') + \
                glob.glob('./mmcv/ops/csrc/common/mps/*.mm') + \
                glob.glob('./mmcv/ops/csrc/pytorch/mps/*.mm')
            extension = CppExtension
            include_dirs.append(os.path.abspath('./mmcv/ops/csrc/common'))
            include_dirs.append(os.path.abspath('./mmcv/ops/csrc/common/mps'))
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
        elif (os.getenv('FORCE_NPU', '0') == '1'):
            print(f'Compiling {ext_name} only with CPU and NPU')
            try:
                from torch_npu.utils.cpp_extension import NpuExtension
                define_macros += [('MMCV_WITH_NPU', None)]
                extension = NpuExtension
            except Exception:
                raise ImportError('can not find any torch_npu')
            # src
            op_files = glob.glob('./mmcv/ops/csrc/pytorch/*.cpp') + \
                glob.glob('./mmcv/ops/csrc/pytorch/cpu/*.cpp') + \
                glob.glob('./mmcv/ops/csrc/common/npu/*.cpp') + \
                glob.glob('./mmcv/ops/csrc/pytorch/npu/*.cpp')
            include_dirs.append(os.path.abspath('./mmcv/ops/csrc/common'))
            include_dirs.append(os.path.abspath('./mmcv/ops/csrc/common/npu'))
344
        else:
345
            print(f'Compiling {ext_name} only with CPU')
346
347
            op_files = glob.glob('./mmcv/ops/csrc/pytorch/*.cpp') + \
                glob.glob('./mmcv/ops/csrc/pytorch/cpu/*.cpp')
Yuanhao Zhu's avatar
Yuanhao Zhu committed
348
            extension = CppExtension
349
            include_dirs.append(os.path.abspath('./mmcv/ops/csrc/common'))
Yuanhao Zhu's avatar
Yuanhao Zhu committed
350

351
352
353
354
355
356
357
358
        # Since the PR (https://github.com/open-mmlab/mmcv/pull/1463) uses
        # c++14 features, the argument ['std=c++14'] must be added here.
        # However, in the windows environment, some standard libraries
        # will depend on c++17 or higher. In fact, for the windows
        # environment, the compiler will choose the appropriate compiler
        # to compile those cpp files, so there is no need to add the
        # argument
        if 'nvcc' in extra_compile_args and platform.system() != 'Windows':
359
360
            extra_compile_args['nvcc'] += ['-std=c++14']

Yuanhao Zhu's avatar
Yuanhao Zhu committed
361
362
363
        ext_ops = extension(
            name=ext_name,
            sources=op_files,
364
            include_dirs=include_dirs,
Yuanhao Zhu's avatar
Yuanhao Zhu committed
365
366
367
            define_macros=define_macros,
            extra_compile_args=extra_compile_args)
        extensions.append(ext_ops)
368
369
    return extensions

370

371
setup(
Zaida Zhou's avatar
Zaida Zhou committed
372
    name='mmcv' if os.getenv('MMCV_WITH_OPS', '1') == '1' else 'mmcv-lite',
373
    version=get_version(),
374
    description='OpenMMLab Computer Vision Foundation',
375
376
    keywords='computer vision',
    packages=find_packages(),
377
    include_package_data=True,
378
379
    classifiers=[
        'Development Status :: 4 - Beta',
Kai Chen's avatar
Kai Chen committed
380
        'License :: OSI Approved :: Apache Software License',
381
382
        'Operating System :: OS Independent',
        'Programming Language :: Python :: 3',
383
        'Programming Language :: Python :: 3.7',
384
        'Programming Language :: Python :: 3.8',
385
        'Programming Language :: Python :: 3.9',
Zaida Zhou's avatar
Zaida Zhou committed
386
        'Programming Language :: Python :: 3.10',
387
388
389
        'Topic :: Utilities',
    ],
    url='https://github.com/open-mmlab/mmcv',
Zaida Zhou's avatar
Zaida Zhou committed
390
    author='MMCV Contributors',
Kai Chen's avatar
Kai Chen committed
391
    author_email='openmmlab@gmail.com',
392
    install_requires=install_requires,
393
394
395
396
397
398
    extras_require={
        'all': parse_requirements('requirements.txt'),
        'tests': parse_requirements('requirements/test.txt'),
        'build': parse_requirements('requirements/build.txt'),
        'optional': parse_requirements('requirements/optional.txt'),
    },
399
    python_requires='>=3.7',
400
    ext_modules=get_extensions(),
401
    cmdclass=cmd_class,
402
    zip_safe=False)