[WIP] add tests for MemoryAsyncPool
@@ -875,3 +875,65 @@ class TestMallocAsync(unittest.TestCase):
s2 = cupy.cuda.Stream()
with s2:
del memptr
+
+
+@testing.gpu
+@pytest.mark.skipIf(cupy.cuda.runtime.is_hip,
+ reason='HIP does not support async allocator')
+@pytest.mark.skipIf(cupy.cuda.driver.get_build_version() < 11020,
+ reason='malloc_async is supported since CUDA 11.2')
+class TestMemoryAsyncPool(unittest.TestCase):
+
+ def setUp(self):
+ self.pool = memory.MemoryAsyncPool()
+
+ def test_zero_size_alloc(self):
+ with cupy.cuda.Device(0):
+ mem = self.pool.malloc(0).mem
+ assert isinstance(mem, memory.MemoryAsync)
+ assert not isinstance(mem, memory.PooledMemory)
+
+ def test_alloc(self):
+ with cupy.cuda.Device(0):
+ mem = self.pool.malloc(100).mem
+ assert isinstance(mem, memory.MemoryAsync)
+ assert not isinstance(mem, memory.PooledMemory)
+
+ def test_alloc_large_chunk(self):
+ self.pool.free_all_blocks()
+ with cupy.cuda.Device(0) as d:
+ _, mem_total = d.mem_info
+ mem = self.pool.malloc(int(0.9 * mem_total)).mem # 90% memory
+ del mem
+ mem = self.pool.malloc(int(0.2 * mem_total)).mem # 20% memory # noqa
+
+ def test_free_all_blocks(self):
+ with cupy.cuda.Device(0):
+ mem = self.pool.malloc(1).mem
+ del mem
+ self.pool.free_all_blocks()
+
+ def test_free_all_blocks_large_chunk(self):
+ # When memory is returned to the async mempool, it is not immediately
+ # visible to normal malloc routines until after a sync happens.
+ default_pool = cupy.get_default_memory_pool()
+ with cupy.cuda.Device(0) as d:
+ _, mem_total = d.mem_info
+ mem = self.pool.malloc(int(0.9 * mem_total)).mem # 90% memory
+ del mem
+ with pytest.raises(memory.OutOfMemoryError):
+ default_pool.malloc(int(0.2 * mem_total)) # 20% memory
+ self.pool.free_all_blocks() # synchronize
+ default_pool.malloc(int(0.2 * mem_total)) # this time it'd work
+
+ def test_interaction_with_CuPy_default_pool(self):
+ # Test saneness of cudaMallocAsync
+ default_pool = cupy.get_default_memory_pool()
+ with cupy.cuda.Device(0) as d:
+ _, mem_total = d.mem_info
+ mem = default_pool.malloc(int(0.9 * mem_total)).mem # 90% memory
+ del mem
+ with pytest.raises(memory.OutOfMemoryError):
+ self.pool.malloc(int(0.2 * mem_total)) # 20% memory
+ default_pool.free_all_blocks()
+ self.pool.malloc(int(0.2 * mem_total)) # this time it'd work
