Merge pull request #5161 from juliantaylor/parallel-distutils

ENH: support parallel compilation of extensions

8 files changed, 153 insertions, 43 deletions

diff --git a/INSTALL.txt b/INSTALL.txt
index 278aab9ef..9c12ba602 100644
--- a/INSTALL.txt
+++ b/INSTALL.txt
@@ -35,6 +35,21 @@ Building NumPy requires the following software installed:
 Python__ http://www.python.org
 nose__ http://somethingaboutorange.com/mrl/projects/nose/
 
+Basic Installation
+==================
+
+To install numpy run:
+
+    python setup.py build -j 4 install --prefix $HOME/.local
+
+This will compile numpy on 4 CPUs and install it into the specified prefix.
+To perform an inplace build that can be run from the source folder run:
+
+   python setup.py build_ext --inplace -j 4
+
+The number of build jobs can also be specified via the environment variable
+NPY_NUM_BUILD_JOBS.
+
 Fortran ABI mismatch
 ====================
 
@@ -65,40 +80,34 @@ this means that g77 has been used. If libgfortran.so is a dependency,
 gfortran has been used. If both are dependencies, this means both have been
 used, which is almost always a very bad idea.
 
-Building with ATLAS support
-===========================
-
-Ubuntu 8.10 (Intrepid)
-----------------------
-
-You can install the necessary packages for optimized ATLAS with this
-command:
-
-    sudo apt-get install libatlas-base-dev
-
-If you have a recent CPU with SIMD support (SSE, SSE2, etc...), you should
-also install the corresponding package for optimal performance. For
-example, for SSE2:
+Building with optimized BLAS support
+====================================
 
-    sudo apt-get install libatlas3gf-sse2
+Ubuntu/Debian
+-------------
 
-*NOTE*: if you build your own atlas, Intrepid changed its default fortran
-compiler to gfortran. So you should rebuild everything from scratch,
-including lapack, to use it on Intrepid.
+In order to build with optimized a BLAS providing development package must be installed.
+Options are for example:
 
-Ubuntu 8.04 and lower
----------------------
+ - libblas-dev
+   reference BLAS not very optimized
+ - libatlas-base-dev
+   generic tuned ATLAS, it is recommended to tune it to the available hardware,
+   see /usr/share/doc/libatlas3-base/README.Debian for instructions
+ - libopenblas-base
+   fast and runtime detected so no tuning required but as of version 2.11 still
+   suffers from correctness issues on some CPUs, test your applications
+   thoughly.
 
-You can install the necessary packages for optimized ATLAS with this
-command:
+The actual implementation can be exchanged also after installation via the
+alternatives mechanism:
 
-    sudo apt-get install atlas3-base-dev
+   update-alternatives --config libblas.so.3
+   update-alternatives --config liblapack.so.3
 
-If you have a recent CPU with SIMD support (SSE, SSE2, etc...), you should
-also install the corresponding package for optimal performance. For
-example, for SSE2:
+Or by preloading a specific BLAS library with
+    LD_PRELOAD=/usr/lib/atlas-base/atlas/libblas.so.3 python ...
 
-    sudo apt-get install atlas3-sse2
 
 Windows 64 bits notes
 =====================
diff --git a/doc/release/1.10.0-notes.rst b/doc/release/1.10.0-notes.rst
index c387f0ec7..cfba0a5cd 100644
--- a/doc/release/1.10.0-notes.rst
+++ b/doc/release/1.10.0-notes.rst
@@ -6,6 +6,8 @@ This release supports Python 2.6 - 2.7 and 3.2 - 3.4.
 
 Highlights
 ==========
+* numpy.distutils now supports parallel compilation via the --jobs/-j argument
+  passed to setup.py build
 
 
 Dropped Support
@@ -35,6 +37,13 @@ New Features
 Compared to `np.power(x, 1./3.)` it is well defined for negative real floats
 and a bit faster.
 
+numpy.distutils now allows parallel compilation
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+By passing `--jobs=n` or `-j n` to `setup.py build` the compilation of
+extensions is now performed in `n` parallel processes.
+The parallelization is limited to files within one extension so projects using
+Cython will not profit because it builds extensions from single files.
+
 
 Improvements
 ============
diff --git a/numpy/distutils/ccompiler.py b/numpy/distutils/ccompiler.py
index 8484685c0..d7fe702a6 100644
--- a/numpy/distutils/ccompiler.py
+++ b/numpy/distutils/ccompiler.py
@@ -16,7 +16,7 @@ from distutils.version import LooseVersion
 from numpy.distutils import log
 from numpy.distutils.exec_command import exec_command
 from numpy.distutils.misc_util import cyg2win32, is_sequence, mingw32, \
-                                      quote_args
+                                      quote_args, get_num_build_jobs
 from numpy.distutils.compat import get_exception
 
 
