program scalapack_example
  use mpi
  implicit none
  integer, parameter :: n = 4
  integer :: ierr, rank, processes, context, grid_rows, grid_columns
  integer :: myrow, mycol, local_rows, local_columns, i, j, column, info
  integer :: descriptor(9)
  integer, external :: numroc
  double precision, allocatable :: a(:,:), b(:,:), c(:,:)
  double precision :: local_error, error

  call MPI_Init(ierr)
  call MPI_Comm_rank(MPI_COMM_WORLD, rank, ierr)
  call MPI_Comm_size(MPI_COMM_WORLD, processes, ierr)
  if (processes /= 2) then
    if (rank == 0) print *, 'This example requires two MPI processes'
    call MPI_Abort(MPI_COMM_WORLD, 1, ierr)
  end if
  call blacs_get(-1, 0, context)
  call blacs_gridinit(context, 'R', 1, processes)
  call blacs_gridinfo(context, grid_rows, grid_columns, myrow, mycol)
  local_rows = numroc(n, 1, myrow, 0, grid_rows)
  local_columns = numroc(n, 1, mycol, 0, grid_columns)
  allocate(a(local_rows,local_columns), b(local_rows,local_columns), &
           c(local_rows,local_columns))
  call descinit(descriptor, n, n, 1, 1, 0, 0, context, max(1,local_rows), info)
  if (info /= 0) call MPI_Abort(MPI_COMM_WORLD, 2, ierr)

  do j = 1, local_columns
    column = mycol + 1 + (j - 1) * grid_columns
    do i = 1, local_rows
      a(i,j) = dble(i + column)
      b(i,j) = 0.0d0
      if (i == column) b(i,j) = 1.0d0
    end do
  end do
  c = 0.0d0
  call pdgemm('N', 'N', n, n, n, 1.0d0, a, 1, 1, descriptor, &
              b, 1, 1, descriptor, 0.0d0, c, 1, 1, descriptor)
  local_error = maxval(abs(c - a))
  call MPI_Allreduce(local_error, error, 1, MPI_DOUBLE_PRECISION, &
                     MPI_MAX, MPI_COMM_WORLD, ierr)
  if (.not. (error <= 1.0d-12)) call MPI_Abort(MPI_COMM_WORLD, 3, ierr)
  if (rank == 0) then
    print *, 'PASS: distributed matrix product with 2 processes'
    print *, 'Maximum absolute error:', error
  end if
  deallocate(a, b, c)
  call blacs_gridexit(context)
  call MPI_Finalize(ierr)
end program scalapack_example