Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
import hashlib
import sys
import numpy as np
from mpi4py import MPI
comm = MPI.COMM_WORLD
UINT32_MAX = 4294967295
INT32_MAX = 2147483647
REPEAT = 10
ROOT = 0
def calc_checksum(array):
h = hashlib.md5()
h.update(array)
return h.hexdigest()
def init_buf(num_elem):
if comm.rank == 0:
return np.random.randint(0, INT32_MAX + 1, num_elem, dtype=np.int32)
else:
return np.zeros(num_elem, dtype=np.int32)
def check_result(checksum, iter):
if comm.rank == ROOT:
peer_checksum = comm.recv(source=1 - ROOT)
if peer_checksum != checksum:
print("Error: checksum mismatch. ")
print(" rank {}: {}".format(ROOT, checksum))
print(" rank {}: {}".format(1 - ROOT, peer_checksum))
else:
print("Iteration {} OK.".format(iter))
else:
comm.send(checksum, dest=ROOT)
def main():
assert comm.size == 2
if len(sys.argv) > 1:
buf_size = int(sys.argv[1])
else:
buf_size = 32 * 1024 * 1024
num_elem = buf_size // 4 # int32
for i in range(REPEAT):
array = init_buf(num_elem)
comm.Bcast(array, root=ROOT)
checksum = calc_checksum(array)
check_result(checksum, i)
if __name__ == '__main__':
main()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment