chundoong-lab-ta/SamsungDS22/submissions/HW4/ssjoong.kim/mat_mul.cpp

#include "mat_mul.h"
#include "util.h"

#include <cstdio>
#include <cstdlib>
#include <mpi.h>

#define ITILESIZE (32)
#define JTILESIZE (1024)
#define KTILESIZE (1024)

static float *A, *B, *C;
static int M, N, K;
static int num_threads;
static int mpi_rank, mpi_world_size;

static void mat_mul_omp( int rows) {
  // TODO: parallelize & optimize matrix multiplication
  // Use num_threads per node
#if 0
  float Aik;
  int bs = 100;
  int k_max;
  
#pragma omp parallel for schedule(guided, 10)
  for (int kk = 0; kk < K; kk += bs)
  {
    if ((kk + bs) < K)
    {
      k_max = kk + bs;
    }
    else
    {
      k_max = K;
    }
    for (int i = 0; i < rows; ++i)
    {
      for (int k = kk; k < k_max; ++k)
      {
        Aik = A[i * K + k];
        for (int j = 0; j < N; ++j)
        {
          C[i * N + j] += Aik * B[k * N + j];
        }
      }
    }
  }

#else
  //printf("rows: %d\n", rows);

#pragma omp parallel for schedule(auto)
  for (int ii = 0; ii < rows; ii += ITILESIZE) {
    int i_min = std::min(ii + ITILESIZE, rows);
    for (int jj = 0; jj < N; jj += JTILESIZE) {
      int j_min = std::min(jj + JTILESIZE, N);
      for (int kk = 0; kk < K; kk += KTILESIZE) {
        int k_min = std::min(kk + KTILESIZE, K);
        for (int k = kk; k < k_min; k++) {
          for (int i = ii; i < i_min; i++) {
            float ar = A[i * K + k];
            for (int j = jj; j < j_min; j+=1) {
              C[i * N + j] += ar * B[k * N + j];
            }
          }
        }

      }
    }
  }
#endif
}

