1# http://mvapich.cse.ohio-state.edu/benchmarks/ 2 3from mpi4py import MPI 4 5def osu_bw( 6 BENCHMARH = "MPI Bandwidth Test", 7 skip = 10, 8 loop = 100, 9 window_size = 64, 10 skip_large = 2, 11 loop_large = 20, 12 window_size_large = 64, 13 large_message_size = 8192, 14 MAX_MSG_SIZE = 1<<22, 15 ): 16 17 comm = MPI.COMM_WORLD 18 myid = comm.Get_rank() 19 numprocs = comm.Get_size() 20 21 if numprocs != 2: 22 if myid == 0: 23 errmsg = "This test requires exactly two processes" 24 else: 25 errmsg = None 26 raise SystemExit(errmsg) 27 28 s_buf = allocate(MAX_MSG_SIZE) 29 r_buf = allocate(MAX_MSG_SIZE) 30 31 if myid == 0: 32 print ('# %s' % (BENCHMARH,)) 33 if myid == 0: 34 print ('# %-8s%20s' % ("Size [B]", "Bandwidth [MB/s]")) 35 36 message_sizes = [2**i for i in range(30)] 37 for size in message_sizes: 38 if size > MAX_MSG_SIZE: 39 break 40 if size > large_message_size: 41 skip = skip_large 42 loop = loop_large 43 window_size = window_size_large 44 45 iterations = list(range(loop+skip)) 46 window_sizes = list(range(window_size)) 47 requests = [MPI.REQUEST_NULL] * window_size 48 # 49 comm.Barrier() 50 if myid == 0: 51 s_msg = [s_buf, size, MPI.BYTE] 52 r_msg = [r_buf, 4, MPI.BYTE] 53 for i in iterations: 54 if i == skip: 55 t_start = MPI.Wtime() 56 for j in window_sizes: 57 requests[j] = comm.Isend(s_msg, 1, 100) 58 MPI.Request.Waitall(requests) 59 comm.Recv(r_msg, 1, 101) 60 t_end = MPI.Wtime() 61 elif myid == 1: 62 s_msg = [s_buf, 4, MPI.BYTE] 63 r_msg = [r_buf, size, MPI.BYTE] 64 for i in iterations: 65 for j in window_sizes: 66 requests[j] = comm.Irecv(r_msg, 0, 100) 67 MPI.Request.Waitall(requests) 68 comm.Send(s_msg, 0, 101) 69 # 70 if myid == 0: 71 MB = size / 1e6 * loop * window_size 72 s = t_end - t_start 73 print ('%-10d%20.2f' % (size, MB/s)) 74 75 76def allocate(n): 77 try: 78 import mmap 79 return mmap.mmap(-1, n) 80 except (ImportError, EnvironmentError): 81 try: 82 from numpy import zeros 83 return zeros(n, 'B') 84 except ImportError: 85 from array import array 86 return array('B', [0]) * n 87 88 89if __name__ == '__main__': 90 osu_bw() 91