/* -*- Mode: C; c-basic-offset:4 ; indent-tabs-mode:nil ; -*- */
/*  
 *  (C) 2007 by Argonne National Laboratory.
 *      See COPYRIGHT in top-level directory.
 */

/* a test to exercise very large extents: on most platforms with 32 bit
 * integers, we'd expect these tests to give unexpected values.  On platforms
 * with 64 bit integers, these tests will be fine.  On BlueGene we're not sure
 * yet :> 
 */ 


#include <mpi.h>
#include <stdint.h>
#include <math.h>
#include <stdio.h>

#define CHECK(fn) {int errcode; errcode = (fn); if (errcode != MPI_SUCCESS) handle_error(errcode, NULL); } 


static void handle_error(int errcode, char *str) 
{
	char msg[MPI_MAX_ERROR_STRING];
	int resultlen;
	MPI_Error_string(errcode, msg, &resultlen);
	fprintf(stderr, "%s: %s\n", str, msg);
	MPI_Abort(MPI_COMM_WORLD, 1);
}

static void typestats(MPI_Datatype type) 
{
    MPI_Aint lb, extent;
    MPI_Count size;

    MPI_Type_get_extent(type, &lb, &extent);
    MPI_Type_size_x(type, &size);

    printf("dtype %d: lb = %ld extent = %ld size = %ld...",
	    type, (long)lb, (long)extent, size);

}

static int verify_type(char *filename, MPI_Datatype type, 
	int64_t expected_extent, int do_coll)
{
    int rank, canary;
    MPI_Count tsize;
    int compare=-1;
    int errs=0, toterrs=0;
    MPI_Status status;
    MPI_File fh;

    MPI_Comm_rank(MPI_COMM_WORLD, &rank);

    CHECK( MPI_File_open(MPI_COMM_WORLD, filename, 
		MPI_MODE_CREATE|MPI_MODE_RDWR, MPI_INFO_NULL, &fh));
    CHECK( MPI_File_set_view(fh, rank*sizeof(int), 
	    MPI_BYTE, type, "native", MPI_INFO_NULL)); 

    MPI_Type_size_x(type, &tsize);

    canary=rank+1000000;

    /* skip over first instance of type */
    if (do_coll) {
	CHECK( MPI_File_write_at_all(fh, tsize, &canary, 1, MPI_INT, &status));
    } else {
	CHECK( MPI_File_write_at(fh, tsize, &canary, 1, MPI_INT, &status));
    }

    CHECK( MPI_File_set_view(fh, 0, MPI_INT, MPI_INT, "native", 
		MPI_INFO_NULL)); 

    if (do_coll) {
	CHECK( MPI_File_read_at_all(fh, expected_extent/sizeof(int)+rank, 
		&compare, 1, MPI_INT, &status));
    } else {
	CHECK( MPI_File_read_at(fh, expected_extent/sizeof(int)+rank, 
		&compare, 1, MPI_INT, &status));
    }

    if (compare != canary)
	errs=1;
    MPI_Allreduce(&errs, &toterrs, 1, MPI_INT, MPI_SUM, MPI_COMM_WORLD);

    MPI_File_close(&fh);

    if (toterrs) {
	printf("%d: got %d expected %d\n", rank, compare, canary);
	/* keep file if there's an error */
    } else {
	if (rank == 0) MPI_File_delete(filename, MPI_INFO_NULL);
    }
	
    return (toterrs);

}

static int testtype(char *filename, MPI_Datatype type, int64_t expected_extent)
{
    int rank, ret, errs=0;
    int collective=1, nocollective=0;

    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    if (!rank) typestats(type);

    ret = verify_type(filename, type, expected_extent, nocollective); 
    if (ret) {
	errs++;
	fprintf(stderr, "type %d failed indep\n", type);
    } else 
	if (!rank) printf("indep: OK ");

    ret = verify_type(filename, type, expected_extent, collective); 
    if (ret) {
	errs++;
	fprintf(stderr, "type %d failed collective\n", type);
    } else
	if (!rank) printf("coll: OK\n");

    return errs;
}

int main(int argc, char **argv)
{
    int count=2;
    int blocks[2];
    int disps[2];

    int ndims=2;
    int sizes[2];
    int subs[2];
    int starts[2];

    MPI_Datatype baseindex, indexed1G, indexed3G, indexed6G; 
    MPI_Datatype subarray1G, subarray3G, subarray6G;
    int ret, rank;

    MPI_Init(&argc, &argv);

    if (argc != 2) {
	fprintf(stderr, "usage: %s <filename>\n", argv[0]);
	MPI_Abort(MPI_COMM_WORLD, 1);
    }

    MPI_Comm_rank(MPI_COMM_WORLD, &rank);

    /* base type: 1MB indexed type of ints*/
    count = 2;
    blocks[0] = 1;
    disps[0] = 0;
    blocks[1] = 1;
    disps[1] = 1024*256-1;

    MPI_Type_indexed(count, blocks, disps, MPI_INT, &baseindex);
    /* simple case: 1GB extent */
    MPI_Type_contiguous(1024, baseindex, &indexed1G);
    MPI_Type_commit(&indexed1G);

    /* a little trickier: 3Gb extent */
    MPI_Type_contiguous(3072, baseindex, &indexed3G);
    MPI_Type_commit(&indexed3G);

    /* and finally 6GB extent */
    MPI_Type_contiguous(6144, baseindex, &indexed6G);
    MPI_Type_commit(&indexed6G);

    /* TODO: 
     * - add a darray test
     * - add a test with crazy extents */
    sizes[0] = 1024*16;
    sizes[1] = 1024*16;
    subs[0] = subs[1] = 256;
    starts[0] = starts[1] = 0;

    MPI_Type_create_subarray(ndims, sizes, subs, starts, 
	    MPI_ORDER_C, MPI_INT, &subarray1G);
    MPI_Type_commit(&subarray1G);

    sizes[1] = 1024*16*3;
    MPI_Type_create_subarray(ndims, sizes, subs, starts, 
	    MPI_ORDER_C, MPI_INT, &subarray3G);
    MPI_Type_commit(&subarray3G);

    sizes[1] = 1024*16*6;
    MPI_Type_create_subarray(ndims, sizes, subs, starts, 
	    MPI_ORDER_C, MPI_INT, &subarray6G);
    MPI_Type_commit(&subarray6G);

    /* assume command line arguments make it out to all processes */
    ret = testtype(argv[1], indexed1G, (int64_t)1024*1024*1024);

    ret = testtype(argv[1], indexed3G, (int64_t)1024*1024*1024*3);

    ret = testtype(argv[1], indexed6G, (int64_t)1024*1024*1024*6);

    ret = testtype(argv[1], subarray1G, (int64_t)1024*1024*1024);

    ret = testtype(argv[1], subarray3G, (int64_t)1024*1024*1024*3);

    ret = testtype(argv[1], subarray6G, (int64_t)1024*1024*1024*6);

    if(!ret && !rank) fprintf(stderr, "  No Errors\n");
    
    MPI_Finalize();
    return (-ret);

}
/* 
 * vim: ts=8 sts=4 sw=4 noexpandtab 
 */