/*
 * Driver for demonstrating MPI_Reduce_scatter performance problems on two 
 * racks of BG/P.
 *
 * Please contact Jack Poulson at <jack.poulson@gmail.com> for more information.
 */
#include "mpi.h"
#include "mpix.h"
#include <stdio.h>

const int num_tests = 20;

int MPIX_Get_last_algorithm(MPI_Comm comm, int *lastalgorithm);

/* Test from 100 doubles to 3000 doubles contributions in multiples of 100 */
void
TestReduceScatter( MPI_Comm comm, char* comm_name )
{
    int num_doubles;
    int comm_size, comm_rank, world_rank;
    int* rcs;
    double* send_buffer;
    double* recv_buffer;
    int result;

    /* Extract the communicator properties and our world rank */
    MPI_Comm_size( comm, &comm_size );
    MPI_Comm_rank( comm, &comm_rank );
    MPI_Comm_rank( MPI_COMM_WORLD, &world_rank );

    rcs = (int*)malloc(comm_size*sizeof(int));

    /* Loop over local contribution sizes */
    if( world_rank == 0 )
    {
        printf("Testing over %s\n", comm_name);
        printf("Local contribution   Megabits/second   Algorithm\n");
        printf("------------------------------------------------\n");
    }
    for( num_doubles=100; num_doubles<=2000; num_doubles+=100 )
    {
        int test;
        int rank;
        double megabits, megabits_per_second;
        double start_time, end_time, average_time;
        int last;

        megabits = num_doubles*comm_size*sizeof(double)*8/1.e6;
        for( rank=0; rank<comm_size; ++rank )
            rcs[rank] = num_doubles;

        send_buffer = (double*)malloc(num_doubles*comm_size*sizeof(double));
        recv_buffer = (double*)malloc(num_doubles*sizeof(double));

        start_time = MPI_Wtime();
        for( test=0; test<num_tests; ++test )
        {
            MPI_Reduce_scatter
            ( send_buffer, recv_buffer, rcs, MPI_DOUBLE, MPI_SUM, comm );
        }
        end_time = MPI_Wtime();
        average_time = (end_time-start_time)/num_tests;
        megabits_per_second = megabits/average_time;

        MPIX_Get_last_algorithm(comm, &last);
        if( world_rank == 0 )
            printf("%8d doubles     %8lf     %8ld \n", num_doubles, megabits_per_second, last);


        free( send_buffer );
        free( recv_buffer );
    }
    if( world_rank == 0 )
        printf("\n");
    free( rcs );
}

int
main( int argc, char* argv[] )
{
    int XYZT_size, XYZT_rank, XY_size, XY_rank, ZT_size, ZT_rank;
    int remainingDimensions[4];
    MPI_Comm XYZT_comm, XY_comm, ZT_comm;

    MPI_Init( &argc, &argv ); 
    MPIX_Cart_comm_create( &XYZT_comm ); 
    MPI_Comm_size( XYZT_comm, &XYZT_size );
    MPI_Comm_rank( XYZT_comm, &XYZT_rank );

    /* Create the XY subcommunicator */
    remainingDimensions[0] = 0;
    remainingDimensions[1] = 0;
    remainingDimensions[2] = 1; /* keep Y */
    remainingDimensions[3] = 1; /* keep X */
    MPI_Cart_sub( XYZT_comm, remainingDimensions, &XY_comm );
    MPI_Comm_rank( XY_comm, &XY_rank );
    MPI_Comm_size( XY_comm, &XY_size );

    /* Create the ZT subcommunicator */
    remainingDimensions[0] = 1; /* keep T */
    remainingDimensions[1] = 1; /* keep Z */
    remainingDimensions[2] = 0;
    remainingDimensions[3] = 0;
    MPI_Cart_sub( XYZT_comm, remainingDimensions, &ZT_comm );
    MPI_Comm_rank( ZT_comm, &ZT_rank );
    MPI_Comm_size( ZT_comm, &ZT_size );

    if( XYZT_rank == 0 )
    {
        printf("Communicator properties:  \n");
        printf("--------------------------\n");
        printf("(X,Y) size:     %d\n", XY_size);
        printf("(Z,T) size:     %d\n", ZT_size);
        printf("(X,Y,Z,T) size: %d\n", XYZT_size);
        printf("\n");
    }

    TestReduceScatter( XY_comm, "(X,Y)" );
    TestReduceScatter( ZT_comm, "(Z,T)" );
    TestReduceScatter( XYZT_comm, "(X,Y,Z,T)" );

    return 0;
}