Indicate that lttng-ust 2.13+ is required

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

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

#include <assert.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.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);
	}

	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;
}