AOT-biomaps 2.9.138__py3-none-any.whl → 2.9.279__py3-none-any.whl

AOT_biomaps/AOT_Recon/AOT_SparseSMatrix/SparseSMatrix_SELL.py

@@ -0,0 +1,328 @@
+import pycuda.driver as drv
+import numpy as np
+from tqdm import trange
+import os
+import gc
+
+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
+        
+        # --- PATH RESOLUTION FIX ---
+        # The cubin file is located in the parent directory (AOT_Recon/)
+        # We use os.path.dirname(os.path.dirname(__file__)) to go up one directory level.
+        cubin_parent_dir = os.path.dirname(os.path.dirname(__file__))
+        self.module_path = os.path.join(cubin_parent_dir, module_path)
+        # --- END FIX ---
+        
+        self.slice_height = slice_height
+
+        # SELL arrays (device) & Size Tracking (CRITICAL FIX: Initialized attributes)
+        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
+        
+        # Attributes to store allocated size in bytes (bypassing the problematic .size attribute)
+        self.sell_values_gpu_size = 0  
+        self.sell_colinds_gpu_size = 0 
+        self.slice_ptr_gpu_size = 0    
+        self.slice_len_gpu_size = 0    
+
+        self.total_storage = 0
+
+        self.norm_factor_inv = None
+        self.norm_factor_inv_gpu = None
+        self.norm_factor_inv_gpu_size = 0
+
+        self.sparse_mod = None
+        self.load_module()
+
+    def load_module(self):
+        """Loads the pre-compiled CUDA module (.cubin file)."""
+        
+        # Check if the file exists at the calculated absolute path
+        if not os.path.exists(self.module_path):
+            # The path is now correctly calculated to the parent directory.
+            raise FileNotFoundError(f"CUDA module {os.path.basename(self.module_path)} not found at path: {self.module_path}")
+            
+        # Try to load the module
+        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. Check compatibility.") from e
+
+    def free(self):
+        try:
+            # Free device allocations
+            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)
+            
+            # Reset stored sizes
+            self.sell_values_gpu_size = 0
+            self.sell_colinds_gpu_size = 0
+            self.slice_ptr_gpu_size = 0
+            self.slice_len_gpu_size = 0
+            self.norm_factor_inv_gpu_size = 0
+            
+            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.
+        NOTE: This is the logic of allocate_sell_c_sigma_direct from the working class.
+        """
+        if self.sparse_mod is None:
+            raise RuntimeError("CUDA module not loaded. Check compilation.")
+            
+        # NOTE: Les noms de kernel (count_nnz_rows_kernel, fill_kernel__SELL) sont utilisés
+        # car ils sont présents dans la classe fonctionnelle.
+        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)
+        
+        # Allocation 1: Dense block GPU memory
+        dense_gpu_size = dense_host.nbytes
+        dense_gpu  = drv.mem_alloc(dense_gpu_size)
+
+        # 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_size = br * np.dtype(np.int32).itemsize
+        row_nnz_gpu_block = drv.mem_alloc(row_nnz_gpu_block_size)
+
+        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_size = total_storage * np.dtype(np.float32).itemsize
+        self.sell_colinds_gpu_size = total_storage * np.dtype(np.uint32).itemsize
+        
+        self.sell_values_gpu = drv.mem_alloc(self.sell_values_gpu_size)
+        self.sell_colinds_gpu = drv.mem_alloc(self.sell_colinds_gpu_size)
+        
+        # allocate slice metadata on device
+        self.slice_ptr = slice_ptr
+        self.slice_len = slice_len
+        
+        self.slice_ptr_gpu_size = self.slice_ptr.nbytes
+        self.slice_len_gpu_size = self.slice_len.nbytes
+        
+        self.slice_ptr_gpu = drv.mem_alloc(self.slice_ptr_gpu_size)
+        self.slice_len_gpu = drv.mem_alloc(self.slice_len_gpu_size)
+        
+        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 allocate new dense_gpu
+        dense_host = np.empty((br, num_cols), dtype=np.float32)
+        
+        dense_gpu_2_size = dense_host.nbytes
+        dense_gpu  = drv.mem_alloc(dense_gpu_2_size)
+        
+        # we also need row_nnz on device per-block; supply global row_nnz on host but the kernel recomputes threshold
+        row_nnz_host_gpu_size = br * np.dtype(np.int32).itemsize
+        row_nnz_host_gpu = drv.mem_alloc(row_nnz_host_gpu_size)
+
+        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()
+
+        # 4) compute norm_factor_inv via GPU accumulate (col sums)
+        self.compute_norm_factor()
+    
+    def apply_apodization_gpu(self, window_vector_gpu):
+        """
+        Applique le fenêtrage directement sur self.sell_values_gpu 
+        en utilisant les indices de colonnes (pixels) pour référencer 
+        la fenêtre. Opération : A_values[i] *= W_vec[A_colinds[i]].
+        """
+        if self.sparse_mod is None:
+            raise RuntimeError("Le module CUDA n'a pas été chargé.")
+            
+        try:
+            apodize_kernel = self.sparse_mod.get_function("apply_apodisation_kernel__SELL")
+        except drv.LogicError as e:
+            raise RuntimeError(
+                f"Le kernel CUDA 'multiply_sell_by_window_kernel' est manquant dans le .cubin. "
+                f"Veuillez le compiler et l'ajouter. Erreur : {e}"
+            )
+
+        # Le total_storage inclut les éléments non-nuls et le padding SELL.
+        threads = 256
+        blocks = (self.total_storage + threads - 1) // threads
+        
+        # Lancement du kernel. Il travaille sur total_storage éléments.
+        apodize_kernel(
+            self.sell_values_gpu, 
+            self.sell_colinds_gpu, 
+            window_vector_gpu, 
+            np.int64(self.total_storage), 
+            block=(threads, 1, 1), 
+            grid=(blocks, 1, 1)
+        )
+        drv.Context.synchronize()
+        print("✅ Multiplication par le fenêtrage effectuée in-place sur GPU (SELL-C-σ).")
+    # --- Ajout de la fonction de normalisation (qui fonctionne) ---
+    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_size = ZX * np.dtype(np.float32).itemsize
+        col_sum_gpu = drv.mem_alloc(col_sum_gpu_size)
+        drv.memset_d32(col_sum_gpu, 0, ZX)
+
+        # FIX: Kernel name is "accumulate_columns_atomic"
+        acc_kernel = self.sparse_mod.get_function("accumulate_columns_atomic")
+        
+        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_size = self.norm_factor_inv.nbytes
+        self.norm_factor_inv_gpu = drv.mem_alloc(self.norm_factor_inv_gpu_size)
+        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 (using .nbytes which works for NumPy 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 (using the stored size attributes instead of the problematic .size)
+        total_bytes += self.sell_values_gpu_size
+        total_bytes += self.sell_colinds_gpu_size
+        total_bytes += self.slice_ptr_gpu_size
+        total_bytes += self.slice_len_gpu_size
+        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 *