MNT: move some thread-unsafe state in thread-unsafe state struct

numpy · ngoldbaum · Jun 19, 2024 · May 30, 2024 · May 21, 2024 · May 30, 2024
commit 402a83ce219c01fd2e0eda8748c31a1e6d352f8c
diff --git a/numpy/_core/src/multiarray/alloc.c b/numpy/_core/src/multiarray/alloc.c
@@ -11,6 +11,7 @@
 #include "numpy/npy_common.h"
 #include "npy_config.h"
 #include "alloc.h"
+#include "multiarraymodule.h"
 
 #include <assert.h>
 #ifdef NPY_OS_LINUX
@@ -35,21 +36,19 @@ typedef struct {
 static cache_bucket datacache[NBUCKETS];
 static cache_bucket dimcache[NBUCKETS_DIM];
 
-static int _madvise_hugepage = 1;
-
 
 /*
  * This function tells whether NumPy attempts to call `madvise` with
  * `MADV_HUGEPAGE`.  `madvise` is only ever used on linux, so the value
- * of `_madvise_hugepage` may be ignored.
+ * of `madvise_hugepage` may be ignored.
  *
  * It is exposed to Python as `np._core.multiarray._get_madvise_hugepage`.
  */
 NPY_NO_EXPORT PyObject *
 _get_madvise_hugepage(PyObject *NPY_UNUSED(self), PyObject *NPY_UNUSED(args))
 {
 #ifdef NPY_OS_LINUX
-    if (_madvise_hugepage) {
+    if (npy_ma_thread_unsafe_state->madvise_hugepage) {
         Py_RETURN_TRUE;
     }
 #endif
@@ -59,20 +58,20 @@ _get_madvise_hugepage(PyObject *NPY_UNUSED(self), PyObject *NPY_UNUSED(args))
 
 /*
  * This function enables or disables the use of `MADV_HUGEPAGE` on Linux
- * by modifying the global static `_madvise_hugepage`.
- * It returns the previous value of `_madvise_hugepage`.
+ * by modifying the global static `madvise_hugepage`.
+ * It returns the previous value of `madvise_hugepage`.
  *
  * It is exposed to Python as `np._core.multiarray._set_madvise_hugepage`.
  */
 NPY_NO_EXPORT PyObject *
 _set_madvise_hugepage(PyObject *NPY_UNUSED(self), PyObject *enabled_obj)
 {
-    int was_enabled = _madvise_hugepage;
+    int was_enabled = npy_ma_thread_unsafe_state->madvise_hugepage;
     int enabled = PyObject_IsTrue(enabled_obj);
     if (enabled < 0) {
         return NULL;
     }
-    _madvise_hugepage = enabled;
+    npy_ma_thread_unsafe_state->madvise_hugepage = enabled;
     if (was_enabled) {
         Py_RETURN_TRUE;
     }
@@ -110,7 +109,8 @@ _npy_alloc_cache(npy_uintp nelem, npy_uintp esz, npy_uint msz,
 #endif
 #ifdef NPY_OS_LINUX
         /* allow kernel allocating huge pages for large arrays */
-        if (NPY_UNLIKELY(nelem * esz >= ((1u<<22u))) && _madvise_hugepage) {
+        if (NPY_UNLIKELY(nelem * esz >= ((1u<<22u))) &&
+            npy_ma_thread_unsafe_state->madvise_hugepage) {
             npy_uintp offset = 4096u - (npy_uintp)p % (4096u);
             npy_uintp length = nelem * esz - offset;
             /**

diff --git a/numpy/_core/src/multiarray/multiarraymodule.c b/numpy/_core/src/multiarray/multiarraymodule.c
@@ -4345,8 +4345,6 @@ _set_numpy_warn_if_no_mem_policy(PyObject *NPY_UNUSED(self), PyObject *arg)
 
 static PyObject *
 _reload_guard(PyObject *NPY_UNUSED(self), PyObject *NPY_UNUSED(args)) {
-    static int initialized = 0;
-
 #if !defined(PYPY_VERSION)
     if (PyThreadState_Get()->interp != PyInterpreterState_Main()) {
         if (PyErr_WarnEx(PyExc_UserWarning,
@@ -4362,19 +4360,19 @@ _reload_guard(PyObject *NPY_UNUSED(self), PyObject *NPY_UNUSED(args)) {
             return NULL;
         }
         /* No need to give the other warning in a sub-interpreter as well... */
-        initialized = 1;
+        npy_ma_thread_unsafe_state->reload_guard_initialized = 1;
         Py_RETURN_NONE;
     }
 #endif
-    if (initialized) {
+    if (npy_ma_thread_unsafe_state->reload_guard_initialized) {
         if (PyErr_WarnEx(PyExc_UserWarning,
                 "The NumPy module was reloaded (imported a second time). "
                 "This can in some cases result in small but subtle issues "
                 "and is discouraged.", 2) < 0) {
             return NULL;
         }
     }
-    initialized = 1;
+    npy_ma_thread_unsafe_state->reload_guard_initialized = 1;
     Py_RETURN_NONE;
 }
 
@@ -4901,6 +4899,7 @@ initialize_static_globals(void)
     // this is module-level global heap allocation, it is currently
     // never freed
     npy_ma_static_data = PyMem_Calloc(1, sizeof(npy_ma_static_data_struct));
+    npy_ma_thread_unsafe_state = PyMem_Calloc(1, sizeof(npy_ma_thread_unsafe_state_struct));
 
     // cached reference to objects defined in python
 

diff --git a/numpy/_core/src/multiarray/multiarraymodule.h b/numpy/_core/src/multiarray/multiarraymodule.h
@@ -194,12 +194,25 @@ typedef struct npy_ma_thread_unsafe_state_struct {
     PyObject *_var;
     PyObject *_view_is_safe;
     PyObject *_void_scalar_to_string;
-    PyObject *array_function_errmsg_formatter;
-    PyObject *array_ufunc_errmsg_formatter;
-    PyObject *internal_gcd_func;
-    PyObject *npy_ctypes_check;
-    PyObject *numpy_matrix;
-    PyObject *NO_NEP50_WARNING;
-} npy_ma_global_data_struct;
+
+    /*
+     * Used to test the internal-only scaled float test dtype
+     */
+    npy_bool get_sfloat_dtype_initialized;
+
+    /*
+     * controls the global madvise hugepage setting
+     */
+    int madvise_hugepage;
+
+    /*
+     * used to detect module reloading in the reload guard
+     */
+    int reload_guard_initialized;
+} npy_ma_thread_unsafe_state_struct;
+
+NPY_VISIBILITY_HIDDEN extern npy_ma_str_struct *npy_ma_str;
+NPY_VISIBILITY_HIDDEN extern npy_ma_static_data_struct *npy_ma_static_data;
+NPY_VISIBILITY_HIDDEN extern npy_ma_thread_unsafe_state_struct *npy_ma_thread_unsafe_state;
 
 #endif  /* NUMPY_CORE_SRC_MULTIARRAY_MULTIARRAYMODULE_H_ */
diff --git a/numpy/_core/src/umath/_scaled_float_dtype.c b/numpy/_core/src/umath/_scaled_float_dtype.c
@@ -867,10 +867,7 @@ sfloat_init_ufuncs(void) {
 NPY_NO_EXPORT PyObject *
 get_sfloat_dtype(PyObject *NPY_UNUSED(mod), PyObject *NPY_UNUSED(args))
 {
-    /* Allow calling the function multiple times. */
-    static npy_bool initialized = NPY_FALSE;
-
    if (initialized) {
+    if (npy_ma_thread_unsafe_state->get_sfloat_dtype_initialized) {
         Py_INCREF(&PyArray_SFloatDType);
         return (PyObject *)&PyArray_SFloatDType;
     }
@@ -899,6 +896,6 @@ get_sfloat_dtype(PyObject *NPY_UNUSED(mod), PyObject *NPY_UNUSED(args))
         return NULL;
     }
 
-    initialized = NPY_TRUE;
+    npy_ma_thread_unsafe_state->get_sfloat_dtype_initialized = NPY_TRUE;
     return (PyObject *)&PyArray_SFloatDType;
 }