feat: parallel compiler (#2521)
diff --git a/docs/compiling.rst b/docs/compiling.rst
index b9bf134..39cde8c 100644
--- a/docs/compiling.rst
+++ b/docs/compiling.rst
@@ -68,6 +68,23 @@
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)
diff --git a/pybind11/setup_helpers.py b/pybind11/setup_helpers.py
index 7a20131..ced8542 100644
--- a/pybind11/setup_helpers.py
+++ b/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 @@
# 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()
diff --git a/tests/extra_setuptools/test_setuphelper.py b/tests/extra_setuptools/test_setuphelper.py
index de0b516..0d8bd0e 100644
--- a/tests/extra_setuptools/test_setuphelper.py
+++ b/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 @@
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",
)