test.c

   1 /*
   2  * SPDX-License-Identifier: MIT
   3  *
   4  * SPDX-FileCopyrightText: 2023 Olivier Dion <odion@efficios.com>
   5  */
   6
   7 #include <assert.h>
   8 #include <stdint.h>
   9 #include <stdio.h>
  10
  11 #include <mpi.h>
  12
  13 static uint64_t sum_of(uint64_t *values, size_t values_count)
  14 {
  15     size_t acc = 0;
  16     for (size_t k=0; k<values_count; ++k) {
  17         acc += values[k];
  18     }
  19     return acc;
  20 }
  21
  22 static void usage()
  23 {
  24     fprintf(stderr, "Usage: test-mpi N\n");
  25     exit(EXIT_FAILURE);
  26 }
  27
  28 static uint64_t *allocate_values(size_t upto)
  29 {
  30     uint64_t *values = (uint64_t*)malloc(sizeof(uint64_t) * upto);
  31     for (size_t k=0; k<upto; ++k) {
  32         values[k] = k + 1;
  33     }
  34     return values;
  35 }
  36
  37 static void send_values(int target, uint64_t *values,
  38                         size_t values_count,
  39                         MPI_Request *request)
  40 {
  41     MPI_Isend(values, values_count, MPI_UINT64_T,
  42               target, 0, MPI_COMM_WORLD, request);
  43 }
  44
  45 static void recv_answer(int target, uint64_t *value,
  46                         MPI_Request *request)
  47 {
  48     MPI_Irecv(value, 1, MPI_UINT64_T,
  49               target, 0, MPI_COMM_WORLD, request);
  50 }
  51
  52 static void send_answer(uint64_t value)
  53 {
  54     MPI_Send(&value, 1, MPI_UINT64_T,
  55              0, 0, MPI_COMM_WORLD);
  56 }
  57
  58 static uint64_t *recv_values(size_t chunk_size)
  59 {
  60     uint64_t *values = (uint64_t*)malloc(sizeof(uint64_t) * chunk_size);
  61     MPI_Recv(values, chunk_size, MPI_UINT64_T,
  62              0, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
  63     return values;
  64 }
  65
  66 int main(int argc, char *argv[])
  67 {
  68     int rank;
  69     int size;
  70     long long upto;
  71     uint64_t *values;
  72
  73     if (argc < 2) {
  74         usage();
  75     }
  76
  77     upto = atoll(argv[1]);
  78
  79     if (upto <= 0) {
  80         fprintf(stderr, "N must be greater than 0\n");
  81         exit(EXIT_FAILURE);
  82     }
  83
  84     cali_init();
  85
  86     MPI_Init(&argc, &argv);
  87
  88     MPI_Comm_set_errhandler(MPI_COMM_WORLD,
  89                             MPI_ERRORS_RETURN);
  90
  91     MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  92     MPI_Comm_size(MPI_COMM_WORLD, &size);
  93
  94     size_t chunk_size;
  95     size_t rest;
  96     uint64_t total;
  97
  98     if (size > 1) {
  99         chunk_size = upto / (size - 1);
 100         rest = upto % (size - 1);
 101     } else {
 102         chunk_size = 0;
 103         rest = upto;
 104     }
 105
 106     if (rank == 0) {
 107         uint64_t sums[size];
 108         MPI_Request requests[size - 1];
 109
 110         values = allocate_values(upto);
 111
 112         for (int k=1; k<size; ++k) {
 113             send_values(k,
 114                         values + (chunk_size * (k - 1)),
 115                         chunk_size,
 116                         &requests[k-1]);
 117         }
 118
 119         sums[0] = sum_of(values + chunk_size * (size - 1),
 120                          rest);
 121
 122         MPI_Waitall(size - 1, requests, MPI_STATUS_IGNORE);
 123
 124         for (int k=1; k<size; ++k) {
 125             recv_answer(k, &sums[k], &requests[k-1]);
 126         }
 127
 128         MPI_Waitall(size - 1, requests, MPI_STATUS_IGNORE);
 129
 130         total = sum_of(sums, size);
 131     } else {
 132         send_answer(sum_of(recv_values(chunk_size),
 133                            chunk_size));
 134     }
 135
 136     MPI_Finalize();
 137
 138     if (rank == 0){
 139         assert(total ==
 140                (((uint64_t)upto * ((uint64_t)upto + 1U)) >> 1U));
 141     }
 142
 143     return 0;
 144 }