[Comm] allow collecting MPI timers from python

src/shamcomm/include/shamcomm/collectives.hpp

@@ -41,6 +41,19 @@ namespace shamcomm {
     void gather_basic_str(
         const std::basic_string<byte> &send_vec, std::basic_string<byte> &recv_vec);
 
+    /**
+     * @brief Allgathers a string from all nodes and concatenates it in a std::string
+     *
+     * This function gathers the string `send_vec` from all nodes and concatenates the
+     * result in `recv_vec` on every rank. The result is ordered by the order of the
+     * nodes in the communicator, i.e. the string is ordered by rank.
+     */
+    void allgather_str(const std::string &send_vec, std::string &recv_vec);
+
+    /// same as allgather_str but with std::basic_string
+    void allgather_basic_str(
+        const std::basic_string<byte> &send_vec, std::basic_string<byte> &recv_vec);
+
     /**
      * @brief Constructs a histogram from a vector of strings, counting occurrences
      *        of each unique string.
@@ -56,8 +69,11 @@ namespace shamcomm {
      * @return An unordered map where keys are unique strings from the input and
      *         values are the counts of their occurrences. (valid only on rank 0)
      */
-
     std::unordered_map<std::string, int> string_histogram(
         const std::vector<std::string> &inputs, std::string delimiter = "\n");
 
+    /// same as string_histogram but with result return on every rank
+    std::unordered_map<std::string, int> all_string_histogram(
+        const std::vector<std::string> &inputs, std::string delimiter = "\n");
+
 } // namespace shamcomm

src/shamcomm/include/shamcomm/wrapper.hpp

@@ -19,6 +19,7 @@
 #include "shambase/aliases_float.hpp"
 #include "shambase/aliases_int.hpp"
 #include "shamcomm/mpi.hpp"
