/*******************************************************************************
 *        MPI Matrix Multiply                                                  *
 *                                                                             *
 *  AUTHOR: Stefano Viola - UNISI                                              *
 * VERSION: 0.0.1                                                              *
 *                                                                             *
 ******************************************************************************/

#include <mpi.h>
#include <stdio.h>
#include <stddef.h>
#include <stdlib.h>
#include <stdint.h>
#include <math.h>
#include <string.h>

/*******************************************************************************
 * Definitions for ROI
 ******************************************************************************/

/* COTSON Base for custom instruction definition */
/* IMPORTANT: keep in sync with AbAeterno: abaeterno/lib/cpuid_call.h */
#define _COTSON_CUSTOM_BASE	0x2DAF0000
#define _COTSON_CUSTOM_TRACER   0xFF
#define _XSMOPCODE(_opc)	(_COTSON_CUSTOM_BASE+(_opc))
#define _XSMOPCODE_I(_opc,_im)	(_COTSON_CUSTOM_BASE+((_im)<<8)+(_opc))
#define _PACKED8(_sw,_zo)       ((((_sw)&1)<<7)|((_zo)&0x7F))
#define _AVOL			__asm__ volatile
#define _RCLB  /*example: #define _RCLB :"%rax" // don't forget the colon! */
#define _XP(str)		#str
#define _AS0(_of)		_XP(prefetchnta _of(%%rax,%%rax))
#define _JHERE			";jmp 1f;1:"
#define _XSM_XZONESTART(_i)	_XSMOPCODE_I(_COTSON_CUSTOM_TRACER,_PACKED8(0,_i))
#define _XSM_XZONESTOP(_i)	_XSMOPCODE_I(_COTSON_CUSTOM_TRACER,_PACKED8(1,_i))
#define XSM_XZONESTART(_z)	_AVOL (_AS0(_XSM_XZONESTART(_z)) _JHERE:: _RCLB)
#define XSM_XZONESTOP(_z)	_AVOL (_AS0(_XSM_XZONESTOP(_z))  _JHERE:: _RCLB)
#define COTSON_ARCH
/******************************************************************************/
/* XSM PERFORMANCE COUNTERS FUNCTIONS */
/******************************************************************************/
#ifdef COTSON_ARCH
    #define xzonestart(_zone) ({  double time; time=MPI_Wtime(); _XSM_XZONESTART(_zone); time;})
    #define xzonestop(_zone)  ({  _XSM_XZONESTOP(_zone); double time; time=MPI_Wtime();time;  })
    #warning "COTSON ARCHITECTURE DETECTED"
#else
    #define xzonestart(_zone) MPI_Wtime();
    #define xzonestop(_zone) MPI_Wtime();
#endif 

/******************************************************************************/

/*******************************************************************************
 * These definitions are modified by filter passed in mydse.sh
 * Plese not change the name!
 ******************************************************************************/
#define N 256
#define BLOCKSZ 8
#define XDATA_INT64

#ifdef XDATA_SINGLE
    #define DATA float
    #define MPI_DATATYPE MPI_FLOAT
#endif
#ifdef XDATA_DOUBLE
    #define DATA double
    #define MPI_DATATYPE MPI_DOUBLE
#endif
#ifdef XDATA_INT64
    #define DATA uint64_t
    #define MPI_DATATYPE MPI_UINT64_T
#endif
/***** mydse.sh parms *********************************************************/

/******************************************
 * MPI TEST 0.8.0
 *****************************************/
#define ERR_BADORDER    255
#define TAG_INIT      31337
#define TAG_RESULT       42
#define DISP_MAXORDER    12

#undef PRINT_MATRIX

/* Functions Prototipe */
int getRowCount(int rowsTotal, int mpiRank, int mpiSize);
int matrixMultiply(DATA *A, DATA *B, DATA *C, int n, int n_local);
void fill_matrix();

/* Global definitions */
#define NODE_MASTER 0
DATA *A, *B, *C;
uint64_t superchecksum;
int n = 0, n_ubound, n_local, n_sq;
int mpiRank = 0, mpiSize = 1;

/*** Standard Templte Functions ***/
void prepare(int *argc, char ***argv);
void compute();
void report();

