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