Avoid threading.local() object overhead
@@ -539,8 +539,7 @@ cpdef MemoryPointer alloc(size):
~cupy.cuda.MemoryPointer: Pointer to the allocated buffer.
"""
- allocator = getattr(_thread_local, 'allocator', _current_allocator)
- return allocator(size)
+ return get_allocator()(size)
cpdef set_allocator(allocator=None):
@@ -559,13 +558,20 @@ cpdef set_allocator(allocator=None):
allocator = _malloc
_current_allocator = allocator
+
cpdef get_allocator():
"""Returns the current allocator for GPU memory.
Returns:
function: CuPy memory allocator.
"""
- allocator = getattr(_thread_local, 'allocator', _current_allocator)
+ try:
+ allocator = _thread_local.allocator
+ except AttributeError:
+ _thread_local.allocator = allocator = None
+ if allocator is None:
+ return _current_allocator
+ else:
return allocator
