MNT: lock initializing the argparse cache

numpy · rgommers · Jul 12, 2024 · May 13, 2024 · May 15, 2024 · Jun 21, 2024
commit 6386423f6a88b35e0ed0c06effe8cd8c3dec82fb
diff --git a/numpy/_core/src/common/npy_argparse.c b/numpy/_core/src/common/npy_argparse.c
@@ -7,11 +7,22 @@
 #include "numpy/ndarraytypes.h"
 #include "numpy/npy_2_compat.h"
 #include "npy_argparse.h"
-
+#include "npy_atomic.h"
 #include "npy_import.h"
 
 #include "arrayfunction_override.h"
 
+static PyThread_type_lock argparse_mutex;
+
+NPY_NO_EXPORT int
+init_argparse_mutex(void) {
+    argparse_mutex = PyThread_allocate_lock();
+    if (argparse_mutex == NULL) {
+        PyErr_NoMemory();
+        return -1;
+    }
+    return 0;
+}
 
 /**
  * Small wrapper converting to array just like CPython does.
@@ -274,14 +285,22 @@ _npy_parse_arguments(const char *funcname,
         /* ... is NULL, NULL, NULL terminated: name, converter, value */
         ...)
 {
-    if (NPY_UNLIKELY(cache->npositional == -1)) {
-        va_list va;
-        va_start(va, kwnames);
-
-        int res = initialize_keywords(funcname, cache, va);
-        va_end(va);
-        if (res < 0) {
-            return -1;
+    if (NPY_UNLIKELY(!cache->initialized)) {
+        // only do a possibly slow atomic load if the cache isn't already initialized
+        if (!npy_atomic_load_uint8(&cache->initialized)) {
+            PyThread_acquire_lock(argparse_mutex, WAIT_LOCK);
+            if (!cache->initialized) {
+                va_list va;
+                va_start(va, kwnames);
+                int res = initialize_keywords(funcname, cache, va);
+                va_end(va);
+                if (res < 0) {
+                    PyThread_release_lock(argparse_mutex);
+                    return -1;
+                }
+                cache->initialized = 1;
+            }
+            PyThread_release_lock(argparse_mutex);
         }
     }
 

diff --git a/numpy/_core/src/common/npy_argparse.h b/numpy/_core/src/common/npy_argparse.h
@@ -20,18 +20,19 @@
 NPY_NO_EXPORT int
 PyArray_PythonPyIntFromInt(PyObject *obj, int *value);
 
-
 #define _NPY_MAX_KWARGS 15
 
 typedef struct {
     int npositional;
     int nargs;
     int npositional_only;
     int nrequired;
+    npy_uint8 initialized;
     /* Null terminated list of keyword argument name strings */
     PyObject *kw_strings[_NPY_MAX_KWARGS+1];
 } _NpyArgParserCache;
 
+NPY_NO_EXPORT int init_argparse_mutex(void);
 
 /*
  * The sole purpose of this macro is to hide the argument parsing cache.

diff --git a/numpy/_core/src/common/npy_atomic.h b/numpy/_core/src/common/npy_atomic.h
@@ -0,0 +1,42 @@
+/*
+ * Provides wrappers around C11 standard library atomics and MSVC intrinsics
+ * to provide basic atomic load and store functionality. This is based on
+ * code in CPython's pyatomic.h, pyatomic_std.h, and pyatomic_msc.h
+ */
+
+#ifndef NUMPY_CORE_SRC_COMMON_NPY_ATOMIC_H_
+#define NUMPY_CORE_SRC_COMMON_NPY_ATOMIC_H_
+
+#include "numpy/npy_common.h"
+
+#if __STDC_VERSION__ >= 201112L && !defined(__STDC_NO_ATOMICS__)
+// TODO: support C++ atomics as well if this header is ever needed in C++
+    #include <stdatomic.h>
+    #include <stdint.h>
+    #define STDC_ATOMICS
+#elif _MSC_VER
+    #include <intrin.h>
+    #define MSC_ATOMICS
+#else
+    #error "no support for missing C11 atomics except with MSVC"
+#endif
+
+
+static inline npy_uint8 npy_atomic_load_uint8(const npy_uint8 *obj) {
+#ifdef STDC_ATOMICS
+    return (npy_uint8)atomic_load((const _Atomic(uint8_t)*)obj);
+#elif defined(MSC_ATOMICS)
+#if defined(_M_X64) || defined(_M_IX86)
+    return *(volatile npy_uint8 *)obj;
+#elif defined(_M_ARM64)
+    return (npy_uint8)__ldar8((unsigned __int8 volatile *)obj);
+#else
+#error "Unsupported MSVC build configuration, neither x86 or ARM"
+#endif
+#endif
+}
+
+#undef MSC_ATOMICS
+#undef STDC_ATOMICS
+
+#endif // NUMPY_CORE_SRC_COMMON_NPY_NPY_ATOMIC_H_
diff --git a/numpy/_core/src/multiarray/_multiarray_tests.c.src b/numpy/_core/src/multiarray/_multiarray_tests.c.src
@@ -2432,6 +2432,9 @@ PyMODINIT_FUNC PyInit__multiarray_tests(void)
         return m;
     }
     import_array();
+    if (init_argparse_mutex() < 0) {
+        return NULL;
+    }
     if (PyErr_Occurred()) {
         PyErr_SetString(PyExc_RuntimeError,
                         "cannot load _multiarray_tests module.");

diff --git a/numpy/_core/src/umath/umathmodule.c b/numpy/_core/src/umath/umathmodule.c
@@ -22,6 +22,7 @@
 #include "numpy/ufuncobject.h"
 #include "numpy/npy_3kcompat.h"
 #include "npy_pycompat.h"
+#include "npy_argparse.h"
 #include "abstract.h"
 
 #include "numpy/npy_math.h"
@@ -321,5 +322,9 @@ int initumath(PyObject *m)
         return -1;
     }
 
+    if (init_argparse_mutex() < 0) {
+        return -1;
+    }
+
     return 0;
 }