void mat_mul(float *_A, float *_B, float *_C, int _M, int _N, int _K,
             int _num_threads, int _mpi_rank, int _mpi_world_size) {
  A = _A, B = _B, C = _C;
  M = _M, N = _N, K = _K;
  num_threads = _num_threads, mpi_rank = _mpi_rank,
  mpi_world_size = _mpi_world_size;

  // TODO: parallelize & optimize matrix multiplication on multi-node
  // You must allocate & initialize A, B, C for non-root processes
  MPI_Status status;

  int source, dest, averow, rows, offset;
  //int i, j, k;
  int i;
  int remainPart;

  if (mpi_rank == 0)
  {
    //printf("rank_count : %d\n", mpi_world_size);
    averow = M / (mpi_world_size);
    remainPart = M % (mpi_world_size);
    offset = (remainPart != 0) ? offset + averow + 1 : offset + averow;

    for (dest = 1; dest < mpi_world_size; dest++)
    {
      //remainPart = (remainPart != 0) ? remainPart + 1 : remainPart;
      rows = (dest < remainPart) ? averow + 1 : averow;

      MPI_Send (&offset, 1, MPI_INT, dest, 1, MPI_COMM_WORLD);
      MPI_Send (&rows, 1, MPI_INT, dest, 1, MPI_COMM_WORLD);
      MPI_Send (A+(offset*K), rows*K, MPI_FLOAT, dest, 1, MPI_COMM_WORLD);
      MPI_Send (B, K*N, MPI_FLOAT, dest, 1, MPI_COMM_WORLD);

      offset = offset + rows;
#if 0
      if (remainPart > 0)
      {
        originalRows = rows;
        ++rows;
        remainPart--;
        MPI_Send (&offset, 1, MPI_INT, dest, 1, MPI_COMM_WORLD);
        MPI_Send (&rows, 1, MPI_INT, dest, 1, MPI_COMM_WORLD);
        MPI_Send (&A[offset*M], rows * M, MPI_FLOAT, dest, 1, MPI_COMM_WORLD);
        MPI_Send (&B, K*N, MPI_FLOAT, dest, 1, MPI_COMM_WORLD);

        offset = offset + rows;
        rows = originalRows;
      }
      else
      {
        MPI_Send (&offset, 1, MPI_INT, dest, 1, MPI_COMM_WORLD);
        MPI_Send (&rows, 1, MPI_INT, dest, 1, MPI_COMM_WORLD);
        MPI_Send (&A[offset*M], rows * M, MPI_FLOAT, dest, 1, MPI_COMM_WORLD);
        MPI_Send (&B, K*N, MPI_FLOAT, dest, 1, MPI_COMM_WORLD);

        offset = offset + rows;
      }
#endif
    }

    rows = (remainPart != 0) ? averow + 1 : averow;
    mat_mul_omp (rows);
  }

  if (mpi_rank > 0)
  {

    source = 0;
    MPI_Recv(&offset, 1, MPI_INT, source, 1, MPI_COMM_WORLD, &status);
    MPI_Recv(&rows, 1, MPI_INT, source, 1, MPI_COMM_WORLD, &status);

    alloc_mat (&A, rows, K);
    alloc_mat (&B, K, N);
    alloc_mat (&C, rows, N);
    zero_mat (C, rows, N);

    MPI_Recv(A, rows*K, MPI_FLOAT, source, 1, MPI_COMM_WORLD, &status);
    MPI_Recv(B, K*N, MPI_FLOAT, source, 1, MPI_COMM_WORLD, &status);  

    mat_mul_omp(rows);
  }

  MPI_Barrier(MPI_COMM_WORLD);

  if(mpi_rank == 0)
    for (i = 1; i < mpi_world_size; i++)
    {
      source = i;
      MPI_Recv (&offset, 1, MPI_INT, source, 2, MPI_COMM_WORLD, &status);
      MPI_Recv (&rows, 1, MPI_INT, source, 2, MPI_COMM_WORLD, &status);
      MPI_Recv (C+(offset*N), rows*N, MPI_INT, source, 2, MPI_COMM_WORLD, &status);
    }

  if(mpi_rank > 0)
  {
    MPI_Send(&offset, 1, MPI_INT, 0, 2, MPI_COMM_WORLD);
    MPI_Send(&rows, 1, MPI_INT, 0, 2, MPI_COMM_WORLD);
    MPI_Send(C, rows*N, MPI_FLOAT, 0, 2, MPI_COMM_WORLD);
  }
#if 0
  data_count = K*N;
  if (mpi_rank == 0) {
    for(int dest = 1; dest < mpi_world_size; dest++) {
      MPI_Send(
          /* data         = */ &B, 
          /* count        = */ data_count,
          /* datatype     = */ MPI_FLOAT, 
          /* destination  = */ dest, 
          /* tag          = */ 1, 
          /* communicator = */ MPI_COMM_WORLD);
    }
    printf("Process %d send data B[%d]: %f\n", mpi_rank, data_count, B[0]);
    printf("Process %d send data B[%d]: %f\n", mpi_rank, data_count, B[data_count-1]);
  } else {
    MPI_Recv(
        /* data         = */ &B, 
        /* count        = */ data_count, 
        /* datatype     = */ MPI_FLOAT, 
        /* source       = */ 0, 
        /* tag          = */ 1, 
        /* communicator = */ MPI_COMM_WORLD, 
        /* status       = */ MPI_STATUS_IGNORE);
    printf("Process %d received data B[%d]: %f\n", mpi_rank, data_count, B[0]);
    printf("Process %d received data B[%d]: %f\n", mpi_rank, data_count, B[data_count-1]);
  }
#endif
  // FIXME: for now, only root process runs the matrix multiplication.
  //if (mpi_rank == 0)
  //  mat_mul_omp();
}