additional installation args, fix #135

CHANGELOG.md

@@ -1,6 +1,6 @@
 # Change Log
 
-## [master] - 2020-05-01
+## [master] - 2020-05-07
 
 ### Changed
 
@@ -19,6 +19,7 @@
 - CoordsManager.getKernelMapGPU returns long type tensors (Issue #125)
 - SyncBatchNorm error fix (Issue #129)
 - Sparse Tensor `dense()` doc update (Issue #126)
+- Installation arguments `--cuda_home=<value>`, `--force_cuda`, `--blas_include_dirs=<comma_separated_values>`, and '--blas_library_dirs=<comma_separated_values>`. (Issue #135)
 
 
 ## [0.4.2] - 2020-03-13

Makefile

@@ -18,6 +18,7 @@ EXTENSION_NAME := minkowski
 # mkl for MKL. For conda, conda install -c intel mkl mkl-include
 # openblas for OpenBlas (default)
 BLAS ?= openblas
+CUDA_HOME ?= $(shell $(PYTHON) -c 'from torch.utils.cpp_extension import _find_cuda_home; print(_find_cuda_home())')
 
 # Custom (MKL/ATLAS/OpenBLAS) include and lib directories.
 # Leave commented to accept the defaults for your choice of BLAS
@@ -57,11 +58,11 @@ endif
 
 ifneq ($(CPU_ONLY), 1)
 	# CUDA ROOT DIR that contains bin/ lib64/ and include/
-	# CUDA_DIR := /usr/local/cuda
-	CUDA_DIR := $(shell $(PYTHON) -c 'from torch.utils.cpp_extension import _find_cuda_home; print(_find_cuda_home())')
+	# CUDA_HOME := /usr/local/cuda
 
-	INCLUDE_DIRS += ./ $(CUDA_DIR)/include
-	LIBRARY_DIRS += $(CUDA_DIR)/lib64
+	NVCC ?= $(CUDA_HOME)/bin/nvcc
+	INCLUDE_DIRS += ./ $(CUDA_HOME)/include
+	LIBRARY_DIRS += $(CUDA_HOME)/lib64
 endif
 
 SRC_DIR := ./src
@@ -163,9 +164,9 @@ $(OBJ_DIR)/%.o: $(SRC_DIR)/%.cpp | $(OBJ_DIR)
 
 $(OBJ_DIR)/cuda/%.o: $(SRC_DIR)/%.cu | $(OBJ_DIR)
 	@ echo NVCC $<
-	$(Q)nvcc $(NVCCFLAGS) $(CUDA_ARCH) -M $< -o ${@:.o=.d} \
+	$(Q)$(NVCC) $(NVCCFLAGS) $(CUDA_ARCH) -M $< -o ${@:.o=.d} \
 		-odir $(@D)
-	$(Q)nvcc $(NVCCFLAGS) $(CUDA_ARCH) -c $< -o $@
+	$(Q)$(NVCC) $(NVCCFLAGS) $(CUDA_ARCH) -c $< -o $@
 
 $(STATIC_LIB): $(ALL_OBJS) | $(OBJ_DIR)
 	$(RM) -f $(STATIC_LIB)

README.md

@@ -238,6 +238,10 @@ If you use the Minkowski Engine, please cite:
 
 ## Projects using Minkowski Engine
 
+Please feel free to update [the wiki page](https://github.com/StanfordVL/MinkowskiEngine/wiki/Usage) to add your projects!
+
+- [A full list of projects using MinkowskiEngine](https://github.com/StanfordVL/MinkowskiEngine/wiki/Usage)
+
 - [3D and 4D Spatio-Temporal Semantic Segmentation, CVPR'19](https://github.com/chrischoy/SpatioTemporalSegmentation)
 - [Fully Convolutional Geometric Features, ICCV'19](https://github.com/chrischoy/FCGF)
 - [Learning multiview 3D point cloud registration, CVPR'20](https://arxiv.org/abs/2001.05119)

setup.py

@@ -1,4 +1,47 @@
+r"""
+Parse additional arguments along with the setup.py arguments such as install, build, distribute, sdist, etc.
+
+
+Usage:
+
+  python setup.py install <additional_flags>..<additional_flags> <additional_arg>=<value>..<additional_arg>=<value>
+
+  export CXX=<C++ compiler>; python setup.py install <additional_flags>..<additional_flags> <additional_arg>=<value>..<additional_arg>=<value>
+
+
+Examples:
+
+  python setup.py install --force_cuda --cuda_home=/usr/local/cuda
+  export CXX=g++7; python setup.py install --force_cuda --cuda_home=/usr/local/cuda
+
+
+Additional flags:
+
+  --cpu_only: Force building only a CPU version. However, if
+      torch.cuda.is_available() is False, it will default to CPU_ONLY.
+
+  --force_cuda: If torch.cuda.is_available() is false, but you have a working
+      nvcc, compile cuda files. --force_cuda will supercede --cpu_only.
+
+
+Additional arguments:
+
+  --blas=<value> : type of blas library to use for CPU matrix multiplications.
+      Options: [openblas, mkl, atlas, blas]. By default, it will use the first
+      numpy blas library it finds.
+
+  --cuda_home=<value> : a directory that contains <value>/bin/nvcc and
+      <value>/lib64/libcudart.so. By default, use
+      `torch.utils.cpp_extension._find_cuda_home()`.
+
+  --blas_include_dirs=<comma_separated_values> : additional include dirs. Only
+      activated when --blas=<value> is set.
+
+  --blas_library_dirs=<comma_separated_values> : additional library dirs. Only
+      activated when --blas=<value> is set.
+"""
 import sys
+
 if sys.version_info < (3, 6):
     sys.stdout.write(
         "Minkowski Engine requires Python 3.6 or higher. Please use anaconda https://www.anaconda.com/distribution/ for an isolated python environment.\n"
@@ -8,7 +51,7 @@
 try:
     import torch
 except ImportError:
-    raise ImportError('Pytorch not found. Please install pytorch first.')
+    raise ImportError("Pytorch not found. Please install pytorch first.")
 
 import codecs
 import os
@@ -20,25 +63,24 @@
 
 from distutils.sysconfig import get_python_inc
 
-if platform == 'win32':
-    raise ImportError('Windows is currently not supported.')
-elif platform == 'darwin':
+if platform == "win32":
+    raise ImportError("Windows is currently not supported.")
+elif platform == "darwin":
     # Set the distutils to use clang instead of g++ for valid std
-    os.environ['CC'] = '/usr/local/opt/llvm/bin/clang'
-    os.environ['CXX'] = '/usr/local/opt/llvm/bin/clang'
+    os.environ["CC"] = "/usr/local/opt/llvm/bin/clang"
+    os.environ["CXX"] = "/usr/local/opt/llvm/bin/clang"
 
 here = os.path.abspath(os.path.dirname(__file__))
 
 
 def read(*parts):
-    with codecs.open(os.path.join(here, *parts), 'r') as fp:
+    with codecs.open(os.path.join(here, *parts), "r") as fp:
         return fp.read()
 
 
 def find_version(*file_paths):
     version_file = read(*file_paths)
-    version_match = re.search(r"^__version__ = ['\"]([^'\"]*)['\"]",
-                              version_file, re.M)
+    version_match = re.search(r"^__version__ = ['\"]([^'\"]*)['\"]", version_file, re.M)
     if version_match:
         return version_match.group(1)
     raise RuntimeError("Unable to find version string.")
@@ -48,129 +90,154 @@ def run_command(*args):
     subprocess.check_call(args)
 
 
+def _argparse(pattern, argv, is_flag=True):
+    if is_flag:
+        found = pattern in argv
+        if found:
+            argv.remove(pattern)
+        return found, argv
+    else:
+        arr = [arg for arg in argv if pattern in arg]
+        if len(arr) == 0:  # not found
+            return False, argv
+        else:
+            assert "=" in arr[0], f"{arr[0]} requires a value."
+            argv.remove(arr[0])
+            return arr[0].split("=")[1], argv
+
+
 # For cpu only build
-CPU_ONLY = '--cpu_only' in argv or not torch.cuda.is_available()
-KEEP_OBJS = '--keep_objs' in argv
-BLAS = [arg for arg in argv if '--blas' in arg]
+CPU_ONLY, argv = _argparse("--cpu_only", argv)
+CPU_ONLY = CPU_ONLY or not torch.cuda.is_available()
+KEEP_OBJS, argv = _argparse("--keep_objs", argv)
+FORCE_CUDA, argv = _argparse("--force_cuda", argv)
+
+# args with return value
+CUDA_HOME, argv = _argparse("--cuda_home", argv, False)
+BLAS, argv = _argparse("--blas", argv, False)
+BLAS_INCLUDE_DIRS, argv = _argparse("--blas_include_dirs", argv, False)
+BLAS_LIBRARY_DIRS, argv = _argparse("--blas_library_dirs", argv, False)
 
 Extension = CUDAExtension
 compile_args = [
-    'make',
-    '-j%d' % min(os.cpu_count(), 12),  # parallel compilation
-    'PYTHON=' + sys.executable,  # curr python
+    "make",
+    "-j%d" % min(os.cpu_count(), 12),  # parallel compilation
+    "PYTHON=" + sys.executable,  # curr python
 ]
-extra_compile_args = ['-Wno-deprecated-declarations']
+
+extra_compile_args = ["-Wno-deprecated-declarations"]
 extra_link_args = []
-libraries = ['minkowski']
+libraries = ["minkowski"]
 
 # extra_compile_args+=['-g']  # Uncomment for debugging
-if CPU_ONLY:
-    print('\nCPU_ONLY build')
-    if '--cpu_only' in argv:
-        argv.remove('--cpu_only')
-    compile_args += ['CPU_ONLY=1']
-    extra_compile_args += ['-DCPU_ONLY']
+if CPU_ONLY and not FORCE_CUDA:
+    print("\nCPU_ONLY build set")
+    compile_args += ["CPU_ONLY=1"]
+    extra_compile_args += ["-DCPU_ONLY"]
     Extension = CppExtension
 else:
     # system python installation
-    libraries.append('cusparse')
+    libraries.append("cusparse")
+
+if not (CUDA_HOME is False):  # False when not set, str otherwise
+    print(f"Using CUDA_HOME={CUDA_HOME}")
+    compile_args += [f"CUDA_HOME={CUDA_HOME}"]
 
 if KEEP_OBJS:
-    print('\nUsing built objects')
-    argv.remove('--keep_objs')
-
-if len(BLAS) > 0:
-    BLAS = BLAS[0]
-    argv.remove(BLAS)
-    BLAS = BLAS.split('=')[1]
-    assert BLAS in ['openblas', 'mkl', 'atlas', 'blas']
+    print("\nUsing built objects")
+
+BLAS_LIST = ["openblas", "mkl", "atlas", "blas"]
+if not (BLAS is False):  # False only when not set, str otherwise
+    assert BLAS in BLAS_LIST
     libraries.append(BLAS)
-    blas_inc_dirs = os.environ.get('BLAS_INCLUDE_DIRS')
-    compile_args += [f'BLAS_INCLUDE_DIRS={blas_inc_dirs}']
-    blas_lib_dirs = os.environ.get('BLAS_LIBRARY_DIRS')
-    if blas_lib_dirs is not None:
-        extra_link_args += [f'-Wl,-rpath,{blas_lib_dirs}']
+    if not (BLAS_INCLUDE_DIRS is False):
+        compile_args += [f"BLAS_INCLUDE_DIRS={BLAS_INCLUDE_DIRS}"]
+    if not (BLAS_LIBRARY_DIRS is False):
+        extra_link_args += [f"-Wl,-rpath,{BLAS_LIBRARY_DIRS}"]
 else:
     # find the default BLAS library
     import numpy.distutils.system_info as sysinfo
+
     # Search blas in this order
-    for blas in ['openblas', 'atlas', 'mkl', 'blas']:
-        if 'libraries' in sysinfo.get_info(blas):
+    for blas in BLAS_LIST:
+        if "libraries" in sysinfo.get_info(blas):
             BLAS = blas
-            libraries += sysinfo.get_info(blas)['libraries']
+            libraries += sysinfo.get_info(blas)["libraries"]
             break
     else:
         # BLAS not found
-        raise ImportError(' \
+        raise ImportError(
+            ' \
 \nBLAS not found from numpy.distutils.system_info.get_info. \
 \nPlease specify BLAS with: python setup.py install --blas=openblas" \
 \nfor more information, please visit https://github.com/StanfordVL/MinkowskiEngine/wiki/Installation'
-                         )
+        )
 
-print(f'\nUsing BLAS={BLAS}')
+print(f"\nUsing BLAS={BLAS}")
 
-compile_args += ['BLAS=' + BLAS]
+compile_args += ["BLAS=" + BLAS]
 
-if 'darwin' in platform:
-    extra_compile_args += ['-stdlib=libc++']
+if "darwin" in platform:
+    extra_compile_args += ["-stdlib=libc++"]
 
 if not KEEP_OBJS:
-    run_command('make', 'clean')
+    run_command("make", "clean")
 
 run_command(*compile_args)
 
 # Python interface
 setup(
-    name='MinkowskiEngine',
-    version=find_version('MinkowskiEngine', '__init__.py'),
-    install_requires=['torch', 'numpy'],
-    packages=[
-        'MinkowskiEngine', 'MinkowskiEngine.utils', 'MinkowskiEngine.modules'
-    ],
-    package_dir={'MinkowskiEngine': './MinkowskiEngine'},
+    name="MinkowskiEngine",
+    version=find_version("MinkowskiEngine", "__init__.py"),
+    install_requires=["torch", "numpy"],
+    packages=["MinkowskiEngine", "MinkowskiEngine.utils", "MinkowskiEngine.modules"],
+    package_dir={"MinkowskiEngine": "./MinkowskiEngine"},
     ext_modules=[
         Extension(
-            name='MinkowskiEngineBackend',
+            name="MinkowskiEngineBackend",
             include_dirs=[here, get_python_inc() + "/.."],
-            library_dirs=['objs'],
-            sources=[
-                'pybind/minkowski.cpp',
-            ],
+            library_dirs=["objs"],
+            sources=["pybind/minkowski.cpp",],
             libraries=libraries,
             extra_compile_args=extra_compile_args,
             extra_link_args=extra_link_args,
         )
     ],
-    cmdclass={'build_ext': BuildExtension},
-    author='Christopher Choy',
-    author_email='chrischoy@ai.stanford.edu',
-    description='a convolutional neural network library for sparse tensors',
-    long_description=read('README.md'),
+    cmdclass={"build_ext": BuildExtension},
+    author="Christopher Choy",
+    author_email="chrischoy@ai.stanford.edu",
+    description="a convolutional neural network library for sparse tensors",
+    long_description=read("README.md"),
     long_description_content_type="text/markdown",
-    url='https://github.com/StanfordVL/MinkowskiEngine',
+    url="https://github.com/StanfordVL/MinkowskiEngine",
     keywords=[
-        'pytorch', 'Minkowski Engine', 'Sparse Tensor',
-        'Convolutional Neural Networks', '3D Vision', 'Deep Learning'
+        "pytorch",
+        "Minkowski Engine",
+        "Sparse Tensor",
+        "Convolutional Neural Networks",
+        "3D Vision",
+        "Deep Learning",
     ],
     zip_safe=False,
     classifiers=[
         # https://pypi.org/classifiers/
-        'Environment :: Console',
-        'Development Status :: 3 - Alpha',
-        'Intended Audience :: Developers',
-        'Intended Audience :: Other Audience',
-        'Intended Audience :: Science/Research',
-        'License :: OSI Approved :: MIT License',
-        'Natural Language :: English',
-        'Programming Language :: C++',
-        'Programming Language :: Python :: 3.6',
-        'Programming Language :: Python :: 3.7',
-        'Programming Language :: Python :: 3.8',
-        'Topic :: Multimedia :: Graphics',
-        'Topic :: Scientific/Engineering',
-        'Topic :: Scientific/Engineering :: Artificial Intelligence',
-        'Topic :: Scientific/Engineering :: Mathematics',
-        'Topic :: Scientific/Engineering :: Physics',
-        'Topic :: Scientific/Engineering :: Visualization',
+        "Environment :: Console",
+        "Development Status :: 3 - Alpha",
+        "Intended Audience :: Developers",
+        "Intended Audience :: Other Audience",
+        "Intended Audience :: Science/Research",
+        "License :: OSI Approved :: MIT License",
+        "Natural Language :: English",
+        "Programming Language :: C++",
+        "Programming Language :: Python :: 3.6",
+        "Programming Language :: Python :: 3.7",
+        "Programming Language :: Python :: 3.8",
+        "Topic :: Multimedia :: Graphics",
+        "Topic :: Scientific/Engineering",
+        "Topic :: Scientific/Engineering :: Artificial Intelligence",
+        "Topic :: Scientific/Engineering :: Mathematics",
+        "Topic :: Scientific/Engineering :: Physics",
+        "Topic :: Scientific/Engineering :: Visualization",
     ],
-    python_requires='>=3.6')
+    python_requires=">=3.6",
+)