Jacobi iteration using MPI

The code below implements Jacobi iteration for solving the linear system arising from the steady state heat equation using MPI. Note that in this code each process, or task, has only a portion of the arrays and must exchange boundary data using message passing.

Compare to:

• Jacobi iteration using OpenMP with parallel do constructs
• Jacobi iteration using OpenMP with coarse-grain parallel block

The code:

! $CLASSHG/codes/mpi/jacobi2_mpi.f90
!
! Domain decomposition version of Jacobi iteration illustrating
! coarse grain parallelism with MPI.
!
! The grid points are split up into ntasks disjoint sets and each task
! is assigned one set that it updates for all iterations.  The tasks 
! correspond to processes.
!
! The task (or process) number is called me in this code for brevity
! rather than proc_num.
!
! Note that each task allocates only as much storage as needed for its 
! portion of the arrays.
!
! Each iteration, boundary values at the edge of each grid must be
! exchanged with the neighbors.


program jacobi2_mpi
    use mpi

    implicit none

    integer, parameter :: maxiter = 10000, nprint = 1000
    real (kind=8), parameter :: alpha = 20.d0, beta = 60.d0

    integer :: i, iter, istart, iend, points_per_task, itask, dummy, n
    integer :: ierr, ntasks, me, req1, req2
    integer, dimension(MPI_STATUS_SIZE) :: mpistatus
    real (kind = 8), dimension(:), allocatable :: f, u, uold
    real (kind = 8) :: x, dumax_task, dumax_global, dx, tol

    ! Initialize MPI; get total number of tasks and ID of this task
    call mpi_init(ierr)
    call mpi_comm_size(MPI_COMM_WORLD, ntasks, ierr)
    call mpi_comm_rank(MPI_COMM_WORLD, me, ierr)

    ! Ask the user for the number of points
    if (me == 0) then
        print *, "Input n ... "
        read *, n
    end if
    ! Broadcast to all tasks; everybody gets the value of n from task 0
    call mpi_bcast(n, 1, MPI_INTEGER, 0, MPI_COMM_WORLD, ierr)

    dx = 1.d0/real(n+1, kind=8)
    tol = 0.1d0*dx**2

    ! Determine how many points to handle with each task
    points_per_task = (n + ntasks - 1)/ntasks
    if (me == 0) then   ! Only one task should print to avoid clutter
        print *, "points_per_task = ", points_per_task
    end if

    ! Determine start and end index for this task's points
    istart = me * points_per_task + 1
    iend = min((me + 1)*points_per_task, n)

    ! Diagnostic: tell the user which points will be handled by which task
    print 201, me, istart, iend
201 format("Task ",i2," will take i = ",i6," through i = ",i6)


    ! Initialize:
    ! -----------

    ! This makes the indices run from istart-1 to iend+1
    ! This is more or less cosmetic, but makes things easier to think about
    allocate(f(istart-1:iend+1), u(istart-1:iend+1), uold(istart-1:iend+1))

    ! Each task sets its own, independent array
    do i = istart, iend
        ! Each task is a single thread with all its variables private
        ! so re-using the scalar variable x from one loop iteration to
        ! the next does not produce a race condition.
        x = dx*real(i, kind=8)
        f(i) = 100.d0*exp(x)               ! Source term
        u(i) = alpha + x*(beta - alpha)    ! Initial guess
    end do
    
    ! Set boundary conditions if this task is keeping track of a boundary
    ! point
    if (me == 0)        u(istart-1) = alpha
    if (me == ntasks-1) u(iend+1)   = beta


    ! Jacobi iteratation:
    ! -------------------

    do iter = 1, maxiter
        uold = u

        ! Send endpoint values to tasks handling neighboring sections of the array
        ! Note that non-blocking sends are used.

        if (me > 0) then
            ! Send left endpoint value to process to the "left"
            call mpi_isend(uold(istart), 1, MPI_DOUBLE_PRECISION, me - 1, &
                1, MPI_COMM_WORLD, req1, ierr)
        end if
        if (me < ntasks-1) then
            ! Send right endpoint value to process on the "right"
            call mpi_isend(uold(iend), 1, MPI_DOUBLE_PRECISION, me + 1, &
                2, MPI_COMM_WORLD, req2, ierr)
        end if

        ! Accept incoming endpoint values from other tasks
        if (me < ntasks-1) then
            ! Receive right endpoint value
            call mpi_recv(uold(iend+1), 1, MPI_DOUBLE_PRECISION, me + 1, &
                1, MPI_COMM_WORLD, mpistatus, ierr)
        end if
        if (me > 0) then
            ! Receive left endpoint value
            call mpi_recv(uold(istart-1), 1, MPI_DOUBLE_PRECISION, me - 1, &
                2, MPI_COMM_WORLD, mpistatus, ierr)
        end if

        dumax_task = 0.d0   ! Max seen by this task

        ! Apply Jacobi iteration on this task's section of the array
        do i = istart, iend
            u(i) = 0.5d0*(uold(i-1) + uold(i+1) + dx**2*f(i))
            dumax_task = max(dumax_task, abs(u(i) - uold(i)))
        end do

        ! Take global maximum of dumax values
        call mpi_allreduce(dumax_task, dumax_global, 1, MPI_DOUBLE_PRECISION, &
            MPI_MAX, MPI_COMM_WORLD, ierr)

        ! Also report progress to the user
        if (me == 0) then
            if (mod(iter, nprint)==0) then
                print 203, iter, dumax_global
203             format("After ",i8," iterations, dumax = ",d16.6,/)
            end if
        end if

        ! All tasks now have dumax_global, and can check for convergence
        if (dumax_global < tol) exit
    end do

    print 212, me, iter, dumax_global
212 format("Task number ",i2," finished after ",i9," iterations, dumax = ",&
            e16.6)


    ! Output result:
    ! --------------

    ! Synchronize to keep the next part from being non-deterministic
    call mpi_barrier(MPI_COMM_WORLD, ierr)

    ! Have each task output to a file in sequence, using messages to
    ! coordinate
    if (me == 0) then    ! Task 0 goes first
        open(unit=20, file="heatsoln.txt", status="replace")
        write(20,*) "          x                  u"
        write(20, 222) 0.d0, u(0)    ! Boundary value at left end
222     format(2e20.10)
        do i = istart, iend
            write(20, 222) i*dx, u(i)
        end do
        close(unit=20)
        ! Closing the file should guarantee that all the ouput will be written to disk

        ! Send go-ahead message to next task
        ! Only the fact that the message was sent is important, not its contents
        call mpi_send(dummy, 1, MPI_INTEGER, 1, 4, MPI_COMM_WORLD, ierr)
    else
        ! Wait for go-ahead message from previous task
        call mpi_recv(dummy, 1, MPI_INTEGER, me - 1, 4, MPI_COMM_WORLD, mpistatus, ierr)
        ! Open file for appending; do not destroy previous contents
        open(unit=20, file="heatsoln.txt", status="old", access="append")
        do i = istart, iend
            write(20, 222) i*dx, u(i)
        end do
        if (me == ntasks - 1) write(20, 222) 1.d0, u(iend+1)    ! Boundary value at right end
        close(unit=20)
        if (me < ntasks - 1) then
            ! Send go-ahead message to next task
            call mpi_send(dummy, 1, MPI_INTEGER, me + 1, 4, MPI_COMM_WORLD, ierr)
        end if
    end if

    if (me == 0) print *, "Solution is in heatsoln.txt"

    ! Close out MPI
    call mpi_finalize(ierr)

end program jacobi2_mpi