File: osu_bw.py

package info (click to toggle)
mpi4py 4.1.0-4
  • links: PTS, VCS
  • area: main
  • in suites: sid
  • size: 4,540 kB
  • sloc: python: 34,465; ansic: 16,475; makefile: 614; sh: 325; cpp: 193; f90: 178
file content (96 lines) | stat: -rw-r--r-- 2,557 bytes parent folder | download | duplicates (2)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
# http://mvapich.cse.ohio-state.edu/benchmarks/

from mpi4py import MPI


def osu_bw(
    BENCHMARH="MPI Bandwidth Test",
    skip=10,
    loop=100,
    window_size=64,
    skip_large=2,
    loop_large=20,
    window_size_large=64,
    large_message_size=8192,
    MAX_MSG_SIZE=1 << 22,
):
    comm = MPI.COMM_WORLD
    myid = comm.Get_rank()
    numprocs = comm.Get_size()

    if numprocs != 2:
        if myid == 0:
            errmsg = "This test requires exactly two processes"
        else:
            errmsg = None
        raise SystemExit(errmsg)

    s_buf = allocate(MAX_MSG_SIZE)
    r_buf = allocate(MAX_MSG_SIZE)

    if myid == 0:
        print(f"# {BENCHMARH}")
    if myid == 0:
        print(f"# {'Size [B]':<8s}{'Latency [us]':>20s}")

    message_sizes = [2**i for i in range(30)]
    for size in message_sizes:
        if size > MAX_MSG_SIZE:
            break
        if size > large_message_size:
            skip = skip_large
            loop = loop_large
            window_size = window_size_large

        iterations = list(range(loop + skip))
        window_sizes = list(range(window_size))
        requests = [MPI.REQUEST_NULL] * window_size
        #
        comm.Barrier()
        if myid == 0:
            s_msg = [s_buf, size, MPI.BYTE]
            r_msg = [r_buf, 4, MPI.BYTE]
            for i in iterations:
                if i == skip:
                    t_start = MPI.Wtime()
                for j in window_sizes:
                    requests[j] = comm.Isend(s_msg, 1, 100)
                MPI.Request.Waitall(requests)
                comm.Recv(r_msg, 1, 101)
            t_end = MPI.Wtime()
        elif myid == 1:
            s_msg = [s_buf, 4, MPI.BYTE]
            r_msg = [r_buf, size, MPI.BYTE]
            for i in iterations:
                if i == skip:
                    pass
                for j in window_sizes:
                    requests[j] = comm.Irecv(r_msg, 0, 100)
                MPI.Request.Waitall(requests)
                comm.Send(s_msg, 0, 101)
        #
        if myid == 0:
            MB = size / 1e6 * loop * window_size
            s = t_end - t_start
            bandwidth = MB / s
            print(f"{size:-10d}{bandwidth:20.2f}")


def allocate(n):
    try:
        import mmap

        return mmap.mmap(-1, n)
    except (ImportError, OSError):
        try:
            from numpy import zeros

            return zeros(n, "B")
        except ImportError:
            from array import array

            return array("B", [0]) * n


if __name__ == "__main__":
    osu_bw()