@@ -70,30 +70,6 @@ def active_grid_comm(base_comm: MPI.Comm, N: int, M: int):
7070
7171 return new_comm , new_rank , new_row , new_col , True
7272
73- def block_distribute (array :NDArray , rank :int , comm : MPI .Comm , pad :bool = False ):
74- size = comm .Get_size ()
75- p_prime = math .isqrt (size )
76- if p_prime * p_prime != size :
77- raise Exception (f"Number of processes must be a square number, provided { size } instead..." )
78-
79- proc_i , proc_j = divmod (rank , p_prime )
80- orig_r , orig_c = array .shape
81-
82- new_r = math .ceil (orig_r / p_prime ) * p_prime
83- new_c = math .ceil (orig_c / p_prime ) * p_prime
84-
85- br , bc = new_r // p_prime , new_c // p_prime
86- i0 , j0 = proc_i * br , proc_j * bc
87- i1 , j1 = min (i0 + br , orig_r ), min (j0 + bc , orig_c )
88-
89- i_end = None if proc_i == p_prime - 1 else i1
90- j_end = None if proc_j == p_prime - 1 else j1
91- block = array [i0 :i_end , j0 :j_end ]
92-
93- pr = (new_r - orig_r ) if proc_i == p_prime - 1 else 0
94- pc = (new_c - orig_c ) if proc_j == p_prime - 1 else 0
95- if pad and (pr or pc ): block = np .pad (block , [(0 , pr ), (0 , pc )], mode = 'constant' )
96- return block , (new_r , new_c )
9773
9874def local_block_spit (global_shape : Tuple [int , int ],
9975 rank : int ,
0 commit comments