ENH: [7/7] enable multi-platform SIMD compiler optimizations

Add testing unit for the utilites of CPU dispatcher
numpy · mattip · Jun 17, 2020 · Jun 13, 2020 · Jun 13, 2020 · Jun 13, 2020
commit 86d25150c03fcddad5b0446342e2b55b38194053
diff --git a/numpy/core/setup.py b/numpy/core/setup.py
@@ -940,8 +940,11 @@ def generate_umath_c(ext, build_dir):
     #                        umath_tests module                           #
     #######################################################################
 
-    config.add_extension('_umath_tests',
-                    sources=[join('src', 'umath', '_umath_tests.c.src')])
+    config.add_extension('_umath_tests', sources=[
+        join('src', 'umath', '_umath_tests.c.src'),
+        join('src', 'umath', '_umath_tests.dispatch.c'),
+        join('src', 'common', 'npy_cpu_features.c.src'),
+    ])
 
     #######################################################################
     #                   custom rational dtype module                      #

diff --git a/numpy/core/src/umath/_umath_tests.c.src b/numpy/core/src/umath/_umath_tests.c.src
@@ -576,13 +576,59 @@ fail:
     return NULL;
 }
 
+// Testing the utilites of the CPU dispatcher
+#ifndef NPY_DISABLE_OPTIMIZATION
+    #include "_umath_tests.dispatch.h"
+#endif
+NPY_CPU_DISPATCH_DECLARE(extern const char *_umath_tests_dispatch_var)
+NPY_CPU_DISPATCH_DECLARE(const char *_umath_tests_dispatch_func, (void))
+NPY_CPU_DISPATCH_DECLARE(void _umath_tests_dispatch_attach, (PyObject *list))
+
+static PyObject *
+UMath_Tests_test_dispatch(PyObject *NPY_UNUSED(dummy), PyObject *NPY_UNUSED(dummy2))
+{
+    const char *highest_func, *highest_var;
+    NPY_CPU_DISPATCH_CALL(highest_func = _umath_tests_dispatch_func, ())
+    NPY_CPU_DISPATCH_CALL(highest_var  = _umath_tests_dispatch_var)
+    const char *highest_func_xb = "nobase", *highest_var_xb = "nobase";
+    NPY_CPU_DISPATCH_CALL_XB(highest_func_xb = _umath_tests_dispatch_func, ())
+    NPY_CPU_DISPATCH_CALL_XB(highest_var_xb  = _umath_tests_dispatch_var)
+
+    PyObject *dict = PyDict_New(), *item;
+    if (dict == NULL) {
+        return NULL;
+    }
+    /**begin repeat
+     * #str = func, var, func_xb, var_xb#
+    */
+    item = PyUnicode_FromString(highest_@str@);
+    if (item == NULL || PyDict_SetItemString(dict, "@str@", item) < 0) {
+        goto err;
+    }
+    /**end repeat**/
+    item = PyList_New(0);
+    if (item == NULL || PyDict_SetItemString(dict, "all", item) < 0) {
+        goto err;
+    }
+    NPY_CPU_DISPATCH_CALL_ALL(_umath_tests_dispatch_attach, (item))
+    if (PyErr_Occurred()) {
+        goto err;
+    }
+    return dict;
+err:
+    Py_XDECREF(item);
+    Py_DECREF(dict);
+    return NULL;
+}
+
 static PyMethodDef UMath_TestsMethods[] = {
     {"test_signature",  UMath_Tests_test_signature, METH_VARARGS,
      "Test signature parsing of ufunc. \n"
      "Arguments: nin nout signature \n"
      "If fails, it returns NULL. Otherwise it returns a tuple of ufunc "
      "internals. \n",
      },
+    {"test_dispatch", UMath_Tests_test_dispatch, METH_NOARGS, NULL},
     {NULL, NULL, 0, NULL}        /* Sentinel */
 };
 
