Indicate that lttng-ust 2.13+ is required

[deliverable/lttng-ust-mpi.git] / test.c

/*
 * SPDX-License-Identifier: MIT
 *
 * SPDX-FileCopyrightText: 2023 Olivier Dion <odion@efficios.com>
 */

#include <assert.h>
#include <stdint.h>
#include <stdio.h>

#include <mpi.h>

static uint64_t sum_of(uint64_t *values, size_t values_count)
{
    size_t acc = 0;
    for (size_t k=0; k<values_count; ++k) {
        acc += values[k];
    }
    return acc;
}

static void usage()
{
    fprintf(stderr, "Usage: test-mpi N\n");
    exit(EXIT_FAILURE);
}

static uint64_t *allocate_values(size_t upto)
{
    uint64_t *values = (uint64_t*)malloc(sizeof(uint64_t) * upto);
    for (size_t k=0; k<upto; ++k) {
        values[k] = k + 1;
    }
    return values;
}

static void send_values(int target, uint64_t *values,
                        size_t values_count,
                        MPI_Request *request)
{
    MPI_Isend(values, values_count, MPI_UINT64_T,
              target, 0, MPI_COMM_WORLD, request);
}

static void recv_answer(int target, uint64_t *value,
                        MPI_Request *request)
{
    MPI_Irecv(value, 1, MPI_UINT64_T,
              target, 0, MPI_COMM_WORLD, request);
}

static void send_answer(uint64_t value)
{
    MPI_Send(&value, 1, MPI_UINT64_T,
             0, 0, MPI_COMM_WORLD);
}

static uint64_t *recv_values(size_t chunk_size)
{
    uint64_t *values = (uint64_t*)malloc(sizeof(uint64_t) * chunk_size);
    MPI_Recv(values, chunk_size, MPI_UINT64_T,
             0, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
    return values;
}

int main(int argc, char *argv[])
{
    int rank;
    int size;
    long long upto;
    uint64_t *values;

    if (argc < 2) {
        usage();
    }

    upto = atoll(argv[1]);

    if (upto <= 0) {
        fprintf(stderr, "N must be greater than 0\n");
        exit(EXIT_FAILURE);
    }

    cali_init();

    MPI_Init(&argc, &argv);

    MPI_Comm_set_errhandler(MPI_COMM_WORLD,
                            MPI_ERRORS_RETURN);

    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);

    size_t chunk_size;
    size_t rest;
    uint64_t total;

    if (size > 1) {
        chunk_size = upto / (size - 1);
        rest = upto % (size - 1);
    } else {
        chunk_size = 0;
        rest = upto;
    }

    if (rank == 0) {
        uint64_t sums[size];
        MPI_Request requests[size - 1];

        values = allocate_values(upto);

        for (int k=1; k<size; ++k) {
            send_values(k,
                        values + (chunk_size * (k - 1)),
                        chunk_size,
                        &requests[k-1]);
        }

        sums[0] = sum_of(values + chunk_size * (size - 1),
                         rest);

        MPI_Waitall(size - 1, requests, MPI_STATUS_IGNORE);

        for (int k=1; k<size; ++k) {
            recv_answer(k, &sums[k], &requests[k-1]);
        }

        MPI_Waitall(size - 1, requests, MPI_STATUS_IGNORE);

        total = sum_of(sums, size);
    } else {
        send_answer(sum_of(recv_values(chunk_size),
                           chunk_size));
    }

    MPI_Finalize();

    if (rank == 0){
        assert(total ==
               (((uint64_t)upto * ((uint64_t)upto + 1U)) >> 1U));
    }

    return 0;
}