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
|
# 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(f"# {'Size [B]':<8s}{'Latency [us]':>20s}")
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:
latency = total_lat / (pairs * 2)
print(f"{size:-10d}{latency: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_multi_lat()
|
