#31: implement global mapping

src/vt-lb/comm/MPI/class_handle.h

@@ -46,6 +46,9 @@
 
 #include <mpi.h>
 
+#include <unordered_map>
+#include <vector>
+
 namespace vt_lb::comm {
 
 struct CommMPI;
@@ -91,6 +94,9 @@ struct ClassHandle {
   template <typename U>
   void broadcast(int root, MPI_Datatype datatype, U buffer, int count);
 
+  template <typename U>
+  std::unordered_map<int, std::vector<U>> allgather(U const* sendbuf, int sendcount);
+
   friend struct ClassHandleRank<T>;
 
   int getIndex() const { return index_; }

src/vt-lb/comm/MPI/class_handle.impl.h

@@ -112,6 +112,12 @@ void ClassHandle<T>::broadcast(int root, MPI_Datatype datatype, U buffer, int co
   comm_->broadcast(root, datatype, buffer, count);
 }
 
+template <typename T>
+template <typename U>
+std::unordered_map<int, std::vector<U>> ClassHandle<T>::allgather(U const* sendbuf, int sendcount) {
+  return comm_->allgather(sendbuf, sendcount);
+}
+
 } // namespace vt_lb::comm
 
 #endif /*INCLUDED_VT_LB_COMM_CLASS_HANDLE_IMPL_H*/

src/vt-lb/comm/MPI/comm_mpi.h

@@ -283,9 +283,103 @@ struct CommMPI {
    */
   bool poll();
 
+  /**
+   * \brief Gather values from all ranks, with varying counts per rank, and remember per-rank contribution
+   *
+   * \tparam T Type of the elements to gather (must be trivially copyable and one of the supported basic types)
+   *
+   * \param sendbuf Pointer to the local data to send
+   * \param sendcount Number of elements to send from this rank
+   *
+   * \return A map from rank -> vector of values contributed by that rank
+   */
+  template <typename T>
+  std::unordered_map<int, std::vector<T>> allgather(T const* sendbuf, int sendcount) {
+    static_assert(std::is_trivially_copyable_v<T>, "CommMPI::allgather requires trivially copyable T");
+    MPI_Datatype datatype = deduceDatatype<T>();
+
+    if (comm_ == MPI_COMM_NULL) {
+      throw std::runtime_error("CommMPI not initialized (MPI_Comm is MPI_COMM_NULL)");
+    }
+
+    const int n = numRanks();
+    std::vector<int> counts(n, 0);
+    // Gather local counts from each rank (block for small metadata)
+    MPI_Allgather(&sendcount, 1, MPI_INT, counts.data(), 1, MPI_INT, comm_);
+
+    std::vector<int> displs(n, 0);
+    int total = 0;
+    for (int i = 0; i < n; ++i) {
+      displs[i] = total;
+      total += counts[i];
+    }
+
+    std::vector<T> recvbuf;
+    recvbuf.resize(static_cast<std::size_t>(total));
+
+    VT_LB_LOG(Communicator, normal, "MPI_Iallgatherv sendcount={} total={} ranks={}\n", sendcount, total, n);
+
+    MPI_Request req;
+    MPI_Iallgatherv(
+      sendbuf, sendcount, datatype,
+      recvbuf.data(), counts.data(), displs.data(), datatype,
+      comm_, &req
+    );
+
+    int flag = 0;
+    while (!flag) {
+      MPI_Status status;
+      MPI_Test(&req, &flag, &status);
+      poll();
+    }
+
+    std::unordered_map<int, std::vector<T>> by_rank;
+    by_rank.reserve(static_cast<std::size_t>(n));
+    for (int r = 0; r < n; ++r) {
+      const int cnt = counts[r];
+      const int off = displs[r];
+      if (cnt > 0) {
+        by_rank.emplace(
+          r,
+          std::vector<T>(recvbuf.begin() + off, recvbuf.begin() + off + cnt)
+        );
+      } else {
+        by_rank.emplace(r, std::vector<T>{});
+      }
+    }
+
+    return by_rank;
+  }
+
 private:
   void initTermination();
 
+  // Helper to trigger static_assert for unsupported types.
+  template <typename> struct AlwaysFalse : std::false_type {};
+
+  // Map a limited set of basic C++ types to MPI_Datatype.
+  template <typename T>
+  static MPI_Datatype deduceDatatype() {
+    if constexpr (std::is_same_v<T, char>) return MPI_CHAR;
+    else if constexpr (std::is_same_v<T, signed char>) return MPI_SIGNED_CHAR;
+    else if constexpr (std::is_same_v<T, unsigned char>) return MPI_UNSIGNED_CHAR;
+    else if constexpr (std::is_same_v<T, short>) return MPI_SHORT;
+    else if constexpr (std::is_same_v<T, unsigned short>) return MPI_UNSIGNED_SHORT;
+    else if constexpr (std::is_same_v<T, int>) return MPI_INT;
+    else if constexpr (std::is_same_v<T, unsigned int>) return MPI_UNSIGNED;
+    else if constexpr (std::is_same_v<T, long>) return MPI_LONG;
+    else if constexpr (std::is_same_v<T, unsigned long>) return MPI_UNSIGNED_LONG;
+    else if constexpr (std::is_same_v<T, long long>) return MPI_LONG_LONG;
+    else if constexpr (std::is_same_v<T, unsigned long long>) return MPI_UNSIGNED_LONG_LONG;
+    else if constexpr (std::is_same_v<T, float>) return MPI_FLOAT;
+    else if constexpr (std::is_same_v<T, double>) return MPI_DOUBLE;
+    else if constexpr (std::is_same_v<T, long double>) return MPI_LONG_DOUBLE;
+    else {
+      static_assert(AlwaysFalse<T>::value, "Unsupported T for CommMPI::allgather");
+      return MPI_DATATYPE_NULL; // unreachable
+    }
+  }
+
   /// @brief Flag indicating if MPI is being used in interop mode
   bool interop_mode_ = false;
   /// @brief MPI communicator

src/vt-lb/comm/vt/collective_handler.h

@@ -72,6 +72,12 @@ struct CollectiveHandler {
     ctx_->done.store(true, std::memory_order_release);
   }
 
+  template <typename DataT>
+  void allgatherValues(DataT const& values) {
+    *static_cast<DataT*>(ctx_->out_ptr) = values;
+    ctx_->done.store(true, std::memory_order_release);
+  }
+
 private:
   CtxT* ctx_ = nullptr;
 };

src/vt-lb/comm/vt/proxy_wrapper.h

@@ -72,6 +72,9 @@ struct ProxyWrapper : ProxyT {
   template <typename T>
   void broadcast(int root, MPI_Datatype datatype, T* buffer, int count);
 
+  template <typename T>
+  std::unordered_map<int, std::vector<T>> allgather(T const* sendbuf, int sendcount);
+
 private:
   enum class VTOp { Plus, Max, Min };
   static VTOp mapOp(MPI_Op mpio);

src/vt-lb/comm/vt/proxy_wrapper.impl.h

@@ -242,6 +242,59 @@ void ProxyWrapper<ProxyT>::broadcast_impl(int root, T* buffer, int count) {
   }
 }
 
+template <typename T>
+struct AllReduceContainer {
+  AllReduceContainer() = default;
+  explicit AllReduceContainer(std::vector<T> const& in) {
+    by_rank[vt::theContext()->getNode()] = in;
+  }
+
+  friend AllReduceContainer<T> operator+(
+    AllReduceContainer<T> lhs,
+    AllReduceContainer<T> const& rhs
+  ) {
+    for (auto const& kv : rhs.by_rank) {
+      auto& vec = lhs.by_rank[kv.first];
+      vec.insert(vec.end(), kv.second.begin(), kv.second.end());
+    }
+    return lhs;
+  }
+
+  template <typename SerializerT>
+  void serialize(SerializerT& s) {
+    s | by_rank;
+  }
+
+  std::unordered_map<int, std::vector<T>> by_rank;
+};
+
+template <typename ProxyT>
+template <typename T>
+std::unordered_map<int, std::vector<T>>
+ProxyWrapper<ProxyT>::allgather(T const* sendbuf, int sendcount) {
+  std::unordered_map<int, std::vector<T>> out;
+  collective_ctx_->out_ptr = static_cast<void*>(&out);
+  collective_ctx_->done.store(false);
+
+  std::vector<T> local_vec;
+  if (sendcount > 0) {
+    local_vec.resize(static_cast<std::size_t>(sendcount));
+    std::memcpy(local_vec.data(), sendbuf, sizeof(T) * static_cast<std::size_t>(sendcount));
+  }
+  AllReduceContainer<T> container{local_vec};
+  collective_proxy_.template allreduce<
+    &CollectiveHandlerType::template allgatherValues<AllReduceContainer<T>>,
+    vt::collective::PlusOp
+  >(
+    container
+  );
+
+  while (!collective_ctx_->done.load(std::memory_order_acquire)) {
+    vt::theSched()->runSchedulerOnceImpl();
+  }
+  return out;
+}
+
 } // namespace vt_lb::comm
 
 #endif /* INCLUDED_VT_LB_COMM_PROXY_WRAPPER_IMPL_H */