/***** main start *************************************************************/
int main(int argc, char *argv[]) {
    double starttime=0,stoptime=0;
    prepare(&argc,&argv);
   
    starttime=xzonestart(1);
 
    compute();
   
    stoptime=xzonestop(1);
  
    if(mpiRank == NODE_MASTER)
    {
     printf("EXECUTION TIME: %f s\n",stoptime-starttime); fflush(stdout);
    }

    report();
    return 0;
}
/***** main end ***************************************************************/

void fill_matrix()
{
  int i,j;
  /****************************************************************/
  if (mpiRank == NODE_MASTER)
  {
    int cr, cc;
    srand(12345); // start always with same numbers in the matrices (optional)
    printf("[INFO]: Fill Matrix B...\n");fflush(stdout);

    for (i = 0; i < n; i++) {
        cr = 0;
        for (j = 0; j < n-1; j++) {
            int val1 = rand()&0xFF;
            cr = (cr + val1) &0xFF;
            B[i*n+j] = (DATA)val1;
        }
        B[i*n+n-1] = (DATA)cr;
    }

    printf("[INFO]: Fill Matrix A...\n");fflush(stdout);
    for (j = 0; j < n; j++) {
        cc = 0;
        for (i = 0; i < n-1; i++) {
            int val2 = rand()&0xFF;
            cc = (cc + val2) &0xFF;
            A[i*n+j] = (DATA)val2;
         }
         A[(n-1)*n+j] = (DATA)cc;
    }
    // Calculate the SuperCheckSum
    superchecksum = 0;
    for (i = 0; i < n; i++) {
        superchecksum =
           (superchecksum + (uint64_t)(A[(n-1)*n+i]) * (uint64_t)(B[i*n+n-1]))&0xFF;
    }
#ifdef VERBOSE
    printf("SUPER CHECKSUM: %ld - %ld\n",superchecksum,(superchecksum<<2)&0xFF);fflush(stdout);
#endif
    superchecksum=(superchecksum<<2)&0xFF;
  }
}
void prepare(int *argc, char ***argv)
{
  n=N;
  int i;
  MPI_Init(argc, argv);
  MPI_Comm_rank(MPI_COMM_WORLD, &mpiRank);
  MPI_Comm_size(MPI_COMM_WORLD, &mpiSize);
  if(mpiRank == NODE_MASTER) { printf("[INFO]: MATRIX N = %d\n",n);fflush(stdout); }
  #ifdef VERBOSE
      printf("[VERBOSE]: MPI RANK: %d\n",mpiRank);fflush(stdout);
      printf("[VERBOSE]: MPI SIZE: %d\n",mpiSize);fflush(stdout);
  #endif
  /*
   * Broadcasts a message from the process with rank 0 (master)
   * to all other processes of the group.
   */
  MPI_Bcast(&n, 1, MPI_INT, 0, MPI_COMM_WORLD);

  n_local  = getRowCount(n, mpiRank, mpiSize);
  n_ubound = n * n_local;  /* slave array's upper bound (partial matrix) */
  n_sq     = n * n;        /* master array's upper bound (total matrix) */

  // Allocate Memory
  A = (DATA *) malloc(sizeof(DATA) * ((mpiRank != NODE_MASTER) ? n_ubound : n_sq));
  B = (DATA *) malloc(sizeof(DATA) * n_sq );
  C = (DATA *) malloc(sizeof(DATA) * ((mpiRank != NODE_MASTER) ? n_ubound : n_sq));
  #ifdef VERBOSE
      printf("--- MEMORY ALLOCATION RANK: %d ---\n",mpiRank);fflush(stdout);
      printf("[VERBOSE] RANK:%d - sizeof(DATA) = %ld\n",mpiRank,sizeof(DATA));fflush(stdout);
      printf("[VERBOSE] RANK:%d A = %ld \t@ %p\n",mpiRank,(sizeof(DATA) * (mpiRank ? n_ubound : n_sq)),A);fflush(stdout);
      printf("[VERBOSE] RANK:%d B = %ld \t@ %p\n",mpiRank,(sizeof(DATA) * n_sq ),B);fflush(stdout);
      printf("[VERBOSE] RANK:%d C = %ld \t@ %p\n",mpiRank,(sizeof(DATA) * (mpiRank ? n_ubound : n_sq)),C);fflush(stdout);
  #endif
  /* Let each process initialize C to zero */
  for (i=0; i<n_ubound; i++) {
      C[i] = 0.0;
  }

  if(mpiRank == NODE_MASTER)
  {
    fill_matrix(mpiRank);
  }
}

