AOT-biomaps 2.9.186__py3-none-any.whl → 2.9.261__py3-none-any.whl

AOT_biomaps/AOT_Recon/AOT_SparseSMatrix/SparseSMatrix_CSR.py

@@ -0,0 +1,252 @@
+# sparse_matrix_gpu.py
+import pycuda.driver as drv
+import numpy as np
+from pycuda.compiler import SourceModule
+from tqdm import trange
+import gc
+import os
+
+class SparseSMatrix_CSR:
+    """Construction d'une matrice CSR à partir d'un objet `manip`.
+    Usage:
+        S = SparseMatrixGPU(manip)
+        S.allocate()
+    Après allocate(), on a: row_ptr (host np.int64 array), row_ptr_gpu (device ptr),
+    h_col_ind, h_values, col_ind_gpu, values_gpu, norm_factor_inv.
+    """
+
+    def __init__(self, manip, block_rows=64, relative_threshold=0.3, device=0):
+        drv.init()
+        self.device = drv.Device(device)
+        self.ctx = self.device.make_context()
+        self.manip = manip
+        self.N = len(manip.AcousticFields)
+        self.T = manip.AcousticFields[0].field.shape[0]
+        self.Z = manip.AcousticFields[0].field.shape[1]
+        self.X = manip.AcousticFields[0].field.shape[2]
+        self.block_rows = block_rows
+        self.relative_threshold = relative_threshold
+        self.h_dense = None
+        self.row_ptr = None
+        self.row_ptr_gpu = None
+        self.h_col_ind = None
+        self.h_values = None
+        self.total_nnz = 0
+        self.norm_factor_inv = None
+        self.sparse_mod = None
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        self.free()
+
+    def load_precompiled_module(self, so_path="AOT_biomaps_kernels.cubin"):
+        try:
+            # If a PTX or cubin is provided via path
+            self.sparse_mod = drv.module_from_file(so_path)
+            print(f"✅ Module CUDA chargé depuis {so_path}")
+        except Exception:
+            # Fallback: try to compile from bundled source (if available)
+            src_path = os.path.join(os.path.dirname(__file__), 'AOT_biomaps_kernels.cu')
+            if os.path.exists(src_path):
+                print("Compilation JIT du kernel CUDA depuis source...")
+                with open(src_path, 'r') as f:
+                    src = f.read()
+                self.sparse_mod = SourceModule(src, no_extern_c=True)
+                print("✅ Module compilé JIT")
+            else:
+                raise
+
+    def estimate_nnz_cpu(self):
+        """Estimation rapide (non-exacte) — utile si tu veux une estimation faible.
+        Recommandé : utiliser la passe GPU exacte (count_nnz_per_row_kernel) à la place.
+        """
+        total = 0
+        for n in range(self.N):
+            field = self.manip.AcousticFields[n].field
+            for t in range(self.T):
+                row = field[t].flatten()
+                row_max = np.max(np.abs(row))
+                thr = row_max * self.relative_threshold
+                total += np.count_nonzero(np.abs(row) > thr)
+        return int(total)
+
+    def allocate(self, kernel_module_path=None):
+        try:
+            # --- 1. Construction bloc par bloc (sans garder tout le dense si possible) ---
+            num_rows = self.N * self.T
+            num_cols = self.Z * self.X
+            bytes_float = np.dtype(np.float32).itemsize
+
+            # Charge module
+            if kernel_module_path:
+                self.load_precompiled_module(kernel_module_path)
+            else:
+                self.load_precompiled_module('AOT_biomaps_kernels.cubin')
+
+            count_nnz_kernel = self.sparse_mod.get_function('count_nnz_per_row_kernel')
+            fill_csr_kernel = self.sparse_mod.get_function('fill_kernel__CSR')
+
+            # allocate host row_ptr
+            self.row_ptr = np.zeros(num_rows + 1, dtype=np.int64)
+
+            # GPU temp buffers
+            dense_block_host = np.empty((self.block_rows, num_cols), dtype=np.float32)
+            dense_block_gpu = drv.mem_alloc(self.block_rows * num_cols * bytes_float)
+            row_nnz_gpu = drv.mem_alloc(self.block_rows * np.dtype(np.int32).itemsize)
+
+            block_size = 128
+
+            # --- Count NNZ per row using GPU kernel to be consistent with filling logic ---
+            for b in trange(0, num_rows, self.block_rows, desc='Comptage NNZ'):
+                current_rows = min(self.block_rows, num_rows - b)
+                # Fill dense_block_host from manip
+                for r in range(current_rows):
+                    global_row = b + r
+                    n_idx = global_row // self.T
+                    t_idx = global_row % self.T
+                    dense_block_host[r, :] = self.manip.AcousticFields[n_idx].field[t_idx].flatten()
+                drv.memcpy_htod(dense_block_gpu, dense_block_host)
+
+                grid = ((current_rows + block_size - 1) // block_size, 1, 1)
+                count_nnz_kernel(dense_block_gpu, row_nnz_gpu,
+                                 np.int32(current_rows), np.int32(num_cols),
+                                 np.float32(self.relative_threshold),
+                                 block=(block_size, 1, 1), grid=grid)
+
+                row_nnz_host = np.empty(current_rows, dtype=np.int32)
+                drv.memcpy_dtoh(row_nnz_host, row_nnz_gpu)
+                self.row_ptr[b + 1:b + current_rows + 1] = self.row_ptr[b] + np.cumsum(row_nnz_host, dtype=np.int64)
+
+            # total nnz
+            self.total_nnz = int(self.row_ptr[-1])
+            print(f"NNZ total : {self.total_nnz}")
+
+            # allocate final arrays
+            self.h_col_ind = np.zeros(self.total_nnz, dtype=np.uint32)
+            self.h_values = np.zeros(self.total_nnz, dtype=np.float32)
+
+            # copy row_ptr to device once
+            self.row_ptr_gpu = drv.mem_alloc(self.row_ptr.nbytes)
+            drv.memcpy_htod(self.row_ptr_gpu, self.row_ptr)
+
+            # allocate device arrays for final csr
+            self.col_ind_gpu = drv.mem_alloc(self.h_col_ind.nbytes)
+            self.values_gpu = drv.mem_alloc(self.h_values.nbytes)
+
+            # --- Fill CSR per-block ---
+            for b in trange(0, num_rows, self.block_rows, desc='Remplissage CSR'):
+                current_rows = min(self.block_rows, num_rows - b)
+                for r in range(current_rows):
+                    global_row = b + r
+                    n_idx = global_row // self.T
+                    t_idx = global_row % self.T
+                    dense_block_host[r, :] = self.manip.AcousticFields[n_idx].field[t_idx].flatten()
+                drv.memcpy_htod(dense_block_gpu, dense_block_host)
+
+                grid = ((current_rows + block_size - 1) // block_size, 1, 1)
+                fill_csr_kernel(dense_block_gpu,
+                                self.row_ptr_gpu,
+                                self.col_ind_gpu,
+                                self.values_gpu,
+                                np.int32(b),
+                                np.int32(current_rows),
+                                np.int32(num_cols),
+                                np.float32(self.relative_threshold),
+                                np.int64(self.total_nnz),
+                                block=(block_size, 1, 1), grid=grid)
+                drv.Context.synchronize()
+
+            # copy back
+            drv.memcpy_dtoh(self.h_col_ind, self.col_ind_gpu)
+            drv.memcpy_dtoh(self.h_values, self.values_gpu)
+            print('CSR généré ✔')
+
+            # compute normalization factor from CSR (sum per column)
+            self.compute_norm_factor_from_csr()
+
+            # free temporaries
+            dense_block_gpu.free(); row_nnz_gpu.free()
+
+        except Exception as e:
+            print(f"❌ Erreur détaillée : {e}")
+            self.free()
+            raise
+
+    def compute_norm_factor_from_csr(self):
+        ZX = self.Z * self.X
+
+        # 1) Allouer un vecteur de somme colonne sur le GPU
+        col_sum_gpu = drv.mem_alloc(ZX * np.dtype(np.float32).itemsize)
+        drv.memset_d32(col_sum_gpu, 0, ZX)
+
+        # 2) Récupérer le kernel
+        acc_kernel = self.sparse_mod.get_function("accumulate_columns_atomic__CSR")
+
+        # 3) Lancer le kernel
+        threads = 256
+        blocks = (self.total_nnz + threads - 1) // threads
+
+        acc_kernel(
+            self.values_gpu,
+            self.col_ind_gpu,
+            np.int64(self.total_nnz),
+            col_sum_gpu,
+            block=(threads,1,1),
+            grid=(blocks,1,1)
+        )
+        drv.Context.synchronize()
+
+        # 4) Récupérer le résultat
+        norm = np.empty(ZX, dtype=np.float32)
+        drv.memcpy_dtoh(norm, col_sum_gpu)
+        col_sum_gpu.free()
+
+        norm = np.maximum(norm.astype(np.float64), 1e-6)
+        self.norm_factor_inv = (1.0 / norm).astype(np.float32)
+
+        self.norm_factor_inv_gpu = drv.mem_alloc(self.norm_factor_inv.nbytes)
+        drv.memcpy_htod(self.norm_factor_inv_gpu, self.norm_factor_inv)
+
+    def getMatrixSize(self):
+        if self.row_ptr is None:
+            return {"error": "La matrice sparse n'est pas encore allouée."}
+        total = (self.row_ptr.nbytes if self.row_ptr is not None else 0) + \
+                (self.h_col_ind.nbytes if self.h_col_ind is not None else 0) + \
+                (self.h_values.nbytes if self.h_values is not None else 0)
+        return total / (1024**3)
+
+    def free(self):
+        try:
+            if hasattr(self, 'col_ind_gpu') and self.col_ind_gpu:
+                self.col_ind_gpu.free()
+            if hasattr(self, 'values_gpu') and self.values_gpu:
+                self.values_gpu.free()
+            if hasattr(self, 'row_ptr_gpu') and self.row_ptr_gpu:
+                self.row_ptr_gpu.free()
+            if hasattr(self, 'norm_factor_inv_gpu') and self.norm_factor_inv_gpu:
+                self.norm_factor_inv_gpu.free()
+            if hasattr(self, 'ctx') and self.ctx:
+                try:
+                    self.ctx.pop()
+                except Exception:
+                    pass
+                self.ctx = None
+            print('✅ Mémoire GPU libérée.')
+        except Exception as e:
+            print(f"❌ Erreur lors de la libération de la mémoire GPU : {e}")
+        
+    def compute_density(self):
+        """
+        Retourne la densité réelle de la CSR = NNZ / (num_rows * num_cols)
+        Nécessite que self.h_values et self.row_ptr existent (host).
+        """
+        if self.row_ptr is None or self.h_values is None:
+            raise RuntimeError("row_ptr et h_values requis pour calculer la densité")
+        num_rows = int(self.N * self.T)
+        num_cols = int(self.Z * self.X)
+        total_nnz = int(self.row_ptr[-1])
+        density = total_nnz / (num_rows * num_cols)
+        return density
+   

AOT_biomaps/AOT_Recon/AOT_SparseSMatrix/SparseSMatrix_SELL.py

@@ -0,0 +1,322 @@
+import pycuda.driver as drv
+import numpy as np
+from tqdm import trange
+import os
+import gc
+import subprocess
+import sys
+
+# ==============================================================================
+# ATTENTION: MANUAL CONFIGURATION REQUIRED
+# ==============================================================================
+# The 'sm_XX' value must correspond to the Compute Capability of your NVIDIA GPU.
+# Examples: Kepler (sm_35), Maxwell (sm_50), Pascal (sm_61), Turing (sm_75), Ampere (sm_86).
+# THIS LINE MUST BE MODIFIED BY THE USER to target their specific architecture.
+GPU_COMPUTE_CAPABILITY = "sm_86"
+# ==============================================================================
+
+class SparseSMatrix_SELL:
+    def __init__(self, manip, block_rows=64, relative_threshold=0.3, device=0,
+                 module_path="AOT_biomaps_kernels.cubin", slice_height=32):
+        drv.init()
+        self.device = drv.Device(device)
+        self.ctx = self.device.make_context()
+        self.manip = manip
+        self.N = len(manip.AcousticFields)
+        self.T = manip.AcousticFields[0].field.shape[0]
+        self.Z = manip.AcousticFields[0].field.shape[1]
+        self.X = manip.AcousticFields[0].field.shape[2]
+        self.block_rows = block_rows
+        self.relative_threshold = relative_threshold
+        # The module_path is relative to this file's directory
+        self.module_path = os.path.join(os.path.dirname(__file__), module_path)
+        self.slice_height = slice_height
+
+        # SELL arrays (device)
+        self.sell_values_gpu = None
+        self.sell_colinds_gpu = None
+        self.slice_ptr = None
+        self.slice_len = None
+        self.slice_ptr_gpu = None
+        self.slice_len_gpu = None
+        self.total_storage = 0
+
+        self.norm_factor_inv = None
+        self.norm_factor_inv_gpu = None
+
+        self.sparse_mod = None
+        self.load_module()
+    
+    def _compile_cubin(self, source_file="AOT_biomaps_kernels.cu"):
+        """
+        Tries to compile the .cu file into .cubin using nvcc.
+        """
+        print("="*60)
+        print("🛠️ CUDA COMPILATION REQUIRED")
+        print(f"Attempting to compile {source_file} to {os.path.basename(self.module_path)}...")
+
+        # The source file is assumed to be in the same directory as this Python file.
+        source_path = os.path.join(os.path.dirname(__file__), source_file)
+        cubin_path = self.module_path
+        
+        if not os.path.exists(source_path):
+            print(f"❌ CRITICAL ERROR: CUDA source file {source_file} not found at {source_path}.")
+            raise FileNotFoundError(f"Could not find source file {source_file} for compilation. AOT_biomaps installation might be incomplete.")
+
+        # Construction of the nvcc command
+        command = [
+            'nvcc',
+            '-cubin',
+            f'-arch={GPU_COMPUTE_CAPABILITY}', # USES THE VARIABLE DEFINED ABOVE
+            source_path,
+            '-o',
+            cubin_path
+        ]
+        
+        print(f"Executing command: {' '.join(command)}")
+
+        try:
+            # Executes the command and waits for completion
+            result = subprocess.run(
+                command,
+                check=True,
+                capture_output=True,
+                text=True
+            )
+            print(f"🎉 Compilation successful! File created: {os.path.basename(cubin_path)}")
+            # print("Output nvcc:\n", result.stdout) # Uncomment to see detailed output
+            print("="*60)
+            return True
+
+        except subprocess.CalledProcessError as e:
+            print("❌ NVCC COMPILATION ERROR:")
+            print(f"Check GPU architecture: {GPU_COMPUTE_CAPABILITY}")
+            print(f"Standard error:\n{e.stderr}")
+            print("="*60)
+            return False
+        
+        except FileNotFoundError:
+            print("❌ ERROR: The 'nvcc' command was not found.")
+            print("Ensure that the CUDA Toolkit is installed and 'nvcc' is in your PATH (or your Conda environment).")
+            print("="*60)
+            return False
+
+    def load_module(self):
+        """Tries to load the CUDA module. If the file is missing, it attempts to compile it."""
+        
+        if not os.path.exists(self.module_path):
+            print(f"CUDA module {os.path.basename(self.module_path)} missing. Attempting compilation...")
+            
+            if not self._compile_cubin():
+                # If compilation fails, re-raise the original error.
+                raise FileNotFoundError(f"{os.path.basename(self.module_path)} not found and compilation failed. Check nvcc and GPU architecture ({GPU_COMPUTE_CAPABILITY}).")
+            
+        # Try to load after compilation (or if the file existed)
+        try:
+            self.sparse_mod = drv.module_from_file(self.module_path)
+            print(f"Loaded CUDA module {os.path.basename(self.module_path)}")
+        except Exception as e:
+            print(f"❌ Error loading CUDA module {os.path.basename(self.module_path)}: {e}")
+            raise RuntimeError(f"File {os.path.basename(self.module_path)} was found, but PyCUDA could not load it.") from e
+
+    def free(self):
+        try:
+            attrs = ["sell_values_gpu","sell_colinds_gpu","slice_ptr_gpu","slice_len_gpu","norm_factor_inv_gpu"]
+            for a in attrs:
+                if hasattr(self, a) and getattr(self, a) is not None:
+                    getattr(self, a).free()
+                    setattr(self, a, None)
+            if hasattr(self, 'ctx') and self.ctx:
+                try: self.ctx.pop()
+                except Exception: pass
+                self.ctx = None
+        except Exception as e:
+            print("Error freeing GPU memory:", e)
+
+    def allocate(self):
+        """
+        Build SELL-C-σ directly from manip AcousticFields in streaming blocks.
+        """
+        # Ensures the module is loaded before attempting to retrieve functions
+        if self.sparse_mod is None:
+             raise RuntimeError("CUDA module not loaded. Check compilation.")
+             
+        count_kernel = self.sparse_mod.get_function("count_nnz_rows_kernel")
+        fill_kernel  = self.sparse_mod.get_function("fill_kernel__SELL")
+
+        num_rows = int(self.N * self.T)
+        num_cols = int(self.Z * self.X)
+        C = int(self.slice_height)
+
+        # host temporary block
+        br = int(self.block_rows)
+        bytes_per_elem = np.dtype(np.float32).itemsize
+        dense_host = np.empty((br, num_cols), dtype=np.float32)
+        dense_gpu  = drv.mem_alloc(dense_host.nbytes)
+
+        # 1) count nnz per row (on host via small blocks with GPU kernel)
+        row_nnz = np.zeros(num_rows, dtype=np.int32)
+        row_nnz_gpu_block = drv.mem_alloc(br * np.dtype(np.int32).itemsize)
+
+        block = 256
+        for b in trange(0, num_rows, br, desc="Count NNZ per row"):
+            R = min(br, num_rows - b)
+            # fill dense_host
+            for i in range(R):
+                rg = b + i
+                n_idx = rg // self.T
+                t_idx = rg % self.T
+                dense_host[i, :] = self.manip.AcousticFields[n_idx].field[t_idx].flatten()
+            # copy only R rows
+            drv.memcpy_htod(dense_gpu, dense_host)
+            grid = ((R + block - 1) // block, 1, 1)
+            count_kernel(dense_gpu, row_nnz_gpu_block, np.int32(R), np.int32(num_cols), np.float32(self.relative_threshold),
+                         block=(block,1,1), grid=grid)
+            tmp = np.empty(R, dtype=np.int32)
+            drv.memcpy_dtoh(tmp, row_nnz_gpu_block)
+            row_nnz[b:b+R] = tmp
+
+        row_nnz_gpu_block.free()
+        dense_gpu.free()
+
+        # 2) compute per-slice maxlen and slice_ptr
+        num_slices = (num_rows + C - 1) // C
+        slice_len = np.zeros(num_slices, dtype=np.int32)
+        for s in range(num_slices):
+            r0 = s * C
+            r1 = min(num_rows, r0 + C)
+            slice_len[s] = int(np.max(row_nnz[r0:r1])) if (r1>r0) else 0
+        # slice_ptr (int64)
+        slice_ptr = np.zeros(num_slices + 1, dtype=np.int64)
+        for s in range(num_slices):
+            slice_ptr[s+1] = slice_ptr[s] + (slice_len[s] * C)
+        total_storage = int(slice_ptr[-1])
+        self.total_storage = total_storage
+        print(f"SELL: num_rows={num_rows}, num_slices={num_slices}, total_storage(padded)={total_storage}")
+
+        # allocate device SELL arrays (values float32, colinds uint32)
+        self.sell_values_gpu = drv.mem_alloc(total_storage * np.dtype(np.float32).itemsize)
+        self.sell_colinds_gpu = drv.mem_alloc(total_storage * np.dtype(np.uint32).itemsize)
+        # allocate slice metadata on device
+        self.slice_ptr = slice_ptr
+        self.slice_len = slice_len
+        self.slice_ptr_gpu = drv.mem_alloc(self.slice_ptr.nbytes)
+        self.slice_len_gpu = drv.mem_alloc(self.slice_len.nbytes)
+        drv.memcpy_htod(self.slice_ptr_gpu, self.slice_ptr)
+        drv.memcpy_htod(self.slice_len_gpu, self.slice_len)
+
+        # 3) fill SELL arrays by streaming blocks again (use GPU fill kernel)
+        # reuse dense_host and dense_gpu
+        dense_host = np.empty((br, num_cols), dtype=np.float32)
+        dense_gpu  = drv.mem_alloc(dense_host.nbytes)
+        # we also need row_nnz on device per-block; supply global row_nnz on host but the kernel recomputes threshold
+        row_nnz_host_gpu = drv.mem_alloc(br * np.dtype(np.int32).itemsize)
+
+        for b in trange(0, num_rows, br, desc="Fill SELL"):
+            R = min(br, num_rows - b)
+            for i in range(R):
+                rg = b + i
+                n_idx = rg // self.T
+                t_idx = rg % self.T
+                dense_host[i, :] = self.manip.AcousticFields[n_idx].field[t_idx].flatten()
+            drv.memcpy_htod(dense_gpu, dense_host)
+            # We pass a dummy row_nnz pointer (not used in this kernel; left for API)
+            # Kernel expects rows_in_block, rows_global_offset to know where to write.
+            grid = ((R + block - 1) // block, 1, 1)
+            fill_kernel(dense_gpu,
+                        np.intp(0), # placeholder for row_nnz pointer (not used)
+                        self.slice_ptr_gpu,
+                        self.slice_len_gpu,
+                        self.sell_colinds_gpu,
+                        self.sell_values_gpu,
+                        np.int32(R),
+                        np.int32(num_cols),
+                        np.int32(b), # rows_global_offset
+                        np.int32(C),
+                        np.float32(self.relative_threshold),
+                        block=(block,1,1), grid=grid)
+        dense_gpu.free()
+        row_nnz_host_gpu.free()
+
+        # At this point sell_values_gpu and sell_colinds_gpu are filled.
+
+        # 4) compute norm_factor_inv via GPU accumulate (col sums)
+        self.compute_norm_factor()
+
+    def compute_norm_factor(self):
+        """
+        Accumulate column sums on GPU using accumulate_columns_atomic, then compute inverse.
+        """
+        if self.total_storage == 0:
+            raise RuntimeError("sell not built")
+        ZX = int(self.Z * self.X)
+
+        # allocate col sum on device
+        col_sum_gpu = drv.mem_alloc(ZX * np.dtype(np.float32).itemsize)
+        drv.memset_d32(col_sum_gpu, 0, ZX)
+
+        acc_kernel = self.sparse_mod.get_function("accumulate_columns_atomic__SELL")
+        threads = 256
+        blocks = (self.total_storage + threads - 1) // threads
+        acc_kernel(self.sell_values_gpu, self.sell_colinds_gpu, np.int64(self.total_storage), col_sum_gpu,
+                   block=(threads,1,1), grid=(blocks,1,1))
+        drv.Context.synchronize()
+
+        # copy back
+        norm_host = np.empty(ZX, dtype=np.float32)
+        drv.memcpy_dtoh(norm_host, col_sum_gpu)
+        col_sum_gpu.free()
+
+        norm = np.maximum(norm_host.astype(np.float64), 1e-6)
+        self.norm_factor_inv = (1.0 / norm).astype(np.float32)
+        if self.norm_factor_inv_gpu is not None:
+            self.norm_factor_inv_gpu.free()
+        self.norm_factor_inv_gpu = drv.mem_alloc(self.norm_factor_inv.nbytes)
+        drv.memcpy_htod(self.norm_factor_inv_gpu, self.norm_factor_inv)
+
+    def compute_density(self):
+        """
+        Returns only the density of the SELL-C-σ matrix.
+        """
+        if not hasattr(self, 'slice_ptr') or self.slice_ptr is None:
+            raise RuntimeError("The SELL-C-σ matrix is not allocated.")
+
+        num_rows = self.N * self.T
+        num_cols = self.Z * self.X
+        total_elements = num_rows * num_cols
+
+        # Conservative estimate of non-zeros (excluding padding)
+        nnz_ell_estimated = int(0.9 * self.total_storage)
+
+        return nnz_ell_estimated / total_elements # Returns only the density
+
+    def getMatrixSize(self):
+        """
+        Returns the total size of the SELL-C-σ matrix in Gigabytes (GB).
+        """
+        if self.sell_values_gpu is None:
+            return {"error": "The SELL-C-σ matrix is not yet allocated."}
+
+        total_bytes = 0
+
+        # Host-side arrays
+        if hasattr(self, 'slice_ptr') and self.slice_ptr is not None:
+            total_bytes += self.slice_ptr.nbytes
+        if hasattr(self, 'slice_len') and self.slice_len is not None:
+            total_bytes += self.slice_len.nbytes
+        if hasattr(self, 'norm_factor_inv') and self.norm_factor_inv is not None:
+            total_bytes += self.norm_factor_inv.nbytes
+
+        # GPU-side arrays
+        if hasattr(self, 'sell_values_gpu') and self.sell_values_gpu:
+            total_bytes += self.sell_values_gpu.size
+        if hasattr(self, 'sell_colinds_gpu') and self.sell_colinds_gpu:
+            total_bytes += self.sell_colinds_gpu.size
+        if hasattr(self, 'slice_ptr_gpu') and self.slice_ptr_gpu:
+            total_bytes += self.slice_ptr_gpu.size
+        if hasattr(self, 'slice_len_gpu') and self.slice_len_gpu:
+            total_bytes += self.slice_len_gpu.size
+        if hasattr(self, 'norm_factor_inv_gpu') and self.norm_factor_inv_gpu:
+            total_bytes += self.norm_factor_inv_gpu.size
+
+        return total_bytes / (1024 ** 3) # Returns only the size in GB

AOT_biomaps/AOT_Recon/AOT_SparseSMatrix/__init__.py

@@ -0,0 +1,2 @@
+from .SparseSMatrix_CSR import *
+from .SparseSMatrix_SELL import *