Merge pull request #22137 from Micky774/enable_simd

ENH allow for specifying CPU features to enable via `NPY_ENABLE_CPU_FEATURES` environment variable

3 files changed, 312 insertions, 53 deletions

diff --git a/numpy/core/src/common/npy_cpu_features.c b/numpy/core/src/common/npy_cpu_features.c
index 92a4e432b..64a85f6fb 100644
--- a/numpy/core/src/common/npy_cpu_features.c
+++ b/numpy/core/src/common/npy_cpu_features.c
@@ -14,13 +14,15 @@ static unsigned char npy__cpu_have[NPY_CPU_FEATURE_MAX];
 static void
 npy__cpu_init_features(void);
 /*
- * Disable CPU dispatched features at runtime if environment variable
- * 'NPY_DISABLE_CPU_FEATURES' is defined.
+ * Enable or disable CPU dispatched features at runtime if the environment variable
+ * `NPY_ENABLE_CPU_FEATURES`  or  `NPY_DISABLE_CPU_FEATURES`
+ * depends on the value of boolean parameter `disable`(toggle).
+ *
  * Multiple features can be present, and separated by space, comma, or tab.
- * Raises an error if parsing fails or if the feature was not enabled
+ * Raises an error if parsing fails or if the feature was not enabled or disabled
 */
 static int
-npy__cpu_try_disable_env(void);
+npy__cpu_check_env(int disable, const char *env);
 
 /* Ensure the build's CPU baseline features are supported at runtime */
 static int
@@ -43,9 +45,22 @@ npy_cpu_init(void)
     if (npy__cpu_validate_baseline() < 0) {
         return -1;
     }
-    if (npy__cpu_try_disable_env() < 0) {
+    char *enable_env = getenv("NPY_ENABLE_CPU_FEATURES");
+    char *disable_env = getenv("NPY_DISABLE_CPU_FEATURES");
+    int is_enable = enable_env && enable_env[0];
+    int is_disable = disable_env && disable_env[0];
+    if (is_enable & is_disable) {
+        PyErr_Format(PyExc_ImportError,
+            "Both NPY_DISABLE_CPU_FEATURES and NPY_ENABLE_CPU_FEATURES "
+            "environment variables cannot be set simultaneously."
+        );
         return -1;
     }
+    if (is_enable | is_disable) {
+        if (npy__cpu_check_env(is_disable, is_disable ? disable_env : enable_env) < 0) {
+            return -1;
+        }
+    }
     return 0;
 }
 
@@ -210,7 +225,8 @@ npy__cpu_validate_baseline(void)
         *(fptr-1) = '\0'; // trim the last space
         PyErr_Format(PyExc_RuntimeError,
             "NumPy was built with baseline optimizations: \n"
-            "(" NPY_WITH_CPU_BASELINE ") but your machine doesn't support:\n(%s).",
+            "(" NPY_WITH_CPU_BASELINE ") but your machine "
+            "doesn't support:\n(%s).",
             baseline_failure
         );
         return -1;
@@ -220,27 +236,31 @@ npy__cpu_validate_baseline(void)
 }
 
 static int
