Support multithreading

bench/large_array_vs_numpy.py

@@ -0,0 +1,152 @@
+#################################################################################
+# To mimic the scenario that computation is i/o bound and constrained by memory
+#
+# It's a much simplified version that the chunk is computed in a loop,
+# and expression is evaluated in a sequence, which is not true in reality.
+# Neverthless, numexpr outperforms numpy.
+#################################################################################
+"""
+Benchmarking Expression 1:
+NumPy time (threaded over 32 chunks with 2 threads): 4.612313 seconds
+numexpr time (threaded with re_evaluate over 32 chunks with 2 threads): 0.951172 seconds
+numexpr speedup: 4.85x
+----------------------------------------
+Benchmarking Expression 2:
+NumPy time (threaded over 32 chunks with 2 threads): 23.862752 seconds
+numexpr time (threaded with re_evaluate over 32 chunks with 2 threads): 2.182058 seconds
+numexpr speedup: 10.94x
+----------------------------------------
+Benchmarking Expression 3:
+NumPy time (threaded over 32 chunks with 2 threads): 20.594895 seconds
+numexpr time (threaded with re_evaluate over 32 chunks with 2 threads): 2.927881 seconds
+numexpr speedup: 7.03x
+----------------------------------------
+Benchmarking Expression 4:
+NumPy time (threaded over 32 chunks with 2 threads): 12.834101 seconds
+numexpr time (threaded with re_evaluate over 32 chunks with 2 threads): 5.392480 seconds
+numexpr speedup: 2.38x
+----------------------------------------
+"""
+
+import os
+
+os.environ["NUMEXPR_NUM_THREADS"] = "16"
+import numpy as np
+import numexpr as ne
+import timeit
+import threading
+
+array_size = 10**8
+num_runs = 10
+num_chunks = 32  # Number of chunks
+num_threads = 2  # Number of threads constrained by how many chunks memory can hold
+
+a = np.random.rand(array_size).reshape(10**4, -1)
+b = np.random.rand(array_size).reshape(10**4, -1)
+c = np.random.rand(array_size).reshape(10**4, -1)
+
+chunk_size = array_size // num_chunks
+
+expressions_numpy = [
+    lambda a, b, c: a + b * c,
+    lambda a, b, c: a**2 + b**2 - 2 * a * b * np.cos(c),
+    lambda a, b, c: np.sin(a) + np.log(b) * np.sqrt(c),
+    lambda a, b, c: np.exp(a) + np.tan(b) - np.sinh(c),
+]
+
+expressions_numexpr = [
+    "a + b * c",
+    "a**2 + b**2 - 2 * a * b * cos(c)",
+    "sin(a) + log(b) * sqrt(c)",
+    "exp(a) + tan(b) - sinh(c)",
+]
+
+
+def benchmark_numpy_chunk(func, a, b, c, results, indices):
+    for index in indices:
+        start = index * chunk_size
+        end = (index + 1) * chunk_size
+        time_taken = timeit.timeit(
+            lambda: func(a[start:end], b[start:end], c[start:end]), number=num_runs
+        )
+        results.append(time_taken)
+
+
+def benchmark_numexpr_re_evaluate(expr, a, b, c, results, indices):
+    for index in indices:
+        start = index * chunk_size
+        end = (index + 1) * chunk_size
+        if index == 0:
+            # Evaluate the first chunk with evaluate
+            time_taken = timeit.timeit(
+                lambda: ne.evaluate(
+                    expr,
+                    local_dict={
+                        "a": a[start:end],
+                        "b": b[start:end],
+                        "c": c[start:end],
+                    },
+                ),
+                number=num_runs,
+            )
+        else:
+            # Re-evaluate subsequent chunks with re_evaluate
+            time_taken = timeit.timeit(
+                lambda: ne.re_evaluate(
+                    local_dict={"a": a[start:end], "b": b[start:end], "c": c[start:end]}
+                ),
+                number=num_runs,
+            )
+        results.append(time_taken)
+
+
+def run_benchmark_threaded():
+    chunk_indices = list(range(num_chunks))
+
+    for i in range(len(expressions_numpy)):
+        print(f"Benchmarking Expression {i+1}:")
+
+        results_numpy = []
+        results_numexpr = []
+
+        threads_numpy = []
+        for j in range(num_threads):
+            indices = chunk_indices[j::num_threads]  # Distribute chunks across threads
+            thread = threading.Thread(
+                target=benchmark_numpy_chunk,
+                args=(expressions_numpy[i], a, b, c, results_numpy, indices),
+            )
+            threads_numpy.append(thread)
+            thread.start()
+
+        for thread in threads_numpy:
+            thread.join()
+
+        numpy_time = sum(results_numpy)
+        print(
+            f"NumPy time (threaded over {num_chunks} chunks with {num_threads} threads): {numpy_time:.6f} seconds"
+        )
+
+        threads_numexpr = []
+        for j in range(num_threads):
+            indices = chunk_indices[j::num_threads]  # Distribute chunks across threads
+            thread = threading.Thread(
+                target=benchmark_numexpr_re_evaluate,
+                args=(expressions_numexpr[i], a, b, c, results_numexpr, indices),
+            )
+            threads_numexpr.append(thread)
+            thread.start()
+
+        for thread in threads_numexpr:
+            thread.join()
+
+        numexpr_time = sum(results_numexpr)
+        print(
+            f"numexpr time (threaded with re_evaluate over {num_chunks} chunks with {num_threads} threads): {numexpr_time:.6f} seconds"
+        )
+        print(f"numexpr speedup: {numpy_time / numexpr_time:.2f}x")
+        print("-" * 40)
+
+
+if __name__ == "__main__":
+    run_benchmark_threaded()

