-
Notifications
You must be signed in to change notification settings - Fork 1
/
mpi_test.py
55 lines (45 loc) · 1.14 KB
/
mpi_test.py
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
from contextlib import contextmanager
from itertools import product
import numpy as np
import time
from mpi4py import MPI
@contextmanager
def benchmark(name):
start = time.time()
yield
end = time.time()
print('{} took {:.2f} ms'.format(name, (end - start) * 1000.0))
comm = MPI.COMM_WORLD
size = comm.Get_size()
rank = comm.Get_rank()
dim = int(1e1 + 1)
dv = np.arange(dim)
dvini = rank * int(dim / size)
if rank == size - 1:
dvend = dim
else:
dvend = dvini + int(dim / size)
print('cpu%d: %d %d' % (rank, dvini, dvend))
mat = np.zeros((dim, dim))
vec = np.zeros(dim**2)
#with benchmark("cpu%d: array looping" % rank):
# r = 0
# for m in dv[dvini:dvend]:
# c = 0
# for n in dv:
# mat[c, r] = m * n
# c += 1
# r += 1
# #print(mat)
dvini = rank * int(dim**2 / size)
if rank == size - 1:
dvend = dim**2
else:
dvend = dvini + int(dim**2 / size)
print('cpu%d: %d %d' % (rank, dvini, dvend))
with benchmark("cpu%d: vector looping" % rank):
c = 0
for m in list(product(dv, dv))[dvini:dvend]:
vec[c] = m[0] * m[1]
c += 1
#print(np.reshape(vec, (dim, dim)))