Add option to build extensions in parallel (#1882)

* build extensions in parallel

* fix: setup.py develop supports --parallel

* update README.md

* Update README.md: Remove an `export` command

Author: yuantailing

README.md

@@ -135,6 +135,12 @@ pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation -
 pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --global-option="--cpp_ext" --global-option="--cuda_ext" ./
 ```
 
+To reduce the build time of APEX, parallel building can be enhanced via
+```bash
+NVCC_APPEND_FLAGS="--threads 4" pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext --cuda_ext --parallel 8" ./
+```
+When CPU cores or memory are limited, the `--parallel` option is generally preferred over `--threads`. See [pull#1882](https://github.com/NVIDIA/apex/pull/1882) for more details.
+
 APEX also supports a Python-only build via
 ```bash
 pip install -v --disable-pip-version-check --no-build-isolation --no-cache-dir ./

setup.py

@@ -1,6 +1,7 @@
 import sys
 import warnings
 import os
+import threading
 import glob
 from packaging.version import parse, Version
 
@@ -859,6 +860,44 @@ def check_cudnn_version_and_warn(global_option: str, required_cudnn_version: int
     )
 
 
+# Patch because `setup.py bdist_wheel` and `setup.py develop` do not support the `parallel` option
+parallel = None
+if "--parallel" in sys.argv:
+    idx = sys.argv.index("--parallel")
+    parallel = int(sys.argv[idx + 1])
+    sys.argv.pop(idx + 1)
+    sys.argv.pop(idx)
+
+
+# Prevent file conflicts when multiple extensions are compiled simultaneously
+class BuildExtensionSeparateDir(BuildExtension):
+    build_extension_patch_lock = threading.Lock()
+    thread_ext_name_map = {}
+
+    def finalize_options(self):
+        if parallel is not None:
+            self.parallel = parallel
+        super().finalize_options()
+
+    def build_extension(self, ext):
+        with self.build_extension_patch_lock:
+            if not getattr(self.compiler, "_compile_separate_output_dir", False):
+                compile_orig = self.compiler.compile
+
+                def compile_new(*args, **kwargs):
+                    return compile_orig(*args, **{
+                        **kwargs,
+                        "output_dir": os.path.join(
+                            kwargs["output_dir"],
+                            self.thread_ext_name_map[threading.current_thread().ident]),
+                    })
+                self.compiler.compile = compile_new
+                self.compiler._compile_separate_output_dir = True
+        self.thread_ext_name_map[threading.current_thread().ident] = ext.name
+        objects = super().build_extension(ext)
+        return objects
+
+
 setup(
     name="apex",
     version="0.1",
@@ -868,6 +907,6 @@ def check_cudnn_version_and_warn(global_option: str, required_cudnn_version: int
     install_requires=["packaging>20.6"],
     description="PyTorch Extensions written by NVIDIA",
     ext_modules=ext_modules,
-    cmdclass={"build_ext": BuildExtension} if ext_modules else {},
+    cmdclass={"build_ext": BuildExtensionSeparateDir} if ext_modules else {},
     extras_require=extras,
 )