/*---------------------------------------------------------------------
parallel.laplace.c               Finite Difference (parallel) program
Author: Robb Newman
Converted to MPI: 11/12/94 by Xianneng Shen
Last revised: 3/14/97 RYL
Last revised: 3/20/97 Saleh Elmohamed

	This program uses a finite difference scheme to solve
Laplace's equation for a square matrix distributed over a square
(logical) processor topology.  A complete description of the algorithm
is found in Fox, et al "Solving problems on concurrent Processors,
Volume 1: General Techniques and Regular Problems, Prentice Hall,
Englewood Cliffs, New Jersey.

	This program works on the SPMD (single program, multiple data)
paradigm.  It illustrates 2-d block decomposition, nodes exchanging
edge values, and convergence checking.

        Each matrix element is updated based on the values of the four
neighboring matrix elements.  This process is repeated until the data
converges -- until the average change in any matrix element (compared
to the value 20 iterations previous) is smaller than a specified value.

        To ensure reproducible results between runs, a red/black
checkerboard algorithm is used. Each process exchanges edge values
with its four neighbors. Then new values are calculated for the upper
left and lower right corners (the "red" corners) of each node's
matrix. The processes exchange edge values again.  The upper right
and lower left corners (the "black" corners) are then calculated.

        The program is currently configured for a 48x48 matrix
distributed over four processors. It can be edited to handle
different matrix sizes or number of processors, as long as the matrix
can be divided evenly between the processors.
-----------------------------------------------------------------------*/

#include <stdio.h>
#include <math.h>
#include "mpi.h"
#define n 250	 	    /* matrix is nxn, excluding boundary values     */
#define nodeedge 125	    /* a task works on a nodeedge x nodeedge matrix */
#define nblock n/nodeedge   /* number of tasks per row of matrix            */
#define nproc nblock*nblock /* total number of tasks (processors)           */
#define w (double)1.2	    /* convergence factor                           */

void setmatrix(double M[][nodeedge+2]);
void setcomm(int rank, int comm[]);
void iterate(double M[][nodeedge+2],double result[][n],int rank,int comm[]);
void setex(double ex[], double M[][nodeedge+2], int which);
void unpack(double M[][nodeedge+2],int where,double in[]);
void exchange(double M[][nodeedge+2], int comm[], int rank);
void dored(double M[][nodeedge+2]);
void doblack(double M[][nodeedge+2]);

main(int argc, char **argv) 
{ 
  double M[nodeedge+2][nodeedge+2], /* my block of the array           */
    result[n][n],              /* will store all final values          */
    start, etime;              /* initial, elapsed wallclock time (sec)*/
    int ntasks,                /* number of tasks started              */
    rank,                      /* rank (process ID)                    */
    comm[4],                   /* directions requiring communication   */
    i,j;
  FILE   *fp;                  /* file pointer to the output file      */

  MPI_Init(&argc, &argv);
  start = MPI_Wtime();
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  MPI_Comm_size(MPI_COMM_WORLD, &ntasks);
  if (ntasks != nproc) {
    if (rank == 0) printf("error: MP_PROCS should be set to %i!\n",nproc);
    exit(1);
  }
  if (rank == 0) printf("Everything's okay so far\n");

  /* set initial values of M */ 
  setmatrix(M);
  if (rank == 0) printf("Matrix is set\n");

  /* figure out who I communicate with */ 
  setcomm(rank, comm);			
  if (rank == 0) 
    {
      printf("Communications are set up\n");
      printf("Beginning to iterate\n");
    }

  /* update M until convergence */

  iterate(M, result, rank, comm);

  /* report results and timing */ 

  etime = (MPI_Wtime() - start);
  printf("task %i took %6.3f seconds\n", rank, etime);
  if (rank == 0) {
    fp = fopen("parallel.laplace.out", "w");
    for (i=0; i<n; i++) {
      for (j=0; j<n; j++) {
	fprintf(fp,"%f \n", result[i][j]);
      }
      fprintf(fp, "\n");	
    }
    close(fp);
  } 
  MPI_Finalize();
}

