File: ring.py

package info (click to toggle)
mpi4py 1.3%2Bhg20120611-3
  • links: PTS, VCS
  • area: main
  • in suites: wheezy
  • size: 2,020 kB
  • sloc: python: 9,503; ansic: 6,296; makefile: 571; f90: 158; sh: 146; cpp: 103
file content (37 lines) | stat: -rw-r--r-- 824 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
#!/usr/bin/env python

# If you want VampirTrace to log MPI calls, you have to add the two
# lines below at the very beginning of your main bootstrap script.
import mpi4py.rc
mpi4py.rc.threaded = False
mpi4py.rc.profile('vt-mpi', logfile='ring')

from mpi4py import MPI

comm = MPI.COMM_WORLD
size = comm.Get_size()
rank = comm.Get_rank()

src  = rank-1
dest = rank+1
if rank == 0:
    src = size-1
if rank == size-1:
    dest = 0

try:
    from numpy import zeros
    a1 = zeros(1000000, 'd')
    a2 = zeros(1000000, 'd')
except ImportError:
    from array import array
    a1 = array('d', [0]*1000); a1 *= 1000
    a2 = array('d', [0]*1000); a2 *= 1000

comm.Sendrecv(sendbuf=a1, recvbuf=a2,
              source=src, dest=dest)

MPI.Request.Waitall([
    comm.Isend(a1, dest=dest),
    comm.Irecv(a2, source=src),
    ])