ENH: mark the dragon4 scratch space as thread-local

Author: ngoldbaum

numpy/_core/src/multiarray/dragon4.c

@@ -163,12 +163,11 @@ typedef struct {
     char repr[16384];
 } Dragon4_Scratch;
 
-static int _bigint_static_in_use = 0;
-static Dragon4_Scratch _bigint_static;
+static NPY_TLS int _bigint_static_in_use = 0;
+static NPY_TLS Dragon4_Scratch _bigint_static;
 
 static Dragon4_Scratch*
 get_dragon4_bigint_scratch(void) {
-    /* this test+set is not threadsafe, but no matter because we have GIL */
     if (_bigint_static_in_use) {
         PyErr_SetString(PyExc_RuntimeError,
             "numpy float printing code is not re-entrant. "

numpy/_core/tests/test_arrayprint.py

@@ -9,6 +9,7 @@
     assert_, assert_equal, assert_raises, assert_warns, HAS_REFCOUNT,
     assert_raises_regex, IS_WASM
     )
+from numpy.testing._private.utils import run_threaded
 from numpy._core.arrayprint import _typelessdata
 import textwrap
 
@@ -1249,3 +1250,10 @@ async def main():
     loop = asyncio.new_event_loop()
     asyncio.run(main())
     loop.close()
+
+@pytest.mark.skipif(IS_WASM, reason="wasm doesn't support threads")
+def test_multithreaded_array_printing():
+    # the dragon4 implementation uses a static scratch space for performance
+    # reasons this test makes sure it is set up in a thread-safe manner
+
+    run_threaded(TestPrintOptions().test_floatmode, 500)

numpy/_core/tests/test_multithreading.py

@@ -1,25 +1,15 @@
-import concurrent.futures
 import threading
 
 import numpy as np
 import pytest
 
 from numpy.testing import IS_WASM
+from numpy.testing._private.utils import run_threaded
 
 if IS_WASM:
     pytest.skip(allow_module_level=True, reason="no threading support in wasm")
 
 
-def run_threaded(func, iters, pass_count=False):
-    with concurrent.futures.ThreadPoolExecutor(max_workers=8) as tpe:
-        if pass_count:
-            futures = [tpe.submit(func, i) for i in range(iters)]
-        else:
-            futures = [tpe.submit(func) for _ in range(iters)]
-        for f in futures:
-            f.result()
-
-
 def test_parallel_randomstate_creation():
     # if the coercion cache is enabled and not thread-safe, creating
     # RandomState instances simultaneously leads to a data race

numpy/testing/_private/utils.py

@@ -17,6 +17,7 @@
 from warnings import WarningMessage
 import pprint
 import sysconfig
+import concurrent.futures
 
 import numpy as np
 from numpy._core import (
@@ -40,7 +41,7 @@
         'HAS_REFCOUNT', "IS_WASM", 'suppress_warnings', 'assert_array_compare',
         'assert_no_gc_cycles', 'break_cycles', 'HAS_LAPACK64', 'IS_PYSTON',
         '_OLD_PROMOTION', 'IS_MUSL', '_SUPPORTS_SVE', 'NOGIL_BUILD',
-        'IS_EDITABLE'
+        'IS_EDITABLE', 'run_threaded',
         ]
 
 
@@ -2697,3 +2698,14 @@ def _get_glibc_version():
 
 _glibcver = _get_glibc_version()
 _glibc_older_than = lambda x: (_glibcver != '0.0' and _glibcver < x)
+
+
+def run_threaded(func, iters, pass_count=False):
+    """Runs a function many times in parallel"""
+    with concurrent.futures.ThreadPoolExecutor(max_workers=8) as tpe:
+        if pass_count:
+            futures = [tpe.submit(func, i) for i in range(iters)]
+        else:
+            futures = [tpe.submit(func) for _ in range(iters)]
+        for f in futures:
+            f.result()