daily pastebin goal
31%
SHARE
TWEET

Untitled

a guest Mar 19th, 2019 54 Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. def AllGatherv_unknown_shape(data, comm, rank, n_threads, mpi_dtype, axis=1):                                        
  2.     """                                                                                                              
  3.     Use MPI to compute shape of result array by asking every thread for their chunk's shape,                          
  4.     then perform the AllGatherv operation (concatenating along the given axis) and return the result.                
  5.     """                                                                                                              
  6.     if axis != 1:                                                                                                    
  7.         data = np.swapaxes(data, axis, 1)                                                                            
  8.                                                                                                                      
  9.     # Require row-major order for concatenation along axis 0                                                          
  10.     data = np.require(data, requirements='C')                                                                        
  11.                                                                                                                      
  12.     dim_1_len = data.shape[1]                                                                                        
  13.                                                                                                                      
  14.     # Gather size of each chunk                                                                                      
  15.     sizes = comm.allgather(data.size)                                                                                
  16.     offsets = np.zeros(len(sizes), dtype=np.int)                                                                      
  17.     offsets[1:] = np.cumsum(sizes)[:-1]                                                                              
  18.                                                                                                                      
  19.     comm.Barrier()                                                                                                    
  20.                                                                                                                      
  21.     # Create big array                                                                                                
  22.     tot_size = sum(sizes)                                                                                            
  23.     full_shape = tot_size/dim_1_len, dim_1_len                                                                        
  24.     all_data = np.zeros(shape=full_shape, dtype=data.dtype, order='C')                                                
  25.                                                                                                                      
  26.     # Perform AllGatherv                                                                                              
  27.     comm.Allgatherv(data, [all_data, sizes, offsets, mpi_dtype])                                                      
  28.                                                                                                                      
  29.     if axis != 1:                                                                                                    
  30.         all_data = np.swapaxes(all_data, axis, 1)                                                                    
  31.                                                                                                                      
  32.     return all_data
RAW Paste Data
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy. OK, I Understand
 
Top