fix macOS compile (#386)

README.rst

@@ -54,6 +54,12 @@ or install from source
     cd mmcv
     pip install -e .
 
+If you are on macOS, replace the last command with
+
+.. code::
+
+   CC=lang CXX=clang++ CFLAGS='-stdlib=libc++' pip install -e .
+
 Note: If you would like to use :code:`opencv-python-headless` instead of :code:`opencv-python`,
 e.g., in a minimum container environment or servers without GUI,
 you can first install it before installing MMCV to skip the installation of :code:`opencv-python`.

mmcv/ops/csrc/pytorch/nms.cpp

@@ -102,7 +102,7 @@ Tensor softnms_cpu(Tensor boxes, Tensor scores, Tensor dets,
 
   int64_t pos = 0;
   Tensor inds_t = at::arange(nboxes, boxes.options().dtype(at::kLong));
-  auto inds = inds_t.data_ptr<long>();
+  auto inds = inds_t.data_ptr<int64_t>();
 
   for (int64_t i = 0; i < nboxes; i++) {
     auto max_score = sc[i];

setup.py

 import glob
 import os
-import platform
 import re
 import setuptools
 from pkg_resources import DistributionNotFound, get_distribution
@@ -10,7 +9,19 @@ dist.Distribution().fetch_build_eggs(['Cython', 'numpy>=1.11.1'])
 
 import numpy  # NOQA: E402  # isort:skip
 from Cython.Build import cythonize  # NOQA: E402  # isort:skip
-from Cython.Distutils import build_ext as build_cmd  # NOQA: E402  # isort:skip
+
+EXT_TYPE = ''
+try:
+    import torch
+    if torch.__version__ == 'parrots':
+        from parrots.utils.build_extension import BuildExtension
+        EXT_TYPE = 'parrots'
+    else:
+        from torch.utils.cpp_extension import BuildExtension
+        EXT_TYPE = 'pytorch'
+except ModuleNotFoundError:
+    from Cython.Distutils import build_ext as BuildExtension
+    print('Skip building ext ops due to the absence of torch.')
 
 
 def choose_requirement(primary, secondary):
@@ -125,13 +136,6 @@ for main, secondary in CHOOSE_INSTALL_REQUIRES:
 def get_extensions():
     extensions = []
 
-    if platform.system() == 'Darwin':
-        extra_compile_args = ['-stdlib=libc++']
-        extra_link_args = ['-stdlib=libc++']
-    else:
-        extra_compile_args = []
-        extra_link_args = []
-
     ext_flow = setuptools.Extension(
         name='mmcv._flow_warp_ext',
         sources=[
@@ -139,16 +143,12 @@ def get_extensions():
             './mmcv/video/optflow_warp/flow_warp_module.pyx'
         ],
         include_dirs=[numpy.get_include()],
-        language='c++',
-        extra_compile_args=extra_compile_args,
-        extra_link_args=extra_link_args)
+        language='c++')
     extensions.extend(cythonize(ext_flow))
 
-    try:
-        import torch
+    if EXT_TYPE == 'parrots':
         ext_name = 'mmcv._ext'
-        if torch.__version__ == 'parrots':
-            from parrots.utils.build_extension import BuildExtension, Extension
+        from parrots.utils.build_extension import Extension
         define_macros = [('MMCV_USE_PARROTS', None)]
         op_files = glob.glob('./mmcv/ops/csrc/parrots/*')
         include_path = os.path.abspath('./mmcv/ops/csrc')
@@ -164,16 +164,15 @@ def get_extensions():
             },
             cuda=True)
         extensions.append(ext_ops)
-        else:
-            from torch.utils.cpp_extension import (BuildExtension,
-                                                   CUDAExtension, CppExtension)
+    elif EXT_TYPE == 'pytorch':
+        ext_name = 'mmcv._ext'
+        from torch.utils.cpp_extension import (CUDAExtension, CppExtension)
         # prevent ninja from using too many resources
         os.environ.setdefault('MAX_JOBS', '4')
         define_macros = []
         extra_compile_args = {'cxx': []}
 
-            if (torch.cuda.is_available()
-                    or os.getenv('FORCE_CUDA', '0') == '1'):
+        if torch.cuda.is_available() or os.getenv('FORCE_CUDA', '0') == '1':
             define_macros += [('MMCV_WITH_CUDA', None)]
             cuda_args = os.getenv('MMCV_CUDA_ARGS')
             extra_compile_args['nvcc'] = [cuda_args] if cuda_args else []
@@ -192,10 +191,6 @@ def get_extensions():
             define_macros=define_macros,
             extra_compile_args=extra_compile_args)
         extensions.append(ext_ops)
-        global build_cmd
-        build_cmd = BuildExtension
-    except ModuleNotFoundError:
-        print('Skip building ext ops due to the absence of torch.')
     return extensions
 
 
@@ -224,5 +219,5 @@ setup(
     tests_require=['pytest'],
     install_requires=install_requires,
     ext_modules=get_extensions(),
-    cmdclass={'build_ext': build_cmd},
+    cmdclass={'build_ext': BuildExtension},
     zip_safe=False)