1# http://mvapich.cse.ohio-state.edu/benchmarks/
2
3from libmpi import ffi, lib
4
5def osu_latency(
6    BENCHMARH = "MPI Latency Test",
7    skip = 1000,
8    loop = 10000,
9    skip_large = 10,
10    loop_large = 100,
11    large_message_size = 8192,
12    MAX_MSG_SIZE = 1<<22,
13    ):
14
15    myid = ffi.new('int*')
16    numprocs = ffi.new('int*')
17    lib.MPI_Comm_rank(lib.MPI_COMM_WORLD, myid)
18    lib.MPI_Comm_size(lib.MPI_COMM_WORLD, numprocs)
19    myid = myid[0]
20    numprocs = numprocs[0]
21
22    if numprocs != 2:
23        if myid == 0:
24            errmsg = "This test requires exactly two processes"
25        else:
26            errmsg = None
27        raise SystemExit(errmsg)
28
29    sbuf = ffi.new('unsigned char[]', MAX_MSG_SIZE)
30    rbuf = ffi.new('unsigned char[]', MAX_MSG_SIZE)
31    dtype = lib.MPI_BYTE
32    tag = 1
33    comm = lib.MPI_COMM_WORLD
34    status = lib.MPI_STATUS_IGNORE
35
36    if myid == 0:
37        print ('# %s' % (BENCHMARH,))
38    if myid == 0:
39        print ('# %-8s%20s' % ("Size [B]", "Latency [us]"))
40
41    message_sizes = [0] + [2**i for i in range(30)]
42    for size in message_sizes:
43        if size > MAX_MSG_SIZE:
44            break
45        if size > large_message_size:
46            skip = skip_large
47            loop = loop_large
48        iterations = list(range(loop+skip))
49        #
50        lib.MPI_Barrier(comm)
51        if myid == 0:
52            for i in iterations:
53                if i == skip:
54                    t_start = lib.MPI_Wtime()
55                lib.MPI_Send(sbuf, size, dtype, 1, tag, comm)
56                lib.MPI_Recv(rbuf, size, dtype, 1, tag, comm, status)
57            t_end = lib.MPI_Wtime()
58        elif myid == 1:
59            for i in iterations:
60                lib.MPI_Recv(rbuf, size, dtype, 0, tag, comm, status)
61                lib.MPI_Send(sbuf, size, dtype, 0, tag, comm)
62        #
63        if myid == 0:
64            latency = (t_end - t_start) * 1e6 / (2 * loop)
65            print ('%-10d%20.2f' % (size, latency))
66
67def main():
68    lib.MPI_Init(ffi.NULL, ffi.NULL)
69    osu_latency()
70    lib.MPI_Finalize()
71
72if __name__ == '__main__':
73    main()
74