Repository: dgasmith/opt_einsum

Issue:
Support multi-threading in shared_intermediates
Following up on https://github.com/dgasmith/opt_einsum/pull/43#issuecomment-414661988, the `shared_intermediates` cache is currently not locked.

One option would be to assign the thread id of a unique owner of the cache, such that at least thread conflicts could be detected and reported, rather than mere silent performance degradation.

Comments:

Patch:
diff --git a/opt_einsum/sharing.py b/opt_einsum/sharing.py
index 54c94bd..038309f 100644
--- a/opt_einsum/sharing.py
+++ b/opt_einsum/sharing.py
@@ -3,11 +3,43 @@
 import contextlib
 import functools
 import numbers
-from collections import Counter
+from collections import Counter, defaultdict
+import threading
 
 from .parser import alpha_canonicalize, parse_einsum_input
 
-_SHARING_STACK = []
+
+_SHARING_STACK = defaultdict(list)
+
+
+try:
+    get_thread_id = threading.get_ident
+except AttributeError:
+    def get_thread_id():
+        return threading.current_thread().ident
+
+
+def currently_sharing():
+    """Check if we are currently sharing a cache -- thread specific.
+    """
+    return get_thread_id() in _SHARING_STACK
+
+
+def get_sharing_cache():
+    """Return the most recent sharing cache -- thread specific.
+    """
+    return _SHARING_STACK[get_thread_id()][-1]
+
+
+def _add_sharing_cache(cache):
+    _SHARING_STACK[get_thread_id()].append(cache)
+
+
+def _remove_sharing_cache():
+    tid = get_thread_id()
+    _SHARING_STACK[tid].pop()
+    if not _SHARING_STACK[tid]:
+        del _SHARING_STACK[tid]
 
 
 @contextlib.contextmanager
@@ -34,11 +66,11 @@ def shared_intermediates(cache=None):
     """
     if cache is None:
         cache = {}
+    _add_sharing_cache(cache)
     try:
-        _SHARING_STACK.append(cache)
         yield cache
     finally:
-        _SHARING_STACK.pop()
+        _remove_sharing_cache()
 
 
 def count_cached_ops(cache):
@@ -52,7 +84,7 @@ def _save_tensors(*tensors):
     """Save tensors in the cache to prevent their ids from being recycled.
     This is needed to prevent false cache lookups.
     """
-    cache = _SHARING_STACK[-1]
+    cache = get_sharing_cache()
     for tensor in tensors:
         cache['tensor', id(tensor)] = tensor
 
@@ -62,7 +94,7 @@ def _memoize(key, fn, *args, **kwargs):
     Results will be stored in the innermost ``cache`` yielded by
     :func:`shared_intermediates`.
     """
-    cache = _SHARING_STACK[-1]
+    cache = get_sharing_cache()
     if key in cache:
         return cache[key]
     result = fn(*args, **kwargs)
@@ -77,7 +109,7 @@ def transpose_cache_wrap(transpose):
 
     @functools.wraps(transpose)
     def cached_transpose(a, axes, backend='numpy'):
-        if not _SHARING_STACK:
+        if not currently_sharing():
             return transpose(a, axes, backend=backend)
 
         # hash by axes
@@ -96,7 +128,7 @@ def tensordot_cache_wrap(tensordot):
 
     @functools.wraps(tensordot)
     def cached_tensordot(x, y, axes=2, backend='numpy'):
-        if not _SHARING_STACK:
+        if not currently_sharing():
             return tensordot(x, y, axes, backend=backend)
 
         # hash based on the (axes_x,axes_y) form of axes
@@ -117,7 +149,7 @@ def einsum_cache_wrap(einsum):
 
     @functools.wraps(einsum)
     def cached_einsum(*args, **kwargs):
-        if not _SHARING_STACK:
+        if not currently_sharing():
             return einsum(*args, **kwargs)
 
         # hash modulo commutativity by computing a canonical ordering and names
@@ -143,7 +175,7 @@ def to_backend_cache_wrap(to_backend):
 
     @functools.wraps(to_backend)
     def cached_to_backend(array):
-        if not _SHARING_STACK:
+        if not currently_sharing():
             return to_backend(array)
 
         # hash by id


Test patch:
diff --git a/opt_einsum/tests/test_sharing.py b/opt_einsum/tests/test_sharing.py
index b0b860a..ea7f00e 100644
--- a/opt_einsum/tests/test_sharing.py
+++ b/opt_einsum/tests/test_sharing.py
@@ -6,11 +6,11 @@ import numpy as np
 import pytest
 
 from opt_einsum import (contract_expression, contract_path, get_symbol,
-                        helpers, shared_intermediates)
+                        contract, helpers, shared_intermediates)
 from opt_einsum.backends import to_cupy, to_torch
 from opt_einsum.contract import _einsum
 from opt_einsum.parser import parse_einsum_input
-from opt_einsum.sharing import count_cached_ops
+from opt_einsum.sharing import count_cached_ops, currently_sharing, get_sharing_cache
 
 try:
     import cupy
@@ -362,3 +362,23 @@ def test_chain_sharing(size, backend):
     print('Without sharing: {} expressions'.format(num_exprs_nosharing))
     print('With sharing: {} expressions'.format(num_exprs_sharing))
     assert num_exprs_nosharing > num_exprs_sharing
+
+
+def test_multithreaded_sharing():
+    from multiprocessing.pool import ThreadPool
+
+    def fn():
+        X, Y, Z = helpers.build_views('ab,bc,cd')
+
+        with shared_intermediates():
+            contract('ab,bc,cd->a', X, Y, Z)
+            contract('ab,bc,cd->b', X, Y, Z)
+
+            return len(get_sharing_cache())
+
+    expected = fn()
+    pool = ThreadPool(8)
+    fs = [pool.apply_async(fn) for _ in range(16)]
+    assert not currently_sharing()
+    assert [f.get() for f in fs] == [expected] * 16
+    pool.close()