fix test_leak for cusz. only complex64 works

danlkv · Dec 15, 2023 · 547c5f5 · 547c5f5
1 parent 4ce7d0c
commit 547c5f5
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 22 deletions.
diff --git a/qtensor/compression/Compressor.py b/qtensor/compression/Compressor.py
@@ -352,8 +352,8 @@ def free_decompressed(self):
             cupy.cuda.runtime.free(x)
             # del x
             cupy.get_default_memory_pool().free_all_blocks()
-            cupy.get_default_pinned_memory_pool().free_all_blocks()
-        torch.cuda.empty_cache()
+            #cupy.get_default_pinned_memory_pool().free_all_blocks()
+        #torch.cuda.empty_cache()
         self.decompressed_own = []
 
     def free_compressed(self, ptr):
@@ -366,8 +366,8 @@ def free_compressed(self, ptr):
         decompressed_int = p_decompressed_int.contents
         cupy.cuda.runtime.free(decompressed_int.value)
         cupy.get_default_memory_pool().free_all_blocks()
-        cupy.get_default_pinned_memory_pool().free_all_blocks()
-        torch.cuda.empty_cache()
+        #cupy.get_default_pinned_memory_pool().free_all_blocks()
+        #torch.cuda.empty_cache()
 
     def compress(self, data):
         import cupy

diff --git a/qtensor/compression/tests/test_memory_leak.py b/qtensor/compression/tests/test_memory_leak.py
@@ -19,34 +19,29 @@ def _get_nvsmi_mem(handle):
     return mem
 
 
-def free_compressed(ptr):
-    cmp_bytes, *_ = ptr
-    p_decompressed_ptr = ctypes.addressof(cmp_bytes[0])
-    # cast to int64 pointer
-    # (effectively converting pointer to pointer to addr to pointer to int64)
-    p_decompressed_int = ctypes.cast(
-        p_decompressed_ptr, ctypes.POINTER(ctypes.c_uint64)
-    )
-    decompressed_int = p_decompressed_int.contents
-    cupy.cuda.runtime.free(decompressed_int.value)
-
-
 def test_leak():
-    N = 1024 * 1024 * 32  # 32MB
-    a = cupy.zeros(N, dtype=float)
+    dtype = cupy.complex64
+    dtype_size = dtype(0).nbytes
+    MB_elems = int(1024 ** 2 / dtype_size)
+    MB_target = 128
+    N = MB_target * MB_elems
+    print(f"== Testing memory leak with {N} elements and {MB_target} MB array ==")
+
+    a = cupy.zeros(N, dtype=dtype)
     a[::1024] = 0.01
+    a[::8] = cupy.random.rand(N // 8) * 0.01
     for i in range(1000):
         a[32 * i] = 0.005 * (i % 5 + 1)
     _nvsmi_handle = _init_nvsmi()
     print(f"Original, [0]={a[0]}, [1024]={a[1024]}")
 
     c = CUSZXCompressor(r2r_error=1e-2, r2r_threshold=1e-2)
-    for i in range(200):
+    for i in range(100):
         out = c.compress(a)
         print(i, "Compression ratio", 4 * N / c.compress_size(out))
         b = c.decompress(out)
-        #a[:] = b
+        a[:] = b
         print(i, f"Decompressed, [0]={b[0]}, [1024]={b[1024]}")
         c.free_decompressed()
-        free_compressed(out)
-        print(f"Memory usage: {_get_nvsmi_mem(_nvsmi_handle) / 1024 ** 3} GB")
+        c.free_compressed(out)
+        print(f"== [{i}] Memory usage: {_get_nvsmi_mem(_nvsmi_handle) / 1024 ** 3} GB ==")