def test_allocator_nested_context_manager(self):
new_pool = memory.MemoryPool()
with memory.using_allocator(new_pool.malloc):
new_pool2 = memory.MemoryPool()
assert memory.get_allocator() == new_pool.malloc
with memory.using_allocator(new_pool2.malloc):
assert memory.get_allocator() == new_pool2.malloc
assert memory.get_allocator() == new_pool.malloc
assert memory.get_allocator() == self.pool.malloc
def thread_body(self):
new_pool = memory.MemoryPool()
with memory.using_allocator(new_pool.malloc):
assert memory.get_allocator() == new_pool.malloc
threading.Barrier(2)
arr = cupy.zeros(128, dtype=cupy.int64)
threading.Barrier(2)
self.assertEqual(arr.data.mem.size, new_pool.used_bytes())
threading.Barrier(2)
assert memory.get_allocator() == self.pool.malloc
def thread_body(self):
cupy.cuda.Device().use()
new_pool = memory.MemoryPool()
with cupy.cuda.using_allocator(new_pool.malloc):
assert memory.get_allocator() == new_pool.malloc
threading.Barrier(2)
arr = cupy.zeros(128, dtype=cupy.int64)
threading.Barrier(2)
assert arr.data.mem.size == new_pool.used_bytes()
threading.Barrier(2)
assert memory.get_allocator() == self.pool.malloc
def test_allocator_thread_local(self):
barrier = threading.Barrier(2)
def thread_body(self):
cupy.cuda.Device().use()
new_pool = memory.MemoryPool()
with cupy.cuda.using_allocator(new_pool.malloc):
assert memory.get_allocator() == new_pool.malloc
barrier.wait()
arr = cupy.zeros(128, dtype=cupy.int64)
barrier.wait()
assert arr.data.mem.size == new_pool.used_bytes()
barrier.wait()
assert memory.get_allocator() == self.pool.malloc
self._success = True
with cupy.cuda.Device():
self._success = False
t = threading.Thread(target=thread_body, args=(self,), daemon=True)
t.start()
barrier.wait()
assert memory.get_allocator() == self.pool.malloc
arr = cupy.ones(256, dtype=cupy.int64)
barrier.wait()
assert arr.data.mem.size == self.pool.used_bytes()
barrier.wait()
t.join()
assert self._success
def test_allocator_thread_local(self):
def thread_body(self):
new_pool = memory.MemoryPool()
with memory.using_allocator(new_pool.malloc):
assert memory.get_allocator() == new_pool.malloc
threading.Barrier(2)
arr = cupy.zeros(128, dtype=cupy.int64)
threading.Barrier(2)
self.assertEqual(arr.data.mem.size, new_pool.used_bytes())
threading.Barrier(2)
assert memory.get_allocator() == self.pool.malloc
with cupy.cuda.Device(0):
t = threading.Thread(target=thread_body, args=(self,))
t.daemon = True
t.start()
threading.Barrier(2)
assert memory.get_allocator() == self.pool.malloc
arr = cupy.ones(256, dtype=cupy.int64)
threading.Barrier(2)
self.assertEqual(arr.data.mem.size, self.pool.used_bytes())
threading.Barrier(2)
t.join()
def test_get_allocator(self):
assert memory.get_allocator() == self.pool.malloc
def test_allocator_context_manager(self):
new_pool = memory.MemoryPool()
with cupy.cuda.using_allocator(new_pool.malloc):
assert memory.get_allocator() == new_pool.malloc
assert memory.get_allocator() == self.pool.malloc
