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

/*

 *  (C) 2003 by Argonne National Laboratory.

 *      See COPYRIGHT in top-level directory.

 */

#include "mpi.h"

#include "mpitestconf.h"

#ifdef HAVE_IOSTREAM

// Not all C++ compilers have iostream instead of iostream.h

#include <iostream>

#ifdef HAVE_NAMESPACE_STD

// Those that do often need the std namespace; otherwise, a bare "cout"

// is likely to fail to compile

using namespace std;

#endif

#else

#include <iostream.h>

#endif

#include "mpitestcxx.h"

#include <assert.h>

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

/* We make the error count global so that we can easily control the output

   of error information (in particular, limiting it after the first 10 

   errors */

int errs = 0;

/* This implements a simple matrix-matrix multiply.  This is an associative

   but not commutative operation.  The matrix size is set in matSize;

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

   in C order, so that

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

 */

#define MAXCOL 256

static int matSize = 0;  /* Must be < MAXCOL */

static int max_offset = 0;

void uop( const void *cinPtr, void *coutPtr, int count, const MPI::Datatype &dtype )

{

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

    int *cout = (int *)coutPtr;

    int i, j, k, nmat;

    int tempcol[MAXCOL];

    int offset1, offset2;

    int matsize2 = matSize*matSize;

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

	for (j=0; j

	    for (i=0; i

		tempcol[i] = 0;

		for (k=0; k

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

		    offset1    = k+i*matSize;

		    offset2    = j+k*matSize;

		    assert(offset1 < max_offset);

		    assert(offset2 < max_offset);

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

		}

	    }

	    for (i=0; i

		offset1       = j+i*matSize;

		assert(offset1 < max_offset);

		cout[offset1] = tempcol[i];

	    }

	}

	cin  += matsize2;

	cout += matsize2;

    }

}

/* Initialize the integer matrix as a permutation of rank with rank+1.

   If we call this matrix P_r, we know that product of P_0 P_1 ... P_{size-2}

   is the the matrix representing the permutation that shifts left by one.

   As the final matrix (in the size-1 position), we use the matrix that

   shifts RIGHT by one

*/   

static void initMat( MPI::Intracomm comm, int mat[] )

{

    int i, j, size, rank;

    int offset;

    rank = comm.Get_rank();

    size = comm.Get_size();

    for (i=0; i

	assert(i < max_offset);

	mat[i] = 0;

    }

    if (rank < size-1) {

	/* Create the permutation matrix that exchanges r with r+1 */

	for (i=0; i

	    if (i == rank) {

		offset = ((i+1)%size) + i * size;

		assert(offset < max_offset);

		mat[offset] = 1;

	    }

	    else if (i == ((rank + 1)%size)) {

		offset = ((i+size-1)%size) + i * size;

		assert(offset < max_offset);

		mat[offset] = 1;

	    }

	    else {

		offset = i+i*size;

		assert(offset < max_offset);

		mat[offset] = 1;

	    }

	}

    }

    else {

	/* Create the permutation matrix that shifts right by one */

	for (i=0; i

	    for (j=0; j

		offset = j + i * size;  /* location of c(i,j) */

		mat[offset] = 0;

		if ( ((j-i+size)%size) == 1 ) mat[offset] = 1;

	    }

	}

    }

}

/* Compare a matrix with the identity matrix */

static int isIdentity( MPI::Intracomm comm, int mat[] )

{

    int i, j, size, rank, lerrs = 0;

    int offset;

    rank = comm.Get_rank();

    size = comm.Get_size();

    for (i=0; i

	for (j=0; j

	    if (i == j) {

		offset = j+i*size;

		assert(offset < max_offset);

		if (mat[offset] != 1) {

		    lerrs++;

		    if (errs + lerrs< 10) {

			cerr << "[" << rank << "] mat[" << i << "," << j << "] = " << mat[offset] << ", expected 1 for comm " << MTestGetIntracommName() << endl;

		    }

		}

	    }

	    else {

		offset = j+i*size;

		assert(offset < max_offset);

		if (mat[offset] != 0) {

		    lerrs++;

		    if (errs + lerrs< 10) {

			cerr << "[" << rank << "] mat[" << i << "," << j << "] = " << mat[offset] << ", expected 0 for comm " << MTestGetIntracommName() << endl;

		    }

		}

	    }

	}

    }

    return lerrs;

}

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

{

    int size;

    int minsize = 2, count; 

    MPI::Intracomm      comm;

    int *buf, *bufout;

    MPI::Op op;

    MPI::Datatype mattype;

    MTest_Init();

    op.Init(uop, false);

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

	if ((MPI::Intracomm)comm == (MPI::Intracomm)MPI_COMM_NULL) {

	    continue;

	}

	size = comm.Get_size();

	matSize = size;

	/* Only one matrix for now */

	count = 1;

	/* A single matrix, the size of the communicator */

	mattype = MPI::INT.Create_contiguous( size*size );

	mattype.Commit();

	max_offset = count * size * size;

	buf = new int[max_offset];

	if (!buf) {

	    MPI::COMM_WORLD.Abort( 1 );

	}

	bufout = new int[max_offset];

	if (!bufout) {

	    MPI::COMM_WORLD.Abort( 1 );

	}

	initMat( comm, buf );

	comm.Allreduce( buf, bufout, count, mattype, op );

	errs += isIdentity( comm, bufout );

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

	initMat( comm, bufout );

	comm.Allreduce( MPI_IN_PLACE, bufout, count, mattype, op );

	errs += isIdentity( comm, bufout );

	delete [] buf;

	delete [] bufout;

	mattype.Free();

	MTestFreeComm( comm );

    }

    op.Free();

    MTest_Finalize( errs );

    MPI::Finalize();

    return 0;

}