File: osu_multi_lat.py

package info (click to toggle)
mpi4py 4.0.3-4
  • links: PTS, VCS
  • area: main
  • in suites: sid, trixie
  • size: 4,196 kB
  • sloc: python: 32,170; ansic: 13,449; makefile: 602; sh: 314; f90: 178; cpp: 148
file content (81 lines) | stat: -rw-r--r-- 2,212 bytes parent folder | download
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
# http://mvapich.cse.ohio-state.edu/benchmarks/

from mpi4py import MPI

def osu_multi_lat(
    BENCHMARH = "MPI Multi Latency Test",
    skip_small = 100,
    loop_small = 10000,
    skip_large = 10,
    loop_large = 1000,
    large_message_size = 8192,
    MAX_MSG_SIZE = 1<<22,
    ):

    comm = MPI.COMM_WORLD
    myid = comm.Get_rank()
    nprocs = comm.Get_size()
    pairs = nprocs/2

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

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

    message_sizes = [0] + [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
        else:
            skip = skip_small
            loop = loop_small
        iterations = list(range(loop+skip))
        s_msg = [s_buf, size, MPI.BYTE]
        r_msg = [r_buf, size, MPI.BYTE]
        #
        comm.Barrier()
        if myid < pairs:
            partner = myid + pairs
            for i in iterations:
                if i == skip:
                    t_start = MPI.Wtime()
                comm.Send(s_msg, partner, 1)
                comm.Recv(r_msg, partner, 1)
            t_end = MPI.Wtime()
        else:
            partner = myid - pairs
            for i in iterations:
                if i == skip:
                    t_start = MPI.Wtime()
                comm.Recv(r_msg, partner, 1)
                comm.Send(s_msg, partner, 1)
            t_end = MPI.Wtime()
        #
        latency = (t_end - t_start) * 1e6 / (2 * loop)
        total_lat = comm.reduce(latency, root=0, op=MPI.SUM)
        if myid == 0:
            avg_lat = total_lat/(pairs * 2)
            print ('%-10d%20.2f' % (size, avg_lat))


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_multi_lat()