testckpt1.c

/* testckpt1.c (2014-04-09) */

/* K-Means program for testing checkpoint restart.

   This program can be run like the following.

   $ mpirun -np 4 ./a.out -c 2 -p 10

   It will be aborted at the end of the second iteration of k-means
   algorithm. There will be two types of checkpoint files.

   1. checkpoint files generated by KMR stored in 'ckptdir0000*'
      directories.
   2. a checkpoint file that stores k-means progress.
      Its file name is 'kmeans_ckpt'.

   You can restart the program from the state stored in the above
   checkpoint files. In this case, you don't need to specify
   parameters.  If you do, they will be ignored.

   $ mpirun -np 4 ./a.out

   The program will be completed after you run it three times.

   When run with the following parameters, this program will perform
   answer checking.

   $ mpirun -np 2 ./a.out -c 4 -p 10
   $ mpirun -np 4 ./a.out -c 2 -p 10
   $ mpirun -np 4 ./a.out -c 4 -p 10
   $ mpirun -np 8 ./a.out -c 4 -p 10

*/

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/time.h>
#include <limits.h>
#include <math.h>
#include <mpi.h>
#include "kmr.h"

#define DEF_NUM_ITERATION 10
#define DEF_NUM_POINTS    10000
#define DEF_NUM_MEANS     100
#define DEF_DIM           3
#define DEF_GRID_SIZE     1000

#define KMEANS_CKPT_FILE     "./kmeans_ckpt"
#define KMEANS_CKPT_FILE_OLD "./kmeans_ckpt.bk"

struct kmr_option kmr_inspect = { .inspect = 1 };

typedef struct {
    // number of processes
    int nprocs;
    // number of k-means iteration
    int n_iteration;
    // current iteration number
    int c_iteration;
    // size of grid
    int grid_size;
    // dimention of points
    int dim;
    // number of points
    int n_points;
    // points
    int *points;
    // number of means
    int n_means;
    // means
    int *means;
} kmeans_t;


/*
 * Ths following code for random number is copied from C standard
 * chapter 7.20.2.
 *
 * http://www.open-std.org/jtc1/sc22/wg14/www/docs/n1256.pdf
 */
static unsigned long int _rand_next = 1;

static int _rand(void)
{
    _rand_next = _rand_next * 1103515245 + 12345;
    return (unsigned int)(_rand_next/65536) % 32768;
}

static void _srand(unsigned int seed)
{
    _rand_next = seed;
}


static int
measure_time(struct timeval *tv)
{
    MPI_Barrier(MPI_COMM_WORLD);
    if (gettimeofday(tv, NULL) == -1) {
        perror("gettimeofday");
        return -1;
    }
    return 0;
}

static double
calc_time_diff(struct timeval *tv_s, struct timeval *tv_e)
{
    return ((double)tv_e->tv_sec - (double)tv_s->tv_sec)
        + ((double)tv_e->tv_usec - (double)tv_s->tv_usec) /1000000.0;
}

static _Bool
is_restart()
{
    _Bool result = 0;
    int ret = access(KMEANS_CKPT_FILE, R_OK);
    if (ret == 0) {
        result = 1;
    }
    return result;
}

/* Parse commandline arguments */
static void
parse_args(int argc, char **argv, kmeans_t *kmeans)
{
    int c;
    extern char *optarg;

    while ((c = getopt(argc, argv, "i:d:c:p:s:")) != EOF) {
        switch (c) {
        case 'i':
            kmeans->n_iteration = atoi(optarg);
            break;
        case 'd':
            kmeans->dim = atoi(optarg);
            break;
        case 'c':
            kmeans->n_means = atoi(optarg);
            break;
        case 'p':
            kmeans->n_points = atoi(optarg);
            break;
        case 's':
            kmeans->grid_size = atoi(optarg);
            break;
        case '?':
            printf("Usage: %s -i <num iteration> -d <vector dimension> "
                   "-c <num clusters> -p <num points> -s <max value>\n",
                   argv[0]);
            MPI_Abort(MPI_COMM_WORLD, 1);
        }
    }

    if (kmeans->n_iteration <= 0 || kmeans->dim <= 0 || kmeans->n_means <= 0
        || kmeans->n_points <= 0 || kmeans->grid_size <= 0) {
        printf("Illegal argument value. All values must be numeric and "
               "greater than 0\n");
        MPI_Abort(MPI_COMM_WORLD, 1);
    }
}

