MNT: move cpu dispatch registry into global data struct

numpy · ngoldbaum · Jun 19, 2024 · May 30, 2024 · May 21, 2024 · May 30, 2024
commit e5c1bd62b4ab92ac8aa47f4b2806c2dc2c109170
diff --git a/numpy/_core/src/common/npy_cpu_dispatch.c b/numpy/_core/src/common/npy_cpu_dispatch.c
@@ -1,11 +1,14 @@
-#include "npy_cpu_dispatch.h"
+#define NPY_NO_DEPRECATED_API NPY_API_VERSION
+#define _MULTIARRAYMODULE
 
-static PyObject *npy__cpu_dispatch_registery = NULL;
+#include "npy_cpu_dispatch.h"
+#include "numpy/ndarraytypes.h"
+#include "multiarraymodule.h"
 
 NPY_VISIBILITY_HIDDEN int
 npy_cpu_dispatch_tracer_init(PyObject *mod)
 {
-    if (npy__cpu_dispatch_registery != NULL) {
+    if (npy_ma_global_data->cpu_dispatch_registry != NULL) {
         PyErr_Format(PyExc_RuntimeError, "CPU dispatcher tracer already initlized");
         return -1;
     }
@@ -22,21 +25,21 @@ npy_cpu_dispatch_tracer_init(PyObject *mod)
     if (err != 0) {
         return -1;
     }
-    npy__cpu_dispatch_registery = reg_dict;
+    npy_ma_global_data->cpu_dispatch_registry = reg_dict;
     return 0;
 }
 
 NPY_VISIBILITY_HIDDEN void
 npy_cpu_dispatch_trace(const char *fname, const char *signature,
                        const char **dispatch_info)
 {
-    PyObject *func_dict = PyDict_GetItemString(npy__cpu_dispatch_registery, fname);
+    PyObject *func_dict = PyDict_GetItemString(npy_ma_global_data->cpu_dispatch_registry, fname);
     if (func_dict == NULL) {
         func_dict = PyDict_New();
         if (func_dict == NULL) {
             return;
         }
-        int err = PyDict_SetItemString(npy__cpu_dispatch_registery, fname, func_dict);
+        int err = PyDict_SetItemString(npy_ma_global_data->cpu_dispatch_registry, fname, func_dict);
         Py_DECREF(func_dict);
         if (err != 0) {
             return;

diff --git a/numpy/_core/src/common/npy_cpu_features.c b/numpy/_core/src/common/npy_cpu_features.c
@@ -5,6 +5,11 @@
 
 /******************** Private Definitions *********************/
 
+// This is initialized during module initialization and thereafter immutable.
+// We don't include it in the global data struct because the definitions in
+// this file are shared by the _simd, _umath_tests, and
+// _multiarray_umath modules
+
 // Hold all CPU features boolean values
 static unsigned char npy__cpu_have[NPY_CPU_FEATURE_MAX];
 

diff --git a/numpy/_core/src/multiarray/multiarraymodule.h b/numpy/_core/src/multiarray/multiarraymodule.h
@@ -54,6 +54,13 @@ typedef struct npy_ma_global_data_struct {
     PyObject *os_PathLike;
     PyObject *os_fspath;
 
+    /*
+     * Used for CPU feature detection and dispatch
+     *
+     * Filled in during module initialization and thereafter immutable
+     */
+    PyObject *cpu_dispatch_registry;
+
     /*
      * The following entries store cached references to object obtained
      * via an import. All of these are initialized at runtime by