-npy__cpu_try_disable_env(void)
-{
-    char *disenv = getenv("NPY_DISABLE_CPU_FEATURES");
-    if (disenv == NULL || disenv[0] == 0) {
-        return 0;
-    }
-    #define NPY__CPU_ENV_ERR_HEAD \
-        "During parsing environment variable 'NPY_DISABLE_CPU_FEATURES':\n"
+npy__cpu_check_env(int disable, const char *env) {
+
+    static const char *names[] = {
+        "enable", "disable",
+        "NPY_ENABLE_CPU_FEATURES", "NPY_DISABLE_CPU_FEATURES",
+        "During parsing environment variable: 'NPY_ENABLE_CPU_FEATURES':\n",
+        "During parsing environment variable: 'NPY_DISABLE_CPU_FEATURES':\n"
+    };
+    disable = disable ? 1 : 0;
+    const char *act_name = names[disable];
+    const char *env_name = names[disable + 2];
+    const char *err_head = names[disable + 4];
 
 #if !defined(NPY_DISABLE_OPTIMIZATION) && NPY_WITH_CPU_DISPATCH_N > 0
     #define NPY__MAX_VAR_LEN 1024 // More than enough for this era
-    size_t var_len = strlen(disenv) + 1;
+    size_t var_len = strlen(env) + 1;
     if (var_len > NPY__MAX_VAR_LEN) {
         PyErr_Format(PyExc_RuntimeError,
-            "Length of environment variable 'NPY_DISABLE_CPU_FEATURES' is %d, only %d accepted",
-            var_len, NPY__MAX_VAR_LEN - 1
+            "Length of environment variable '%s' is %d, only %d accepted",
+            env_name, var_len, NPY__MAX_VAR_LEN
         );
         return -1;
     }
-    char disable_features[NPY__MAX_VAR_LEN];
-    memcpy(disable_features, disenv, var_len);
+    char features[NPY__MAX_VAR_LEN];
+    memcpy(features, env, var_len);
 
     char nexist[NPY__MAX_VAR_LEN];
     char *nexist_cur = &nexist[0];
@@ -250,17 +270,19 @@ npy__cpu_try_disable_env(void)
 
     //comma and space including (htab, vtab, CR, LF, FF)
     const char *delim = ", \t\v\r\n\f";
-    char *feature = strtok(disable_features, delim);
+    char *feature = strtok(features, delim);
     while (feature) {
-        if (npy__cpu_baseline_fid(feature) > 0) {
-            PyErr_Format(PyExc_RuntimeError,
-                NPY__CPU_ENV_ERR_HEAD
-                "You cannot disable CPU feature '%s', since it is part of "
-                "the baseline optimizations:\n"
-                "(" NPY_WITH_CPU_BASELINE ").",
-                feature
-            );
-            return -1;
+        if (npy__cpu_baseline_fid(feature) > 0){
+            if (disable) {
+                PyErr_Format(PyExc_RuntimeError,
+                    "%s"
+                    "You cannot disable CPU feature '%s', since it is part of "
+                    "the baseline optimizations:\n"
+                    "(" NPY_WITH_CPU_BASELINE ").",
+                    err_head, feature
+                );
+                return -1;
+            } goto next;
         }
         // check if the feature is part of dispatched features
         int feature_id = npy__cpu_dispatch_fid(feature);
@@ -277,47 +299,58 @@ npy__cpu_try_disable_env(void)
             notsupp_cur[flen] = ' '; notsupp_cur += flen + 1;
             goto next;
         }
-        // Finally we can disable it
-        npy__cpu_have[feature_id] = 0;
+        // Finally we can disable or mark for enabling
+        npy__cpu_have[feature_id] = disable ? 0:2;
     next:
         feature = strtok(NULL, delim);
     }
+    if (!disable){
+        // Disables any unmarked dispatched feature.
+        #define NPY__CPU_DISABLE_DISPATCH_CB(FEATURE, DUMMY) \
+            if(npy__cpu_have[NPY_CAT(NPY_CPU_FEATURE_, FEATURE)] != 0)\
+            {npy__cpu_have[NPY_CAT(NPY_CPU_FEATURE_, FEATURE)]--;}\
+
+        NPY_WITH_CPU_DISPATCH_CALL(NPY__CPU_DISABLE_DISPATCH_CB, DUMMY) // extra arg for msvc
+    }
 
     *nexist_cur = '\0';
     if (nexist[0] != '\0') {
         *(nexist_cur-1) = '\0'; // trim the last space
-        if (PyErr_WarnFormat(PyExc_RuntimeWarning, 1,
-                NPY__CPU_ENV_ERR_HEAD
-                "You cannot disable CPU features (%s), since "
-                "they are not part of the dispatched optimizations\n"
-                "(" NPY_WITH_CPU_DISPATCH ").",
-                nexist
+        if (PyErr_WarnFormat(PyExc_ImportWarning, 1,
+            "%sYou cannot %s CPU features (%s), since "
+            "they are not part of the dispatched optimizations\n"
+            "(" NPY_WITH_CPU_DISPATCH ").",
+            err_head, act_name, nexist
         ) < 0) {
             return -1;
         }
+        return 0;
     }
 
+    #define NOTSUPP_BODY \
+                "%s" \
+                "You cannot %s CPU features (%s), since " \
+                "they are not supported by your machine.", \
+                err_head, act_name, notsupp
+
     *notsupp_cur = '\0';
     if (notsupp[0] != '\0') {
         *(notsupp_cur-1) = '\0'; // trim the last space
-        if (PyErr_WarnFormat(PyExc_RuntimeWarning, 1,
-                NPY__CPU_ENV_ERR_HEAD
-                "You cannot disable CPU features (%s), since "
-                "they are not supported by your machine.",
-                notsupp
-        ) < 0) {
+        if (!disable){
+            PyErr_Format(PyExc_RuntimeError, NOTSUPP_BODY);
             return -1;
         }
     }
 #else
-    if (PyErr_WarnFormat(PyExc_RuntimeWarning, 1,
-            NPY__CPU_ENV_ERR_HEAD
-            "You cannot use environment variable 'NPY_DISABLE_CPU_FEATURES', since "
+    if (PyErr_WarnFormat(PyExc_ImportWarning, 1,
+            "%s"
+            "You cannot use environment variable '%s', since "
         #ifdef NPY_DISABLE_OPTIMIZATION
-            "the NumPy library was compiled with optimization disabled."
+            "the NumPy library was compiled with optimization disabled.",
         #else
-            "the NumPy library was compiled without any dispatched optimizations."
+            "the NumPy library was compiled without any dispatched optimizations.",
         #endif
+        err_head, env_name, act_name
     ) < 0) {
         return -1;
     }
diff --git a/numpy/core/src/common/npy_cpu_features.h b/numpy/core/src/common/npy_cpu_features.h
index b49aea247..9180670c4 100644
--- a/numpy/core/src/common/npy_cpu_features.h
+++ b/numpy/core/src/common/npy_cpu_features.h
@@ -106,13 +106,25 @@ enum npy_cpu_features
  *  - detects runtime CPU features
  *  - check that baseline CPU features are present
  *  - uses 'NPY_DISABLE_CPU_FEATURES' to disable dispatchable features
+ *  - uses 'NPY_ENABLE_CPU_FEATURES' to enable dispatchable features
  *
  * It will set a RuntimeError when 
  *  - CPU baseline features from the build are not supported at runtime
  *  - 'NPY_DISABLE_CPU_FEATURES' tries to disable a baseline feature
- * and will warn if 'NPY_DISABLE_CPU_FEATURES' tries to disable a feature that
- * is not disabled (the machine or build does not support it, or the project was
- * not built with any feature optimization support)
+ *  - 'NPY_DISABLE_CPU_FEATURES' and 'NPY_ENABLE_CPU_FEATURES' are
+ *    simultaneously set
+ *  - 'NPY_ENABLE_CPU_FEATURES' tries to enable a feature that is not supported
+ *    by the machine or build
+ *  - 'NPY_ENABLE_CPU_FEATURES' tries to enable a feature when the project was
+ *    not built with any feature optimization support
+ *  
+ * It will set an ImportWarning when:
+ *  - 'NPY_DISABLE_CPU_FEATURES' tries to disable a feature that is not supported
+ *    by the machine or build
+ *  - 'NPY_DISABLE_CPU_FEATURES' or 'NPY_ENABLE_CPU_FEATURES' tries to
+ *    disable/enable a feature when the project was not built with any feature
+ *    optimization support
+ * 
  * return 0 on success otherwise return -1
  */
 NPY_VISIBILITY_HIDDEN int
diff --git a/numpy/core/tests/test_cpu_features.py b/numpy/core/tests/test_cpu_features.py
index 4718615b7..2fad4dfd9 100644
--- a/numpy/core/tests/test_cpu_features.py
+++ b/numpy/core/tests/test_cpu_features.py
@@ -1,5 +1,14 @@
 import sys, platform, re, pytest
-from numpy.core._multiarray_umath import __cpu_features__
+from numpy.core._multiarray_umath import (
+    __cpu_features__,
+    __cpu_baseline__,
+    __cpu_dispatch__,
+)
+import numpy as np
+import subprocess
+import pathlib
+import os
+import re
 
 def assert_features_equal(actual, desired, fname):
     __tracebackhide__ = True  # Hide traceback for py.test
@@ -48,6 +57,10 @@ def assert_features_equal(actual, desired, fname):
         "%s"
     ) % (fname, actual, desired, error_report))
 
+def _text_to_list(txt):
+    out = txt.strip("][\n").replace("'", "").split(', ')
+    return None if out[0] == "" else out
+
 class AbstractTest:
     features = []
     features_groups = {}
@@ -92,7 +105,6 @@ class AbstractTest:
         return values
 
     def load_flags_auxv(self):
-        import subprocess
         auxv = subprocess.check_output(['/bin/true'], env=dict(LD_SHOW_AUXV="1"))
         for at in auxv.split(b'\n'):
             if not at.startswith(b"AT_HWCAP"):
@@ -103,6 +115,208 @@ class AbstractTest:
                     hwcap_value[1].upper().decode().split()
                 )
 
+@pytest.mark.skipif(
+    sys.platform == 'emscripten',
+    reason= (
+        "The subprocess module is not available on WASM platforms and"
+        " therefore this test class cannot be properly executed."
+    ),
+)
+class TestEnvPrivation:
+    cwd = pathlib.Path(__file__).parent.resolve()
+    env = os.environ.copy()
+    _enable = os.environ.pop('NPY_ENABLE_CPU_FEATURES', None)
+    _disable = os.environ.pop('NPY_DISABLE_CPU_FEATURES', None)
+    SUBPROCESS_ARGS = dict(cwd=cwd, capture_output=True, text=True, check=True)
+    unavailable_feats = [
+        feat for feat in __cpu_dispatch__ if not __cpu_features__[feat]
+    ]
+    UNAVAILABLE_FEAT = (
+        None if len(unavailable_feats) == 0
+        else unavailable_feats[0]
+    )
+    BASELINE_FEAT = None if len(__cpu_baseline__) == 0 else __cpu_baseline__[0]
+    SCRIPT = """
+def main():
+    from numpy.core._multiarray_umath import __cpu_features__, __cpu_dispatch__
+
+    detected = [feat for feat in __cpu_dispatch__ if __cpu_features__[feat]]
+    print(detected)
+
+if __name__ == "__main__":
+    main()
+    """
+
+    @pytest.fixture(autouse=True)
+    def setup_class(self, tmp_path_factory):
+        file = tmp_path_factory.mktemp("runtime_test_script")
+        file /= "_runtime_detect.py"
+        file.write_text(self.SCRIPT)
+        self.file = file
+        return
+
+    def _run(self):
+        return subprocess.run(
+            [sys.executable, self.file],
+            env=self.env,
+            **self.SUBPROCESS_ARGS,
+            )
+
+    # Helper function mimicing pytest.raises for subprocess call
+    def _expect_error(
+        self,
+        msg,
+        err_type,
+        no_error_msg="Failed to generate error"
+    ):
+        try:
+            self._run()
+        except subprocess.CalledProcessError as e:
+            assertion_message = f"Expected: {msg}\nGot: {e.stderr}"
+            assert re.search(msg, e.stderr), assertion_message
+
+            assertion_message = (
+                f"Expected error of type: {err_type}; see full "
+                f"error:\n{e.stderr}"
+            )
+            assert re.search(err_type, e.stderr), assertion_message
+        else:
+            assert False, no_error_msg
+
+    def setup_method(self):
+        """Ensure that the environment is reset"""
+        self.env = os.environ.copy()
+        return
+
+    def test_runtime_feature_selection(self):
+        """
+        Ensure that when selecting `NPY_ENABLE_CPU_FEATURES`, only the
+        features exactly specified are dispatched.
+        """
+
+        # Capture runtime-enabled features
+        out = self._run()
+        non_baseline_features = _text_to_list(out.stdout)
+
+        if non_baseline_features is None:
+            pytest.skip(
+                "No dispatchable features outside of baseline detected."
+            )
+        feature = non_baseline_features[0]
+
+        # Capture runtime-enabled features when `NPY_ENABLE_CPU_FEATURES` is
+        # specified
+        self.env['NPY_ENABLE_CPU_FEATURES'] = feature
+        out = self._run()
+        enabled_features = _text_to_list(out.stdout)
+
+        # Ensure that only one feature is enabled, and it is exactly the one
+        # specified by `NPY_ENABLE_CPU_FEATURES`
+        assert set(enabled_features) == {feature}
+
+        if len(non_baseline_features) < 2:
+            pytest.skip("Only one non-baseline feature detected.")
+        # Capture runtime-enabled features when `NPY_ENABLE_CPU_FEATURES` is
+        # specified
+        self.env['NPY_ENABLE_CPU_FEATURES'] = ",".join(non_baseline_features)
+        out = self._run()
+        enabled_features = _text_to_list(out.stdout)
+
+        # Ensure that both features are enabled, and they are exactly the ones
+        # specified by `NPY_ENABLE_CPU_FEATURES`
+        assert set(enabled_features) == set(non_baseline_features)
+        return
+
+    @pytest.mark.parametrize("enabled, disabled",
+    [
+        ("feature", "feature"),
+        ("feature", "same"),
+    ])
+    def test_both_enable_disable_set(self, enabled, disabled):
+        """
+        Ensure that when both environment variables are set then an
+        ImportError is thrown
+        """
+        self.env['NPY_ENABLE_CPU_FEATURES'] = enabled
+        self.env['NPY_DISABLE_CPU_FEATURES'] = disabled
+        msg = "Both NPY_DISABLE_CPU_FEATURES and NPY_ENABLE_CPU_FEATURES"
+        err_type = "ImportError"
+        self._expect_error(msg, err_type)
+
+    @pytest.mark.skipif(
+        not __cpu_dispatch__,
+        reason=(
+            "NPY_*_CPU_FEATURES only parsed if "
+            "`__cpu_dispatch__` is non-empty"
+        )
+    )
+    @pytest.mark.parametrize("action", ["ENABLE", "DISABLE"])
+    def test_variable_too_long(self, action):
+        """
+        Test that an error is thrown if the environment variables are too long
+        to be processed. Current limit is 1024, but this may change later.
+        """
+        MAX_VAR_LENGTH = 1024
+        # Actual length is MAX_VAR_LENGTH + 1 due to null-termination
+        self.env[f'NPY_{action}_CPU_FEATURES'] = "t" * MAX_VAR_LENGTH
+        msg = (
+            f"Length of environment variable 'NPY_{action}_CPU_FEATURES' is "
+            f"{MAX_VAR_LENGTH + 1}, only {MAX_VAR_LENGTH} accepted"
+        )
+        err_type = "RuntimeError"
+        self._expect_error(msg, err_type)
+
+    @pytest.mark.skipif(
+        not __cpu_dispatch__,
+        reason=(
+            "NPY_*_CPU_FEATURES only parsed if "
+            "`__cpu_dispatch__` is non-empty"
+        )
+    )
+    def test_impossible_feature_disable(self):
+        """
+        Test that a RuntimeError is thrown if an impossible feature-disabling
+        request is made. This includes disabling a baseline feature.
+        """
+
+        if self.BASELINE_FEAT is None:
+            pytest.skip("There are no unavailable features to test with")
+        bad_feature = self.BASELINE_FEAT
+        self.env['NPY_DISABLE_CPU_FEATURES'] = bad_feature
+        msg = (
+            f"You cannot disable CPU feature '{bad_feature}', since it is "
+            "part of the baseline optimizations"
+        )
+        err_type = "RuntimeError"
+        self._expect_error(msg, err_type)
+
+    def test_impossible_feature_enable(self):
+        """
+        Test that a RuntimeError is thrown if an impossible feature-enabling
+        request is made. This includes enabling a feature not supported by the
+        machine, or disabling a baseline optimization.
+        """
+
+        if self.UNAVAILABLE_FEAT is None:
+            pytest.skip("There are no unavailable features to test with")
+        bad_feature = self.UNAVAILABLE_FEAT
+        self.env['NPY_ENABLE_CPU_FEATURES'] = bad_feature
+        msg = (
+            f"You cannot enable CPU features \\({bad_feature}\\), since "
+            "they are not supported by your machine."
+        )
+        err_type = "RuntimeError"
+        self._expect_error(msg, err_type)
+
+        # Ensure that only the bad feature gets reported
+        feats = f"{bad_feature}, {self.BASELINE_FEAT}"
+        self.env['NPY_ENABLE_CPU_FEATURES'] = feats
+        msg = (
+            f"You cannot enable CPU features \\({bad_feature}\\), since they "
+            "are not supported by your machine."
+        )
+        self._expect_error(msg, err_type)
+
 is_linux = sys.platform.startswith('linux')
 is_cygwin = sys.platform.startswith('cygwin')
 machine  = platform.machine()