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