Merge pull request #11957 from juliantaylor/hugepages

ENH: mark that large allocations can use huge pages

3 files changed, 25 insertions, 12 deletions

diff --git a/numpy/core/setup.py b/numpy/core/setup.py
index 1588a2634..b41580030 100644
--- a/numpy/core/setup.py
+++ b/numpy/core/setup.py
@@ -153,7 +153,8 @@ def check_math_capabilities(config, moredefs, mathlibs):
 
     for h in OPTIONAL_HEADERS:
         if config.check_func("", decl=False, call=False, headers=[h]):
-            moredefs.append((fname2def(h).replace(".", "_"), 1))
+            h = h.replace(".", "_").replace(os.path.sep, "_")
+            moredefs.append((fname2def(h), 1))
 
     for tup in OPTIONAL_INTRINSICS:
         headers = None
diff --git a/numpy/core/setup_common.py b/numpy/core/setup_common.py
index 356482b07..e637dbc20 100644
--- a/numpy/core/setup_common.py
+++ b/numpy/core/setup_common.py
@@ -110,7 +110,7 @@ OPTIONAL_STDFUNCS = ["expm1", "log1p", "acosh", "asinh", "atanh",
         "rint", "trunc", "exp2", "log2", "hypot", "atan2", "pow",
         "copysign", "nextafter", "ftello", "fseeko",
         "strtoll", "strtoull", "cbrt", "strtold_l", "fallocate",
-        "backtrace"]
+        "backtrace", "madvise"]
 
 
 OPTIONAL_HEADERS = [
@@ -120,6 +120,7 @@ OPTIONAL_HEADERS = [
                 "features.h",  # for glibc version linux
                 "xlocale.h",  # see GH#8367
                 "dlfcn.h", # dladdr
+                "sys/mman.h", #madvise
 ]
 
 # optional gcc compiler builtins and their call arguments and optional a
diff --git a/numpy/core/src/multiarray/alloc.c b/numpy/core/src/multiarray/alloc.c
index fe957bf16..6755095d7 100644
--- a/numpy/core/src/multiarray/alloc.c
+++ b/numpy/core/src/multiarray/alloc.c
@@ -22,8 +22,16 @@
 #include "npy_config.h"
 #include "alloc.h"
 
+
 #include <assert.h>
 
+#ifdef HAVE_SYS_MMAN_H
+#include <sys/mman.h>
+#if defined MADV_HUGEPAGE && defined HAVE_MADVISE
+#define HAVE_MADV_HUGEPAGE
+#endif
+#endif
+
 #define NBUCKETS 1024 /* number of buckets for data*/
 #define NBUCKETS_DIM 16 /* number of buckets for dimensions/strides */
 #define NCACHE 7 /* number of cache entries per bucket */
@@ -52,6 +60,7 @@ static NPY_INLINE void *
 _npy_alloc_cache(npy_uintp nelem, npy_uintp esz, npy_uint msz,
                  cache_bucket * cache, void * (*alloc)(size_t))
 {
+    void * p;
     assert((esz == 1 && cache == datacache) ||
            (esz == sizeof(npy_intp) && cache == dimcache));
     assert(NPY_CHECK_GIL_HELD());
@@ -60,19 +69,21 @@ _npy_alloc_cache(npy_uintp nelem, npy_uintp esz, npy_uint msz,
             return cache[nelem].ptrs[--(cache[nelem].available)];
         }
     }
+    p = alloc(nelem * esz);
+    if (p) {
 #ifdef _PyPyGC_AddMemoryPressure
-    {
-        size_t size = nelem * esz;
-        void * ret = alloc(size);
-        if (ret != NULL)
-        {
-            _PyPyPyGC_AddMemoryPressure(size);
+        _PyPyPyGC_AddMemoryPressure(nelem * esz);
+#endif
+#ifdef HAVE_MADV_HUGEPAGE
+        /* allow kernel allocating huge pages for large arrays */
+        if (NPY_UNLIKELY(nelem * esz >= ((1u<<22u)))) {
+            npy_uintp offset = 4096u - (npy_uintp)p % (4096u);
+            npy_uintp length = nelem * esz - offset;
+            madvise((void*)((npy_uintp)p + offset), length, MADV_HUGEPAGE);
         }
-        return ret;
-    }
-#else
-     return alloc(nelem * esz);
 #endif
+    }
+    return p;
 }
 
 /*