From 44ebd7fcd615e6463aaae59851c1bddbdc46f9e0 Mon Sep 17 00:00:00 2001 From: Jonathan Lifflander Date: Tue, 6 Jan 2026 16:04:24 -0800 Subject: [PATCH 1/2] #31: mpi: implement a simple allgather --- src/vt-lb/comm/MPI/class_handle.h | 6 ++ src/vt-lb/comm/MPI/class_handle.impl.h | 6 ++ src/vt-lb/comm/MPI/comm_mpi.h | 94 ++++++++++++++++++++++++++ 3 files changed, 106 insertions(+) diff --git a/src/vt-lb/comm/MPI/class_handle.h b/src/vt-lb/comm/MPI/class_handle.h index b881e3e..28a482f 100644 --- a/src/vt-lb/comm/MPI/class_handle.h +++ b/src/vt-lb/comm/MPI/class_handle.h @@ -46,6 +46,9 @@ #include +#include +#include + namespace vt_lb::comm { struct CommMPI; @@ -91,6 +94,9 @@ struct ClassHandle { template void broadcast(int root, MPI_Datatype datatype, U buffer, int count); + template + std::unordered_map> allgather(U const* sendbuf, int sendcount); + friend struct ClassHandleRank; int getIndex() const { return index_; } diff --git a/src/vt-lb/comm/MPI/class_handle.impl.h b/src/vt-lb/comm/MPI/class_handle.impl.h index 1cc43b9..ff90ba1 100644 --- a/src/vt-lb/comm/MPI/class_handle.impl.h +++ b/src/vt-lb/comm/MPI/class_handle.impl.h @@ -112,6 +112,12 @@ void ClassHandle::broadcast(int root, MPI_Datatype datatype, U buffer, int co comm_->broadcast(root, datatype, buffer, count); } +template +template +std::unordered_map> ClassHandle::allgather(U const* sendbuf, int sendcount) { + return comm_->allgather(sendbuf, sendcount); +} + } // namespace vt_lb::comm #endif /*INCLUDED_VT_LB_COMM_CLASS_HANDLE_IMPL_H*/ \ No newline at end of file diff --git a/src/vt-lb/comm/MPI/comm_mpi.h b/src/vt-lb/comm/MPI/comm_mpi.h index 1597025..9d93d65 100644 --- a/src/vt-lb/comm/MPI/comm_mpi.h +++ b/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 + std::unordered_map> allgather(T const* sendbuf, int sendcount) { + static_assert(std::is_trivially_copyable_v, "CommMPI::allgather requires trivially copyable T"); + MPI_Datatype datatype = deduceDatatype(); + + if (comm_ == MPI_COMM_NULL) { + throw std::runtime_error("CommMPI not initialized (MPI_Comm is MPI_COMM_NULL)"); + } + + const int n = numRanks(); + std::vector 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 displs(n, 0); + int total = 0; + for (int i = 0; i < n; ++i) { + displs[i] = total; + total += counts[i]; + } + + std::vector recvbuf; + recvbuf.resize(static_cast(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> by_rank; + by_rank.reserve(static_cast(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(recvbuf.begin() + off, recvbuf.begin() + off + cnt) + ); + } else { + by_rank.emplace(r, std::vector{}); + } + } + + return by_rank; + } + private: void initTermination(); + // Helper to trigger static_assert for unsupported types. + template struct AlwaysFalse : std::false_type {}; + + // Map a limited set of basic C++ types to MPI_Datatype. + template + static MPI_Datatype deduceDatatype() { + if constexpr (std::is_same_v) return MPI_CHAR; + else if constexpr (std::is_same_v) return MPI_SIGNED_CHAR; + else if constexpr (std::is_same_v) return MPI_UNSIGNED_CHAR; + else if constexpr (std::is_same_v) return MPI_SHORT; + else if constexpr (std::is_same_v) return MPI_UNSIGNED_SHORT; + else if constexpr (std::is_same_v) return MPI_INT; + else if constexpr (std::is_same_v) return MPI_UNSIGNED; + else if constexpr (std::is_same_v) return MPI_LONG; + else if constexpr (std::is_same_v) return MPI_UNSIGNED_LONG; + else if constexpr (std::is_same_v) return MPI_LONG_LONG; + else if constexpr (std::is_same_v) return MPI_UNSIGNED_LONG_LONG; + else if constexpr (std::is_same_v) return MPI_FLOAT; + else if constexpr (std::is_same_v) return MPI_DOUBLE; + else if constexpr (std::is_same_v) return MPI_LONG_DOUBLE; + else { + static_assert(AlwaysFalse::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 From 1242653e6a282604f709185e1783fb4bb7813917 Mon Sep 17 00:00:00 2001 From: Jonathan Lifflander Date: Tue, 13 Jan 2026 16:41:07 -0800 Subject: [PATCH 2/2] #31: allgather: implement VT comm allgather and test --- src/vt-lb/comm/vt/collective_handler.h | 6 +++ src/vt-lb/comm/vt/proxy_wrapper.h | 3 ++ src/vt-lb/comm/vt/proxy_wrapper.impl.h | 53 ++++++++++++++++++++++++++ tests/unit/comm/test_comm.cc | 29 ++++++++++++++ 4 files changed, 91 insertions(+) diff --git a/src/vt-lb/comm/vt/collective_handler.h b/src/vt-lb/comm/vt/collective_handler.h index 6afed41..eda618d 100644 --- a/src/vt-lb/comm/vt/collective_handler.h +++ b/src/vt-lb/comm/vt/collective_handler.h @@ -72,6 +72,12 @@ struct CollectiveHandler { ctx_->done.store(true, std::memory_order_release); } + template + void allgatherValues(DataT const& values) { + *static_cast(ctx_->out_ptr) = values; + ctx_->done.store(true, std::memory_order_release); + } + private: CtxT* ctx_ = nullptr; }; diff --git a/src/vt-lb/comm/vt/proxy_wrapper.h b/src/vt-lb/comm/vt/proxy_wrapper.h index 2ca3dd6..349216a 100644 --- a/src/vt-lb/comm/vt/proxy_wrapper.h +++ b/src/vt-lb/comm/vt/proxy_wrapper.h @@ -72,6 +72,9 @@ struct ProxyWrapper : ProxyT { template void broadcast(int root, MPI_Datatype datatype, T* buffer, int count); + template + std::unordered_map> allgather(T const* sendbuf, int sendcount); + private: enum class VTOp { Plus, Max, Min }; static VTOp mapOp(MPI_Op mpio); diff --git a/src/vt-lb/comm/vt/proxy_wrapper.impl.h b/src/vt-lb/comm/vt/proxy_wrapper.impl.h index 08d0002..c393387 100644 --- a/src/vt-lb/comm/vt/proxy_wrapper.impl.h +++ b/src/vt-lb/comm/vt/proxy_wrapper.impl.h @@ -242,6 +242,59 @@ void ProxyWrapper::broadcast_impl(int root, T* buffer, int count) { } } +template +struct AllReduceContainer { + AllReduceContainer() = default; + explicit AllReduceContainer(std::vector const& in) { + by_rank[vt::theContext()->getNode()] = in; + } + + friend AllReduceContainer operator+( + AllReduceContainer lhs, + AllReduceContainer 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 + void serialize(SerializerT& s) { + s | by_rank; + } + + std::unordered_map> by_rank; +}; + +template +template +std::unordered_map> +ProxyWrapper::allgather(T const* sendbuf, int sendcount) { + std::unordered_map> out; + collective_ctx_->out_ptr = static_cast(&out); + collective_ctx_->done.store(false); + + std::vector local_vec; + if (sendcount > 0) { + local_vec.resize(static_cast(sendcount)); + std::memcpy(local_vec.data(), sendbuf, sizeof(T) * static_cast(sendcount)); + } + AllReduceContainer container{local_vec}; + collective_proxy_.template allreduce< + &CollectiveHandlerType::template allgatherValues>, + 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 */ diff --git a/tests/unit/comm/test_comm.cc b/tests/unit/comm/test_comm.cc index 0db5ca5..5ca139c 100644 --- a/tests/unit/comm/test_comm.cc +++ b/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 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(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(i)); + } + } +} + }}} // end namespace vt_lb::tests::unit