add more info for compression profiling

qtensor/compression/Compressor.py

@@ -101,6 +101,9 @@ def get_profile_stats(self):
         compress_ratios = np.mean([x.size_in/x.size_out for x in compress])
         compress_size = sum([x.size_out for x in compress])
         return compress_time, decompress_time, compress_size, compress_ratios
+
+    def compress_size(self, ptr):
+        return self.compressor.compress_size(ptr)
 # --
 
 class NumpyCompressor(Compressor):

qtensor/contraction_backends/performance_measurement_decorator.py

@@ -1,7 +1,7 @@
 import numpy as np
 from dataclasses import dataclass
 from qtensor.contraction_backends import ContractionBackend, NumpyBackend
-from qtensor.contraction_backends.compression import CompressionBackend
+from qtensor.contraction_backends.compression import CompressionBackend, CompressedTensor
 from pyrofiler import timing
 from qtensor.tools.lazy_import import torch, pandas
 import string
@@ -40,6 +40,14 @@ def max_mem(self):
     def nvsmi_max_mem(self):
         mems = [m['nvmem'] for m in self.mem_history]
         return max(mems)
+    @property
+    def cupy_buffer_max_mem(self):
+        mems = [m['cupy_bufsize'] for m in self.mem_history]
+        return max(mems)
+    @property
+    def object_max_mem(self):
+        mems = [m['objmem'] for m in self.mem_history]
+        return max(mems)
 
     def check_store(self):
         import cupy
@@ -54,6 +62,8 @@ def check_store(self):
                 continue
             else:
                 size = self.tensor_size(tensor)
+                if isinstance(tensor, CompressedTensor):
+                    print("Tensor", tensor, "size", size)
                 total_mem += size
         for key in deleted_keys:
             self.object_keys.remove(key)
@@ -70,22 +80,25 @@ def check_store(self):
             mem=gpu_mem,
             cupy_bufsize=mempool.total_bytes(),
             nvmem = self._get_nvsmi_mem(),
+            cupybuf=mempool.total_bytes(),
+            objmem=total_mem,
             tensors_sizes=[len(tensor.indices) for tensor in self.object_store.values()]
         ))
         # --
+        print('MH', self.mem_history[-1])
         if cupy_mem>1024**2:
             self._print("CuPy memory usage", cupy_mem/1024/1024, "MB. Total MB:", mempool.total_bytes()/1024**2)
 
     def tensor_size(self, tensor)->int:
         from qtensor.compression import Tensor, CompressedTensor
         if tensor.data is None:
             return 0
-        if isinstance(tensor, Tensor):
-            return tensor.data.nbytes
-        elif isinstance(tensor, CompressedTensor):
-            chunks = tensor.data
+        if isinstance(tensor, CompressedTensor):
+            chunks = tensor._data
             sizes = [tensor.compressor.compress_size(x) for x in chunks]
             return sum(sizes)
+        elif isinstance(tensor, Tensor):
+            return tensor.data.nbytes
         else:
             raise ValueError("Unknown tensor type")