static int
save_kmeans_state(kmeans_t *kmeans)
{
    const char *tmp_file = KMEANS_CKPT_FILE ".tmp";
    FILE *fp = fopen(tmp_file, "w+");
    if (fp == NULL) {
        fprintf(stderr, "Failed to open kmeans state file.\n");
        return 1;
    }
    size_t size = sizeof(kmeans_t);
    size_t ret = fwrite(&size, sizeof(size_t), 1, fp);
    if (ret != 1) {
        fprintf(stderr, "Failed to write size of kmeans data.\n");
        return 1;
    }
    ret = fwrite(kmeans, size, 1, fp);
    if (ret != 1) {
        fprintf(stderr, "Failed to write kmeans data.\n");
        return 1;
    }
    size = sizeof(int) * (size_t)kmeans->n_means * (size_t)kmeans->dim;
    ret = fwrite(&size, sizeof(size_t), 1, fp);
    if (ret != 1) {
        fprintf(stderr, "Failed to write size of means array.\n");
        return 1;
    }
    ret = fwrite(kmeans->means, size, 1, fp);
    if (ret != 1) {
        fprintf(stderr, "Failed to write means array.\n");
        return 1;
    }
    fflush(fp);
    fsync(fileno(fp));
    fclose(fp);

    rename(KMEANS_CKPT_FILE, KMEANS_CKPT_FILE_OLD);
    rename(tmp_file, KMEANS_CKPT_FILE);
    return 0;
}

static int
load_kmeans_state(kmeans_t *kmeans)
{
    FILE *fp = fopen(KMEANS_CKPT_FILE, "r");
    if (fp == NULL) {
        fprintf(stderr, "Failed to open kmeans state file.\n");
        return 1;
    }
    size_t size;
    size_t ret = fread(&size, sizeof(size_t), 1, fp);
    if (ret != 1) {
        fprintf(stderr, "Failed to read size of kmeans data.\n");
        return 1;
    }
    ret = fread(kmeans, size, 1, fp);
    if (ret != 1) {
        fprintf(stderr, "Failed to read kmeans data.\n");
        return 1;
    }
    kmeans->points = NULL;
    ret = fread(&size, sizeof(size_t), 1, fp);
    if (ret != 1) {
        fprintf(stderr, "Failed to read size of means array.\n");
        return 1;
    }
    if (size != (size_t)kmeans->n_means * (size_t)kmeans->dim * sizeof(int)) {
        fprintf(stderr, "The size of means array is mismatch.\n");
        return 1;
    }
    kmeans->means = (int*)malloc(size);
    ret = fread(kmeans->means, size, 1, fp);
    if (ret != 1) {
        fprintf(stderr, "Failed to read kmeans data.\n");
        return 1;
    }
    return 0;
}

static void
delete_kmeans_state()
{
    remove(KMEANS_CKPT_FILE_OLD);
    remove(KMEANS_CKPT_FILE);
}

/* Generate the points. */
static void
generate_randoms(int *pts, int size, int max_val)
{
    int i;
    for (i = 0; i < size; i++)
        pts[i] = (_rand() % max_val);
}

/* KMR map function
   It copies multi-dimension points to KVS.
*/
static int
generate_points(const struct kmr_kv_box kv,
                const KMR_KVS *kvi, KMR_KVS *kvo, void *p, long i_)
{
    int i;
    kmeans_t *kmeans = (kmeans_t *)p;
    int *points = kmeans->points;

    for (i = 0; i < kmeans->n_points * kmeans->dim; i += kmeans->dim) {
        struct kmr_kv_box nkv = { .klen = (int)sizeof(long),
                                  .vlen = kmeans->dim * (int)sizeof(int),
                                  .k.i  = i,
                                  .v.p  = (void *)&points[i] };
        kmr_add_kv(kvo, nkv);
    }

    return MPI_SUCCESS;
}

