/* -*- Mode: C; c-basic-offset:4 ; indent-tabs-mode:nil ; -*- */
/*
 *  (C) 2012 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>
#include <string.h>

#define LAST_TEST 14
#define RMA_SIZE  2048
#define OFFSET_1  7
#define OFFSET_2  83
#define OFFSET_3  157

#define PUT_VAL 0xdcba97
#define ACC_VAL 10771134

/*
 * Additional tests for lock contention.  These are designed to exercise
 * some of the optimizations within MPICH, but all are valid MPI programs.
 * Tests structure includes
 *    lock local (must happen at this time since application can use load
 *                store after the lock)
 *    send message to partner
 *                                  receive message
 *                                  send ack
 *    receive ack
 *    Provide a delay so that
 *      the partner will see the
 *      conflict
 *                                  partner executes:
 *                                  lock         // Note: this may block
 *                                     rma operations (see below)
 *                                  unlock
 *
 *    unlock                        send back to partner
 *    receive from partner
 *    check for correct data
 *
 * The delay may be implemented as a ring of message communication; this
 * is likely to automatically scale the time to what is needed
 */

/* Define a datatype to be used with */
int stride = 11;
int veccount = 7;
MPI_Datatype vectype;
/* Define long RMA ops size */
int longcount = 512;
int medcount = 127;
int mednum = 4;

void RMATest(int i, MPI_Win win, int master, int *srcbuf, int srcbufsize, int *getbuf,
             int getbufsize);
int RMACheck(int i, int *buf, MPI_Aint bufsize);
int RMACheckGet(int i, MPI_Win win, int *getbuf, MPI_Aint getsize);
void RMATestInit(int i, int *buf, MPI_Aint bufsize);

