Addressing changes

Author: astroC86

examples/plot_matrixmult.py

@@ -27,9 +27,9 @@
     exit(-1)
 
 # matrix dims
-M = 32
-K = 35
-N = 37
+M = 5
+K = 4
+N = 3
 
 A = np.random.rand(M * K).astype(dtype=np.float32).reshape(M, K)
 B = np.random.rand(K * N).astype(dtype=np.float32).reshape(K, N)
@@ -47,21 +47,16 @@
 #      - :math:`B_{p} \in \mathbb{R}^{K\times \text{my\_own\_cols}}`
 #    where
 blk_rows = int(math.ceil(M / P_prime))
-row_start = my_group * blk_rows
-row_end = min(M, row_start + blk_rows)
-my_own_rows = row_end - row_start
-
 blk_cols = int(math.ceil(N / P_prime))
-col_start = my_layer * blk_cols
-col_end = min(N, col_start + blk_cols)
-my_own_cols = col_end - col_start
-
 
-rs = (rank % P_prime) * blk_rows
+rs = my_group * blk_rows
 re = min(M, rs + blk_rows)
+my_own_rows = re - rs
 
-cs = (rank // P_prime) * blk_cols
+cs = my_layer * blk_cols
 ce = min(N, cs + blk_cols)
+my_own_cols = ce - cs
+
 A_p, B_p = A[rs:re, :].copy(), B[:, cs:ce].copy()
 
 Aop = MPIMatrixMult(A_p, N, dtype="float32")
@@ -81,19 +76,19 @@
 xadj_loc = (A.T.dot(y_loc.conj())).conj()
 
 
-expected_y_loc = y_loc[:, col_start:col_end].flatten().astype(np.float32)
-expected_xadj_loc = xadj_loc[:, col_start:col_end].flatten().astype(np.float32)
+expected_y_loc = y_loc[:, cs:ce].flatten().astype(np.float32)
+expected_xadj_loc = xadj_loc[:, cs:ce].flatten().astype(np.float32)
 
 xadj = Aop.H @ y
 if not np.allclose(y.local_array, expected_y_loc, rtol=1e-6):
     print(f"RANK {rank}: FORWARD VERIFICATION FAILED")
-    print(f'{rank} local: {y.local_array}, expected: {y_loc[:, col_start:col_end]}')
+    print(f'{rank} local: {y.local_array}, expected: {y_loc[:, cs:ce]}')
 else:
     print(f"RANK {rank}: FORWARD VERIFICATION PASSED")
 
 if not np.allclose(xadj.local_array, expected_xadj_loc, rtol=1e-6):
     print(f"RANK {rank}: ADJOINT VERIFICATION FAILED")
-    print(f'{rank} local: {xadj.local_array}, expected: {xadj_loc[:, col_start:col_end]}')
+    print(f'{rank} local: {xadj.local_array}, expected: {xadj_loc[:, cs:ce]}')
 else:
     print(f"RANK {rank}: ADJOINT VERIFICATION PASSED")
 

pylops_mpi/basicoperators/MatrixMult.py

@@ -16,6 +16,7 @@ def __init__(
             self,
             A: NDArray,
             N: int,
+            saveAt: bool = False,
             base_comm: MPI.Comm = MPI.COMM_WORLD,
             dtype: DTypeLike = "float64",
     ) -> None:
@@ -25,113 +26,91 @@ def __init__(
         # Determine grid dimensions (P_prime × C) such that P_prime * C ≥ size
         self._P_prime = int(math.ceil(math.sqrt(size)))
         self._C = int(math.ceil(size / self._P_prime))
-        if self._P_prime * self._C < size:
+        if self._P_prime * self._C != size:
             raise Exception("Number of Procs must be a square number")
 
         # Compute this process's group and layer indices
         self._group_id = rank % self._P_prime
         self._layer_id = rank // self._P_prime
 
         # Split communicators by layer (rows) and by group (columns)
-        self.base_comm   = base_comm
+        self.base_comm = base_comm
         self._layer_comm = base_comm.Split(color=self._layer_id, key=self._group_id)
         self._group_comm = base_comm.Split(color=self._group_id, key=self._layer_id)
         self.A = A.astype(np.dtype(dtype))
+        if saveAt: self.At = A.T.conj()
 
         self.M = self._layer_comm.allreduce(self.A.shape[0], op=MPI.SUM)
         self.K = A.shape[1]
         self.N = N
 
         # Determine how many columns each group holds
         block_cols = int(math.ceil(self.N / self._P_prime))
-        local_col_start = self._group_id * block_cols
-        local_col_end = min(self.N, local_col_start + block_cols)
-        local_ncols = local_col_end - local_col_start
+        blk_rows = int(math.ceil(self.M / self._P_prime))
 
-        # Sum up the total number of input columns across all processes
-        total_ncols = base_comm.allreduce(local_ncols, op=MPI.SUM)
-        self.dims = (self.K, total_ncols)
+        self._row_start = self._group_id * blk_rows
+        self._row_end = min(self.M, self._row_start + blk_rows)
+
+        self._col_start = self._layer_id * block_cols
+        self._col_end = min(self.N, self._col_start + block_cols)
 
-        # Recompute how many output columns each layer holds
-        layer_col_start  = self._layer_id * block_cols
-        layer_col_end    = min(self.N, layer_col_start + block_cols)
-        layer_ncols      = layer_col_end - layer_col_start
-        total_layer_cols = self.base_comm.allreduce(layer_ncols, op=MPI.SUM)
+        self._local_ncols = self._col_end - self._col_start
+        self._rank_col_lens = self.base_comm.allgather(self._local_ncols)
+        total_ncols = np.sum(self._rank_col_lens)
 
-        self.dimsd = (self.M, total_layer_cols)
+        self.dims = (self.K, total_ncols)
+        self.dimsd = (self.M, total_ncols)
         shape = (int(np.prod(self.dimsd)), int(np.prod(self.dims)))
         super().__init__(shape=shape, dtype=np.dtype(dtype), base_comm=base_comm)
-        
+
     def _matvec(self, x: DistributedArray) -> DistributedArray:
         ncp = get_module(x.engine)
         if x.partition != Partition.SCATTER:
             raise ValueError(f"x should have partition={Partition.SCATTER} Got {x.partition} instead...")
-        blk_cols    = int(math.ceil(self.N / self._P_prime))
-        col_start   = self._layer_id * blk_cols
-        col_end     = min(self.N, col_start + blk_cols)
-        my_own_cols = max(0, col_end - col_start)
-        x = x.local_array.reshape((self.dims[0], my_own_cols))
-        x = x.astype(self.dtype)
-
-        B_block = self._layer_comm.bcast(x if self._group_id == self._layer_id else None, root=self._layer_id)
-        C_local = ncp.vstack(
+
+        my_own_cols = self._rank_col_lens[self.rank]
+        x_arr = x.local_array.reshape((self.dims[0], my_own_cols))
+        x_arr = x_arr.astype(self.dtype)
+
+        X_local = self._layer_comm.bcast(x_arr if self._group_id == self._layer_id else None, root=self._layer_id)
+        Y_local = ncp.vstack(
             self._layer_comm.allgather(
-                ncp.matmul(self.A, B_block)
+                ncp.matmul(self.A, X_local)
             )
         )
 
-        layer_col_start = self._layer_id * blk_cols
-        layer_col_end   = min(self.N, layer_col_start + blk_cols)
-        layer_ncols     = max(0, layer_col_end - layer_col_start)
-        layer_col_lens  = self.base_comm.allgather(layer_ncols)
-        mask = [i // self._P_prime for i in range(self.size)]
-
-        y = DistributedArray(global_shape= (self.M * self.dimsd[1]),
-                             local_shapes=[(self.M * c) for c in layer_col_lens],
-                             mask=mask,
+        y = DistributedArray(global_shape=(self.M * self.dimsd[1]),
+                             local_shapes=[(self.M * c) for c in self._rank_col_lens],
+                             mask=x.mask,
                              partition=Partition.SCATTER,
                              dtype=self.dtype)
-        y[:] = C_local.flatten()
+        y[:] = Y_local.flatten()
         return y
 
     def _rmatvec(self, x: DistributedArray) -> DistributedArray:
         ncp = get_module(x.engine)
         if x.partition != Partition.SCATTER:
             raise ValueError(f"x should have partition={Partition.SCATTER}. Got {x.partition} instead.")
 
-        # Determine local column block for this layer
-        blk_cols        = int(math.ceil(self.N / self._P_prime))
-        layer_col_start = self._layer_id * blk_cols
-        layer_col_end   = min(self.N, layer_col_start + blk_cols)
-        layer_ncols     = layer_col_end - layer_col_start
-        layer_col_lens  = self.base_comm.allgather(layer_ncols)
-        x               = x.local_array.reshape((self.M, layer_ncols)).astype(self.dtype)
-
-        # Determine local row block for this process group
-        blk_rows  = int(math.ceil(self.M / self._P_prime))
-        row_start = self._group_id * blk_rows
-        row_end   = min(self.M, row_start + blk_rows)
-
-        B_tile = x[row_start:row_end, :].astype(self.dtype)
-        A_local = self.A.T.conj().astype(self.dtype)
-
-        m, b    = A_local.shape
-        pad     = (-m) % self._P_prime
-        r       = (m + pad) // self._P_prime
-        A_pad   = np.pad(A_local, ((0, pad), (0, 0)),  mode='constant', constant_values=self.dtype.type(0.0))
+        x_arr = x.local_array.reshape((self.M, self._local_ncols)).astype(self.dtype)
+        X_tile = x_arr[self._row_start:self._row_end, :]
+
+        A_local = self.At if hasattr(self, "At") else self.A.T.conj()
+        m, b = A_local.shape
+        pad = (-m) % self._P_prime
+        r = (m + pad) // self._P_prime
+        A_pad = np.pad(A_local, ((0, pad), (0, 0)), mode='constant', constant_values=self.dtype.type(0.0))
         A_batch = A_pad.reshape(self._P_prime, r, b)
 
-        # Perform local matmul and unpad
-        Y_batch = ncp.matmul(A_batch, B_tile).astype(self.dtype)
-        Y_pad   = Y_batch.reshape(r * self._P_prime, -1)
+        Y_batch = ncp.matmul(A_batch, X_tile)
+        Y_pad = Y_batch.reshape(r * self._P_prime, -1)
         y_local = Y_pad[:m, :]
         y_layer = self._layer_comm.allreduce(y_local, op=MPI.SUM)
 
-        mask = [i // self._P_prime for i in range(self.size)]
         y = DistributedArray(
             global_shape=(self.K * self.dimsd[1]),
-            local_shapes=[self.K * c for c in layer_col_lens],
-            mask=mask,
+            local_shapes=[self.K * c for c in self._rank_col_lens],
+            mask=x.mask,
             partition=Partition.SCATTER,
             dtype=self.dtype,
         )

tests/test_matrixmult.py

@@ -53,13 +53,10 @@ def test_SUMMAMatrixMult(M, K, N, dtype_str):
     my_own_rows_A = max(0, row_end_A - row_start_A)
 
     blk_cols_BC = int(math.ceil(N / P_prime))
-    col_start_B = my_group * blk_cols_BC
+    col_start_B = my_layer * blk_cols_BC
     col_end_B = min(N, col_start_B + blk_cols_BC)
     my_own_cols_B = max(0, col_end_B - col_start_B)
 
-    # Initialize local matrices
-    A_p = np.empty((my_own_rows_A, K), dtype=dtype)
-    B_p = np.empty((K, my_own_cols_B), dtype=dtype)
 
     A_glob_real = np.arange(M * K, dtype=base_float_dtype).reshape(M, K)
     A_glob_imag = np.arange(M * K, dtype=base_float_dtype).reshape(M, K) * 0.5
@@ -69,53 +66,19 @@ def test_SUMMAMatrixMult(M, K, N, dtype_str):
     B_glob_imag = np.arange(K * N, dtype=base_float_dtype).reshape(K, N) * 0.7
     B_glob = (B_glob_real + cmplx * B_glob_imag).astype(dtype)
 
-    if rank == 0:
-        # Distribute matrix blocks to all ranks
-        for dest_rank in range(size):
-            dest_my_group = dest_rank % P_prime
-
-            # Calculate destination rank's block dimensions
-            dest_row_start_A = dest_my_group * blk_rows_A
-            dest_row_end_A = min(M, dest_row_start_A + blk_rows_A)
-            dest_my_own_rows_A = max(0, dest_row_end_A - dest_row_start_A)
-
-            dest_col_start_B = dest_my_group * blk_cols_BC
-            dest_col_end_B = min(N, dest_col_start_B + blk_cols_BC)
-            dest_my_own_cols_B = max(0, dest_col_end_B - dest_col_start_B)
-
-            A_block_send = A_glob[dest_row_start_A:dest_row_end_A, :].copy()
-            B_block_send = B_glob[:, dest_col_start_B:dest_col_end_B].copy()
-
-            # Validate block shapes
-            assert A_block_send.shape == (dest_my_own_rows_A, K)
-            assert B_block_send.shape == (K, dest_my_own_cols_B)
-
-            if dest_rank == 0:
-                A_p, B_p = A_block_send, B_block_send
-            else:
-                if A_block_send.size > 0:
-                    comm.Send(A_block_send, dest=dest_rank, tag=100 + dest_rank)
-                if B_block_send.size > 0:
-                    comm.Send(B_block_send, dest=dest_rank, tag=200 + dest_rank)
-    else:
-        if A_p.size > 0:
-            comm.Recv(A_p, source=0, tag=100 + rank)
-        if B_p.size > 0:
-            comm.Recv(B_p, source=0, tag=200 + rank)
-
-    comm.Barrier()
+    A_p = A_glob[row_start_A:row_end_A,:]
+    B_p = B_glob[:,col_start_B:col_end_B]
 
     # Create SUMMAMatrixMult operator
     Aop = MPIMatrixMult(A_p, N, base_comm=comm, dtype=dtype_str)
 
     # Create DistributedArray for input x (representing B flattened)
     all_my_own_cols_B = comm.allgather(my_own_cols_B)
-    total_cols = sum(all_my_own_cols_B)
-    local_shapes_x = [(K * cl_b,) for cl_b in all_my_own_cols_B]
+    total_cols = np.sum(all_my_own_cols_B)
 
     x_dist = DistributedArray(
         global_shape=(K * total_cols),
-        local_shapes=local_shapes_x,
+        local_shapes=[K * cl_b for cl_b in all_my_own_cols_B],
         partition=Partition.SCATTER,
         base_comm=comm,
         dtype=dtype