add thread-local storage and helpers from scipy/scipy#12512
# distutils: language = c++
-from libcpp.list cimport list as cpplist
-from libcpp cimport vector
+import threading
from cupy.cuda import cufft
+cdef object _thread_local = threading.local()
+
+
cdef class _Node:
# Unfortunately cython cdef class cannot be nested, so the node class
# has to live outside of the linked list...
@@ -93,8 +95,8 @@ cdef class PlanCache:
cdef dict cache
# for keeping track of least recently used plans
- # lru.head: most recent used
- # lru.tail: least recent used
+ # lru.head: least recent used
+ # lru.tail: most recent used
cdef _LinkedList lru
cdef inline void _reset(self):
@@ -104,10 +106,17 @@ cdef class PlanCache:
self.lru = _LinkedList()
def __init__(self, Py_ssize_t size=4, Py_ssize_t memsize=-1):
+ # TODO(leofang): use stream as part of cache key?
self._validate_size_memsize(size, memsize)
self._set_size_memsize(size, memsize)
self._reset()
+ if hasattr(_thread_local, '_current_plan_cache'):
+ import warnings
+ warnings.warn('invalidating the existing thread-local '
+ 'plan cache...')
+ _thread_local._current_plan_cache = self
+
def __getitem__(self, tuple key):
cdef _Node node = self.cache.get(key)
if node is not None:
@@ -231,6 +240,7 @@ cdef class PlanCache:
'(unlimited)' if self.memsize == -1 else self.memsize)
output += '\ncached plans (least used first):\n'
+ # TODO(leofang): maybe traverse from the end?
cdef _Node node = self.lru.head
cdef size_t count = 0
while node.next is not self.lru.tail:
@@ -243,3 +253,42 @@ cdef class PlanCache:
cpdef show_info(self):
print(self)
+
+
+cpdef Py_ssize_t get_plan_cache_size(size):
+ if not hasattr(_thread_local, '_current_plan_cache'):
+ raise RuntimeError('cache not found')
+ cdef PlanCache cache = _thread_local._current_plan_cache
+ return cache.get_size()
+
+
+cpdef set_plan_cache_size(size):
+ if not hasattr(_thread_local, '_current_plan_cache'):
+ raise RuntimeError('cache not found')
+ cdef PlanCache cache = _thread_local._current_plan_cache
+ cache.set_size(size)
+
+
+cpdef Py_ssize_t get_plan_cache_max_memsize(size):
+ if not hasattr(_thread_local, '_current_plan_cache'):
+ raise RuntimeError('cache not found')
+ cdef PlanCache cache = _thread_local._current_plan_cache
+ return cache.get_memsize()
+
+
+cpdef set_plan_cache_max_memsize(size):
+ if not hasattr(_thread_local, '_current_plan_cache'):
+ raise RuntimeError('cache not found')
+ cdef PlanCache cache = _thread_local._current_plan_cache
+ cache.set_memsize(size)
+
+
+cpdef clear_plan_cache():
+ if not hasattr(_thread_local, '_current_plan_cache'):
+ raise RuntimeError('cache not found')
+ cdef PlanCache cache = _thread_local._current_plan_cache
+ cache.clear()
+
+
+# enable cache by default
+plan_cache = PlanCache(size=10)
