/* -*- Mode: C; c-basic-offset:4 ; indent-tabs-mode:nil ; -*- */

/*

 *  (C) 2004 by Argonne National Laboratory.

 *      See COPYRIGHT in top-level directory.

 */

#include "mpi.h"

#include <stdio.h>

#include <stdlib.h>

#include "mpitest.h"

#include <assert.h>

/*

static char MTEST_Descrip[] = "Test MPI_Allreduce with non-commutative user-defined operations using matrix rotations";

*/

/* This example is similar to allred3.c, but uses only 3x3 matrics with

   integer-valued entries.  This is an associative but not commutative

   operation.

   The number of matrices is the count argument. The matrix is stored

   in C order, so that

     c(i,j) is cin[j+i*3]

   Three different matrices are used:

   I = identity matrix

   A = (1 0 0    B = (0 1 0

        0 0 1         1 0 0

        0 1 0)        0 0 1)

   The product

         I^k A I^(p-2-k-j) B I^j

   is

   (0 1 0

     0 0 1

     1 0 0)

   for all values of k, p, and j.

 */

void matmult(void *cinPtr, void *coutPtr, int *count, MPI_Datatype * dtype);

void matmult(void *cinPtr, void *coutPtr, int *count, MPI_Datatype * dtype)

{

    const int *cin = (const int *) cinPtr;

    int *cout = (int *) coutPtr;

    int i, j, k, nmat;

    int tempcol[3];

    int offset1, offset2;

    for (nmat = 0; nmat < *count; nmat++) {

        for (j = 0; j < 3; j++) {

            for (i = 0; i < 3; i++) {

                tempcol[i] = 0;

                for (k = 0; k < 3; k++) {

                    /* col[i] += cin(i,k) * cout(k,j) */

                    offset1 = k + i * 3;

                    offset2 = j + k * 3;

                    tempcol[i] += cin[offset1] * cout[offset2];

                }

            }

            for (i = 0; i < 3; i++) {

                offset1 = j + i * 3;

                cout[offset1] = tempcol[i];

            }

        }

        /* Advance to the next matrix */

        cin += 9;

        cout += 9;

    }

}

/* Initialize the integer matrix as one of the

   above matrix entries, as a function of count.

   We guarantee that both the A and B matrices are included.

*/

static void initMat(int rank, int size, int nmat, int mat[])

{

    int i, kind;

    /* Zero the matrix */

    for (i = 0; i < 9; i++) {

        mat[i] = 0;

    }

    /* Decide which matrix to create (I, A, or B) */

    if (size == 2) {

        /* rank 0 is A, 1 is B */

        kind = 1 + rank;

    } else {

        int tmpA, tmpB;

        /* Most ranks are identity matrices */

        kind = 0;

        /* Make sure exactly one rank gets the A matrix

         * and one the B matrix */

        tmpA = size / 4;

        tmpB = (3 * size) / 4;

        if (rank == tmpA)

            kind = 1;

        if (rank == tmpB)

            kind = 2;

    }

    switch (kind) {

        case 0:        /* Identity */

            mat[0] = 1;

            mat[4] = 1;

            mat[8] = 1;

            break;

        case 1:        /* A */

            mat[0] = 1;

            mat[5] = 1;

            mat[7] = 1;

            break;

        case 2:        /* B */

            mat[1] = 1;

            mat[3] = 1;

            mat[8] = 1;

            break;

    }

}

/* Compare a matrix with the known result */

static int checkResult(int nmat, int mat[], const char *msg)

{

    int n, k, errs = 0, wrank;

    static int solution[9] = { 0, 1, 0,

        0, 0, 1,

        1, 0, 0

    };

    MPI_Comm_rank(MPI_COMM_WORLD, &wrank);

    for (n = 0; n < nmat; n++) {

        for (k = 0; k < 9; k++) {

            if (mat[k] != solution[k]) {

                errs++;

                if (errs == 1) {

                    printf("Errors for communicators %s\n", MTestGetIntracommName());

                    fflush(stdout);

                }

                if (errs < 10) {

                    printf("[%d]matrix #%d(%s): Expected mat[%d,%d] = %d, got %d\n",

                           wrank, n, msg, k / 3, k % 3, solution[k], mat[k]);

                    fflush(stdout);

                }

            }

        }

        /* Advance to the next matrix */

        mat += 9;

    }

    return errs;

}

int main(int argc, char *argv[])

{

    int errs = 0;

    int size, rank;

    int minsize = 2, count;

    MPI_Comm comm;

    int *buf, *bufout;

    MPI_Op op;

    MPI_Datatype mattype;

    int i;

    MTest_Init(&argc, &argv);

    MPI_Op_create(matmult, 0, &op);

    /* A single rotation matrix (3x3, stored as 9 consequetive elements) */

    MPI_Type_contiguous(9, MPI_INT, &mattype);

    MPI_Type_commit(&mattype);

    /* Sanity check: test that our routines work properly */

    {

        int one = 1;

        buf = (int *) malloc(4 * 9 * sizeof(int));

        initMat(0, 4, 0, &buf[0]);

        initMat(1, 4, 0, &buf[9]);

        initMat(2, 4, 0, &buf[18]);

        initMat(3, 4, 0, &buf[27]);

        matmult(&buf[0], &buf[9], &one, &mattype);

        matmult(&buf[9], &buf[18], &one, &mattype);

        matmult(&buf[18], &buf[27], &one, &mattype);

        checkResult(1, &buf[27], "Sanity Check");

        free(buf);

    }

    while (MTestGetIntracommGeneral(&comm, minsize, 1)) {

        if (comm == MPI_COMM_NULL)

            continue;

        MPI_Comm_size(comm, &size);

        MPI_Comm_rank(comm, &rank;;

        for (count = 1; count < size; count++) {

            /* Allocate the matrices */

            buf = (int *) malloc(count * 9 * sizeof(int));

            if (!buf) {

                MPI_Abort(MPI_COMM_WORLD, 1);

            }

            bufout = (int *) malloc(count * 9 * sizeof(int));

            if (!bufout) {

                MPI_Abort(MPI_COMM_WORLD, 1);

            }

            for (i = 0; i < count; i++) {

                initMat(rank, size, i, &buf[i * 9]);

            }

            MPI_Allreduce(buf, bufout, count, mattype, op, comm);

            errs += checkResult(count, bufout, "");

            /* Try the same test, but using MPI_IN_PLACE */

            for (i = 0; i < count; i++) {

                initMat(rank, size, i, &bufout[i * 9]);

            }

            MPI_Allreduce(MPI_IN_PLACE, bufout, count, mattype, op, comm);

            errs += checkResult(count, bufout, "IN_PLACE");

            free(buf);

            free(bufout);

        }

        MTestFreeComm(&comm);

    }

    MPI_Op_free(&op);

    MPI_Type_free(&mattype);

    MTest_Finalize(errs);

    return MTestReturnValue(errs);

}