/* calculate squared distance
 */
static unsigned int
calc_sq_dist(int *v1, int *v2, int dim)
{
    int i;
    unsigned int sum = 0;
    for (i = 0; i < dim; i++)
        sum += (unsigned int)((v1[i] - v2[i]) * (v1[i] - v2[i]));
    return sum;
}

/* KMR map function
   It calculates cluster of a point stored in kv.
   It emits a Key-Value whose key is cluster id (index of array
   "kmeans.means") and value is the point.
*/
static int
calc_cluster(const struct kmr_kv_box kv,
             const KMR_KVS *kvi, KMR_KVS *kvo, void *p, long i_)
{
    int i;
    kmeans_t *kmeans = (kmeans_t *)p;
    int dim    = kmeans->dim;
    int *means = kmeans->means;
    int n_means = kmeans->n_means;
    int *point = (int *)kv.v.p;
    int min_idx = 0;
    unsigned int min_dist = calc_sq_dist(point, &means[0], dim);

    for (i = 1; i < n_means; i++) {
        unsigned int dist = calc_sq_dist(point, &means[i * dim], dim);
        if (dist < min_dist) {
            min_idx  = i;
            min_dist = dist;
        }
    }
    struct kmr_kv_box nkv = { .klen = (int)sizeof(long),
                              .vlen = dim * (int)sizeof(int),
                              .k.i  = min_idx,
                              .v.p  = (void *)point };
    kmr_add_kv(kvo, nkv);

    return MPI_SUCCESS;
}

/* KMR reduce function
   It calculates center of clusters.
   It emits a Key-Value whose key is cluster id and value is the new center.
*/
static int
update_cluster(const struct kmr_kv_box kv[], const long n,
               const KMR_KVS *kvi, KMR_KVS *kvo, void *p)
{
    int i, j;
    int cid = (int)kv[0].k.i;
    kmeans_t *kmeans = (kmeans_t *)p;
    int dim = kmeans->dim;
    long sum[dim];
    int average[dim];

    for (i = 0; i < dim; i++)
        sum[i] = 0;
    for (i = 0; i < n; i++) {
        int *point = (int *)kv[i].v.p;
        for (j = 0; j < dim; j++) {
            sum[j] += point[j];
        }
    }
    for (i = 0; i < dim; i++)
        average[i] = (int)(sum[i] / n);

    struct kmr_kv_box nkv = { .klen = sizeof(long),
                              .vlen = dim * (int)sizeof(int),
                              .k.i  = cid,
                              .v.p  = (void *)average };
    kmr_add_kv(kvo, nkv);

    return MPI_SUCCESS;
}

/* KMR map function
   It copies centers of clusters stored in kv to "kmeans.means" array.
*/
static int
copy_center(const struct kmr_kv_box kv,
            const KMR_KVS *kvi, KMR_KVS *kvo, void *p, long i_)
{
    int i;
    int cid = (int)kv.k.i;
    int *center = (int *)kv.v.p;
    kmeans_t *kmeans = (kmeans_t *)p;
    int dim = kmeans->dim;
    int *target = &kmeans->means[cid * dim];

    for (i = 0; i < dim; i++)
        target[i] = center[i];

    return MPI_SUCCESS;
}

static void
print_means(int *means, int size, int dim)
{
    int i, j;
    for (i = 0; i < size; i++) {
        int *mean = &means[i * dim];
        fprintf(stderr, "( ");
        for (j = 0; j < dim; j++)
            fprintf(stderr, "%d ", mean[j]);
        fprintf(stderr, ")\n");
    }
}