tests/unit/comm/test_comm.cc

@@ -221,4 +221,33 @@ TYPED_TEST(TestCommBasic, test_broadcast_int_array) {
   EXPECT_EQ(buf[3], 4);
 }
 
+TYPED_TEST(TestCommBasic, test_allgather_int_array) {
+  auto& the_comm = this->comm;
+  auto rank = the_comm.getRank();
+
+  SET_MIN_NUM_NODES_CONSTRAINT(2);
+
+  typename TestFixture::TestObject obj{};
+  auto handle = this->makeHandle(&obj);
+
+  int const root = 0;
+  int offset = (rank + 1) * 4;
+  std::array<int,4> buf{{offset,offset+1,offset+2,offset+3}};
+
+  auto res = handle.allgather(buf.data(), int(buf.size()));
+
+  // Check that we received from all ranks
+  EXPECT_EQ(res.size(), static_cast<size_t>(the_comm.numRanks()));
+  for (int r = 0; r < the_comm.numRanks(); ++r) {
+    auto it = res.find(r);
+    ASSERT_NE(it, res.end());
+    const auto& vec = it->second;
+    EXPECT_EQ(vec.size(), buf.size());
+    int expected_offset = (r + 1) * 4;
+    for (size_t i = 0; i < buf.size(); ++i) {
+      EXPECT_EQ(vec[i], expected_offset + static_cast<int>(i));
+    }
+  }
+}
+
 }}} // end namespace vt_lb::tests::unit