+#include <unordered_map>
 #include <string>
 
 namespace shamcomm::mpi {
@@ -29,6 +30,9 @@ namespace shamcomm::mpi {
     /// get a timer value
     f64 get_timer(std::string timername);
 
+    /// return all internal timers
+    const std::unordered_map<std::string, f64> &get_timers();
+
     /// MPI wrapper for MPI_Allreduce
     void Allreduce(
         const void *sendbuf,

src/shamcomm/src/collectives.cpp

@@ -81,6 +81,56 @@ namespace {
         recv_vec = result;
     }
 
+    /**
+     * @brief Allgather a vector of characters from all MPI ranks into a single string
+     *
+     * The resulting string is concatenated in rank order and is returned on every rank.
+     */
+    template<class Tchar>
+    inline void _internal_allgather_str(
+        const std::basic_string<Tchar> &send_vec, std::basic_string<Tchar> &recv_vec) {
+        StackEntry stack_loc{};
+
+        if (shamcomm::world_size() == 1) {
+            recv_vec = send_vec;
+            return;
+        }
+
+        i32 wsize       = shamcomm::world_size();
+        size_t wsize_sz = static_cast<size_t>(wsize);
+
+        // counts/displacements are expressed in number of characters.
+        std::vector<int> counts(wsize_sz);
+        std::vector<int> disps(wsize_sz);
+
+        // MPI counts/displacements use `int`.
+        int local_count = static_cast<int>(send_vec.size());
+
+        shamcomm::mpi::Allgather(
+            &local_count, 1, MPI_INT, counts.data(), 1, MPI_INT, MPI_COMM_WORLD);
+
+        for (size_t i = 0; i < wsize_sz; i++) {
+            disps[i] = (i > 0) ? (disps[i - 1] + counts[i - 1]) : 0;
+        }
+
+        int global_len = disps[wsize_sz - 1] + counts[wsize_sz - 1];
+
+        std::basic_string<Tchar> result;
+        result.resize(static_cast<size_t>(global_len));
+
+        shamcomm::mpi::Allgatherv(
+            send_vec.data(),
+            local_count,
+            MPI_CHAR,
+            result.data(),
+            counts.data(),
+            disps.data(),
+            MPI_CHAR,
+            MPI_COMM_WORLD);
+
+        recv_vec = result;
+    }
+
 } // namespace
 
 void shamcomm::gather_str(const std::string &send_vec, std::string &recv_vec) {
@@ -94,6 +144,17 @@ void shamcomm::gather_basic_str(
     _internal_gather_str(send_vec, recv_vec);
 }
 
+void shamcomm::allgather_str(const std::string &send_vec, std::string &recv_vec) {
+    StackEntry stack_loc{};
+    _internal_allgather_str(send_vec, recv_vec);
+}
+
+void shamcomm::allgather_basic_str(
+    const std::basic_string<byte> &send_vec, std::basic_string<byte> &recv_vec) {
+    StackEntry stack_loc{};
+    _internal_allgather_str(send_vec, recv_vec);
+}
+
 std::unordered_map<std::string, int> shamcomm::string_histogram(
     const std::vector<std::string> &inputs, std::string delimiter) {
     std::string accum_loc = "";
@@ -119,3 +180,24 @@ std::unordered_map<std::string, int> shamcomm::string_histogram(
 
     return {};
 }
+
+std::unordered_map<std::string, int> shamcomm::all_string_histogram(
+    const std::vector<std::string> &inputs, std::string delimiter) {
+    std::string accum_loc = "";
+    for (auto &s : inputs) {
+        accum_loc += s + delimiter;
+    }
+
+    std::string recv = "";
+    allgather_str(accum_loc, recv);
+
+    std::vector<std::string> splitted = shambase::split_str(recv, delimiter);
+
+    std::unordered_map<std::string, int> histogram;
+
+    for (size_t i = 0; i < splitted.size(); i++) {
+        histogram[splitted[i]] += 1;
+    }
+
+    return histogram;
+}

src/shamcomm/src/wrapper.cpp

@@ -43,6 +43,8 @@ namespace shamcomm::mpi {
 
     f64 get_timer(std::string timername) { return mpi_timers[timername]; }
 
+    const std::unordered_map<std::string, f64> &get_timers() { return mpi_timers; }
+
 } // namespace shamcomm::mpi
 
 namespace {

src/shampylib/src/pyShamcomm.cpp

@@ -0,0 +1,59 @@
+// -------------------------------------------------------//
+//
+// SHAMROCK code for hydrodynamics
+// Copyright (c) 2021-2026 Timothée David--Cléris <tim.shamrock@proton.me>
+// SPDX-License-Identifier: CeCILL Free Software License Agreement v2.1
+// Shamrock is licensed under the CeCILL 2.1 License, see LICENSE for more information
+//
+// -------------------------------------------------------//
+
+/**
+ * @file pyShamcomm.cpp
+ * @author Timothée David--Cléris (tim.shamrock@proton.me)
+ * @brief
+ */
+
+#include "shamalgs/collective/reduction.hpp"
+#include "shambindings/pybind11_stl.hpp"
+#include "shambindings/pybindaliases.hpp"
+#include "shambindings/pytypealias.hpp"
+#include "shamcomm/collectives.hpp"
+#include "shamcomm/logs.hpp"
+#include "shamcomm/wrapper.hpp"
+#include <pybind11/pytypes.h>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+Register_pymod(shamcommlibinit) {
+
+    py::module shamcomm_module = m.def_submodule("comm", "comm library");
+
+    shamcomm_module.def("get_timer", [](std::string name) {
+        return shamcomm::mpi::get_timer(std::move(name));
+    });
+
+    shamcomm_module.def("get_timers", []() {
+        return shamcomm::mpi::get_timers();
+    });
+
+    shamcomm_module.def(
+        "mpi_timers_delta",
+        [](std::unordered_map<std::string, f64> start, std::unordered_map<std::string, f64> end) {
+            std::vector<std::string> keys{};
+
+            for (auto &[k, v] : end) {
+                keys.push_back(k);
+            }
+
+            auto key_histo = shamcomm::all_string_histogram(keys);
+
+            std::unordered_map<std::string, f64> deltas{};
+
+            for (auto &[k, c] : key_histo) {
+                deltas[k] = shamalgs::collective::allreduce_max(end[k] - start[k]);
+            }
+
+            return deltas;
+        });
+}

src/tests/shamcomm/collectivesTests.cpp

@@ -37,3 +37,26 @@ TestStart(Unittest, "shamcomm/collectives::gather_str", test_gather_str, 4) {
 
     REQUIRE_EQUAL(recv, result);
 }
+
+TestStart(Unittest, "shamcomm/collectives::allgather_str", test_allgather_str, 4) {
+
+    std::array<std::string, 4> ref_base{
+        "I'm a very important string",
+        "But I'm a very important string",
+        "Listen, I'm a very important string",
+        "The most importantest string",
+    };
+
+    std::string result = "";
+    for (u32 i = 0; i < ref_base.size(); i++) {
+        result += ref_base[i];
+    }
+
+    std::string send = ref_base[shamcomm::world_rank()];
+
+    std::string recv = "random string"; // Just to check that it is overwritten
+
+    shamcomm::allgather_str(send, recv);
+
+    REQUIRE_EQUAL(recv, result);
+}