added handling for padding

Author: astroC86

examples/matrixmul.py

@@ -24,28 +24,30 @@
 B_data = np.arange(int(B_shape[0] * B_shape[1])).reshape(B_shape)
 
 i, j = divmod(rank, p_prime)
-
 A_local, (N_new, K_new) = MPIMatrixMult.block_distribute(A_data, i, j,comm)
 B_local, (K_new, M_new) = MPIMatrixMult.block_distribute(B_data, i, j,comm)
 
-B_dist = pylops_mpi.DistributedArray(global_shape=(K_new*M_new),
+B_dist = pylops_mpi.DistributedArray(global_shape=(K * M),
                                      local_shapes=comm.allgather(B_local.shape[0] * B_local.shape[1]),
                                      base_comm=comm,
                                      partition=pylops_mpi.Partition.SCATTER)
 B_dist.local_array[:] = B_local.flatten()
 
-Aop = MPIMatrixMult(A_local, M_new, base_comm=comm)
+Aop = MPIMatrixMult(A_local, M, base_comm=comm)
 C_dist = Aop @ B_dist
 Z_dist = Aop.H @ C_dist
 
-C = MPIMatrixMult.block_gather(C_dist, (N_new,M_new), (N,M), comm)
-Z = MPIMatrixMult.block_gather(Z_dist, (K_new,M_new), (K,M), comm)
+C = MPIMatrixMult.block_gather(C_dist, (N,M), (N,M), comm)
+Z = MPIMatrixMult.block_gather(Z_dist, (K,M), (K,M), comm)
 if rank == 0 :
-    print("expected:\n", np.allclose((A_data.T.dot((A_data @ B_data).conj())).conj(), Z.astype(np.int32)))
-    # print("expected:\n", (A_data.T.dot((A_data @ B_data).conj())).conj())
-    # print("calculated:\n",Z.astype(np.int32))
-    # print("calculated:\n", (A_data.T.dot((A_data @ B_data).conj())).conj() == Z.astype(np.int32))
-
-    # print("expected:\n",np.allclose(A_data @ B_data, C))
-    # print("expected:\n", A_data @ B_data)
-    # print("calculated:\n",C)
+    C_correct = np.allclose(A_data @ B_data, C)
+    print("C expected: ", C_correct)
+    if not C_correct:
+        print("expected:\n", A_data @ B_data)
+        print("calculated:\n",C)
+
+    Z_correct = np.allclose((A_data.T.dot((A_data @ B_data).conj())).conj(), Z.astype(np.int32))
+    print("Z expected: ", Z_correct)
+    if not Z_correct:
+        print("expected:\n", (A_data.T.dot((A_data @ B_data).conj())).conj())
+        print("calculated:\n", Z.astype(np.int32))

pylops_mpi/basicoperators/MatrixMult.py

@@ -136,13 +136,32 @@ def __init__(
         self._col_comm = base_comm.Split(color=self._col_id, key=self._row_id)
 
         self.A = A.astype(np.dtype(dtype))
-        if saveAt:
-            self.At = A.T.conj()
 
         self.N = self._col_comm.allreduce(A.shape[0])
         self.K = self._row_comm.allreduce(A.shape[1])
         self.M = M
 
+        self._N_padded = math.ceil(self.N / self._P_prime) * self._P_prime
+        self._K_padded = math.ceil(self.K / self._P_prime) * self._P_prime
+        self._M_padded = math.ceil(self.M / self._P_prime) * self._P_prime
+
+        bn = self._N_padded // self._P_prime
+        bk = self._K_padded // self._P_prime
+        bm = self._M_padded // self._P_prime
+
+        pr = (bn - A.shape[0]) if self._row_id == self._P_prime - 1 else 0
+        pc = (bk - A.shape[1]) if self._col_id == self._P_prime - 1 else 0
+
+        if pr < 0 or pc < 0:
+            raise Exception(f"Improper distribution of A expected local shape "
+                            f"( ≤ {bn}, ≤ {bk}) but got ({A.shape[0]},{A.shape[1]})")
+
+        if pr > 0 or pc > 0:
+            self.A = np.pad(self.A, [(0, pr), (0, pc)], mode='constant')
+
+        if saveAt:
+            self.At = self.A.T.conj()
+
         self.dims  = (self.K, self.M)
         self.dimsd = (self.N, self.M)
         shape = (int(np.prod(self.dimsd)), int(np.prod(self.dims)))
@@ -218,65 +237,185 @@ def block_distribute(array, proc_i, proc_j, comm):
         i0, j0 = proc_i * br, proc_j * bc
         i1, j1 = min(i0 + br, orig_r), min(j0 + bc, orig_c)
 
-        block = array[i0:i1, j0:j1]
+        i_end = None if proc_i == p_prime - 1 else i1
+        j_end = None if proc_j == p_prime - 1 else j1
+        block = array[i0:i_end, j0:j_end]
+
         pr = (new_r - orig_r) if proc_i == p_prime - 1 else 0
         pc = (new_c - orig_c) if proc_j == p_prime - 1 else 0
-        if pr or pc:
-            block = np.pad(block, [(0, pr), (0, pc)], mode='constant')
-
+        #comment the padding to get the block as unpadded
+        # if pr or pc: block = np.pad(block, [(0, pr), (0, pc)], mode='constant')
         return block, (new_r, new_c)
 
     @staticmethod
     def block_gather(x, new_shape, orig_shape, comm):
         ncp = get_module(x.engine)
         p_prime = math.isqrt(comm.Get_size())
         all_blks = comm.allgather(x.local_array)
-        nr, nc   = new_shape
+
+        nr, nc = new_shape
         orr, orc = orig_shape
-        br, bc = nr // p_prime, nc // p_prime
-        C = ncp.array(all_blks).reshape(p_prime, p_prime, br, bc).transpose(0, 2, 1, 3).reshape(nr, nc)
+
+        # Calculate base block sizes
+        br_base = nr // p_prime
+        bc_base = nc // p_prime
+
+        # Calculate remainder rows/cols that need to be distributed
+        r_remainder = nr % p_prime
+        c_remainder = nc % p_prime
+
+        # Create the output matrix
+        C = ncp.zeros((nr, nc), dtype=all_blks[0].dtype)
+
+        # Place each block in the correct position
+        for rank in range(p_prime * p_prime):
+            # Convert linear rank to 2D grid position
+            proc_row = rank // p_prime
+            proc_col = rank % p_prime
+
+            # Calculate this process's block dimensions
+            block_rows = br_base + (1 if proc_row < r_remainder else 0)
+            block_cols = bc_base + (1 if proc_col < c_remainder else 0)
+
+            # Calculate starting position in global matrix
+            start_row = proc_row * br_base + min(proc_row, r_remainder)
+            start_col = proc_col * bc_base + min(proc_col, c_remainder)
+
+            # Place the block
+            block = all_blks[rank]
+            if block.ndim == 1:
+                block = block.reshape(block_rows, block_cols)
+            C[start_row:start_row + block_rows, start_col:start_col + block_cols] = block
         return C[:orr, :orc]
 
     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...")
-        local_shape = ((self.N  *  self.M) // self.size)
-        y = DistributedArray(global_shape=(self.N  * self.M),
+
+        # Calculate local shapes for block distribution
+        bn = self._N_padded // self._P_prime  # block size in N dimension
+        bm = self._M_padded // self._P_prime  # block size in M dimension
+
+        # Calculate actual local shape for this process (considering original dimensions)
+        local_n = bn
+        local_m = bm
+
+        # Adjust for edge/corner processes that might have smaller blocks
+        if self._row_id == self._P_prime - 1:
+            local_n = self.N - (self._P_prime - 1) * bn
+        if self._col_id == self._P_prime - 1:
+            local_m = self.M - (self._P_prime - 1) * bm
+
+        local_shape = local_n * local_m
+
+        # Create local_shapes array for all processes
+        local_shapes = []
+        for rank in range(self.size):
+            row_id, col_id = divmod(rank, self._P_prime)
+            proc_n = bn if row_id != self._P_prime - 1 else self.N - (self._P_prime - 1) * bn
+            proc_m = bm if col_id != self._P_prime - 1 else self.M - (self._P_prime - 1) * bm
+            local_shapes.append(proc_n * proc_m)
+
+        y = DistributedArray(global_shape=(self.N * self.M),
                              mask=x.mask,
-                             local_shapes=[local_shape] * self.size,
+                             local_shapes=local_shapes,
                              partition=Partition.SCATTER,
-                             dtype=self.dtype)
+                             dtype=self.dtype,
+                             base_comm=self.base_comm
+                             )
+
+        # Calculate expected padded dimensions for x
+        bk = self._K_padded // self._P_prime  # block size in K dimension
+
+        # The input x corresponds to blocks from matrix B (K x M)
+        # This process should receive a block of size (local_k x local_m)
+        local_k = bk
+        if self._row_id == self._P_prime - 1:
+            local_k = self.K - (self._P_prime - 1) * bk
+
+        # Reshape x.local_array to its 2D block form
+        x_block = x.local_array.reshape((local_k, local_m))
+
+        # Pad the block to the full padded size if necessary
+        pad_k = bk - local_k
+        pad_m = bm - local_m
+
+        if pad_k > 0 or pad_m > 0:
+            x_block = np.pad(x_block, [(0, pad_k), (0, pad_m)], mode='constant')
+
+        Y_local = np.zeros((self.A.shape[0], bm))
 
-        x = x.local_array.reshape((self.A.shape[1], -1))
-        Y_local = np.zeros((self.A.shape[0], x.shape[1]))
         for k in range(self._P_prime):
             Atemp = self.A.copy() if self._col_id == k else np.empty_like(self.A)
-            Xtemp = x.copy() if self._row_id == k else np.empty_like(x)
+            Xtemp = x_block.copy() if self._row_id == k else np.empty_like(x_block)
             self._row_comm.Bcast(Atemp, root=k)
             self._col_comm.Bcast(Xtemp, root=k)
             Y_local += ncp.dot(Atemp, Xtemp)
-        y[:] = Y_local.flatten()
-        return y
 
+        Y_local_unpadded = Y_local[:local_n, :local_m]
+        y[:] = Y_local_unpadded.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.")
 
-        local_shape = ((self.K  * self.M ) // self.size)
+        # Calculate local shapes for block distribution
+        bk = self._K_padded // self._P_prime  # block size in K dimension
+        bm = self._M_padded // self._P_prime  # block size in M dimension
+
+        # Calculate actual local shape for this process (considering original dimensions)
+        local_k = bk
+        local_m = bm
+
+        # Adjust for edge/corner processes that might have smaller blocks
+        if self._row_id == self._P_prime - 1:
+            local_k = self.K - (self._P_prime - 1) * bk
+        if self._col_id == self._P_prime - 1:
+            local_m = self.M - (self._P_prime - 1) * bm
+
+        local_shape = local_k * local_m
+
+        # Create local_shapes array for all processes
+        local_shapes = []
+        for rank in range(self.size):
+            row_id, col_id = divmod(rank, self._P_prime)
+            proc_k = bk if row_id != self._P_prime - 1 else self.K - (self._P_prime - 1) * bk
+            proc_m = bm if col_id != self._P_prime - 1 else self.M - (self._P_prime - 1) * bm
+            local_shapes.append(proc_k * proc_m)
+
         y = DistributedArray(
             global_shape=(self.K * self.M),
             mask=x.mask,
-            local_shapes=[local_shape] * self.size,
+            local_shapes=local_shapes,
             partition=Partition.SCATTER,
             dtype=self.dtype,
             base_comm=self.base_comm
         )
-        x_reshaped = x.local_array.reshape((self.A.shape[0], -1))
+
+        # Calculate expected padded dimensions for x
+        bn = self._N_padded // self._P_prime  # block size in N dimension
+
+        # The input x corresponds to blocks from the result (N x M)
+        # This process should receive a block of size (local_n x local_m)
+        local_n = bn
+        if self._row_id == self._P_prime - 1:
+            local_n = self.N - (self._P_prime - 1) * bn
+
+        # Reshape x.local_array to its 2D block form
+        x_block = x.local_array.reshape((local_n, local_m))
+
+        # Pad the block to the full padded size if necessary
+        pad_n = bn - local_n
+        pad_m = bm - local_m
+
+        if pad_n > 0 or pad_m > 0:
+            x_block = np.pad(x_block, [(0, pad_n), (0, pad_m)], mode='constant')
+
         A_local = self.At if hasattr(self, "At") else self.A.T.conj()
-        Y_local = np.zeros((self.A.shape[1], x_reshaped.shape[1]))
+        Y_local = np.zeros((self.A.shape[1], bm))
 
         for k in range(self._P_prime):
             requests = []
@@ -289,10 +428,12 @@ def _rmatvec(self, x: DistributedArray) -> DistributedArray:
                 for moving_col in range(self._P_prime):
                     destA = fixed_col * self._P_prime + moving_col
                     tagA = (100 + k) * 1000 + destA
-                    requests.append(self.base_comm.Isend(A_local, dest=destA,tag=tagA))
-            Xtemp = x_reshaped.copy() if self._row_id == k else np.empty_like(x_reshaped)
+                    requests.append(self.base_comm.Isend(A_local, dest=destA, tag=tagA))
+            Xtemp = x_block.copy() if self._row_id == k else np.empty_like(x_block)
             requests.append(self._col_comm.Ibcast(Xtemp, root=k))
             MPI.Request.Waitall(requests)
             Y_local += ncp.dot(ATtemp, Xtemp)
-        y[:] = Y_local.flatten()
+
+        Y_local_unpadded = Y_local[:local_k, :local_m]
+        y[:] = Y_local_unpadded.flatten()
         return y