From b9d00273eefd4c07f7c9e27fbcf04fba590b798f Mon Sep 17 00:00:00 2001 From: Henry Schreiner Date: Fri, 2 Oct 2020 10:03:35 -0400 Subject: [PATCH] feat: parallel compiler (#2521) --- docs/compiling.rst | 17 ++++ pybind11/setup_helpers.py | 106 +++++++++++++++++++++ tests/extra_setuptools/test_setuphelper.py | 10 +- 3 files changed, 131 insertions(+), 2 deletions(-) diff --git a/docs/compiling.rst b/docs/compiling.rst index b9bf134b..39cde8c5 100644 --- a/docs/compiling.rst +++ b/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) diff --git a/pybind11/setup_helpers.py b/pybind11/setup_helpers.py index 7a201312..ced8542e 100644 --- a/pybind11/setup_helpers.py +++ b/pybind11/setup_helpers.py @@ -49,6 +49,7 @@ except ImportError: from distutils.extension import Extension as _Extension import distutils.errors +import distutils.ccompiler WIN = sys.platform.startswith("win32") @@ -279,3 +280,108 @@ class build_ext(_build_ext): # noqa: N801 # 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 de0b516a..0d8bd0e4 100644 --- a/tests/extra_setuptools/test_setuphelper.py +++ b/tests/extra_setuptools/test_setuphelper.py @@ -10,8 +10,9 @@ DIR = os.path.abspath(os.path.dirname(__file__)) 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", )