Do tree reductions instead of allgather in distributed quantile construction (#12061)

Author: RAMitchell

src/collective/allreduce.h

@@ -9,6 +9,7 @@
 
 #include "../common/type.h"             // for EraseType, RestoreType
 #include "../data/array_interface.h"    // for ToDType, ArrayInterfaceHandler
+#include "broadcast.h"                  // for Broadcast
 #include "comm.h"                       // for Comm, RestoreType
 #include "comm_group.h"                 // for GlobalCommGroup
 #include "xgboost/collective/result.h"  // for Result
@@ -75,4 +76,142 @@ template <typename T>
 Allreduce(Context const* ctx, T* data, Op op) {
   return Allreduce(ctx, linalg::MakeVec(data, 1), op);
 }
+
+/**
+ * @brief Allreduce a variable-length vector over `comm`.
+ *
+ * The method performs a tree reduction rooted at rank 0 using `redop`, then broadcasts
+ * the result so every rank ends with the same reduced payload in `data`.
+ *
+ * `redop` must have the signature
+ * `void(Fn(const Span<T const>& lhs, const Span<T const>& rhs, std::vector<T>* out))` and must
+ * write the combined result into `out`.
+ */
+template <typename T, typename Fn>
+std::enable_if_t<
+    std::is_invocable_v<Fn, common::Span<T const>, common::Span<T const>, std::vector<T>*>, Result>
+AllreduceV(Comm const& comm, std::vector<T>* data, Fn redop) {
+  static_assert(std::is_standard_layout_v<T> && std::is_trivially_copyable_v<T>,
+                "AllreduceV supports only standard-layout trivially-copyable types.");
+  CHECK(data);
+  if (!comm.IsDistributed() || comm.World() == 1) {
+    return Success();
+  }
+
+  auto const world = comm.World();
+  auto const rank = comm.Rank();
+  auto constexpr kRoot = 0;
+
+  auto send = [&](std::int32_t peer, std::vector<T> const& vec) {
+    std::int64_t n = static_cast<std::int64_t>(vec.size());
+    auto n_bytes =
+        common::Span<std::int8_t const>{reinterpret_cast<std::int8_t const*>(&n), sizeof(n)};
+    return Success() << [&] {
+      return comm.Chan(peer)->SendAll(n_bytes);
+    } << [&] {
+      if (n == 0) {
+        return Success();
+      }
+      auto payload_bytes = static_cast<std::size_t>(n) * sizeof(T);
+      auto bytes = common::Span<std::int8_t const>{reinterpret_cast<std::int8_t const*>(vec.data()),
+                                                   payload_bytes};
+      return comm.Chan(peer)->SendAll(bytes);
+    } << [&] {
+      return comm.Chan(peer)->Block();
+    };
+  };
+
+  auto recv = [&](std::int32_t peer, std::vector<T>* out) {
+    std::int64_t n = 0;
+    auto n_bytes = common::Span<std::int8_t>{reinterpret_cast<std::int8_t*>(&n), sizeof(n)};
+    auto rc = Success() << [&] {
+      return comm.Chan(peer)->RecvAll(n_bytes);
+    } << [&] {
+      return comm.Chan(peer)->Block();
+    };
+    if (!rc.OK()) {
+      return rc;
+    }
+    CHECK_GE(n, 0);
+    out->resize(static_cast<std::size_t>(n));
+    if (n == 0) {
+      return Success();
+    }
+    auto payload_bytes = static_cast<std::size_t>(n) * sizeof(T);
+    auto bytes =
+        common::Span<std::int8_t>{reinterpret_cast<std::int8_t*>(out->data()), payload_bytes};
+    return Success() << [&] {
+      return comm.Chan(peer)->RecvAll(bytes);
+    } << [&] {
+      return comm.Chan(peer)->Block();
+    };
+  };
+
+  auto shifted_rank = rank;
+  std::vector<T> incoming;
+  std::vector<T> out;
+  bool continue_reduce = true;
+  for (std::int32_t step = 1; step < world; step <<= 1) {
+    if (!continue_reduce) {
+      continue;
+    }
+    if (shifted_rank % (step * 2) == step) {
+      auto parent = shifted_rank - step;
+      auto rc = send(parent, *data);
+      if (!rc.OK()) {
+        return Fail("AllreduceV failed to send data to parent.", std::move(rc));
+      }
+      continue_reduce = false;
+      continue;
+    }
+    if (shifted_rank % (step * 2) == 0 && shifted_rank + step < world) {
+      auto child = shifted_rank + step;
+      auto rc = recv(child, &incoming);
+      if (!rc.OK()) {
+        return Fail("AllreduceV failed to receive data from child.", std::move(rc));
+      }
+      out.clear();
+      redop(common::Span<T const>{data->data(), data->size()},
+            common::Span<T const>{incoming.data(), incoming.size()}, &out);
+      data->swap(out);
+    }
+  }
+
+  std::int64_t reduced_size = static_cast<std::int64_t>(rank == kRoot ? data->size() : 0);
+  auto rc = Broadcast(comm, common::Span<std::int64_t>{&reduced_size, 1}, kRoot);
+  if (!rc.OK()) {
+    return Fail("AllreduceV failed to broadcast reduced size.", std::move(rc));
+  }
+  if (reduced_size == 0) {
+    data->clear();
+    return Success();
+  }
+  if (rank != kRoot) {
+    data->resize(static_cast<std::size_t>(reduced_size));
+  }
+  auto reduced = common::Span<T>{data->data(), static_cast<std::size_t>(reduced_size)};
+  rc = Broadcast(comm, reduced, kRoot);
+  if (!rc.OK()) {
+    return Fail("AllreduceV failed to broadcast reduced payload.", std::move(rc));
+  }
+  return Success();
+}
+
+template <typename T, typename Fn>
+std::enable_if_t<
+    std::is_invocable_v<Fn, common::Span<T const>, common::Span<T const>, std::vector<T>*>, Result>
+AllreduceV(Context const* ctx, CommGroup const& comm, std::vector<T>* data, Fn redop) {
+  if (!comm.IsDistributed()) {
+    return Success();
+  }
+  auto const& cctx = comm.Ctx(ctx, DeviceOrd::CPU());
+  return AllreduceV(cctx, data, redop);
+}
+
+template <typename T, typename Fn>
+std::enable_if_t<
+    std::is_invocable_v<Fn, common::Span<T const>, common::Span<T const>, std::vector<T>*>, Result>
+AllreduceV(Context const* ctx, std::vector<T>* data, Fn redop) {
+  return AllreduceV(ctx, *GlobalCommGroup(), data, redop);
+}
 }  // namespace xgboost::collective