int main(int argc, char *argv[])
{
    int errs = 0;
    MPI_Win win;
    int *rmabuffer = 0, *getbuf = 0;
    MPI_Aint bufsize = 0, getbufsize = 0;
    int master, partner, next, wrank, wsize, i;
    int ntest = LAST_TEST;
    int *srcbuf;

    MTest_Init(&argc, &argv);

    /* Determine who is responsible for each part of the test */
    MPI_Comm_rank(MPI_COMM_WORLD, &wrank);
    MPI_Comm_size(MPI_COMM_WORLD, &wsize);
    if (wsize < 3) {
        fprintf(stderr, "This test requires at least 3 processes\n");
        MPI_Abort(MPI_COMM_WORLD, 1);
    }

    master = 0;
    partner = 1;
    next = wrank + 1;
    if (next == partner)
        next++;
    if (next >= wsize) {
        next = 0;
        if (next == partner)
            next++;
    }

    /* Determine the last test to run (by default, run them all) */
    for (i = 1; i < argc; i++) {
        if (strcmp("-ntest", argv[i]) == 0) {
            i++;
            if (i < argc) {
                ntest = atoi(argv[i]);
            }
            else {
                fprintf(stderr, "Missing value for -ntest\n");
                MPI_Abort(MPI_COMM_WORLD, 1);
            }
        }
    }

    MPI_Type_vector(veccount, 1, stride, MPI_INT, &vectype);
    MPI_Type_commit(&vectype);

    /* Create the RMA window */
    bufsize = 0;
    if (wrank == master) {
        bufsize = RMA_SIZE;
        MPI_Alloc_mem(bufsize * sizeof(int), MPI_INFO_NULL, &rmabuffer);
    }
    else if (wrank == partner) {
        getbufsize = RMA_SIZE;
        getbuf = (int *) malloc(getbufsize * sizeof(int));
        if (!getbuf) {
            fprintf(stderr, "Unable to allocated %d bytes for getbuf\n", (int) getbufsize);
            MPI_Abort(MPI_COMM_WORLD, 1);
        }
    }
    srcbuf = malloc(RMA_SIZE * sizeof(*srcbuf));
    assert(srcbuf);

    MPI_Win_create(rmabuffer, bufsize * sizeof(int), sizeof(int), MPI_INFO_NULL,
                   MPI_COMM_WORLD, &win);

    /* Run a sequence of tests */
    for (i = 0; i <= ntest; i++) {
        if (wrank == master) {
            MTestPrintfMsg(0, "Test %d\n", i);
            /* Because this lock is local, it must return only when the
             * lock is acquired */
            MPI_Win_lock(MPI_LOCK_EXCLUSIVE, 0, master, win);
            RMATestInit(i, rmabuffer, bufsize);
            MPI_Send(MPI_BOTTOM, 0, MPI_INT, partner, i, MPI_COMM_WORLD);
            MPI_Send(MPI_BOTTOM, 0, MPI_INT, next, i, MPI_COMM_WORLD);
            MPI_Recv(MPI_BOTTOM, 0, MPI_INT, MPI_ANY_SOURCE, i, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            MPI_Win_unlock(master, win);
            MPI_Recv(MPI_BOTTOM, 0, MPI_INT, partner, i, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            errs += RMACheck(i, rmabuffer, bufsize);
        }
        else if (wrank == partner) {
            MPI_Recv(MPI_BOTTOM, 0, MPI_INT, master, i, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            MPI_Win_lock(MPI_LOCK_EXCLUSIVE, 0, master, win);
            RMATest(i, win, master, srcbuf, RMA_SIZE, getbuf, getbufsize);
            MPI_Win_unlock(master, win);
            errs += RMACheckGet(i, win, getbuf, getbufsize);
            MPI_Send(MPI_BOTTOM, 0, MPI_INT, master, i, MPI_COMM_WORLD);
        }
        else {
            MPI_Recv(MPI_BOTTOM, 0, MPI_INT, MPI_ANY_SOURCE, i, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            MPI_Send(MPI_BOTTOM, 0, MPI_INT, next, i, MPI_COMM_WORLD);
        }
    }

    if (rmabuffer) {
        MPI_Free_mem(rmabuffer);
    }
    if (getbuf) {
        free(getbuf);
    }
    free(srcbuf);
    MPI_Win_free(&win);
    MPI_Type_free(&vectype);

    MTest_Finalize(errs);
    MPI_Finalize();
    return MTestReturnValue(errs);
}

/* Perform the tests.
 *
 * The srcbuf must be passed in because the buffer must remain valid
 * until the subsequent unlock call. */
void RMATest(int i, MPI_Win win, int master, int *srcbuf, int srcbufsize, int *getbuf,
             int getbufsize)
{
    int j, k;
    int *source = srcbuf;
    assert(srcbufsize == RMA_SIZE);

    for (j = 0; j < srcbufsize; j++)
        source[j] = -j;

    switch (i) {
    case 0:    /* Single short put (1 word at OFFSET_1) */
        source[0] = PUT_VAL;
        MPI_Put(source, 1, MPI_INT, master, OFFSET_1, 1, MPI_INT, win);
        break;
    case 1:    /* Single short accumulate (1 word of value 17 at OFFSET_2) */
        source[0] = ACC_VAL;
        MPI_Accumulate(source, 1, MPI_INT, master, OFFSET_2, 1, MPI_INT, MPI_SUM, win);
        break;
    case 2:    /* Single short get (1 word at OFFSET_3) */
        getbuf[0] = -1;
        MPI_Get(getbuf, 1, MPI_INT, master, OFFSET_3, 1, MPI_INT, win);
        break;
    case 3:    /* Datatype single put (strided put) */
        for (j = 0; j < veccount; j++) {
            source[j * stride] = PUT_VAL + j;
        }
        MPI_Put(source, 1, vectype, master, OFFSET_1, 1, vectype, win);
        break;
    case 4:    /* Datatype single accumulate (strided acc) */
        for (j = 0; j < veccount; j++) {
            source[j * stride] = ACC_VAL + j;
        }
        MPI_Accumulate(source, 1, vectype, master, OFFSET_2, 1, vectype, MPI_SUM, win);
        break;
    case 5:    /* Datatype single get (strided get) */
        for (j = 0; j < veccount; j++) {
            getbuf[j] = -j;
        }
        MPI_Get(getbuf, 1, vectype, master, OFFSET_3, 1, vectype, win);
        break;
    case 6:    /* a few small puts (like strided put, but 1 word at a time) */
        for (j = 0; j < veccount; j++) {
            source[j * stride] = PUT_VAL + j;
        }
        for (j = 0; j < veccount; j++) {
            MPI_Put(source + j * stride, 1, MPI_INT, master,
                    OFFSET_1 + j * stride, 1, MPI_INT, win);
        }
        break;
    case 7:    /* a few small accumulates (like strided acc, but 1 word at a time) */
        for (j = 0; j < veccount; j++) {
            source[j * stride] = ACC_VAL + j;
        }
        for (j = 0; j < veccount; j++) {
            MPI_Accumulate(source + j * stride, 1, MPI_INT, master,
                           OFFSET_2 + j * stride, 1, MPI_INT, MPI_SUM, win);
        }
        break;
    case 8:    /* a few small gets (like strided get, but 1 word at a time) */
        for (j = 0; j < veccount; j++) {
            getbuf[j * stride] = -j;
        }
        for (j = 0; j < veccount; j++) {
            MPI_Get(getbuf + j * stride, 1, MPI_INT, master,
                    OFFSET_3 + j * stride, 1, MPI_INT, win);
        }
        break;
    case 9:    /* Single long put (OFFSET_1) */
        for (j = 0; j < longcount; j++)
            source[j] = j;
        MPI_Put(source, longcount, MPI_INT, master, OFFSET_1, longcount, MPI_INT, win);
        break;
    case 10:   /* Single long accumulate (OFFSET_2) */
        for (j = 0; j < longcount; j++)
            source[j] = j;
        MPI_Accumulate(source, longcount, MPI_INT, master,
                       OFFSET_2, longcount, MPI_INT, MPI_SUM, win);
        break;
    case 11:   /* Single long get (OFFSET_3) */
        for (j = 0; j < longcount; j++)
            getbuf[j] = -j;
        MPI_Get(getbuf, longcount, MPI_INT, master, OFFSET_3, longcount, MPI_INT, win);
        break;
    case 12:   /* a few long puts (start at OFFSET_1, medcount) */
        for (j = 0; j < mednum; j++) {
            for (k = 0; k < medcount; k++) {
                source[j * medcount + k] = j * 2 * medcount + k;
            }
            MPI_Put(source + j * medcount, medcount, MPI_INT, master,
                    OFFSET_1 + j * 2 * medcount, medcount, MPI_INT, win);
        }
        break;
    case 13:   /* a few long accumulates (start at OFFSET_2, medcount) */
        for (j = 0; j < mednum; j++) {
            for (k = 0; k < medcount; k++) {
                source[j * medcount + k] = ACC_VAL + j * 2 * medcount + k;
            }
            MPI_Accumulate(source + j * medcount, medcount, MPI_INT, master,
                           OFFSET_2 + j * 2 * medcount, medcount, MPI_INT, MPI_SUM, win);
        }
        break;
    case 14:   /* a few long gets (start at OFFSET_3, medcount) */
        for (j = 0; j < mednum; j++) {
            for (k = 0; k < medcount; k++) {
                getbuf[j * medcount + k] = -(j * medcount + k);
            }
            MPI_Get(getbuf + j * medcount, medcount, MPI_INT, master,
                    OFFSET_3 + j * 2 * medcount, medcount, MPI_INT, win);
        }
        break;
    }
}

int RMACheck(int i, int *buf, MPI_Aint bufsize)
{
    int j, k;
    int errs = 0;

    switch (i) {
    case 0:    /* Single short put (1 word at OFFSET_1) */
        if (buf[OFFSET_1] != PUT_VAL) {
            errs++;
            printf("case 0: value is %d should be %d\n", buf[OFFSET_1], PUT_VAL);
        }
        break;
    case 1:    /* Single short accumulate (1 word of value 17 at OFFSET_2) */
        if (buf[OFFSET_2] != ACC_VAL + OFFSET_2) {
            errs++;
            printf("case 1: value is %d should be %d\n", buf[OFFSET_2], ACC_VAL + OFFSET_2);
        }
        break;
    case 2:    /* Single short get (1 word at OFFSET_3) */
        /* See RMACheckGet */
        break;
    case 3:    /* Datatype single put (strided put) */
    case 6:    /* a few small puts (like strided put, but 1 word at a time) */
        /* FIXME: The conditional and increment are reversed below.  This looks
         * like a bug, and currently prevents the following test from running. */
        for (j = 0; j++; j < veccount) {
            if (buf[j * stride] != PUT_VAL + j) {
                errs++;
                printf("case %d: value is %d should be %d\n", i, buf[j * stride], PUT_VAL + j);
            }
        }
        break;
    case 4:    /* Datatype single accumulate (strided acc) */
    case 7:    /* a few small accumulates (like strided acc, but 1 word at a time) */
        /* FIXME: The conditional and increment are reversed below.  This looks
         * like a bug, and currently prevents the following test from running. */
        for (j = 0; j++; j < veccount) {
            if (buf[j * stride] != ACC_VAL + j + OFFSET_2 + j * stride) {
                errs++;
                printf("case %d: value is %d should be %d\n", i,
                       buf[j * stride], ACC_VAL + j + OFFSET_2 + j * stride);
            }
        }
        break;
    case 5:    /* Datatype single get (strided get) */
    case 8:    /* a few small gets (like strided get, but 1 word at a time) */
        /* See RMACheckGet */
        break;
    case 9:    /* Single long put (OFFSET_1) */
        for (j = 0; j < longcount; j++) {
            if (buf[OFFSET_1 + j] != j) {
                errs++;
                printf("case 9: value is %d should be %d\n", buf[OFFSET_1 + j], OFFSET_1 + j);
            }
        }
        break;
    case 10:   /* Single long accumulate (OFFSET_2) */
        for (j = 0; j < longcount; j++) {
            if (buf[OFFSET_2 + j] != OFFSET_2 + j + j) {
                errs++;
                printf("case 10: value is %d should be %d\n", buf[OFFSET_2 + j], OFFSET_2 + j + j);
            }
        }
        break;
    case 11:   /* Single long get (OFFSET_3) */
        /* See RMACheckGet */
        break;
    case 12:   /* a few long puts (start at OFFSET_1, medcount) */
        for (j = 0; j < mednum; j++) {
            for (k = 0; k < medcount; k++) {
                if (buf[OFFSET_1 + j * 2 * medcount + k] != j * 2 * medcount + k) {
                    errs++;
                    printf("case 12: value is %d should be %d\n",
                           buf[OFFSET_1 + j * 2 * medcount + k], j * 2 * medcount + k);
                }
            }
        }
        break;
    case 13:   /* a few long accumulates (start at OFFSET_2, medcount) */
        for (j = 0; j < mednum; j++) {
            for (k = 0; k < medcount; k++) {
                if (buf[OFFSET_2 + j * 2 * medcount + k] !=
                    OFFSET_2 + 2 * j * 2 * medcount + 2 * k + ACC_VAL) {
                    errs++;
                    printf("case 13: value is %d should be %d\n",
                           buf[OFFSET_2 + j * 2 * medcount + k],
                           OFFSET_2 + 2 * j * 2 * medcount + k + ACC_VAL);
                }
            }
        }
        break;
    case 14:   /* a few long gets (start at OFFSET_3, medcount) */
        /* See RMACheckGet */
        break;
    default:
        fprintf(stderr, "Unrecognized case %d\n", i);
        errs++;
        break;
    }
    return errs;
}

int RMACheckGet(int i, MPI_Win win, int *getbuf, MPI_Aint getsize)
{
    int errs = 0;
    int j, k;

    /* */
    switch (i) {
    case 0:    /* Single short put (1 word at OFFSET_1) */
        break;
    case 1:    /* Single short accumulate (1 word of value 17 at OFFSET_2) */
        break;
    case 2:    /* Single short get (1 word at OFFSET_3) */
        if (getbuf[0] != OFFSET_3) {
            errs++;
            printf("case 2: value is %d should be %d\n", getbuf[0], OFFSET_3);
        }
        break;
    case 3:    /* Datatype single put (strided put) */
        break;
    case 4:    /* Datatype single accumulate (strided acc) */
        break;
    case 5:    /* Datatype single get (strided get) */
    case 8:    /* a few small gets (like strided get, but 1 word at a time) */
        for (j = 0; j < veccount; j++) {
            if (getbuf[j * stride] != OFFSET_3 + j * stride) {
                errs++;
                printf("case %d: value is %d should be %d\n", i,
                       getbuf[j * stride], OFFSET_3 + j * stride);
            }
        }

        break;
    case 6:    /* a few small puts (like strided put, but 1 word at a time) */
        break;
    case 7:    /* a few small accumulates (like strided acc, but 1 word at a time) */
        break;
    case 9:    /* Single long put (OFFSET_1) */
        break;
    case 10:   /* Single long accumulate (OFFSET_2) */
        break;
    case 11:   /* Single long get (OFFSET_3) */
        for (j = 0; j < longcount; j++) {
            if (getbuf[j] != OFFSET_3 + j) {
                errs++;
                printf("case 11: value is %d should be %d\n", getbuf[j], OFFSET_3 + j);
            }
        }
        break;
    case 12:   /* a few long puts (start at OFFSET_1, medcount) */
        break;
    case 13:   /* a few long accumulates (start at OFFSET_2, medcount) */
        break;
    case 14:   /* a few long gets (start at OFFSET_3, medcount) */
        for (j = 0; j < mednum; j++) {
            for (k = 0; k < medcount; k++) {
                if (getbuf[j * medcount + k] != OFFSET_3 + j * 2 * medcount + k) {
                    errs++;
                    printf("case 14: buf[%d] value is %d should be %d\n",
                           j * medcount + k,
                           getbuf[j * medcount + k], OFFSET_3 + j * 2 * medcount + k);
                }
            }
        }
        break;
    default:
        fprintf(stderr, "Unrecognized case %d\n", i);
        errs++;
        break;
    }
    return errs;
}


void RMATestInit(int i, int *buf, MPI_Aint bufsize)
{
    int j;
    for (j = 0; j < bufsize; j++) {
        buf[j] = j;
    }
}