static void
print_progress(const char *msg, int rank, int n_itr)
{
    MPI_Barrier(MPI_COMM_WORLD);
    if (rank == 0) {
        fprintf(stderr, "ITR[%d]: %s\n", n_itr, msg);
    }
    MPI_Barrier(MPI_COMM_WORLD);
}

static void
check_answer(kmeans_t *kmeans)
{
    int count   = 0;
    int *answer = NULL;
    if (kmeans->nprocs      == 2 &&
        kmeans->dim         == DEF_DIM &&
        kmeans->grid_size   == DEF_GRID_SIZE &&
        kmeans->n_points    == 10 &&
        kmeans->n_means     == 4 &&
        kmeans->n_iteration >= DEF_NUM_ITERATION) {
        // 2 nodes, 10 points, 4 centers
        printf("\n"
               "Case of 2 nodes, 10 points and 4 centers.\n"
               "Check the answer...");
        count = 4 * kmeans->dim;
        answer = (int*)malloc((size_t)count * sizeof(int));
        answer[0] = 906;
        answer[1] = 733;
        answer[2] = 301;
        answer[3] = 426;
        answer[4] = 179;
        answer[5] = 279;
        answer[6] = 121;
        answer[7] = 702;
        answer[8] = 107;
        answer[9] = 349;
        answer[10] = 727;
        answer[11] = 751;
    } else if (kmeans->nprocs      == 4 &&
               kmeans->dim         == DEF_DIM &&
               kmeans->grid_size   == DEF_GRID_SIZE &&
               kmeans->n_points    == 10 &&
               kmeans->n_means     == 2 &&
               kmeans->n_iteration >= DEF_NUM_ITERATION) {
        // 4 nodes, 10 points, 2 centers
        printf("\n"
               "Case of 4 nodes, 10 points and 2 centers.\n"
               "Check the answer...");
        count = 2 * kmeans->dim;
        answer = (int*)malloc((size_t)count * sizeof(int));
        answer[0] = 830;
        answer[1] = 787;
        answer[2] = 389;
        answer[3] = 346;
        answer[4] = 425;
        answer[5] = 448;
    } else if (kmeans->nprocs      == 4 &&
               kmeans->dim         == DEF_DIM &&
               kmeans->grid_size   == DEF_GRID_SIZE &&
               kmeans->n_points    == 10 &&
               kmeans->n_means     == 4 &&
               kmeans->n_iteration >= DEF_NUM_ITERATION) {
        // 4 nodes, 10 points, 4 centers
        printf("\n"
               "Case of 4 nodes, 10 points and 4 centers.\n"
               "Check the answer...");
        count = 4 * kmeans->dim;
        answer = (int*)malloc((size_t)count * sizeof(int));
        answer[0] = 870;
        answer[1] = 756;
        answer[2] = 434;
        answer[3] = 654;
        answer[4] = 220;
        answer[5] = 435;
        answer[6] = 194;
        answer[7] = 501;
        answer[8] = 167;
        answer[9] = 230;
        answer[10] = 609;
        answer[11] = 729;
    } else if (kmeans->nprocs      == 8 &&
               kmeans->dim         == DEF_DIM &&
               kmeans->grid_size   == DEF_GRID_SIZE &&
               kmeans->n_points    == 10 &&
               kmeans->n_means     == 4 &&
               kmeans->n_iteration >= DEF_NUM_ITERATION) {
        // 8 nodes, 10 points, 4 centers
        printf("\n"
               "Case of 8 nodes, 10 points and 4 centers.\n"
               "Check the answer...");
        count = 4 * kmeans->dim;
        answer = (int*)malloc((size_t)count * sizeof(int));
        answer[0] = 809;
        answer[1] = 753;
        answer[2] = 467;
        answer[3] = 668;
        answer[4] = 260;
        answer[5] = 495;
        answer[6] = 194;
        answer[7] = 438;
        answer[8] = 235;
        answer[9] = 220;
        answer[10] = 620;
        answer[11] = 792;
    }

    if (count != 0) {
        _Bool ok = 1;
        for (int i = 0; i < count; i++) {
            if (kmeans->means[i] != answer[i]) {
                ok = 0;
                break;
            }
        }
        if (ok) {
            printf("The answer is correct.\n");
        } else {
            printf("The answer is wrong.\n");
        }
        free(answer);
    }
}

