#include <stdio.h>
#include <stdlib.h>
#include <mpi.h>
#include <assert.h>

#define NITER (100000000)
#define NKV (20)

int main(int argc, char *argv[])
{
    int wrank, wsize;
    int i;
    int kv = MPI_KEYVAL_INVALID;
    int kvs[NKV] = { MPI_KEYVAL_INVALID };
    int x, *y;
    int flag;
    double start, t;

    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &wrank);
    MPI_Comm_size(MPI_COMM_WORLD, &wsize);

    /* put an attribute per MPI-2.2, example 16.17 (p.506) */
    MPI_Keyval_create(MPI_NULL_COPY_FN, MPI_NULL_DELETE_FN, &kv, NULL);
    x = 60439;
    MPI_Attr_put(MPI_COMM_WORLD, kv, &x);

    /* sanity check correctness */
    y = NULL;
    flag = 0;
    MPI_Attr_get(MPI_COMM_WORLD, kv, &y, &flag);
    assert(flag);
    assert(*y == x);

    /* test performance */
    start = MPI_Wtime();
    for (i = 0; i < NITER; i++) {
        y = NULL;
        flag = 0;
        MPI_Attr_get(MPI_COMM_WORLD, kv, &y, &flag);
    }
    t = MPI_Wtime() - start;

    printf("SINGLE ATTRIBUTE\n");
    printf("t= %f (us/iteration)\n", 1000000*(t/NITER));
    printf("t= %f (rough cycles/iteration @ 2.5 GHz)\n", ((t/NITER)*2.5e9));

    MPI_Attr_delete(MPI_COMM_WORLD, kv);
    MPI_Keyval_free(&kv);

    /* -------------------------------------------------------------------------- */

    /* now let's do it again but with NKV attributes attached to the comm */
    for (i = 0; i < NKV; i++) {
        MPI_Keyval_create(MPI_NULL_COPY_FN, MPI_NULL_DELETE_FN, &kvs[i], NULL);
        x = 60439;
        MPI_Attr_put(MPI_COMM_WORLD, kvs[i], &x);

        /* sanity check correctness */
        y = NULL;
        flag = 0;
        MPI_Attr_get(MPI_COMM_WORLD, kvs[i], &y, &flag);
        assert(flag);
        assert(*y == x);
    }

    /* now test performance of gets from kvs[0], which should be slowest
     * under MPICH2's linked list implementation */
    start = MPI_Wtime();
    for (i = 0; i < NITER; i++) {
        y = NULL;
        flag = 0;
        MPI_Attr_get(MPI_COMM_WORLD, kvs[0], &y, &flag);
    }
    t = MPI_Wtime() - start;

    printf("\n");
    printf("%d ATTRIBUTES\n", NKV);
    printf("t= %f (us/iteration)\n", 1000000*(t/NITER));
    printf("t= %f (rough cycles/iteration @ 2.5 GHz)\n", ((t/NITER)*2.5e9));

    for (i = 0; i < NKV; i++) {
        MPI_Attr_delete(MPI_COMM_WORLD, kvs[i]);
        MPI_Keyval_free(&kvs[i]);
    }

    MPI_Finalize();

    return 0;
}