source: /cluster/svnroot/bccd-ng/bw-institute/2016/day06/matmul/mpi.c @ 5713

/*
 * $Id: mpi.c 5713 2016-05-27 01:06:37Z skylar $
 */

#include "matmul.h"
#include "mpi.h"

#define FIRST_RANK 0

int main(int argc,char **argv) {
    int rank,size,c;
    unsigned int start_row,stop_row,stride;
    bool print;
    struct matrix m1,m2,local_dst_m,dst_m;

    MPI_Init(&argc,&argv);
    MPI_Comm_size(MPI_COMM_WORLD,&size);
    MPI_Comm_rank(MPI_COMM_WORLD,&rank);

#ifdef DEBUG
    fprintf(stderr,"Rank %d/%d - Hello World\n",rank,size);
#endif

    m1.rows = m1.cols = m2.rows = m2.cols = 0;

    while((c = getopt(argc,argv, "x:y:a:b:p")) != -1) {
        switch(c) {
            case 'x':
                m1.cols = atoi(optarg);
                break;
            case 'y':
                m1.rows = atoi(optarg);
                break;
            case 'a':
                m2.cols = atoi(optarg);
                break;
            case 'b':
                m2.rows = atoi(optarg);
                break;
            case 'p':
                print = true;
                break;
            case '?':
                usage();
                exit(EXIT_FAILURE);
        }
    }

    if(m1.rows == 0 || m1.cols == 0 || m2.rows == 0 || m2.cols == 0) {
        fprintf(stderr,"Supply row and column counts!\n");
        usage();
        exit(EXIT_FAILURE);
    }

    if(m1.cols != m2.rows) {
        fprintf(stderr,"Invalid matrix dimensions!\n");
        exit(EXIT_FAILURE);
    }


    m1.matrix       = safe_malloc_int(
            m1.rows*m1.cols,
            "Allocating first matrix"
            );
    m2.matrix       = safe_malloc_int(
            m2.rows*m2.cols,
            "Allocating second matrix"
            );

    if(rank == FIRST_RANK) {
        // Each thread will get a separate random seed
        unsigned int *random_seeds = init_random_seeds();
        init_matrix(&m1,random_seeds);
        init_matrix(&m2,random_seeds);

        // Allocate full destination matrix on first rank
        // Will be populated via MPI_Gather
        dst_m.rows = m1.rows;
        dst_m.cols = m2.cols;
        dst_m.matrix    = safe_malloc_int(
                dst_m.rows*dst_m.cols,
                "Allocating destination matrix"
                );

        if(print) {
            puts("Matrix 1\n");
            print_matrix(&m1);
            puts("");
            puts("Matrix 2\n");
            print_matrix(&m2);
            puts("");
        }
    }

    // Broadcast each element of the structs, to avoid the complexity
    // of creating custom data types
    MPI_Bcast(&m1.rows,1,MPI_INT,FIRST_RANK,MPI_COMM_WORLD);
    MPI_Bcast(&m1.cols,1,MPI_INT,FIRST_RANK,MPI_COMM_WORLD);
    MPI_Bcast(&m2.rows,1,MPI_INT,FIRST_RANK,MPI_COMM_WORLD);
    MPI_Bcast(&m2.cols,1,MPI_INT,FIRST_RANK,MPI_COMM_WORLD);
    MPI_Bcast(m1.matrix,m1.rows*m1.cols,MPI_INT,FIRST_RANK,MPI_COMM_WORLD);
    MPI_Bcast(m2.matrix,m2.rows*m2.cols,MPI_INT,FIRST_RANK,MPI_COMM_WORLD);

    // Calculate row offset in product matrix to start and stop calculation
    stride                  = m1.rows/size;
    start_row               = rank*stride;

    // Assign an even number of rows to each rank, except for the last rank
    // whiich gets the remainder
    if(rank<(size-1)) {
        stop_row = stride*(rank+1);
    }
    else {
        stop_row = m1.rows;
    }

    // Only need to allocate enough for the local computation
    local_dst_m.cols        = m2.cols;
    local_dst_m.rows        = (stop_row-start_row),
    local_dst_m.matrix      = safe_malloc_int(
            local_dst_m.rows*local_dst_m.cols,
            "Allocating local destination matrix"
            );

#ifdef DEBUG
    fprintf(stderr,"Rank %d - Stride %u, start_row %u, stop_row %u\n",
            rank,stride,start_row,stop_row);
#endif
    matmul(&m1,&m2,&local_dst_m,start_row,stop_row);

#ifdef DEBUG
    fprintf(stderr,"Rank %d local destination matrix:\n",rank);
    print_matrix(&local_dst_m);
#endif

    MPI_Gather(
            &local_dst_m.matrix,
            (local_dst_m.cols*local_dst_m.rows),
            MPI_INT,
            &dst_m.matrix,
            (dst_m.cols*dst_m.rows),
            MPI_INT,
            FIRST_RANK,
            MPI_COMM_WORLD
            );

    MPI_Finalize();

    free(m1.matrix);
    free(m2.matrix);
    free(local_dst_m.matrix);

    exit(EXIT_SUCCESS);
}