Videre

linedance-app/venv/lib/python3.12/site-packages/numba/tests/test_threadsafety.py

@@ -0,0 +1,97 @@
+"""
+Test threadsafety for compiler.
+These tests will cause segfault if fail.
+"""
+import threading
+import random
+
+import numpy as np
+
+from numba import jit, vectorize, guvectorize
+
+from numba.tests.support import temp_directory, override_config
+from numba.core import config
+import unittest
+
+
+def foo(n, v):
+    return np.ones(n)
+
+
+def ufunc_foo(a, b):
+    return a + b
+
+
+def gufunc_foo(a, b, out):
+    out[0] = a + b
+
+
+
+class TestThreadSafety(unittest.TestCase):
+
+    def run_jit(self, **options):
+        def runner():
+            cfunc = jit(**options)(foo)
+
+            return cfunc(4, 10)
+        return runner
+
+    def run_compile(self, fnlist):
+        self._cache_dir = temp_directory(self.__class__.__name__)
+        with override_config('CACHE_DIR', self._cache_dir):
+            def chooser():
+                for _ in range(10):
+                    fn = random.choice(fnlist)
+                    fn()
+
+            ths = [threading.Thread(target=chooser)
+                   for i in range(4)]
+            for th in ths:
+                th.start()
+            for th in ths:
+                th.join()
+
+    def test_concurrent_jit(self):
+        self.run_compile([self.run_jit(nopython=True)])
+
+    def test_concurrent_jit_cache(self):
+        self.run_compile([self.run_jit(nopython=True, cache=True)])
+
+    def run_vectorize(self, **options):
+        def runner():
+            cfunc = vectorize(['(f4, f4)'], **options)(ufunc_foo)
+            a = b = np.random.random(10).astype(np.float32)
+            return cfunc(a, b)
+        return runner
+
+    def test_concurrent_vectorize(self):
+        self.run_compile([self.run_vectorize(nopython=True)])
+
+    def test_concurrent_vectorize_cache(self):
+        self.run_compile([self.run_vectorize(nopython=True, cache=True)])
+
+    def run_guvectorize(self, **options):
+        def runner():
+            sig = ['(f4, f4, f4[:])']
+            cfunc = guvectorize(sig, '(),()->()', **options)(gufunc_foo)
+            a = b = np.random.random(10).astype(np.float32)
+            return cfunc(a, b)
+        return runner
+
+    def test_concurrent_guvectorize(self):
+        self.run_compile([self.run_guvectorize(nopython=True)])
+
+    def test_concurrent_guvectorize_cache(self):
+        self.run_compile([self.run_guvectorize(nopython=True, cache=True)])
+
+    def test_concurrent_mix_use(self):
+        self.run_compile([self.run_jit(nopython=True, cache=True),
+                          self.run_jit(nopython=True),
+                          self.run_vectorize(nopython=True, cache=True),
+                          self.run_vectorize(nopython=True),
+                          self.run_guvectorize(nopython=True, cache=True),
+                          self.run_guvectorize(nopython=True)])
+
+
+if __name__ == '__main__':
+    unittest.main()