python_parallel_computing/virtualTopology.py at master · qiweihan/python_parallel_computing · GitHub

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
from mpi4py import MPI
import numpy as np

UP = 0
DOWN = 1
LEFT = 2
RIGHT = 3
neighbour_processes = [0, 0, 0, 0]

if __name__ == "__main__":
    comm = MPI.COMM_WORLD
    rank = comm.rank
    size = comm.size

    grid_rows = int(np.floor(np.sqrt(comm.size)))
    grid_column = comm.size // grid_rows

    if grid_rows * grid_column > size:
        grid_column -= 1
    if grid_rows * grid_column > size:
        grid_rows -= 1

    if (rank == 0):
        print("Building a %d x %d grid topology:" % (grid_rows, grid_column))

    cartesian_communicator = comm.Create_cart((grid_rows, grid_column),
                                              periods=(True, True), reorder=True)
    my_mpi_row, my_mpi_col = cartesian_communicator.Get_coords(cartesian_communicator.rank)

    neighbour_processes[UP], neighbour_processes[DOWN] = cartesian_communicator.Shift(0, 1)
    neighbour_processes[LEFT], neighbour_processes[RIGHT] = cartesian_communicator.Shift(1, 1)

    print("Process = %s row = %s column = %s ----> \n "
          "neighbour_processes[UP] = %s neighbour_processes[DOWN] = %s \
          neighbour_processes[LEFT] =%s neighbour_processes[RIGHT]=%s" % (rank, my_mpi_row,
                                                                    my_mpi_col, neighbour_processes[UP],
                                                                    neighbour_processes[DOWN],
                                                                    neighbour_processes[LEFT],
                                                                    neighbour_processes[RIGHT]))