From cde5583abda2bd3da410058b2e1ea4108073b4d3 Mon Sep 17 00:00:00 2001
From: Julian Taylor <jtaylor.debian@googlemail.com>
Date: Tue, 27 Mar 2018 22:29:17 +0200
Subject: [PATCH] BUG: verify the OS supports avx instruction

On some systems you can disable avx registers but the gcc builtin does
only checks if the cpu has the feature.
Before using avx functions check the OS support with xgetbv.
Closes gh-10787
Closes gh-9534
---
 numpy/core/code_generators/generate_umath.py |  4 +-
 numpy/core/include/numpy/npy_common.h        | 16 ------
 numpy/core/setup.py                          |  1 +
 numpy/core/setup_common.py                   |  1 +
 numpy/core/src/umath/cpuid.c                 | 56 ++++++++++++++++++++
 numpy/core/src/umath/cpuid.h                 |  9 ++++
 6 files changed, 69 insertions(+), 18 deletions(-)
 create mode 100644 numpy/core/src/umath/cpuid.c
 create mode 100644 numpy/core/src/umath/cpuid.h

diff --git a/numpy/core/code_generators/generate_umath.py b/numpy/core/code_generators/generate_umath.py
index ebcf864ea4c9..1d3550e065ad 100644
--- a/numpy/core/code_generators/generate_umath.py
+++ b/numpy/core/code_generators/generate_umath.py
@@ -966,7 +966,7 @@ def make_arrays(funcdict):
                     for vt in t.simd:
                         code2list.append(textwrap.dedent("""\
                         #ifdef HAVE_ATTRIBUTE_TARGET_{ISA}
-                        if (NPY_CPU_SUPPORTS_{ISA}) {{
+                        if (npy_cpu_supports("{ISA}")) {{
                             {fname}_functions[{idx}] = {type}_{fname}_{isa};
                         }}
                         #endif
@@ -1073,7 +1073,7 @@ def make_code(funcdict, filename):
 
         Please make changes to the code generator program (%s)
     **/
-
+    #include "cpuid.h"
     %s
 
     static int
diff --git a/numpy/core/include/numpy/npy_common.h b/numpy/core/include/numpy/npy_common.h
index b7634a9306e5..5faff43851fe 100644
--- a/numpy/core/include/numpy/npy_common.h
+++ b/numpy/core/include/numpy/npy_common.h
@@ -101,22 +101,6 @@
 #endif
 #endif
 
-#ifdef HAVE___BUILTIN_CPU_SUPPORTS
-  #ifdef HAVE_ATTRIBUTE_TARGET_AVX2
-    #define NPY_CPU_SUPPORTS_AVX2 __builtin_cpu_supports("avx2")
-  #else
-    #define NPY_CPU_SUPPORTS_AVX2 0
-  #endif
-  #ifdef HAVE_ATTRIBUTE_TARGET_AVX
-    #define NPY_CPU_SUPPORTS_AVX __builtin_cpu_supports("avx")
-  #else
-    #define NPY_CPU_SUPPORTS_AVX 0
-  #endif
-#else
-  #define NPY_CPU_SUPPORTS_AVX 0
-  #define NPY_CPU_SUPPORTS_AVX2 0
-#endif
-
 #if defined(_MSC_VER)
         #define NPY_INLINE __inline
 #elif defined(__GNUC__)
diff --git a/numpy/core/setup.py b/numpy/core/setup.py
index d519e0eb817c..15f6e1522695 100644
--- a/numpy/core/setup.py
+++ b/numpy/core/setup.py
@@ -888,6 +888,7 @@ def generate_umath_c(ext, build_dir):
             join('src', 'umath', 'loops.c.src'),
             join('src', 'umath', 'ufunc_object.c'),
             join('src', 'umath', 'extobj.c'),
+            join('src', 'umath', 'cpuid.c'),
             join('src', 'umath', 'scalarmath.c.src'),
             join('src', 'umath', 'ufunc_type_resolution.c'),
             join('src', 'umath', 'override.c'),
diff --git a/numpy/core/setup_common.py b/numpy/core/setup_common.py
index 1fe953910591..f36d61f558ca 100644
--- a/numpy/core/setup_common.py
+++ b/numpy/core/setup_common.py
@@ -146,6 +146,7 @@ def check_api_version(apiversion, codegen_dir):
                         "stdio.h", "LINK_AVX"),
                        ("__asm__ volatile", '"vpand %ymm1, %ymm2, %ymm3"',
                         "stdio.h", "LINK_AVX2"),
+                       ("__asm__ volatile", '"xgetbv"', "stdio.h", "XGETBV"),
                        ]
 
 # function attributes
diff --git a/numpy/core/src/umath/cpuid.c b/numpy/core/src/umath/cpuid.c
new file mode 100644
index 000000000000..912d51eeb9aa
--- /dev/null
+++ b/numpy/core/src/umath/cpuid.c
@@ -0,0 +1,56 @@
+#define _UMATHMODULE
+#define NPY_NO_DEPRECATED_API NPY_API_VERSION
+
+#include <Python.h>
+
+#include "npy_config.h"
+
+#define PY_ARRAY_UNIQUE_SYMBOL _npy_umathmodule_ARRAY_API
+#define NO_IMPORT_ARRAY
+
+#include "cpuid.h"
+
+#define XCR_XFEATURE_ENABLED_MASK 0x0
+#define XSTATE_SSE 0x2
+#define XSTATE_YMM 0x4
+
+/*
+ * verify the OS supports avx instructions
+ * it can be disabled in some OS, e.g. with the nosavex boot option of linux
+ */
+static NPY_INLINE
+int os_avx_support(void)
+{
+#if HAVE_XGETBV
+    /*
+     * use bytes for xgetbv to avoid issues with compiler not knowing the
+     * instruction
+     */
+    unsigned int eax, edx;
+    unsigned int ecx = XCR_XFEATURE_ENABLED_MASK;
+    __asm__("xgetbv" : "=a" (eax), "=d" (edx) : "c" (ecx));
+    return (eax & (XSTATE_SSE | XSTATE_YMM)) == (XSTATE_SSE | XSTATE_YMM);
+#else
+    return 0;
+#endif
+}
+
+
+/*
+ * Primitive cpu feature detect function
+ * Currently only supports checking for avx on gcc compatible compilers.
+ */
+NPY_NO_EXPORT int
+npy_cpu_supports(const char * feature)
+{
+#ifdef HAVE___BUILTIN_CPU_SUPPORTS
+    if (strcmp(feature, "avx2") == 0) {
+        return __builtin_cpu_supports("avx2") && os_avx_support();
+    }
+    else if (strcmp(feature, "avx") == 0) {
+        return __builtin_cpu_supports("avx") && os_avx_support();
+    }
+#endif
+
+    return 0;
+}
diff --git a/numpy/core/src/umath/cpuid.h b/numpy/core/src/umath/cpuid.h
new file mode 100644
index 000000000000..33702ed41f48
--- /dev/null
+++ b/numpy/core/src/umath/cpuid.h
@@ -0,0 +1,9 @@
+#ifndef _NPY_PRIVATE__CPUID_H_
+#define _NPY_PRIVATE__CPUID_H_
+
+#include <numpy/ndarraytypes.h>  /* for NPY_NO_EXPORT */
+
+NPY_NO_EXPORT int
+npy_cpu_supports(const char * feature);
+
+#endif