/*--------------------------------------------------------------------
  setmatrix sets the initial values of my piece of the matrix 
  --------------------------------------------------------------------*/
void setmatrix(double M[][nodeedge+2])
{
  int i,j;
  double bv[4], avg;

  bv[0] = 100.0;
  bv[1] = 0.0;
  bv[2] = 0.0;
  bv[3] = 0.0;
 
  /* put boundary values in edge vectors of M */ 
  for (i=1; i<=nodeedge; i++) {	
    M[0][i]=         bv[0];
    M[i][nodeedge+1]=bv[1];
    M[nodeedge+1][i]=bv[2];
    M[i][0]=         bv[3];
  }

  /* set all interior values to be the average of the boundary values */ 
  avg = (bv[0]+bv[1]+bv[2]+bv[3])/4.0;
  for (i=1; i<=nodeedge; i++)
    for (j=1; j<=nodeedge; j++)
      M[i][j] = avg;
}


/*-------------------------------------------------------------------------
  setcomm figures out what directions I must communicate with (1=yes, 0=no)
  comm[0]=up, comm[1]=right, comm[2]=down, comm[3]=left
  -------------------------------------------------------------------------*/
void setcomm(int rank, int comm[])
{
  int i;

  for (i=0; i<4; i++) comm[i] = 1;
  if (rank < nblock) comm[0] = 0;	
  if (fmod((double)(rank+1), (double)nblock) == 0.0) comm[1] = 0;
  if (rank > (nblock*(nblock-1)-1)) comm[2] = 0;
  if (fmod((double)rank, (double)(nblock)) == 0.0) comm[3] = 0;
}


/*-------------------------------------------------------------------------
  iterate takes care of everything to do with iteration, including the 
  convergence checking
  -------------------------------------------------------------------------*/
