Merge pull request pyFFTW#241 from grlee77/default_num_threads_v2

Allow run-time choice of the default number of threads and planning effort

docs/source/pyfftw/pyfftw.rst

@@ -98,3 +98,35 @@ Utility Functions
 .. autofunction:: pyfftw.is_n_byte_aligned
 
 .. autofunction:: pyfftw.next_fast_len
+
+.. _configuration_variables:
+
+FFTW Configuration
+------------------
+
+.. data:: pyfftw.config.NUM_THREADS
+
+   This variable controls the default number of threads used by the functions
+   in :mod:`pyfftw.builders` and :mod:`pyfftw.interfaces`.
+
+   The default value is read from the environment variable
+   ``PYFFTW_NUM_THREADS``. If this variable is undefined and the user's
+   underlying FFTW library was built using OpenMP threading, the number of
+   threads will be read from the environment variable ``OMP_NUM_THREADS``
+   instead. If neither environment variable is defined, the default value is 1.
+
+   If the specified value is ``<= 0``, the library will use
+   :func:`multiprocessing.cpu_count` to determine the number of threads.
+
+   The user can modify the value at run time by assigning to this variable.
+
+.. data:: pyfftw.config.PLANNER_EFFORT
+
+   This variable controls the default planning effort used by the functions
+   in :mod:`pyfftw.builders` and :mod:`pyfftw.interfaces`.
+
+   The default value of is determined by reading from the environment variable
+   ``PYFFTW_PLANNER_EFFORT``. If this environment variable is undefined, it
+   defaults to ``'FFTW_ESTIMATE'``.
+
+   The user can modify the value at run time by assigning to this variable.

docs/source/tutorial.rst

@@ -102,14 +102,19 @@ Monkey patching 3rd party libraries
 Since :mod:`pyfftw.interfaces.numpy_fft` and
 :mod:`pyfftw.interfaces.scipy_fftpack` are drop-in replacements for their
 :mod:`numpy.fft` and :mod:`scipy.fftpack` libraries respectively, it is
-possible use them as replacements at run-time through monkey patching.
+possible to use them as replacements at run-time through monkey patching. Note
+that the interfaces (and builders) all currently default to a single thread.
+The number of threads to use can be configured by assigning a positive integer
+to `pyfftw.config.NUM_THREADS` (see more details under
+:ref:configuration <interfaces_tutorial>).
 
 The following code demonstrates :func:`scipy.signal.fftconvolve` being monkey
 patched in order to speed it up.
 
 .. testcode::
 
    import pyfftw
+   import multiprocessing
    import scipy.signal
    import numpy
    from timeit import Timer
@@ -123,7 +128,10 @@ patched in order to speed it up.
    t = Timer(lambda: scipy.signal.fftconvolve(a, b))
 
    print('Time with scipy.fftpack: %1.3f seconds' % t.timeit(number=100))
-
+
+   # Configure PyFFTW to use all cores (the default is single-threaded)
+   pyfftw.config.NUM_THREADS = multiprocessing.cpu_count()
+
    # Monkey patch fftpack with pyfftw.interfaces.scipy_fftpack
    scipy.fftpack = pyfftw.interfaces.scipy_fftpack
    scipy.signal.fftconvolve(a, b) # We cheat a bit by doing the planning first
@@ -464,6 +472,35 @@ Understanding :mod:`numpy.fft`, these functions are largely
 self-explanatory. We point the reader to the :mod:`API docs <pyfftw.builders>`
 for more information.
 
+.. _configuration:
+
+Configuring FFTW planning effort and number of threads
+------------------------------------------------------
+The user may set the default number of threads used by the interfaces and
+builders at run time by assigning to ``pyfftw.config.NUM_THREADS``. Similarly
+the default
+`planning effort <http://www.fftw.org/fftw3_doc/Planner-Flags.html>`_
+may be set by assigning a string such as ``'FFTW_ESTIMATE'`` or
+``'FFTW_MEASURE'`` to ``pyfftw.config.PLANNER_EFFORT``.
+
+For example, to change the effort to ``'FFTW_MEASURE'`` and specify 4 threads:
+
+.. testcode::
+
+   import pyfftw
+
+   pyfftw.config.NUM_THREADS = 4
+
+   pyfftw.config.PLANNER_EFFORT = 'FFTW_MEASURE'
+
+All functions in :mod:`pyfftw.interfaces` and :mod:`pyfftw.builders` use the
+values from :mod:`pyfftw.config` when determining the default number of threads
+and planning effort.
+
+The initial values in pyfftw.config at import time can be controlled via the
+environment variables as detailed in the
+:ref:`configuration <_configuration_variables>` documentation.
+
 .. rubric:: Footnotes
 
 .. [#caveat] :mod:`pyfftw.interfaces` deals with repeated values in the

pyfftw/__init__.py

@@ -8,9 +8,13 @@
 
 This module represents the full interface to the underlying `FFTW
 library <http://www.fftw.org/>`_. However, users may find it easier to
-use the helper routines provided in :mod:`pyfftw.builders`.
+use the helper routines provided in :mod:`pyfftw.builders`. Default values
+used by the helper routines can be controlled as via
+:ref:`configuration variables <configuration_variables>`.
 '''
 
+import os
+
 from .pyfftw import (
         FFTW,
         export_wisdom,
@@ -30,12 +34,15 @@
         _supported_nptypes_complex,
         _supported_nptypes_real,
         _all_types_human_readable,
-        _all_types_np
+        _all_types_np,
+        _threading_type
 )
 
+from . import config
 from . import builders
 from . import interfaces
 
+
 # clean up the namespace
 del builders.builders
 

pyfftw/builders/_utils.py

@@ -50,9 +50,12 @@
 and may change in future versions.
 '''
 
+import multiprocessing
 import pyfftw
 import numpy
 import warnings
+from .. import _threading_type
+from .. import config
 
 __all__ = ['_FFTWWrapper', '_rc_dtype_pairs', '_default_dtype', '_Xfftn',
         '_setup_input_slicers', '_compute_array_shapes', '_precook_1d_args',
@@ -92,6 +95,27 @@
     raise NotImplementedError("No default precision available")
 
 
+def _default_effort(effort):
+    if effort is None:
+        return config.PLANNER_EFFORT
+    else:
+        return effort
+
+
+def _default_threads(threads):
+    if threads is None:
+        if config.NUM_THREADS <= 0:
+            return multiprocessing.cpu_count()
+        return config.NUM_THREADS
+    else:
+        if threads > 1 and _threading_type is None:
+            raise ValueError("threads > 1 requested, but pyFFTW was not built "
+                             "with multithreaded FFTW.")
+        elif threads <= 0:
+            return multiprocessing.cpu_count()
+        return threads
+
+
 def _unitary(norm):
     """_unitary() utility copied from numpy"""
     if norm not in (None, "ortho"):