Aluode/PerceptionLabPortable

1import os
2from joblib import cpu_count
3
4
5# Module level cache for cpu_count as we do not expect this to change during
6# the lifecycle of a Python program. This dictionary is keyed by
7# only_physical_cores.
8_CPU_COUNTS = {}
9
10
11def _openmp_parallelism_enabled():
12    """Determines whether scikit-learn has been built with OpenMP
13
14    It allows to retrieve at runtime the information gathered at compile time.
15    """
16    # SKLEARN_OPENMP_PARALLELISM_ENABLED is resolved at compile time and defined
17    # in _openmp_helpers.pxd as a boolean. This function exposes it to Python.
18    return SKLEARN_OPENMP_PARALLELISM_ENABLED
19
20
21cpdef _openmp_effective_n_threads(n_threads=None, only_physical_cores=True):
22    """Determine the effective number of threads to be used for OpenMP calls
23
24    - For ``n_threads = None``,
25      - if the ``OMP_NUM_THREADS`` environment variable is set, return
26        ``openmp.omp_get_max_threads()``
27      - otherwise, return the minimum between ``openmp.omp_get_max_threads()``
28        and the number of cpus, taking cgroups quotas into account. Cgroups
29        quotas can typically be set by tools such as Docker.
30      The result of ``omp_get_max_threads`` can be influenced by environment
31      variable ``OMP_NUM_THREADS`` or at runtime by ``omp_set_num_threads``.
32
33    - For ``n_threads > 0``, return this as the maximal number of threads for
34      parallel OpenMP calls.
35
36    - For ``n_threads < 0``, return the maximal number of threads minus
37      ``|n_threads + 1|``. In particular ``n_threads = -1`` will use as many
38      threads as there are available cores on the machine.
39
40    - Raise a ValueError for ``n_threads = 0``.
41
42    Passing the `only_physical_cores=False` flag makes it possible to use extra
43    threads for SMT/HyperThreading logical cores. It has been empirically
44    observed that using as many threads as available SMT cores can slightly
45    improve the performance in some cases, but can severely degrade
46    performance other times. Therefore it is recommended to use
47    `only_physical_cores=True` unless an empirical study has been conducted to
48    assess the impact of SMT on a case-by-case basis (using various input data
49    shapes, in particular small data shapes).
50
51    If scikit-learn is built without OpenMP support, always return 1.
52    """
53    if n_threads == 0:
54        raise ValueError("n_threads = 0 is invalid")
55
56    if not SKLEARN_OPENMP_PARALLELISM_ENABLED:
57        # OpenMP disabled at build-time => sequential mode
58        return 1
59
60    if os.getenv("OMP_NUM_THREADS"):
61        # Fall back to user provided number of threads making it possible
62        # to exceed the number of cpus.
63        max_n_threads = omp_get_max_threads()
64    else:
65        try:
66            n_cpus = _CPU_COUNTS[only_physical_cores]
67        except KeyError:
68            n_cpus = cpu_count(only_physical_cores=only_physical_cores)
69            _CPU_COUNTS[only_physical_cores] = n_cpus
70        max_n_threads = min(omp_get_max_threads(), n_cpus)
71
72    if n_threads is None:
73        return max_n_threads
74    elif n_threads < 0:
75        return max(1, max_n_threads + n_threads + 1)
76
77    return n_threads
78