numexpr/necompiler.py

@@ -19,7 +19,7 @@
 
 is_cpu_amd_intel = False # DEPRECATION WARNING: WILL BE REMOVED IN FUTURE RELEASE
 from numexpr import interpreter, expressions, use_vml
-from numexpr.utils import CacheDict
+from numexpr.utils import CacheDict, ContextDict
 
 # Declare a double type that does not exist in Python space
 double = numpy.double
@@ -776,11 +776,9 @@ def getArguments(names, local_dict=None, global_dict=None, _frame_depth: int=2):
 # Dictionaries for caching variable names and compiled expressions
 _names_cache = CacheDict(256)
 _numexpr_cache = CacheDict(256)
-_numexpr_last = {}
+_numexpr_last = ContextDict()
 evaluate_lock = threading.Lock()
 
-# MAYBE: decorate this function to add attributes instead of having the 
-# _numexpr_last dictionary?
 def validate(ex: str, 
              local_dict: Optional[Dict] = None, 
              global_dict: Optional[Dict] = None,
@@ -887,7 +885,7 @@ def validate(ex: str,
             compiled_ex = _numexpr_cache[numexpr_key] = NumExpr(ex, signature, sanitize=sanitize, **context)
         kwargs = {'out': out, 'order': order, 'casting': casting,
                 'ex_uses_vml': ex_uses_vml}
-        _numexpr_last = dict(ex=compiled_ex, argnames=names, kwargs=kwargs)
+        _numexpr_last.set(ex=compiled_ex, argnames=names, kwargs=kwargs)
     except Exception as e:
         return e
     return None

numexpr/tests/test_numexpr.py

@@ -1201,6 +1201,7 @@ def run(self):
         test.join()
 
     def test_multithread(self):
+
         import threading
 
         # Running evaluate() from multiple threads shouldn't crash
@@ -1218,6 +1219,77 @@ def work(n):
         for t in threads:
             t.join()
 
+    def test_thread_safety(self):
+        """
+        Expected output
+
+        When not safe (before the pr this test is commited)
+        AssertionError: Thread-0 failed: result does not match expected
+
+        When safe (after the pr this test is commited)
+        Should pass without failure
+        """
+        import threading
+        import time
+
+        barrier = threading.Barrier(4)
+
+        # Function that each thread will run with different expressions
+        def thread_function(a_value, b_value, expression, expected_result, results, index):
+            validate(expression, local_dict={"a": a_value, "b": b_value})
+            # Wait for all threads to reach this point
+            # such that they all set _numexpr_last
+            barrier.wait()
+
+            # Simulate some work or a context switch delay
+            time.sleep(0.1)
+
+            result = re_evaluate(local_dict={"a": a_value, "b": b_value})
+            results[index] = np.array_equal(result, expected_result)
+
+        def test_thread_safety_with_numexpr():
+            num_threads = 4
+            array_size = 1000000
+
+            expressions = [
+                "a + b",
+                "a - b",
+                "a * b",
+                "a / b"
+            ]
+
+            a_value = [np.full(array_size, i + 1) for i in range(num_threads)]
+            b_value = [np.full(array_size, (i + 1) * 2) for i in range(num_threads)]
+
+            expected_results = [
+                a_value[i] + b_value[i] if expr == "a + b" else
+                a_value[i] - b_value[i] if expr == "a - b" else
+                a_value[i] * b_value[i] if expr == "a * b" else
+                a_value[i] / b_value[i] if expr == "a / b" else None
+                for i, expr in enumerate(expressions)
+            ]
+
+            results = [None] * num_threads
+            threads = []
+
+            # Create and start threads with different expressions
+            for i in range(num_threads):
+                thread = threading.Thread(
+                    target=thread_function,
+                    args=(a_value[i], b_value[i], expressions[i], expected_results[i], results, i)
+                )
+                threads.append(thread)
+                thread.start()
+
+            for thread in threads:
+                thread.join()
+
+            for i in range(num_threads):
+                if not results[i]:
+                    self.fail(f"Thread-{i} failed: result does not match expected")
+
+        test_thread_safety_with_numexpr()
+
 
 # The worker function for the subprocess (needs to be here because Windows
 # has problems pickling nested functions with the multiprocess module :-/)

numexpr/utils.py

@@ -13,6 +13,7 @@
 
 import os
 import subprocess
+import contextvars
 
 from numexpr.interpreter import _set_num_threads, _get_num_threads, MAX_THREADS
 from numexpr import use_vml
@@ -226,3 +227,83 @@ def __setitem__(self, key, value):
                 super(CacheDict, self).__delitem__(k)
         super(CacheDict, self).__setitem__(key, value)
 
+
+class ContextDict:
+    """
+    A context aware version dictionary
+    """
+    def __init__(self):
+        self._context_data = contextvars.ContextVar('context_data', default={})
+
+    def set(self, key=None, value=None, **kwargs):
+        data = self._context_data.get().copy()
+
+        if key is not None:
+            data[key] = value
+
+        for k, v in kwargs.items():
+            data[k] = v
+
+        self._context_data.set(data)
+
+    def get(self, key, default=None):
+        data = self._context_data.get()
+        return data.get(key, default)
+
+    def delete(self, key):
+        data = self._context_data.get().copy()
+        if key in data:
+            del data[key]
+        self._context_data.set(data)
+
+    def clear(self):
+        self._context_data.set({})
+
+    def all(self):
+        return self._context_data.get()
+
+    def update(self, *args, **kwargs):
+        data = self._context_data.get().copy()
+
+        if args:
+            if len(args) > 1:
+                raise TypeError(f"update() takes at most 1 positional argument ({len(args)} given)")
+            other = args[0]
+            if isinstance(other, dict):
+                data.update(other)
+            else:
+                for k, v in other:
+                    data[k] = v
+
+        data.update(kwargs)
+        self._context_data.set(data)
+
+    def keys(self):
+        return self._context_data.get().keys()
+
+    def values(self):
+        return self._context_data.get().values()
+
+    def items(self):
+        return self._context_data.get().items()
+
+    def __getitem__(self, key):
+        return self.get(key)
+
+    def __setitem__(self, key, value):
+        self.set(key, value)
+
+    def __delitem__(self, key):
+        self.delete(key)
+
+    def __contains__(self, key):
+        return key in self._context_data.get()
+
+    def __len__(self):
+        return len(self._context_data.get())
+
+    def __iter__(self):
+        return iter(self._context_data.get())
+
+    def __repr__(self):
+        return repr(self._context_data.get())