Scalapack solution

C

#include "stdio.h"
#include "stdlib.h"
#include "mpi.h"
#include "math.h"
#include "wrapper.h"
#define N 10
#define NB 2

#ifndef max
#define max( a, b ) ( ((a) > (b)) ? (a) : (b) )
#endif

int main( int argc, char** argv)
{
  int i, j, zero = 0, uno = 1, sei = 6;
  // MPI
  int myrank, nprocs, ndims = 2, dims[2] = {0,0};
  MPI_Status status;
  // BLACS/SCALAPACK
  int nprow, npcol, info, ic = -1, context, myrow, mycol, desca[9], descb[9];
  char order = 'R';
  // MATRIX
  int n = N, nb = NB, nrhs = 1, nr, nc, lld;
  double *A, *b, *work, alpha;
  int *ipiv;

  // Initialize MPI environment
  MPI_Init(&argc, &argv);
  MPI_Comm_size(MPI_COMM_WORLD, &nprocs);
  MPI_Comm_rank(MPI_COMM_WORLD, &myrank);

  // Initialize a default BLACS context and the processes grid
  MPI_Dims_create(nprocs, ndims, dims);
  nprow = dims[0];
  npcol = dims[1];
  blacs_get_( &ic, &zero, &context );
  blacs_gridinit_( &context, &order, &nprow, &npcol );
  blacs_gridinfo_( &context, &nprow, &npcol, &myrow, &mycol );

  // Computation of local matrix size
  nr = numroc_( &n, &nb, &myrow, &zero, &nprow );
  nc = numroc_( &n, &nb, &mycol, &zero, &npcol );
  lld = max( 1 , nr );
  A = malloc(nr*nc*sizeof(double));
  b = malloc(nr*sizeof(double));
  work = malloc(nb*sizeof(double));
  ipiv = malloc((lld+nb)*sizeof(int));

  // Descriptors
  descinit_( descb, &n, &uno, &nb, &nb, &zero, &zero, &context, &lld, &info );
  descinit_( desca, &n, &n, &nb, &nb, &zero, &zero, &context, &lld, &info );

  // Vector and matrix elements generation
  for( i = 1; i <= n; i++ ){
    alpha = (double)(207 - i);
    pdelset_( b, &i, &uno, descb, &alpha );
  }
  for( i = 1; i <= n; i++ ){
    for( j = 1; j <= n; j++ ){
      if( i == j ) {
	alpha = (double)(10000);
	pdelset_( A, &i, &j, desca, &alpha );
      } else {
	alpha = ((double)(i)+(double)(j)/(double)(2));
	pdelset_( A, &i, &j, desca, &alpha );
      }
    }
  }

  // Linear system equations solver
  pdgesv_( &n, &uno, A, &uno, &uno, desca, ipiv, b, &uno, &uno, descb, &info );

  // Print the results vector
  Cpdlaprnt( &n, &uno, b, &uno, &uno, descb, &zero, &zero, "x", &sei, work );

  free(A);
  free(b);
  free(ipiv);
  free(work);

  // Close BLACS environment
  blacs_gridexit_( &context );
  blacs_exit_( &zero );

  // Close MPI environment if blacs_exit paramater is not equal zero
  //MPI_Finalize();
  return 0;
}

FORTRAN

program es_scalapack
use mpi
implicit none

integer :: i, j, stdout
parameter ( stdout=6 )
! MPI
integer :: ierr, nprocs, myrank
! BLACS/SCALAPACK
integer :: context, nprow, npcol, myrow, mycol, info, ndims, dims(2)
parameter ( ndims=2 )
external :: blacs_exit, blacs_gridexit, blacs_gridinfo, blacs_get
external :: blacs_gridinit, descinit, numroc, pdelset, pdlaprnt
integer :: numroc
! Matrix
integer :: M, N, MB, NB, Nloc, Mloc, aval
!parameter ( M=16384, N=16384, MB=64, NB=64 )
parameter ( M=10, N=10, MB=2, NB=2 )
real*8 , allocatable :: A(:,:), B(:), work(:)
integer, allocatable :: ipiv(:)
integer :: descA(9), descB(9)
! Benchmarks
real*8 :: t1, t2, dt

! Initialize MPI environment
call MPI_Init(ierr)
call MPI_Comm_Size(mpi_comm_world,nprocs,ierr)
call MPI_Comm_Rank(mpi_comm_world,myrank,ierr)

! Initialize a default BLACS context and the processes grid
dims = 0
call MPI_Dims_Create( nprocs, ndims, dims, ierr)
nprow = dims(1)
npcol = dims(2)
call blacs_get( -1, 0, context )
call blacs_gridinit( context, 'Row-major', nprow, npcol)
call blacs_gridinfo( context, nprow, npcol, myrow, mycol )

! Computation of local matrix size
Mloc = numroc( M, MB, myrow, 0, nprow )
Nloc = numroc( N, NB, mycol, 0, npcol )
write(*,*) "myrank= ",myrank,"     Mloc= ",Mloc,"     Nloc= ",Nloc 
allocate( A( Mloc, Nloc ) )
allocate( B( Mloc ) )
allocate( ipiv( Mloc+MB ) )
allocate( work( M ) )

! Descriptos
call descinit(descA, M, N, MB, NB, 0, 0, context, max(1,Mloc), info)
call descinit(descB, M, 1, MB,  1, 0, 0, context, max(1,Mloc), info)

! Vector and matrix elements generation
do j=1,N
  do i=1,M
    if(i.eq.j) then
      call pdelset( A, i, j, descA, dble(10000) )
   else
      call pdelset( A, i, j, descA, (dble(i)+dble(j)/dble(2.0)) )
   endif
  enddo
enddo
do i=1,M
  call pdelset( B, i, 1, descB, dble(207-i) )
enddo

!!! for benchmarking !!!
call MPI_Barrier( mpi_comm_world, ierr)
t1 = MPI_wtime()
!!!!!!!!!!!!!!!!!!!!!!!!

! Linear system equations solver
call pdgesv( N, 1, A, 1, 1, descA, ipiv, B, 1, 1, descB, info )

!!! for benchmarking !!!
call MPI_Barrier( mpi_comm_world, ierr)
t2 = MPI_wtime()
dt = t2-t1
if(myrank.eq.0) write(*,*) nprocs, dt
!!!!!!!!!!!!!!!!!!!!!!!!

! Print the results vector
call pdlaprnt(M, 1, B, 1, 1, descB, 0, 0, 'X', stdout, work)

deallocate( A )
deallocate( B )
deallocate( ipiv )
deallocate( work )

! Close BLACS environment
call blacs_gridexit( context )
call blacs_exit( 0 )

! Close MPI environment if blacs_exit paramater is not equal zero
!call MPI_Finalize(ierr)

end program es_scalapack