# Jacobi computation, using persitent requests from mpi4py import MPI try: import numpy except ImportError: raise SystemExit n = 5 * MPI.COMM_WORLD.Get_size() # compute number of processes and myrank p = MPI.COMM_WORLD.Get_size() myrank = MPI.COMM_WORLD.Get_rank() # compute size of local block m = n/p if myrank < (n - p * m): m = m + 1 #compute neighbors if myrank == 0: left = MPI.PROC_NULL else: left = myrank - 1 if myrank == p - 1: right = MPI.PROC_NULL else: right = myrank + 1 # allocate local arrays A = numpy.empty((n+2, m+2), dtype=float, order='fortran') B = numpy.empty((n, m), dtype=float, order='fortran') A.fill(1) A[0, :] = A[-1, :] = 0 A[:, 0] = A[:, -1] = 0 # create persintent requests tag = 0 sreq1 = MPI.COMM_WORLD.Send_init((B[:, 0], MPI.DOUBLE), left, tag) sreq2 = MPI.COMM_WORLD.Send_init((B[:, -1], MPI.DOUBLE), right, tag) rreq1 = MPI.COMM_WORLD.Recv_init((A[:, 0], MPI.DOUBLE), left, tag) rreq2 = MPI.COMM_WORLD.Recv_init((A[:, -1], MPI.DOUBLE), right, tag) reqlist = [sreq1, sreq2, rreq1, rreq2] for req in reqlist: assert req != MPI.REQUEST_NULL # main loop converged = False while not converged: # compute boundary columns N, S = A[ :-2, 1], A[2:, 1] E, W = A[1:-1, 0], A[1:-1, 2] C = B[:, 0] numpy.add(N, S, C) numpy.add(C, E, C) numpy.add(C, W, C) C *= 0.25 N, S = A[ :-2, -2], A[2:, -2] E, W = A[1:-1, -3], A[1:-1, -1] C = B[:, -1] numpy.add(N, S, C) numpy.add(C, E, C) numpy.add(C, W, C) C *= 0.25 # start communication #MPI.Prequest.Startall(reqlist) for r in reqlist: r.Start() # compute interior N, S = A[ :-2, 2:-2], A[2, 2:-2] E, W = A[1:-1, 2:-2], A[1:-1, 2:-2] C = B[:, 1:-1] numpy.add(N, S, C) numpy.add(E, C, C) numpy.add(W, C, C) C *= 0.25 A[1:-1, 1:-1] = B # complete communication MPI.Prequest.Waitall(reqlist) # convergence myconv = numpy.allclose(B, 0) loc_conv = numpy.asarray(myconv, dtype='i') glb_conv = numpy.asarray(0, dtype='i') MPI.COMM_WORLD.Allreduce([loc_conv, MPI.INT], [glb_conv, MPI.INT], op=MPI.LAND) converged = bool(glb_conv) # free persintent requests for req in reqlist: req.Free()