void iterate(double M[][nodeedge+2], double result[][n],
	     int rank, int comm[])
{
  int  it,             /* current iteration                             */
    i,j,k,l,index,     /* index variables                               */
    done,              /* loop control variable                         */
    count;             /* length (in elements) of messages              */
  double mold[nodeedge+2][nodeedge+2], /* my matrix, 20 iterations ago  */
                                       /* used to check convergence     */
    diff,              /* the average absolute difference in elements   */
                       /*   of M and mold                               */
    in;                /* sum of diff from each task                    */
  double MM[n*n],      /* initially holds results from all tasks        */
    ediff,             /* the difference in two elements of M and mold  */
    send[nodeedge][nodeedge];  /* the result I will send to task 0      */

  it   = 0;
  done = 0;
  for (i=1; i<=nodeedge; i++)
    for (j=1; j<=nodeedge; j++)
      mold[i][j] = M[i][j];
 
  while (done == 0) 
    {
      it++;

    /* exchange values with neighbors, update red squares, exchange values */
    /* with neighbors, update black squares */

      exchange(M, comm, rank);
      dored(M);
      exchange(M, comm, rank);
      doblack(M);

    /* check for convergence every 20 iterations */
    /* find the average absolute change in elements of M */
    /* maximum iterations is 50000 */

      if (it>50000) done=1;	
      if ((fmod((double)it, (double)20.0) == 0.0) && (done != 1)) { 
	diff = 0.0;
	for (i=1; i<=nodeedge; i++)
	  for (j=1; j<=nodeedge; j++) {
	    ediff = M[i][j] - mold[i][j];
	    if (ediff < 0.0) 
	      ediff = -ediff;
	    diff += ediff;
	    mold[i][j] = M[i][j];
	  }
	diff = diff/((double)(nodeedge*nodeedge));
	MPI_Allreduce(&diff, &in, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
	if (in < (double)nproc*.001) 
	  done = 1;
      }
    }

  /* send results to task 0 */
  /* don't need to include boundary points */ 

  for (i=0; i<nodeedge; i++)
    for (j=0; j<nodeedge; j++)
      send[i][j] = M[i+1][j+1];
  count = nodeedge*nodeedge;
  MPI_Gather(&send, count, MPI_DOUBLE, &MM,
	     count, MPI_DOUBLE, 0, MPI_COMM_WORLD);
  printf("I finished gather\n");

  /* storage on task 0 has to be consistent with a nblock x nblock */
  /* decomposition */
  if (rank == 0) 
    {
      printf("did %i iterations\n",it);
      index=0;
      for (k=0; k<nblock; k++)
	for (l=0; l<nblock; l++)
	  for (i=k*nodeedge; i<(k+1)*nodeedge; i++)
	    for (j=l*nodeedge; j<(l+1)*nodeedge; j++) {
	      result[i][j] = MM[index];
	      index++;
	    }
    }
}


/*-------------------------------------------------------------------------
  pulls off the edge values of M to send to another task
  -------------------------------------------------------------------------*/
void setex(double ex[], double M[][nodeedge+2], int which)
{
  int i;
  switch (which) 
    {
    case 0: {
      for (i=1; i<=nodeedge; i++) ex[i-1] = M[1][i];
      break; }
    case 1: {
      for (i=1; i<=nodeedge; i++) ex[i-1] = M[i][nodeedge];
      break; }
    case 2: {
      for (i=1; i<=nodeedge; i++) ex[i-1] = M[nodeedge][i];
      break; }
    case 3: {
      for (i=1; i<=nodeedge; i++) ex[i-1] = M[i][1];
      break; }
    }
}


/*-------------------------------------------------------------------------
  unpack puts the vector of new edge values I just received into the
  edges of M
  -------------------------------------------------------------------------*/
void unpack(double M[][nodeedge+2],int where,double in[])
{
  int i;
  switch (where)
    {
    case 0: {
      for (i=0; i<nodeedge; i++) M[0][i+1] = in[i];
      break; }
    case 1: {
      for (i=0; i<nodeedge; i++) M[i+1][nodeedge+1] = in[i];
      break; }
    case 2: {
      for (i=0; i<nodeedge; i++) M[nodeedge+1][i+1] = in[i];
      break; }
    case 3: {
      for (i=0; i<nodeedge; i++) M[i+1][0] = in[i];
      break; }
    }
}


/*-----------------------------------------------------------------------
  makes the communication calls for each task to exchange edge values 
  with up to four neighbors
  -----------------------------------------------------------------------*/
void exchange(double M[][nodeedge+2], int comm[], int rank)
{
  double ex0[nodeedge], ex1[nodeedge],  /* vectors being sent            */
    ex2[nodeedge], ex3[nodeedge];
  double in0[nodeedge], in1[nodeedge],  /* vectors being received        */
    in2[nodeedge], in3[nodeedge];
  int    tag,                           /* message tag                   */
                                        /* 0=up, 1=right, 2=down, 3=left */
    partner,                            /* who I am exchanging with      */
    i;
  MPI_Request requests[8];              /* communication identifier      */
  MPI_Status status[8];                 /* status of completed communication */

  /* initialize requests */
  for (i = 0; i < 8; i++) 
    requests[i] = MPI_REQUEST_NULL; 

  /* receive incoming messages */
  if (comm[0] == 1) 
    {				/* receive from above */
      partner = rank-nblock;
      tag = 0;
      MPI_Irecv(&in0, nodeedge, MPI_DOUBLE, partner, tag, MPI_COMM_WORLD, 
		&requests[0]);
    }
  if (comm[1] == 1) 
    {				/* receive from right */
      partner = rank+1;
      tag = 1;
      MPI_Irecv(&in1, nodeedge, MPI_DOUBLE, partner, tag, MPI_COMM_WORLD,
		&requests[1]);
    }
  if (comm[2] == 1) 
    {				/* receive from below */
      partner = rank + nblock;
      tag = 2;
      MPI_Irecv(&in2, nodeedge, MPI_DOUBLE, partner, tag, MPI_COMM_WORLD,
		&requests[2]);
    }
  if (comm[3] == 1) 
    {				/* receive from left */
      partner = rank-1;
      tag = 3;
      MPI_Irecv(&in3, nodeedge, MPI_DOUBLE, partner, tag, MPI_COMM_WORLD,
		&requests[3]);
    }
 
  /* send messages to my partners */
  if (comm[0] == 1) {				/* send up */
    partner = rank-nblock;
    setex(ex0,M,0);
    MPI_Isend(&ex0, nodeedge, MPI_DOUBLE, partner, 2, MPI_COMM_WORLD,
	      &requests[4]);
  }
  if (comm[1] == 1) {				/* send right */
    partner = rank+1;
    setex(ex1,M,1);
    MPI_Isend(&ex1, nodeedge, MPI_DOUBLE, partner, 3, MPI_COMM_WORLD,
	      &requests[5]);
  }
  if (comm[2] == 1) {				/* send down */
    partner = rank + nblock;
    setex(ex2,M,2);
    MPI_Isend(&ex2, nodeedge, MPI_DOUBLE, partner, 0, MPI_COMM_WORLD,
	      &requests[6]);
  }
  if (comm[3] == 1) {				/* send left */
    partner = rank - 1;
    setex(ex3,M,3);
    MPI_Isend(&ex3, nodeedge, MPI_DOUBLE, partner, 1,  MPI_COMM_WORLD,
	      &requests[7]);
  }

  /* wait for all communication to complete */ 
  MPI_Waitall(8, requests, status);
   
  if (comm[0] == 1) unpack(M, 0, in0);   /* put msg from above into matrix */
  if (comm[1] == 1) unpack(M, 1, in1);   /* put msg from right into matrix */
  if (comm[2] == 1) unpack(M, 2, in2);   /* put msg from below into matrix */
  if (comm[3] == 1) unpack(M, 3, in3);   /* put msg from left  into matrix */
}


/*-----------------------------------------------------------------------
  iterates on the upper left and lower right corners of my matrix
  -----------------------------------------------------------------------*/
void dored (double M[][nodeedge+2])
{
  int i,j;
 
  for (i=1; i<=nodeedge/2; i++) 		  /* upper left */
    for (j=1; j<=nodeedge/2; j++) 
      M[i][j] = w/4.0*(M[i-1][j]+M[i][j-1]+M[i+1][j]+M[i][j+1])+(1.0-w)*M[i][j];
  for (i=nodeedge/2+1; i<=nodeedge; i++)	  /* lower right */
    for (j=nodeedge/2+1; j<=nodeedge; j++)
      M[i][j] = w/4.0*(M[i-1][j]+M[i][j-1]+M[i+1][j]+M[i][j+1])+(1.0-w)*M[i][j];
}


/*----------------------------------------------------------------------
  iterates on the upper right and lower left corners of my matrix
  ----------------------------------------------------------------------*/
void doblack(double M[][nodeedge+2])
{
  int i,j;
  for (i=1; i<=nodeedge/2; i++)		          /* upper right */
    for (j=nodeedge/2+1; j<=nodeedge; j++)
      M[i][j] = w/4.0*(M[i-1][j]+M[i][j-1]+M[i+1][j]+M[i][j+1])+(1.0-w)*M[i][j];
  for (i=nodeedge/2+1; i<=nodeedge; i++)	  /* lower left */
    for (j=1; j<=nodeedge/2; j++)
      M[i][j] = w/4.0*(M[i-1][j]+M[i][j-1]+M[i+1][j]+M[i][j+1])+(1.0-w)*M[i][j];
}