diff options
Diffstat (limited to 'numpy')
| -rw-r--r-- | numpy/__init__.py | 21 | ||||
| -rw-r--r-- | numpy/core/src/multiarray/alloc.c | 21 | ||||
| -rw-r--r-- | numpy/core/src/multiarray/alloc.h | 3 | ||||
| -rw-r--r-- | numpy/core/src/multiarray/multiarraymodule.c | 4 |
4 files changed, 48 insertions, 1 deletions
diff --git a/numpy/__init__.py b/numpy/__init__.py index 1d8570f71..42adfc085 100644 --- a/numpy/__init__.py +++ b/numpy/__init__.py @@ -286,3 +286,24 @@ else: error_message)) raise RuntimeError(msg) del _mac_os_check + + # We usually use madvise hugepages support, but on some old kernels it + # is slow and thus better avoided. + # Specifically kernel version 4.6 had a bug fix which probably fixed this: + # https://github.com/torvalds/linux/commit/7cf91a98e607c2f935dbcc177d70011e95b8faff + import os + use_hugepage = os.environ.get("NUMPY_MADVISE_HUGEPAGE", None) + if sys.platform == "linux" and use_hugepage is None: + use_hugepage = 1 + kernel_version = os.uname().release.split(".")[:2] + kernel_version = tuple(int(v) for v in kernel_version) + if kernel_version < (4, 6): + use_hugepage = 0 + elif use_hugepage is None: + # This is not Linux, so it should not matter, just enable anyway + use_hugepage = 1 + else: + use_hugepage = int(use_hugepage) + + # Note that this will currently only make a difference on Linux + core.multiarray._multiarray_umath._set_madvise_hugepage(use_hugepage) diff --git a/numpy/core/src/multiarray/alloc.c b/numpy/core/src/multiarray/alloc.c index c2b7e9ca7..9d11608f9 100644 --- a/numpy/core/src/multiarray/alloc.c +++ b/numpy/core/src/multiarray/alloc.c @@ -47,6 +47,25 @@ typedef struct { static cache_bucket datacache[NBUCKETS]; static cache_bucket dimcache[NBUCKETS_DIM]; +static int _madvise_hugepage = 1; + + +NPY_NO_EXPORT PyObject * +_set_madvise_hugepage(PyObject *NPY_UNUSED(self), PyObject *enabled_obj) +{ + int was_enabled = _madvise_hugepage; + int enabled = PyObject_IsTrue(enabled_obj); + if (enabled < 0) { + return NULL; + } + _madvise_hugepage = enabled; + if (was_enabled) { + Py_RETURN_TRUE; + } + Py_RETURN_FALSE; +} + + /* as the cache is managed in global variables verify the GIL is held */ /* @@ -75,7 +94,7 @@ _npy_alloc_cache(npy_uintp nelem, npy_uintp esz, npy_uint msz, #endif #ifdef NPY_OS_LINUX /* allow kernel allocating huge pages for large arrays */ - if (NPY_UNLIKELY(nelem * esz >= ((1u<<22u)))) { + if (NPY_UNLIKELY(nelem * esz >= ((1u<<22u))) && _madvise_hugepage) { npy_uintp offset = 4096u - (npy_uintp)p % (4096u); npy_uintp length = nelem * esz - offset; /** diff --git a/numpy/core/src/multiarray/alloc.h b/numpy/core/src/multiarray/alloc.h index 2b69efc35..15e31ebb5 100644 --- a/numpy/core/src/multiarray/alloc.h +++ b/numpy/core/src/multiarray/alloc.h @@ -6,6 +6,9 @@ #define NPY_TRACE_DOMAIN 389047 +NPY_NO_EXPORT PyObject * +_set_madvise_hugepage(PyObject *NPY_UNUSED(self), PyObject *enabled_obj); + NPY_NO_EXPORT void * npy_alloc_cache(npy_uintp sz); diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c index 7792fcdcb..6411ac0a3 100644 --- a/numpy/core/src/multiarray/multiarraymodule.c +++ b/numpy/core/src/multiarray/multiarraymodule.c @@ -34,6 +34,7 @@ NPY_NO_EXPORT int NPY_NUMUSERTYPES = 0; /* Internal APIs */ +#include "alloc.h" #include "arrayfunction_override.h" #include "arraytypes.h" #include "arrayobject.h" @@ -3971,6 +3972,7 @@ normalize_axis_index(PyObject *NPY_UNUSED(self), PyObject *args, PyObject *kwds) return PyInt_FromLong(axis); } + static struct PyMethodDef array_module_methods[] = { {"_get_implementing_args", (PyCFunction)array__get_implementing_args, @@ -4159,6 +4161,8 @@ static struct PyMethodDef array_module_methods[] = { METH_VARARGS, NULL}, {"_add_newdoc_ufunc", (PyCFunction)add_newdoc_ufunc, METH_VARARGS, NULL}, + {"_set_madvise_hugepage", (PyCFunction)_set_madvise_hugepage, + METH_O, "Toggle and return madvise hugepage (no OS support check)."}, {NULL, NULL, 0, NULL} /* sentinel */ }; |