6. Several devices with MPI

A node often has several GPUs, and a cluster many nodes. The usual layout is one MPI rank per device: the domain is split among the ranks, each rank advances its part on its device, and neighbouring ranks exchange the cells next to their boundary (the halo) at each step. mpih_object, in the module fundal_mpih_object (built with MPI=1 or an mpi FoBiS mode), initializes MPI and the device together:

call mpih%initialize(do_mpi_init=.true., do_device_init=.true.) ! MPI_Init, then dev_init(local_rank=...)
first = mpih%myrank * n / mpih%procs_number + 1                   ! this rank owns the cells first:last
last  = (mpih%myrank + 1) * n / mpih%procs_number
left  = mpih%myrank - 1 ; if (mpih%myrank == 0) left = MPI_PROC_NULL
right = mpih%myrank + 1 ; if (mpih%myrank == mpih%procs_number - 1) right = MPI_PROC_NULL

initialize(do_mpi_init=.true.) calls MPI_Init; do_device_init=.true. splits MPI_COMM_WORLD by node (MPI_COMM_TYPE_SHARED, stored in local_comm) and calls dev_init(local_rank=...), so that the ranks of a node take the devices mod(local_rank, devs_number). The 15 cells are split into contiguous blocks, and the neighbours at the physical boundaries are MPI_PROC_NULL.

Each rank allocates only its cells and their two ghost cells, with the global indexes as bounds, so the kernel and the initial condition keep the indexing of the serial code:

call dev_alloc(fptr_dev=t_dev,  lbounds=[first-1], ubounds=[last+1], ierr=ierr, init_value=0._R8P, label='t')
if (ierr /= 0) call mpih%abort(msg='device allocation failed')
call dev_alloc(fptr_dev=tn_dev, lbounds=[first-1], ubounds=[last+1], ierr=ierr, init_value=0._R8P, label='tn')
if (ierr /= 0) call mpih%abort(msg='device allocation failed')
t = [(sin(pi*i*dx), i=0, n+1)]
t(0) = 0._R8P ; t(n+1) = 0._R8P
call dev_memcpy_to_device(dst=t_dev, src=t(first-1:last+1)) ! only the cells of this rank, and its ghosts

The halo exchange

Before each step, each rank sends its first and last cells to its neighbours and receives their values into its ghost cells. MPI cannot read device memory here, so the values go through host buffers:

call dev_memcpy_from_device(dst=edge(1:1), src=t_dev(first:first)) ! device -> host buffer
call dev_memcpy_from_device(dst=edge(2:2), src=t_dev(last:last))
call MPI_SENDRECV(edge(2), 1, MPI_DOUBLE_PRECISION, right, 0, ghost(1), 1, MPI_DOUBLE_PRECISION, left,  0, &
                  MPI_COMM_WORLD, MPI_STATUS_IGNORE, ierr)
call MPI_SENDRECV(edge(1), 1, MPI_DOUBLE_PRECISION, left,  1, ghost(2), 1, MPI_DOUBLE_PRECISION, right, 1, &
                  MPI_COMM_WORLD, MPI_STATUS_IGNORE, ierr)
if (left  /= MPI_PROC_NULL) call dev_memcpy_to_device(dst=t_dev(first-1:first-1), src=ghost(1:1)) ! host -> device
if (right /= MPI_PROC_NULL) call dev_memcpy_to_device(dst=t_dev(last+1:last+1),   src=ghost(2:2))

The copies take array sections of the device pointer: a section is contiguous here, and only contiguous arrays may be copied. MPI_SENDRECV pairs every send with a receive, so no rank waits for another, and a transfer to or from MPI_PROC_NULL does nothing: the physical ghost cells keep the boundary value 0 set by init_value.

Gathering the result

Only rank 0 prints, after collecting the whole field (each rank contributes its cells, zero elsewhere, so a sum reassembles it):

allocate(t_loc(first-1:last+1))
call dev_memcpy_from_device(dst=t_loc, src=t_dev)
t = 0._R8P
t(first:last) = t_loc(first:last) ! each rank fills its cells, zero elsewhere: the sum is the whole field
call MPI_REDUCE(t, t_all, n+2, MPI_DOUBLE_PRECISION, MPI_SUM, 0, MPI_COMM_WORLD, ierr)
if (mpih%myrank == 0) then
   do p=0, mpih%procs_number - 1
      print '(A,I0,A,I0,A,I0)', 'rank ', p, ' owns the cells ', p * n / mpih%procs_number + 1, ' to ', &
                                (p + 1) * n / mpih%procs_number
   enddo
   exact = (1._R8P - 4._R8P * r * sin(pi * dx / 2._R8P)**2)**steps * [(sin(pi*i*dx), i=0, n+1)]
   exact(0) = 0._R8P ; exact(n+1) = 0._R8P
   print '(A,I0,A,F7.5)', 'temperature at the centre after ', steps, ' steps: ', t_all((n+1)/2)
   print '(A,L1)', 'equal to the exact discrete solution: ', maxval(abs(t_all - exact)) < 1.e-12_R8P
endif

$ mpirun --oversubscribe -np 2 heat_6
rank 0 owns the cells 1 to 7
rank 1 owns the cells 8 to 15
temperature at the centre after 50 steps: 0.61712
equal to the exact discrete solution: T

The same temperature as the serial chapters: the domain decomposition is exact.

MPI and device memory

Passing a device pointer to MPI works only with a GPU-aware MPI library and the device address of the buffer; staging through host buffers, as here, works with every MPI. On a node, each rank must use its own device: initialize through mpih_object (or dev_init(local_rank=...)), not by setting mydev by hand, so that the host fallback and the device count are checked.

What you learned

mpih_object%initialize, one device per rank from the local rank, global bounds per rank, a halo exchange through host buffers, gathering before printing. Reference: MPI handler, dev_init.

Next: 7. Production runs.