Improve threading defaults, docs, and sys_info (#14103)

Author: larsoner

.gitignore

@@ -103,3 +103,4 @@ venv/
 .hypothesis/
 .ruff_cache/
 .ipynb_checkpoints/
+/.claude/

doc/changes/dev/14103.other.rst

@@ -0,0 +1 @@
+Decomposition-heavy operations such as :func:`mne.preprocessing.maxwell_filter`, :func:`mne.compute_covariance` and :meth:`mne.preprocessing.ICA.fit` now limit how many BLAS threads they use, which is often much faster; see :ref:`faq_cpu` to control this yourself, by `Eric Larson`_.

doc/help/faq.rst

@@ -174,27 +174,43 @@ support multithreading:
 - `OpenBLAS <http://www.openblas.net/>`_
 - `Intel Math Kernel Library (MKL) <https://www.intel.com/content/www/us/en/developer/tools/oneapi/onemkl.html>`_,
   which uses `OpenMP <https://www.openmp.org/>`_
+- `Accelerate <https://developer.apple.com/accelerate/>`_, used by the NumPy and
+  SciPy wheels on PyPI for macOS on Apple silicon
 
 To control how many cores are used for linear-algebra-heavy functions like
-:func:`mne.preprocessing.maxwell_filter`, you can set the ``OMP_NUM_THREADS``
-or ``OPENBLAS_NUM_THREADS`` environment variable to the desired number of cores
-for MKL or OpenBLAS, respectively.
+:func:`mne.preprocessing.maxwell_filter`, you can set an environment variable
+to the desired number of cores. **Which variable works depends on the threading
+layer your BLAS was built against, not just on which BLAS it is.**
+:func:`mne.sys_info` reports both, for example
+``OpenBLAS 0.3.33 with 10 threads (openmp threading layer)``.
+
+- OpenBLAS with the ``pthreads`` threading layer, typically shipped by the
+  NumPy and SciPy wheels on PyPI: use ``OPENBLAS_NUM_THREADS``.
+- OpenBLAS with the ``openmp`` threading layer, typically shipped by
+  ``conda-forge`` and by the :ref:`standalone installers <installers>`: use
+  ``OMP_NUM_THREADS``. On these builds ``OPENBLAS_NUM_THREADS`` is silently
+  ignored, because OpenBLAS delegates threading to OpenMP.
+- MKL: use ``MKL_NUM_THREADS``, or ``OMP_NUM_THREADS`` for any OpenMP-based
+  threading layer.
+- Accelerate: there is no effective setting, and none is needed. Accelerate
+  manages its own threads and ignores all of the variables above;
+  ``VECLIB_MAXIMUM_THREADS`` is advisory at best. Its performance is also
+  essentially flat with respect to thread settings, so there is nothing to tune.
 
 Using more threads is not always faster. In particular, decomposition-heavy
 operations such as :func:`mne.preprocessing.maxwell_filter` and
-:func:`mne.compute_covariance` with ``method="shrunk"`` can be slower when
-OpenBLAS uses all logical CPU cores on Linux. If one of these operations is
-unexpectedly slow, use :func:`mne.sys_info` to check the BLAS library and its
-current thread count. Compare a few settings, for example 1, 2, 4, and the
-number of physical CPU cores, using a representative part of your analysis.
-The best setting depends on the operation, CPU, and BLAS library.
+:func:`mne.compute_covariance` with ``method="shrunk"`` can be much slower when
+OpenBLAS or MKL uses all logical CPU cores. If one of these operations is
+unexpectedly slow, compare a few settings, for example 1, 2, 3, 4, and the
+number of physical CPU cores, using a representative part of your analysis. The
+best setting depends on the operation, CPU, and BLAS library.
 
 Set the environment variable before importing MNE-Python, NumPy, or SciPy. For
 example, start the analysis from the shell with:
 
 .. code-block:: console
 
-    $ OPENBLAS_NUM_THREADS=4 python analysis.py
+    $ OMP_NUM_THREADS=4 python analysis.py
 
 Changes made after the linear algebra library has been loaded might have no
 effect in the same Python session.

mne/cov.py

@@ -58,6 +58,7 @@
     _check_fname,
     _check_on_missing,
     _check_option,
+    _limit_blas_threads,
     _on_missing,
     _pl,
     _scaled_array,
@@ -1259,6 +1260,7 @@ def _compute_rank_raw_array(
     )
 
 
+@_limit_blas_threads()
 def _compute_covariance_auto(
     data,
     method,

mne/preprocessing/ica.py

@@ -68,6 +68,7 @@
     _check_preload,
     _ensure_int,
     _get_inst_data,
+    _limit_blas_threads,
     _on_missing,
     _pl,
     _reject_data_segments,
@@ -887,6 +888,7 @@ def _pre_whiten(self, data):
             data = self.pre_whitener_ @ data
         return data
 
+    @_limit_blas_threads()
     def _fit(self, data, fit_type):
         """Aux function."""
         if not np.isfinite(data).all():

mne/preprocessing/maxwell.py

@@ -51,6 +51,7 @@
     _check_option,
     _clean_names,
     _ensure_int,
+    _limit_blas_threads,
     _pl,
     _time_mask,
     _validate_type,
@@ -713,6 +714,7 @@ def _prep_maxwell_filter(
     return params
 
 
+@_limit_blas_threads()
 def _run_maxwell_filter(
     raw,
     skip_by_annotation,

mne/utils/__init__.pyi

@@ -89,6 +89,7 @@ __all__ = [
     "_is_numeric",
     "_is_vtk",
     "_julian_to_date",
+    "_limit_blas_threads",
     "_mask_to_onsets_offsets",
     "_on_missing",
     "_open_lock",
@@ -317,6 +318,7 @@ from .docs import (
 from .fetching import _url_to_local_path
 from .linalg import (
     _get_blas_funcs,
+    _limit_blas_threads,
     _repeated_svd,
     _svd_lwork,
     _sym_mat_pow,

mne/utils/config.py

@@ -633,25 +633,51 @@ def _get_root_dir():
     return root_dir
 
 
+_blas_rename = dict(
+    openblas="OpenBLAS",
+    mkl="MKL",
+    accelerate="Accelerate",
+)
+
+
 def _get_numpy_libs():
     bad_lib = "unknown linalg bindings"
     try:
         from threadpoolctl import threadpool_info
     except Exception as exc:
         return bad_lib + f" (threadpoolctl module not found: {exc})"
     pools = threadpool_info()
-    rename = dict(
-        openblas="OpenBLAS",
-        mkl="MKL",
-    )
     for pool in pools:
         if pool["internal_api"] in ("openblas", "mkl"):
+            layer = pool.get("threading_layer")
+            layer = f" via {layer}" if layer else ""
+            name = pool["internal_api"]
+            name = _blas_rename.get(name, name)
             return (
-                f"{rename[pool['internal_api']]} "
+                f"{name} "
                 f"{pool['version']} with "
                 f"{pool['num_threads']} thread{_pl(pool['num_threads'])}"
+                f"{layer}"
             )
-    return bad_lib
+    return _get_numpy_build_blas() or bad_lib
+
+
+def _get_numpy_build_blas():
+    """Name the BLAS from the build config, for backends threadpoolctl can't see."""
+    # Accelerate has no threadpoolctl controller, so macOS wheels otherwise report only
+    # "unknown linalg bindings"
+    import numpy as np
+
+    try:
+        blas = np.show_config(mode="dicts")["Build Dependencies"]["blas"]
+        name = blas["name"].lower()
+    except Exception:
+        return None
+    version = blas.get("version", "")
+    name = _blas_rename.get(name, name)
+    if version and version != "unknown":  # Accelerate reports a literal "unknown"
+        name = f"{name} {version}"
+    return f"{name}, threads not introspectable"
 
 
 _gpu_cmd = """\

mne/utils/linalg.py

@@ -22,14 +22,90 @@
 # License: BSD-3-Clause
 # Copyright the MNE-Python contributors.
 
+import contextlib
 import functools
+import os
 
 import numpy as np
 from scipy import linalg
 from scipy._lib._util import _asarray_validated
 
 from ..fixes import _safe_svd
 
+###############################################################################
+# BLAS thread limiting
+
+# Setting any of these means the user has already expressed a preference
+_BLAS_THREAD_ENV_VARS = (
+    "OMP_NUM_THREADS",
+    "OPENBLAS_NUM_THREADS",
+    "MKL_NUM_THREADS",
+    "BLIS_NUM_THREADS",
+    "VECLIB_MAXIMUM_THREADS",
+)
+# Measured optima across the machines in gh-13766 range from 1 to 8 depending on CPU and
+# BLAS, but 3 is within ~11% of the best on all of them, and unlike a larger cap it
+# leaves a core free on 4-core machines.
+_MAX_BLAS_THREADS = 3
+
+
+@functools.cache
+def _blas_thread_controller():
+    """Get a cached controller, or None if we should not touch thread counts."""
+    # Constructing one rescans the process's shared libraries (~120 us); reusing it
+    # drops enter/exit to ~3 us, hence the cache. It still reports live thread counts,
+    # so caching does not stale the checks in _limit_blas_threads.
+    try:
+        from threadpoolctl import ThreadpoolController
+    except Exception:  # not a required dependency
+        return None
+    controller = ThreadpoolController()
+    # Accelerate has no working control, so there is nothing to limit
+    if not any(
+        pool["internal_api"] in ("openblas", "mkl") for pool in controller.info()
+    ):
+        return None
+    return controller
+
+
+def _n_available_cpus():
+    """Get the number of CPUs this process may actually use."""
+    # TODO VERSION: drop both fallbacks once Python 3.13 is the minimum
+    if (process_cpu_count := getattr(os, "process_cpu_count", None)) is not None:
+        # preferred: affinity-aware everywhere, and honors -X cpu_count and
+        # PYTHON_CPU_COUNT, which gives users a documented override
+        return process_cpu_count() or 1
+    if (sched_getaffinity := getattr(os, "sched_getaffinity", None)) is not None:
+        return len(sched_getaffinity(0))  # Linux only
+    return os.cpu_count() or 1
+
+
+@contextlib.contextmanager
+def _limit_blas_threads():
+    """Cap BLAS threads around decomposition-heavy code (gh-13766).
+
+    Works as a context manager or as a decorator. Yields immediately, leaving thread
+    counts untouched, whenever the user has expressed their own preference or the BLAS
+    in use cannot be controlled.
+    """
+    controller = _blas_thread_controller()
+    if controller is None or any(os.getenv(var) for var in _BLAS_THREAD_ENV_VARS):
+        yield
+        return
+    n_cpus = _n_available_cpus()
+    # a pool already below the machine size means something else is managing threads
+    if any(pool["num_threads"] < n_cpus for pool in controller.info()):
+        yield
+        return
+    # limit every API rather than just user_api="blas": OpenBLAS built against
+    # OpenMP silently ignores the BLAS-level call before OpenBLAS 0.3.34
+    with controller.limit(limits=min(_MAX_BLAS_THREADS, n_cpus)):
+        yield
+
+
+###############################################################################
+# Fast linalg helpers
+
 # For efficiency, names should be str or tuple of str, dtype a builtin
 # NumPy dtype
 

mne/utils/tests/test_linalg.py

@@ -4,6 +4,8 @@
 # License: BSD-3-Clause
 # Copyright the MNE-Python contributors.
 
+from contextlib import contextmanager
+
 import numpy as np
 import pytest
 from numpy.testing import assert_allclose, assert_array_equal
@@ -101,3 +103,76 @@ def test_pos_semidef_inv(ndim, dtype, n, deficient, reduce_rank, psdef, func):
         want = np.repeat(want[np.newaxis], n_extra, axis=0)
     assert_allclose(np.matmul(mat_symv, mat), want, **kwargs)
     assert_allclose(np.matmul(mat, mat_symv), want, **kwargs)
+
+
+class _FakeController:
+    """Minimal stand-in for threadpoolctl.ThreadpoolController."""
+
+    def __init__(self, num_threads):
+        self.num_threads = num_threads
+        self.applied = []
+
+    def info(self):
+        return [dict(internal_api="openblas", num_threads=self.num_threads)]
+
+    @contextmanager
+    def limit(self, *, limits):
+        self.applied.append(limits)
+        yield
+
+
+def test_limit_blas_threads_logic(monkeypatch):
+    """Test when BLAS thread limiting engages and when it defers."""
+    pytest.importorskip("threadpoolctl")
+    from mne.utils import linalg
+
+    controller = _FakeController(8)
+    monkeypatch.setattr(linalg, "_blas_thread_controller", lambda: controller)
+    monkeypatch.setattr(linalg, "_n_available_cpus", lambda: 8)
+    for var in linalg._BLAS_THREAD_ENV_VARS:
+        monkeypatch.delenv(var, raising=False)
+    with linalg._limit_blas_threads():
+        pass
+    assert controller.applied == [3]
+
+    # defer to an explicit user preference
+    monkeypatch.setenv("OMP_NUM_THREADS", "2")
+    with linalg._limit_blas_threads():
+        pass
+    assert controller.applied == [3]
+
+    # defer when something already limited us (outer limits, joblib worker, ...)
+    monkeypatch.delenv("OMP_NUM_THREADS")
+    controller.num_threads = 2
+    with linalg._limit_blas_threads():
+        pass
+    assert controller.applied == [3]
+
+    # but do not oversubscribe a machine smaller than the cap
+    monkeypatch.setattr(linalg, "_n_available_cpus", lambda: 2)
+    with linalg._limit_blas_threads():
+        pass
+    assert controller.applied == [3, 2]
+
+    # no-op when threads cannot be controlled at all (no threadpoolctl, Accelerate)
+    monkeypatch.setattr(linalg, "_blas_thread_controller", lambda: None)
+    with linalg._limit_blas_threads():
+        pass
+    assert controller.applied == [3, 2]
+
+
+def test_limit_blas_threads(monkeypatch):
+    """Test that BLAS threads are actually limited and restored."""
+    threadpool_info = pytest.importorskip("threadpoolctl").threadpool_info
+    from mne.utils import linalg
+
+    if linalg._blas_thread_controller() is None:
+        pytest.skip("BLAS threads are not controllable (e.g. Accelerate)")
+    for var in linalg._BLAS_THREAD_ENV_VARS:
+        monkeypatch.delenv(var, raising=False)
+    # 1 so that no ambient limit (e.g. from conftest) counts as already-limited
+    monkeypatch.setattr(linalg, "_n_available_cpus", lambda: 1)
+    before = [pool["num_threads"] for pool in threadpool_info()]
+    with linalg._limit_blas_threads():
+        assert [pool["num_threads"] for pool in threadpool_info()] == [1] * len(before)
+    assert [pool["num_threads"] for pool in threadpool_info()] == before