@@ -604,6 +650,11 @@ PyMODINIT_FUNC PyInit__umath_tests(void) {
     PyObject *d;
     PyObject *version;
 
+    // Initialize CPU features
+    if (npy_cpu_init() < 0) {
+        return NULL;
+    }
+
     m = PyModule_Create(&moduledef);
     if (m == NULL) {
         return NULL;
@@ -632,6 +683,5 @@ PyMODINIT_FUNC PyInit__umath_tests(void) {
                         "cannot load _umath_tests module.");
         return NULL;
     }
-
     return m;
 }
diff --git a/numpy/core/src/umath/_umath_tests.dispatch.c b/numpy/core/src/umath/_umath_tests.dispatch.c
@@ -0,0 +1,33 @@
+/**
+ * Testing the utilites of the CPU dispatcher
+ *
+ * @targets $werror baseline
+ * SSE2 SSE41 AVX2
+ * VSX VSX2 VSX3
+ * NEON ASIMD ASIMDHP
+ */
+#include <Python.h>
+#include "npy_cpu_dispatch.h"
+
+#ifndef NPY_DISABLE_OPTIMIZATION
+    #include "_umath_tests.dispatch.h"
+#endif
+
+NPY_CPU_DISPATCH_DECLARE(const char *_umath_tests_dispatch_func, (void))
+NPY_CPU_DISPATCH_DECLARE(extern const char *_umath_tests_dispatch_var)
+NPY_CPU_DISPATCH_DECLARE(void _umath_tests_dispatch_attach, (PyObject *list))
+
+const char *NPY_CPU_DISPATCH_CURFX(_umath_tests_dispatch_var) = NPY_TOSTRING(NPY_CPU_DISPATCH_CURFX(var));
+const char *NPY_CPU_DISPATCH_CURFX(_umath_tests_dispatch_func)(void)
+{
+    static const char *current = NPY_TOSTRING(NPY_CPU_DISPATCH_CURFX(func));
+    return current;
+}
+
+void NPY_CPU_DISPATCH_CURFX(_umath_tests_dispatch_attach)(PyObject *list)
+{
+    PyObject *item = PyUnicode_FromString(NPY_TOSTRING(NPY_CPU_DISPATCH_CURFX(func)));
+    if (item) {
+        PyList_Append(list, item);
+    }
+}
diff --git a/numpy/core/tests/test_cpu_dispatcher.py b/numpy/core/tests/test_cpu_dispatcher.py
@@ -0,0 +1,42 @@
+from numpy.core._multiarray_umath import __cpu_features__, __cpu_baseline__, __cpu_dispatch__
+from numpy.core import _umath_tests
+from numpy.testing import assert_equal
+
+def test_dispatcher():
+    """
+    Testing the utilites of the CPU dispatcher
+    """
+    targets = (
+        "SSE2", "SSE41", "AVX2",
+        "VSX", "VSX2", "VSX3",
+        "NEON", "ASIMD", "ASIMDHP"
+    )
+    highest_sfx = "" # no suffix for the baseline
+    all_sfx = []
+    for feature in reversed(targets):
+        # skip baseline features, by the default `CCompilerOpt` do not generate separated objects
+        # for the baseline,  just one object combined all of them via 'baseline' option
+        # within the configuration statments.
+        if feature in __cpu_baseline__:
+            continue
+        # check compiler and running machine support
+        if feature not in __cpu_dispatch__ or not __cpu_features__[feature]:
+            continue
+
+        if not highest_sfx:
+            highest_sfx = "_" + feature
+        all_sfx.append("func" + "_" + feature)
+
+    test = _umath_tests.test_dispatch()
+    assert_equal(test["func"], "func" + highest_sfx)
+    assert_equal(test["var"], "var"  + highest_sfx)
+
+    if highest_sfx:
+        assert_equal(test["func_xb"], "func" + highest_sfx)
+        assert_equal(test["var_xb"], "var"  + highest_sfx)
+    else:
+        assert_equal(test["func_xb"], "nobase")
+        assert_equal(test["var_xb"], "nobase")
+
+    all_sfx.append("func") # add the baseline
+    assert_equal(test["all"], all_sfx)