From 8411edf44f43d5549ab66e180f359860ae4cd031 Mon Sep 17 00:00:00 2001
From: Dmitri Naumov <github@naumov.de>
Date: Fri, 1 Nov 2024 11:15:44 +0100
Subject: [PATCH] [BL/MPI] allgatherv for variable size data exchg
---
BaseLib/MPI.h | 28 ++++++++++++++++++++++++++++
1 file changed, 28 insertions(+)
diff --git a/BaseLib/MPI.h b/BaseLib/MPI.h
index 891b81f98b3..3d60e44a61b 100644
--- a/BaseLib/MPI.h
+++ b/BaseLib/MPI.h
@@ -11,6 +11,7 @@
#include <algorithm>
+#include "Algorithm.h"
#include "Error.h"
#ifdef USE_PETSC
@@ -136,6 +137,33 @@ static void allreduceInplace(std::vector<T>& vector,
mpi_op,
mpi.communicator);
}
+
+/// Allgather for variable data. Returns offsets in the receive buffer.
+/// The receive buffer is resized to accommodate the gathered data.
+template <typename T>
+static std::vector<int> allgatherv(
+ std::span<T> const send_buffer,
+ std::vector<std::remove_const_t<T>>& receive_buffer,
+ Mpi const& mpi)
+{
+ // Determine the number of elements to send
+ int const size = static_cast<int>(send_buffer.size());
+
+ // Gather sizes from all ranks
+ std::vector<int> const sizes = allgather(size, mpi);
+
+ // Compute offsets based on counts
+ std::vector<int> const offsets = BaseLib::sizesToOffsets(sizes);
+
+ // Resize receive buffer to hold all gathered data
+ receive_buffer.resize(offsets.back());
+
+ MPI_Allgatherv(send_buffer.data(), size, mpiType<T>(),
+ receive_buffer.data(), sizes.data(), offsets.data(),
+ mpiType<T>(), mpi.communicator);
+
+ return offsets;
+}
#endif
/// The reduction is implemented transparently for with and without MPI. In the
--
GitLab