feat: parallel compiler (#2521)

Author: henryiii

docs/compiling.rst

@@ -68,6 +68,23 @@ that is supported via a ``build_ext`` command override; it will only affect
         ext_modules=ext_modules
     )
 
+Since pybind11 does not require NumPy when building, a light-weight replacement
+for NumPy's parallel compilation distutils tool is included. Use it like this:
+
+    from pybind11.setup_helpers import ParallelCompile
+
+    # Optional multithreaded build
+    ParallelCompile("NPY_NUM_BUILD_JOBS").install()
+
+    setup(...
+
+The argument is the name of an environment variable to control the number of
+threads, such as ``NPY_NUM_BUILD_JOBS`` (as used by NumPy), though you can set
+something different if you want. You can also pass ``default=N`` to set the
+default number of threads (0 will take the number of threads available) and
+``max=N``, the maximum number of threads; if you have a large extension you may
+want set this to a memory dependent number.
+
 .. _setup_helpers-pep518:
 
 PEP 518 requirements (Pip 10+ required)

pybind11/setup_helpers.py

@@ -49,6 +49,7 @@
     from distutils.extension import Extension as _Extension
 
 import distutils.errors
+import distutils.ccompiler
 
 
 WIN = sys.platform.startswith("win32")
@@ -279,3 +280,108 @@ def build_extensions(self):
         # Python 2 doesn't allow super here, since distutils uses old-style
         # classes!
         _build_ext.build_extensions(self)
+
+
+# Optional parallel compile utility
+# inspired by: http://stackoverflow.com/questions/11013851/speeding-up-build-process-with-distutils
+# and: https://github.com/tbenthompson/cppimport/blob/stable/cppimport/build_module.py
+# and NumPy's parallel distutils module:
+#              https://github.com/numpy/numpy/blob/master/numpy/distutils/ccompiler.py
+class ParallelCompile(object):
+    """
+    Make a parallel compile function. Inspired by
+    numpy.distutils.ccompiler.CCompiler_compile and cppimport.
+
+    This takes several arguments that allow you to customize the compile
+    function created:
+
+    envvar: Set an environment variable to control the compilation threads, like NPY_NUM_BUILD_JOBS
+    default: 0 will automatically multithread, or 1 will only multithread if the envvar is set.
+    max: The limit for automatic multithreading if non-zero
+
+    To use:
+        ParallelCompile("NPY_NUM_BUILD_JOBS").install()
+    or:
+        with ParallelCompile("NPY_NUM_BUILD_JOBS"):
+            setup(...)
+    """
+
+    __slots__ = ("envvar", "default", "max", "old")
+
+    def __init__(self, envvar=None, default=0, max=0):
+        self.envvar = envvar
+        self.default = default
+        self.max = max
+        self.old = []
+
+    def function(self):
+        """
+        Builds a function object usable as distutils.ccompiler.CCompiler.compile.
+        """
+
+        def compile_function(
+            compiler,
+            sources,
+            output_dir=None,
+            macros=None,
+            include_dirs=None,
+            debug=0,
+            extra_preargs=None,
+            extra_postargs=None,
+            depends=None,
+        ):
+
+            # These lines are directly from distutils.ccompiler.CCompiler
+            macros, objects, extra_postargs, pp_opts, build = compiler._setup_compile(
+                output_dir, macros, include_dirs, sources, depends, extra_postargs
+            )
+            cc_args = compiler._get_cc_args(pp_opts, debug, extra_preargs)
+
+            # The number of threads; start with default.
+            threads = self.default
+
+            # Determine the number of compilation threads, unless set by an environment variable.
+            if self.envvar is not None:
+                threads = int(os.environ.get(self.envvar, self.default))
+
+            def _single_compile(obj):
+                try:
+                    src, ext = build[obj]
+                except KeyError:
+                    return
+                compiler._compile(obj, src, ext, cc_args, extra_postargs, pp_opts)
+
+            try:
+                import multiprocessing
+                from multiprocessing.pool import ThreadPool
+            except ImportError:
+                threads = 1
+
+            if threads == 0:
+                try:
+                    threads = multiprocessing.cpu_count()
+                    threads = self.max if self.max and self.max < threads else threads
+                except NotImplementedError:
+                    threads = 1
+
+            if threads > 1:
+                for _ in ThreadPool(threads).imap_unordered(_single_compile, objects):
+                    pass
+            else:
+                for ob in objects:
+                    _single_compile(ob)
+
+            return objects
+
+        return compile_function
+
+    def install(self):
+        distutils.ccompiler.CCompiler.compile = self.function()
+        return self
+
+    def __enter__(self):
+        self.old.append(distutils.ccompiler.CCompiler.compile)
+        return self.install()
+
+    def __exit__(self, *args):
+        distutils.ccompiler.CCompiler.compile = self.old.pop()

tests/extra_setuptools/test_setuphelper.py

@@ -10,8 +10,9 @@
 MAIN_DIR = os.path.dirname(os.path.dirname(DIR))
 
 
+@pytest.mark.parametrize("parallel", [False, True])
 @pytest.mark.parametrize("std", [11, 0])
-def test_simple_setup_py(monkeypatch, tmpdir, std):
+def test_simple_setup_py(monkeypatch, tmpdir, parallel, std):
     monkeypatch.chdir(tmpdir)
     monkeypatch.syspath_prepend(MAIN_DIR)
 
@@ -39,13 +40,18 @@ def test_simple_setup_py(monkeypatch, tmpdir, std):
                 cmdclass["build_ext"] = build_ext
 
 
+            parallel = {parallel}
+            if parallel:
+                from pybind11.setup_helpers import ParallelCompile
+                ParallelCompile().install()
+
             setup(
                 name="simple_setup_package",
                 cmdclass=cmdclass,
                 ext_modules=ext_modules,
             )
             """
-        ).format(MAIN_DIR=MAIN_DIR, std=std),
+        ).format(MAIN_DIR=MAIN_DIR, std=std, parallel=parallel),
         encoding="ascii",
     )