void compute()
{
  int sizeSent, sizeToBeSent;
  int i;
  /* Send A by splitting it in row-wise parts */
  /* MASTER NODE */
  if (mpiRank == NODE_MASTER) {
      sizeSent = n_ubound;
      for (i=1; i<mpiSize; i++) {
          sizeToBeSent = n * getRowCount(n, i, mpiSize);
          MPI_Send(A + sizeSent, sizeToBeSent, MPI_DATATYPE, i, TAG_INIT,
                   MPI_COMM_WORLD);
          sizeSent += sizeToBeSent;
      }
  }
  /* SLAVE NODES */
  else { /* Receive parts of A */
      MPI_Recv(A, n_ubound, MPI_DATATYPE, 0, TAG_INIT, MPI_COMM_WORLD,
               MPI_STATUS_IGNORE);
  }

  /* Send B completely to each process */
  MPI_Bcast(B, n*n, MPI_DATATYPE, 0, MPI_COMM_WORLD);
  
  /* Let each process perform its own multiplications */
#ifdef VERBOSE
  printf("[DEBUG] RANK:%d make matrixMultiply() - n_local = %d\n",mpiRank,n_local);fflush(stdout);
#endif
  matrixMultiply(A, B, C, n, n_local);
  /* Receive partial results from each slave */
  if (!mpiRank) {
      sizeSent = n_ubound;
      for (i=1; i<mpiSize; i++) {
          sizeToBeSent = n * getRowCount(n, i, mpiSize);
          MPI_Recv(C + sizeSent, sizeToBeSent, MPI_DATATYPE, i, TAG_RESULT,
                   MPI_COMM_WORLD, MPI_STATUS_IGNORE);
          sizeSent += sizeToBeSent;
      }
  }
  else { /* Send partial results to master */
      MPI_Send(C, n_ubound, MPI_DATATYPE, 0, TAG_RESULT, MPI_COMM_WORLD);
  }
  MPI_Barrier(MPI_COMM_WORLD);
}

int matrixMultiply(DATA *a, DATA *b, DATA *c, int n, int n_local) {
    int i, j, k;
    for (i=0; i<n_local; i++) {
        for (j=0; j<n; j++) {
            for (k=0; k<n; k++) {
                c[i*n + j] += a[i*n + k] * b[k*n + j];
            }
        }
    }
    return 0;
}

void report()
{
  /* Print out the final results matrix (root process only) */
  if (mpiRank == NODE_MASTER)
  {
    uint64_t xchecksum = 0L;
    int i,j,n=N;
  #ifdef PRINT_MATRIX
      n_sq = n * n;
      for (i=0; i<n_sq; i++) {
          printf("%5.1lf\t", C[i]);fflush(stdout);
          if (i%n == 0) { printf("\n");fflush(stdout); }
      }
  #endif
    // Verify the SuperCheckSum
    printf("Check the result of Matrix Multiply...\n");fflush(stdout);
    for (i = 0; i < n; i++) {   // i = row pointer
        for (j = 0; j < n; j++) { // j = column pointer
            xchecksum = (xchecksum + (int)(C[i*n + j])) &0xFF;
        }
    }
  #ifdef VERBOSE
    printf("checksum: %ld - superchecksum: %ld\n", xchecksum, superchecksum);fflush(stdout);
  #endif
    if(xchecksum == superchecksum)
    {
      printf("*** %s ***\n","SUCCESS");fflush(stdout);
      printf("All workers done, goodbye\n");fflush(stdout);
    }
    else
    {
      printf("*** %s ***\n","FAILURE");fflush(stdout);
    }
  }
  /* Goodbye, world */
  MPI_Finalize();
}

int getRowCount(int rowsTotal, int mpiRank, int mpiSize) {
    /* Adjust slack of rows in case rowsTotal is not exactly divisible */
#ifdef VERBOSE
    printf("%s:%d %s rowTotal = %d - mpiRank = %d - mpiSize = %d - result = %d\n",fflush(stdout);
    __FILE__, __LINE__, __FUNCTION__,rowsTotal,mpiRank,mpiSize,(rowsTotal / mpiSize) + (rowsTotal % mpiSize > mpiRank));
#endif
    return (rowsTotal / mpiSize) + ((rowsTotal % mpiSize > mpiRank)?1:0);
}