Add threading tests for shared input

Author: crusaderky

pyproject.toml

@@ -63,3 +63,7 @@ addopts = "--strict-markers --strict-config -v -r sxfE --color=yes --durations=1
 python_files = ["test_*.py"]
 testpaths = ["tests"]
 filterwarnings = ["error"]
+markers = [
+  # Needed when pytest-run-parallel is not installed
+  "thread_unsafe: Run test in single thread in pytest-run-parallel"
+]

tests/blis_tests_common.py

@@ -1,4 +1,5 @@
 # Copyright ExplosionAI GmbH, released under BSD.
+from concurrent.futures import ThreadPoolExecutor
 from hypothesis import settings
 from hypothesis import strategies as st
 from hypothesis.extra.numpy import arrays
@@ -23,3 +24,12 @@ def ndarrays(shape, lo, hi, dtype):
     width = int(dtype[5:])
     elements = st.floats(min_value=lo, max_value=hi, width=width)
     return arrays(dtype, shape=shape, elements=elements)
+
+
+def run_threaded(func, max_workers=8, outer_iterations=2):
+    """Runs a function many times in parallel"""
+    with ThreadPoolExecutor(max_workers=max_workers) as tpe:
+        for _ in range(outer_iterations):
+            futures = [tpe.submit(func) for _ in range(max_workers)]
+            for f in futures:
+                f.result()

tests/test_dotv.py

@@ -5,7 +5,7 @@
 from hypothesis import strategies as st
 from numpy.testing import assert_allclose
 
-from blis_tests_common import dtypes_tols, ndarrays
+from blis_tests_common import dtypes_tols, ndarrays, run_threaded
 from blis.py import dotv
 
 
@@ -29,3 +29,17 @@ def test_memoryview_noconj(arrays):
     numpy_result = A.dot(B)
     result = dotv(A, B)
     assert_allclose(result, numpy_result, atol=atol, rtol=rtol)
+
+
+@given(dotv_arrays())
+@pytest.mark.thread_unsafe(reason="Uses run_threaded")
+def test_threads_share_input(arrays):
+    """Test when multiple threads share the same input arrays."""
+    A, B, atol, rtol = arrays
+    numpy_result = A.dot(B)
+
+    def f():
+        result = dotv(A, B)
+        assert_allclose(result, numpy_result, atol=atol, rtol=rtol)
+
+    run_threaded(f)

tests/test_gemm.py

@@ -5,7 +5,7 @@
 from hypothesis import strategies as st
 from numpy.testing import assert_allclose
 
-from blis_tests_common import dtypes_tols, ndarrays
+from blis_tests_common import dtypes_tols, ndarrays, run_threaded
 from blis.py import gemm
 
 
@@ -40,3 +40,18 @@ def test_memoryview_notrans(arrays):
     C = np.zeros_like(numpy_result)  # (l, n)
     gemm(A, B, out=C)
     assert_allclose(C, numpy_result, atol=atol, rtol=rtol)
+
+
+@given(gemm_arrays())
+@pytest.mark.thread_unsafe(reason="Uses run_threaded")
+def test_threads_share_input(arrays):
+    """Test when multiple threads share the same input arrays."""
+    A, B, atol, rtol = arrays
+    numpy_result = A.dot(B)
+
+    def f():
+        C = np.zeros_like(numpy_result)  # (l, n)
+        gemm(A, B, out=C)
+        assert_allclose(C, numpy_result, atol=atol, rtol=rtol)
+
+    run_threaded(f)