////////////////////////////////////////////////////////////////////////////////

int
main(int argc, char **argv)
{
    kmeans_t kmeans;
    int nprocs, rank;
    int p_iteration = -1;

    MPI_Init(&argc, &argv);
    MPI_Comm_size(MPI_COMM_WORLD, &nprocs);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);

    kmr_init();
    MPI_Info info;
    MPI_Info_create(&info);
    MPI_Info_set(info, "ckpt_enable", "1");
    KMR *mr = kmr_create_context(MPI_COMM_WORLD, info, 0);
    MPI_Info_free(&info);

    // Initialize using MPI functions
    _srand((unsigned int)(rank + 1));
    //srand((rank + 1) * getpid());
    if (rank == 0){
        if (is_restart()) {
            load_kmeans_state(&kmeans);
            p_iteration = kmeans.c_iteration;
            kmeans.c_iteration = 0;
            printf("*** Restarted ***\n\n");
        } else {
            kmeans.nprocs      = nprocs;
            kmeans.n_iteration = DEF_NUM_ITERATION;
            kmeans.c_iteration = 0;
            kmeans.grid_size   = DEF_GRID_SIZE;
            kmeans.dim         = DEF_DIM;
            kmeans.n_points    = DEF_NUM_POINTS;
            kmeans.n_means     = DEF_NUM_MEANS;

            parse_args(argc, argv, &kmeans);
            // set initial centers randomly on rank 0
            kmeans.means = (int *) malloc((size_t)kmeans.n_means * (size_t)kmeans.dim * sizeof(int));
            generate_randoms(kmeans.means, kmeans.n_means * kmeans.dim,
                             kmeans.grid_size);
        }

        printf("#### Configuration ###########################\n");
        printf("Number of processes = %d\n", nprocs);
        printf("Iteration           = %d\n", kmeans.n_iteration);
        printf("Dimension           = %d\n", kmeans.dim);
        printf("Number of clusters  = %d\n", kmeans.n_means);
        printf("Number of points    = %d\n", kmeans.n_points);
        printf("##############################################\n");
    }
    MPI_Bcast(&(kmeans.n_iteration), 1, MPI_INT, 0, MPI_COMM_WORLD);
    MPI_Bcast(&(kmeans.c_iteration), 1, MPI_INT, 0, MPI_COMM_WORLD);
    MPI_Bcast(&(kmeans.grid_size), 1, MPI_INT, 0, MPI_COMM_WORLD);
    MPI_Bcast(&(kmeans.dim), 1, MPI_INT, 0, MPI_COMM_WORLD);
    MPI_Bcast(&(kmeans.n_points), 1, MPI_INT, 0, MPI_COMM_WORLD);
    MPI_Bcast(&(kmeans.n_means), 1, MPI_INT, 0, MPI_COMM_WORLD);
    MPI_Bcast(&p_iteration, 1, MPI_INT, 0, MPI_COMM_WORLD);
    if (rank != 0) {
        kmeans.means = (int *)malloc((size_t)kmeans.n_means * (size_t)kmeans.dim * sizeof(int));
    }
    MPI_Bcast(kmeans.means, kmeans.n_means * kmeans.dim, MPI_INT,
              0, MPI_COMM_WORLD);
    if (is_restart() != 1) {
        // set points randomly on each rank
        kmeans.points = (int *)malloc((size_t)kmeans.n_points * (size_t)kmeans.dim * sizeof(int));
        generate_randoms(kmeans.points, kmeans.n_points * kmeans.dim,
                         kmeans.grid_size);
    }
    KMR_KVS *kvs_points = kmr_create_kvs(mr, KMR_KV_INTEGER, KMR_KV_OPAQUE);
    kmr_map_once(kvs_points, (void *)&kmeans, kmr_noopt, 0, generate_points);
    print_progress("map_onece done", rank, 0);

    // kernel //
    struct timeval tv_s, tv_e;
    double time_diff, total_time = 0.0;
    while (kmeans.c_iteration < kmeans.n_iteration) {
        // measure iteration start time
        if (kmeans.c_iteration > p_iteration) {
            if (measure_time(&tv_s) == -1) {
                MPI_Abort(MPI_COMM_WORLD, 1);
            }
        }

        // call kmr_map to calculate cluster
        KMR_KVS *kvs_c2p = kmr_create_kvs(mr, KMR_KV_INTEGER, KMR_KV_OPAQUE);
        kmr_map(kvs_points, kvs_c2p, (void *)&kmeans, kmr_inspect,
                calc_cluster);
        print_progress("map done", rank, kmeans.c_iteration);

        // call kmr_shuffle to gather points which belong to the same cluster
        KMR_KVS *kvs_c2p_s = kmr_create_kvs(mr, KMR_KV_INTEGER, KMR_KV_OPAQUE);
        kmr_shuffle(kvs_c2p, kvs_c2p_s, kmr_noopt);
        print_progress("shuffle done", rank, kmeans.c_iteration);

        // call kmr_reduce to update cluster center
        KMR_KVS *kvs_cluster = kmr_create_kvs(mr,
                                              KMR_KV_INTEGER, KMR_KV_OPAQUE);
        kmr_reduce(kvs_c2p_s, kvs_cluster, (void *)&kmeans, kmr_noopt,
                   update_cluster);
        print_progress("reduce done", rank, kmeans.c_iteration);

        if (mr->ckpt_enable && kmeans.c_iteration > (p_iteration + 1)
            && kmeans.c_iteration == 4) {
            if (rank == 0) {
                fprintf(stderr, "Aborted on rank 0 for testing purpose.\n");
                MPI_Abort(MPI_COMM_WORLD, 1);
            }
        }

        // cal kmr_replicate to share new cluster centers
        KMR_KVS *kvs_all_clusters =
            kmr_create_kvs(mr, KMR_KV_INTEGER, KMR_KV_OPAQUE);
        kmr_replicate(kvs_cluster, kvs_all_clusters, kmr_noopt);
        print_progress("replicate done", rank, kmeans.c_iteration);

        // call kmr_map to copy new cluster centers to kmeans.means array.
        kmr_map(kvs_all_clusters, NULL, (void *)&kmeans, kmr_noopt,
                copy_center);

        // measure iteration end time
        if (kmeans.c_iteration > p_iteration) {
            if (measure_time(&tv_e) == -1) {
                MPI_Abort(MPI_COMM_WORLD, 1);
            }
            if (rank == 0) {
                time_diff = calc_time_diff(&tv_s, &tv_e);
                total_time += time_diff;
                //printf("Iteration[%2d]: Elapse time: %f\n", itr, time_diff);
                print_means(kmeans.means, kmeans.n_means, kmeans.dim);
                if (save_kmeans_state(&kmeans)) {
                    MPI_Abort(MPI_COMM_WORLD, 1);
                }
            }
        }

        if (mr->ckpt_enable && kmeans.c_iteration > p_iteration
            && kmeans.c_iteration == 1) {
            if (rank == 0) {
                fprintf(stderr, "Aborted on rank 0 for testing purpose.\n");
                MPI_Abort(MPI_COMM_WORLD, 1);
            }
        }

        kmeans.c_iteration += 1;
    }
    if (rank == 0) {
        printf("Total elapse time: %f\n", total_time);
        check_answer(&kmeans);
        delete_kmeans_state();
    }

    free(kmeans.means);
    if (is_restart() != 1) {
        free(kmeans.points);
    }
    kmr_free_kvs(kvs_points);
    kmr_free_context(mr);
    kmr_fin();

    MPI_Finalize();
    return 0;
}