wehub-resource-sync
28558dca80
Build / Build and test on ubuntu-24.04-arm for aarch64 (push) Waiting to run
Build / Build and test on windows-11-arm for aarch64 (push) Waiting to run
Build / Build and test on macos-latest for x86_64 (push) Waiting to run
Build / Build and test on windows-latest for x86_64 (push) Waiting to run
Build / Build and test on ubuntu-latest for x86_64 (push) Failing after 1s
Build / Build and test on ubuntu-latest (numpy 1.26) for x86_64 (push) Failing after 0s
312 lines
9.6 KiB
Python
312 lines
9.6 KiB
Python
###################################################################
|
|
# Numexpr - Fast numerical array expression evaluator for NumPy.
|
|
#
|
|
# License: MIT
|
|
# Author: See AUTHORS.txt
|
|
#
|
|
# See LICENSE.txt and LICENSES/*.txt for details about copyright and
|
|
# rights to use.
|
|
####################################################################
|
|
|
|
import logging
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
import contextvars
|
|
import os
|
|
import subprocess
|
|
|
|
from numexpr import use_vml
|
|
from numexpr.interpreter import MAX_THREADS, _get_num_threads, _set_num_threads
|
|
|
|
from . import version
|
|
|
|
if use_vml:
|
|
from numexpr.interpreter import (_get_vml_num_threads, _get_vml_version,
|
|
_set_vml_accuracy_mode,
|
|
_set_vml_num_threads)
|
|
|
|
|
|
def get_vml_version():
|
|
"""
|
|
Get the VML/MKL library version.
|
|
"""
|
|
if use_vml:
|
|
return _get_vml_version()
|
|
else:
|
|
return None
|
|
|
|
|
|
def set_vml_accuracy_mode(mode):
|
|
"""
|
|
Set the accuracy mode for VML operations.
|
|
|
|
The `mode` parameter can take the values:
|
|
- 'high': high accuracy mode (HA), <1 least significant bit
|
|
- 'low': low accuracy mode (LA), typically 1-2 least significant bits
|
|
- 'fast': enhanced performance mode (EP)
|
|
- None: mode settings are ignored
|
|
|
|
This call is equivalent to the `vmlSetMode()` in the VML library.
|
|
See:
|
|
|
|
http://www.intel.com/software/products/mkl/docs/webhelp/vml/vml_DataTypesAccuracyModes.html
|
|
|
|
for more info on the accuracy modes.
|
|
|
|
Returns old accuracy settings.
|
|
"""
|
|
if use_vml:
|
|
acc_dict = {None: 0, 'low': 1, 'high': 2, 'fast': 3}
|
|
acc_reverse_dict = {1: 'low', 2: 'high', 3: 'fast'}
|
|
if mode not in list(acc_dict.keys()):
|
|
raise ValueError(
|
|
"mode argument must be one of: None, 'high', 'low', 'fast'")
|
|
retval = _set_vml_accuracy_mode(acc_dict.get(mode, 0))
|
|
return acc_reverse_dict.get(retval)
|
|
else:
|
|
return None
|
|
|
|
|
|
def set_vml_num_threads(nthreads):
|
|
"""
|
|
Suggests a maximum number of threads to be used in VML operations.
|
|
|
|
This function is equivalent to the call
|
|
`mkl_domain_set_num_threads(nthreads, MKL_DOMAIN_VML)` in the MKL
|
|
library. See:
|
|
|
|
http://www.intel.com/software/products/mkl/docs/webhelp/support/functn_mkl_domain_set_num_threads.html
|
|
|
|
for more info about it.
|
|
"""
|
|
if use_vml:
|
|
_set_vml_num_threads(nthreads)
|
|
pass
|
|
|
|
def get_vml_num_threads():
|
|
"""
|
|
Gets the maximum number of threads to be used in VML operations.
|
|
|
|
This function is equivalent to the call
|
|
`mkl_domain_get_max_threads (MKL_DOMAIN_VML)` in the MKL
|
|
library. See:
|
|
|
|
http://software.intel.com/en-us/node/522118
|
|
|
|
for more info about it.
|
|
"""
|
|
if use_vml:
|
|
return _get_vml_num_threads()
|
|
return None
|
|
|
|
def set_num_threads(nthreads):
|
|
"""
|
|
Sets a number of threads to be used in operations.
|
|
|
|
DEPRECATED: returns the previous setting for the number of threads.
|
|
|
|
During initialization time NumExpr sets this number to the number
|
|
of detected cores in the system (see `detect_number_of_cores()`).
|
|
"""
|
|
old_nthreads = _set_num_threads(nthreads)
|
|
return old_nthreads
|
|
|
|
def get_num_threads():
|
|
"""
|
|
Gets the number of threads currently in use for operations.
|
|
"""
|
|
return _get_num_threads()
|
|
|
|
def _init_num_threads():
|
|
"""
|
|
Detects the environment variable 'NUMEXPR_MAX_THREADS' to set the threadpool
|
|
size, and if necessary the slightly redundant 'NUMEXPR_NUM_THREADS' or
|
|
'OMP_NUM_THREADS' env vars to set the initial number of threads used by
|
|
the virtual machine.
|
|
"""
|
|
# Any platform-specific short-circuits
|
|
if 'sparc' in version.platform_machine:
|
|
log.warning('The number of threads have been set to 1 because problems related '
|
|
'to threading have been reported on some sparc machine. '
|
|
'The number of threads can be changed using the "set_num_threads" '
|
|
'function.')
|
|
set_num_threads(1)
|
|
return 1
|
|
|
|
env_configured = False
|
|
n_cores = detect_number_of_cores()
|
|
if ('NUMEXPR_MAX_THREADS' in os.environ and os.environ['NUMEXPR_MAX_THREADS'] != '' or
|
|
'OMP_NUM_THREADS' in os.environ and os.environ['OMP_NUM_THREADS'] != ''):
|
|
# The user has configured NumExpr in the expected way, so suppress logs.
|
|
env_configured = True
|
|
n_cores = MAX_THREADS
|
|
else:
|
|
# The use has not set 'NUMEXPR_MAX_THREADS', so likely they have not
|
|
# configured NumExpr as desired, so we emit info logs.
|
|
if n_cores > MAX_THREADS:
|
|
log.info('Note: detected %d virtual cores but NumExpr set to maximum of %d, check "NUMEXPR_MAX_THREADS" environment variable.'%(n_cores, MAX_THREADS))
|
|
if n_cores > 16:
|
|
# Back in 2019, 8 threads would be considered safe for performance. We are in 2024 now, so adjusting.
|
|
log.info('Note: NumExpr detected %d cores but "NUMEXPR_MAX_THREADS" not set, so enforcing safe limit of 16.'%n_cores)
|
|
n_cores = 16
|
|
|
|
# Now we check for 'NUMEXPR_NUM_THREADS' or 'OMP_NUM_THREADS' to set the
|
|
# actual number of threads used.
|
|
if 'NUMEXPR_NUM_THREADS' in os.environ and os.environ['NUMEXPR_NUM_THREADS'] != '':
|
|
requested_threads = int(os.environ['NUMEXPR_NUM_THREADS'])
|
|
elif 'OMP_NUM_THREADS' in os.environ and os.environ['OMP_NUM_THREADS'] != '':
|
|
# Empty string is commonly used to unset the variable
|
|
requested_threads = int(os.environ['OMP_NUM_THREADS'])
|
|
else:
|
|
requested_threads = n_cores
|
|
if not env_configured:
|
|
log.info('NumExpr defaulting to %d threads.'%n_cores)
|
|
|
|
# The C-extension function performs its own checks against `MAX_THREADS`
|
|
set_num_threads(requested_threads)
|
|
return requested_threads
|
|
|
|
|
|
def detect_number_of_cores():
|
|
"""
|
|
Detects the number of cores on a system. Cribbed from pp.
|
|
"""
|
|
# Linux, Unix and MacOS:
|
|
if hasattr(os, "sysconf"):
|
|
if "SC_NPROCESSORS_ONLN" in os.sysconf_names:
|
|
# Linux & Unix:
|
|
ncpus = os.sysconf("SC_NPROCESSORS_ONLN")
|
|
if isinstance(ncpus, int) and ncpus > 0:
|
|
return ncpus
|
|
else: # OSX:
|
|
return int(subprocess.check_output(["sysctl", "-n", "hw.ncpu"]))
|
|
# Windows:
|
|
try:
|
|
ncpus = int(os.environ.get("NUMBER_OF_PROCESSORS", ""))
|
|
if ncpus > 0:
|
|
return ncpus
|
|
except ValueError:
|
|
pass
|
|
return 1 # Default
|
|
|
|
|
|
def detect_number_of_threads():
|
|
"""
|
|
DEPRECATED: use `_init_num_threads` instead.
|
|
If this is modified, please update the note in: https://github.com/pydata/numexpr/wiki/Numexpr-Users-Guide
|
|
"""
|
|
log.warning('Deprecated, use `init_num_threads` instead.')
|
|
try:
|
|
nthreads = int(os.environ.get('NUMEXPR_NUM_THREADS', ''))
|
|
except ValueError:
|
|
try:
|
|
nthreads = int(os.environ.get('OMP_NUM_THREADS', ''))
|
|
except ValueError:
|
|
nthreads = detect_number_of_cores()
|
|
|
|
# Check that we don't surpass the MAX_THREADS in interpreter.cpp
|
|
if nthreads > MAX_THREADS:
|
|
nthreads = MAX_THREADS
|
|
return nthreads
|
|
|
|
|
|
class CacheDict(dict):
|
|
"""
|
|
A dictionary that prevents itself from growing too much.
|
|
"""
|
|
|
|
def __init__(self, maxentries):
|
|
self.maxentries = maxentries
|
|
super(CacheDict, self).__init__(self)
|
|
|
|
def __setitem__(self, key, value):
|
|
# Protection against growing the cache too much
|
|
if len(self) > self.maxentries:
|
|
# Remove a 10% of (arbitrary) elements from the cache
|
|
entries_to_remove = self.maxentries // 10
|
|
for k in list(self.keys())[:entries_to_remove]:
|
|
super(CacheDict, self).__delitem__(k)
|
|
super(CacheDict, self).__setitem__(key, value)
|
|
|
|
|
|
class ContextDict:
|
|
"""
|
|
A context aware version dictionary
|
|
"""
|
|
def __init__(self):
|
|
self._context_data = contextvars.ContextVar('context_data', default={})
|
|
|
|
def set(self, key=None, value=None, **kwargs):
|
|
data = self._context_data.get().copy()
|
|
|
|
if key is not None:
|
|
data[key] = value
|
|
|
|
for k, v in kwargs.items():
|
|
data[k] = v
|
|
|
|
self._context_data.set(data)
|
|
|
|
def get(self, key, default=None):
|
|
data = self._context_data.get()
|
|
return data.get(key, default)
|
|
|
|
def delete(self, key):
|
|
data = self._context_data.get().copy()
|
|
if key in data:
|
|
del data[key]
|
|
self._context_data.set(data)
|
|
|
|
def clear(self):
|
|
self._context_data.set({})
|
|
|
|
def all(self):
|
|
return self._context_data.get()
|
|
|
|
def update(self, *args, **kwargs):
|
|
data = self._context_data.get().copy()
|
|
|
|
if args:
|
|
if len(args) > 1:
|
|
raise TypeError(f"update() takes at most 1 positional argument ({len(args)} given)")
|
|
other = args[0]
|
|
if isinstance(other, dict):
|
|
data.update(other)
|
|
else:
|
|
for k, v in other:
|
|
data[k] = v
|
|
|
|
data.update(kwargs)
|
|
self._context_data.set(data)
|
|
|
|
def keys(self):
|
|
return self._context_data.get().keys()
|
|
|
|
def values(self):
|
|
return self._context_data.get().values()
|
|
|
|
def items(self):
|
|
return self._context_data.get().items()
|
|
|
|
def __getitem__(self, key):
|
|
return self.get(key)
|
|
|
|
def __setitem__(self, key, value):
|
|
self.set(key, value)
|
|
|
|
def __delitem__(self, key):
|
|
self.delete(key)
|
|
|
|
def __contains__(self, key):
|
|
return key in self._context_data.get()
|
|
|
|
def __len__(self):
|
|
return len(self._context_data.get())
|
|
|
|
def __iter__(self):
|
|
return iter(self._context_data.get())
|
|
|
|
def __repr__(self):
|
|
return repr(self._context_data.get())
|