/*
  mpi example for matrix matrix product
  
  -- row-column partition
  -- dim(A) is an integral multiple of np 
*/
#include "mpi.h"
#include <stdio.h>
#include <stdlib.h>
#include <time.h>

#define NP  4
#define N  1024
#define NL N/NP

int main(int argc, char *argv[])
{
  	int myid, np;
  	int i, j, k, ii, ij, ik, shift, src, dest;
  	double t, (*pA)[N], (*pB)[NL], (*pC)[N];
  	MPI_Status status;

  	MPI_Init(&argc, &argv);
  	MPI_Comm_rank(MPI_COMM_WORLD, &myid);
  	MPI_Comm_size(MPI_COMM_WORLD, &np);

  	if ( (np != NP) && (myid == 0) )
	{   
  		printf("Error: np not equal NP\n");
      	MPI_Abort(MPI_COMM_WORLD, 1);
  	}
  
  	pA = malloc(NL*N*sizeof(double));
  	pB = malloc(N*NL*sizeof(double));
  	pC = malloc(NL*N*sizeof(double));
  	if ( !pA | !pB | !pC ) 
	{ 
    	printf("Out of memory!\n"); 
    	MPI_Abort(MPI_COMM_WORLD, 2); 
  	}

  	t = MPI_Wtime(); // time

	// set initial value 
  	shift = myid*NL;
  	for(i=0; i<NL; i++) {
    	for(j=0; j<N; j++) {
      		pA[i][j] = 1.0/(double)(i+shift+j+1);
      		pC[i][j] = 0.0;
    	}
  	}
  	for(i=0; i<N; i++) 
    	for(j=0; j<NL; j++) 
      		pB[i][j] = (double)(i+shift+j+1);

// ********************************** //
// computations
  	for (k=0; k<np; k++)
  	{
		// compute C_{ij} = A_i*B_j
    	j = ((myid + k) % np) * NL;
    	for(ij=0; ij<NL; ij++)
      		for(ik=0; ik<N; ik++)
        		for(ii=0; ii<NL; ii++)
          			pC[ii][ij+j] += pA[ii][ik]*pB[ik][ij];

//  send B to the left processor and 
//  get new B from the right processor
    	src  = ((myid+1+np) % np);
    	dest = ((myid-1+np) % np);
    	if ( k != np-1)
      		MPI_Sendrecv_replace(pB, N*NL, MPI_DOUBLE, dest, 111, 
                           src, 111, MPI_COMM_WORLD, &status);
  	}

  	t = MPI_Wtime() - t; // time

  	printf("myid=%d, C[0][0]=%.6f, Time=%.4e\n", myid, pC[0][0], t);

  	free(pA); free(pB); free(pC);
  	MPI_Finalize();  	
  	return 0;
}