@@ -165,9 +165,10 @@ def CCompiler_compile(self, sources, output_dir=None, macros=None,
         return []
     # FIXME:RELATIVE_IMPORT
     if sys.version_info[0] < 3:
-        from .fcompiler import FCompiler
+        from .fcompiler import FCompiler, is_f_file, has_f90_header
     else:
-        from numpy.distutils.fcompiler import FCompiler
+        from numpy.distutils.fcompiler import (FCompiler, is_f_file,
+                                               has_f90_header)
     if isinstance(self, FCompiler):
         display = []
         for fc in ['f77', 'f90', 'fix']:
@@ -189,20 +190,45 @@ def CCompiler_compile(self, sources, output_dir=None, macros=None,
         display += "\nextra options: '%s'" % (' '.join(extra_postargs))
     log.info(display)
 
-    # build any sources in same order as they were originally specified
-    #   especially important for fortran .f90 files using modules
+    def single_compile(args):
+        obj, (src, ext) = args
+        self._compile(obj, src, ext, cc_args, extra_postargs, pp_opts)
+
     if isinstance(self, FCompiler):
         objects_to_build = list(build.keys())
+        f77_objects, other_objects = [], []
         for obj in objects:
             if obj in objects_to_build:
                 src, ext = build[obj]
                 if self.compiler_type=='absoft':
                     obj = cyg2win32(obj)
                     src = cyg2win32(src)
-                self._compile(obj, src, ext, cc_args, extra_postargs, pp_opts)
+                if is_f_file(src) and not has_f90_header(src):
+                    f77_objects.append((obj, (src, ext)))
+                else:
+                    other_objects.append((obj, (src, ext)))
+
+        # f77 objects can be built in parallel
+        build_items = f77_objects
+        # build f90 modules serial, module files are generated during
+        # compilation and may be used by files later in the list so the
+        # ordering is important
+        for o in other_objects:
+            single_compile(o)
+    else:
+        build_items = build.items()
+
+    jobs = get_num_build_jobs()
+    if len(build) > 1 and jobs > 1:
+        # build parallel
+        import multiprocessing.pool
+        pool = multiprocessing.pool.ThreadPool(jobs)
+        pool.map(single_compile, build_items)
+        pool.close()
     else:
-        for obj, (src, ext) in build.items():
-            self._compile(obj, src, ext, cc_args, extra_postargs, pp_opts)
+        # build serial
+        for o in build_items:
+            single_compile(o)
 
     # Return *all* object filenames, not just the ones we just built.
     return objects
diff --git a/numpy/distutils/command/build.py b/numpy/distutils/command/build.py
index b6912be15..f7249ae81 100644
--- a/numpy/distutils/command/build.py
+++ b/numpy/distutils/command/build.py
@@ -16,6 +16,8 @@ class build(old_build):
     user_options = old_build.user_options + [
         ('fcompiler=', None,
          "specify the Fortran compiler type"),
+        ('jobs=', 'j',
+         "number of parallel jobs"),
         ]
 
     help_options = old_build.help_options + [
@@ -26,8 +28,14 @@ class build(old_build):
     def initialize_options(self):
         old_build.initialize_options(self)
         self.fcompiler = None
+        self.jobs = None
 
     def finalize_options(self):
+        if self.jobs:
+            try:
+                self.jobs = int(self.jobs)
+            except ValueError:
+                raise ValueError("--jobs/-j argument must be an integer")
         build_scripts = self.build_scripts
         old_build.finalize_options(self)
         plat_specifier = ".%s-%s" % (get_platform(), sys.version[0:3])
diff --git a/numpy/distutils/command/build_clib.py b/numpy/distutils/command/build_clib.py
index 84ca87250..6e65a3bfb 100644
--- a/numpy/distutils/command/build_clib.py
+++ b/numpy/distutils/command/build_clib.py
@@ -30,6 +30,8 @@ class build_clib(old_build_clib):
         ('fcompiler=', None,
          "specify the Fortran compiler type"),
         ('inplace', 'i', 'Build in-place'),
+        ('jobs=', 'j',
+         "number of parallel jobs"),
         ]
 
     boolean_options = old_build_clib.boolean_options + ['inplace']
@@ -38,7 +40,16 @@ class build_clib(old_build_clib):
         old_build_clib.initialize_options(self)
         self.fcompiler = None
         self.inplace = 0
-        return
+        self.jobs = None
+
+    def finalize_options(self):
+        if self.jobs:
+            try:
+                self.jobs = int(self.jobs)
+            except ValueError:
+                raise ValueError("--jobs/-j argument must be an integer")
+        old_build_clib.finalize_options(self)
+        self.set_undefined_options('build', ('jobs', 'jobs'))
 
     def have_f_sources(self):
         for (lib_name, build_info) in self.libraries:
diff --git a/numpy/distutils/command/build_ext.py b/numpy/distutils/command/build_ext.py
index 68aab21eb..59c453607 100644
--- a/numpy/distutils/command/build_ext.py
+++ b/numpy/distutils/command/build_ext.py
@@ -34,6 +34,8 @@ class build_ext (old_build_ext):
     user_options = old_build_ext.user_options + [
         ('fcompiler=', None,
          "specify the Fortran compiler type"),
+        ('jobs=', 'j',
+         "number of parallel jobs"),
         ]
 
     help_options = old_build_ext.help_options + [
@@ -44,12 +46,19 @@ class build_ext (old_build_ext):
     def initialize_options(self):
         old_build_ext.initialize_options(self)
         self.fcompiler = None
+        self.jobs = None
 
     def finalize_options(self):
+        if self.jobs:
+            try:
+                self.jobs = int(self.jobs)
+            except ValueError:
+                raise ValueError("--jobs/-j argument must be an integer")
         incl_dirs = self.include_dirs
         old_build_ext.finalize_options(self)
         if incl_dirs is not None:
             self.include_dirs.extend(self.distribution.include_dirs or [])
+        self.set_undefined_options('build', ('jobs', 'jobs'))
 
     def run(self):
         if not self.extensions:
diff --git a/numpy/distutils/misc_util.py b/numpy/distutils/misc_util.py
index c146178f0..cab0a4cba 100644
--- a/numpy/distutils/misc_util.py
+++ b/numpy/distutils/misc_util.py
@@ -13,6 +13,10 @@ import shutil
 
 import distutils
 from distutils.errors import DistutilsError
+try:
+    from threading import local as tlocal
+except ImportError:
+    from dummy_threading import local as tlocal
 
 try:
     set
@@ -31,7 +35,8 @@ __all__ = ['Configuration', 'get_numpy_include_dirs', 'default_config_dict',
            'get_script_files', 'get_lib_source_files', 'get_data_files',
            'dot_join', 'get_frame', 'minrelpath', 'njoin',
            'is_sequence', 'is_string', 'as_list', 'gpaths', 'get_language',
-           'quote_args', 'get_build_architecture', 'get_info', 'get_pkg_info']
+           'quote_args', 'get_build_architecture', 'get_info', 'get_pkg_info',
+           'get_num_build_jobs']
 
 class InstallableLib(object):
     """
@@ -60,6 +65,36 @@ class InstallableLib(object):
         self.build_info = build_info
         self.target_dir = target_dir
 
+
+def get_num_build_jobs():
+    """
+    Get number of parallel build jobs set by the --jobs command line argument
+    of setup.py
+    If the command did not receive a setting the environment variable
+    NPY_NUM_BUILD_JOBS checked and if that is unset it returns 1.
+
+    Returns
+    -------
+    out : int
+        number of parallel jobs that can be run
+
+    """
+    from numpy.distutils.core import get_distribution
+    envjobs = int(os.environ.get("NPY_NUM_BUILD_JOBS", 1))
+    dist = get_distribution()
+    # may be None during configuration
+    if dist is None:
+        return envjobs
+
+    # any of these three may have the job set, take the largest
+    cmdattr = (getattr(dist.get_command_obj('build'), 'jobs'),
+               getattr(dist.get_command_obj('build_ext'), 'jobs'),
+               getattr(dist.get_command_obj('build_clib'), 'jobs'))
+    if all(x is None for x in cmdattr):
+        return envjobs
+    else:
+        return max(x for x in cmdattr if x is not None)
+
 def quote_args(args):
     # don't used _nt_quote_args as it does not check if
     # args items already have quotes or not.
@@ -249,9 +284,9 @@ def gpaths(paths, local_path='', include_non_existing=True):
     return _fix_paths(paths, local_path, include_non_existing)
 
 
-_temporary_directory = None
 def clean_up_temporary_directory():
-    global _temporary_directory
+    tdata = tlocal()
+    _temporary_directory = getattr(tdata, 'tempdir', None)
     if not _temporary_directory:
         return
     try:
@@ -261,13 +296,13 @@ def clean_up_temporary_directory():
     _temporary_directory = None
 
 def make_temp_file(suffix='', prefix='', text=True):
-    global _temporary_directory
-    if not _temporary_directory:
-        _temporary_directory = tempfile.mkdtemp()
+    tdata = tlocal()
+    if not hasattr(tdata, 'tempdir'):
+        tdata.tempdir = tempfile.mkdtemp()
         atexit.register(clean_up_temporary_directory)
     fid, name = tempfile.mkstemp(suffix=suffix,
                                  prefix=prefix,
-                                 dir=_temporary_directory,
+                                 dir=tdata.tempdir,
                                  text=text)
     fo = os.fdopen(fid, 'w')
     return fo, name
diff --git a/tools/travis-test.sh b/tools/travis-test.sh
index 17d520891..3981c3b58 100755
--- a/tools/travis-test.sh
+++ b/tools/travis-test.sh
@@ -1,6 +1,9 @@
 #!/bin/sh
 set -ex
 
+# travis boxes give you 1.5 cpus
+export NPY_NUM_BUILD_JOBS=2
+
 # setup env
 if [ -r /usr/lib/libeatmydata/libeatmydata.so ]; then
   # much faster package installation