diff --git a/qtensor/compression/Compressor.py b/qtensor/compression/Compressor.py index 71302263..ced72ad0 100644 --- a/qtensor/compression/Compressor.py +++ b/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): diff --git a/qtensor/contraction_backends/performance_measurement_decorator.py b/qtensor/contraction_backends/performance_measurement_decorator.py index cedb3e45..39efffa7 100644 --- a/qtensor/contraction_backends/performance_measurement_decorator.py +++ b/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,9 +80,12 @@ 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) @@ -80,12 +93,12 @@ 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")