[coll] Avoid all-to-all connection.

R-package/src/Makevars.in

@@ -112,6 +112,7 @@ OBJECTS= \
     $(PKGROOT)/src/collective/tracker.o \
     $(PKGROOT)/src/collective/in_memory_handler.o \
     $(PKGROOT)/src/collective/loop.o \
+    $(PKGROOT)/src/collective/topo.o \
     $(PKGROOT)/src/collective/socket.o \
     $(PKGROOT)/src/common/charconv.o \
     $(PKGROOT)/src/common/column_matrix.o \

R-package/src/Makevars.win

@@ -112,6 +112,7 @@ OBJECTS= \
     $(PKGROOT)/src/collective/tracker.o \
     $(PKGROOT)/src/collective/in_memory_handler.o \
     $(PKGROOT)/src/collective/loop.o \
+    $(PKGROOT)/src/collective/topo.o \
     $(PKGROOT)/src/collective/socket.o \
     $(PKGROOT)/src/common/charconv.o \
     $(PKGROOT)/src/common/column_matrix.o \

include/xgboost/collective/socket.h

@@ -3,6 +3,7 @@
  */
 #pragma once
 
+#include <algorithm>     // for max
 #include <cerrno>        // errno, EINTR, EBADF
 #include <climits>       // HOST_NAME_MAX
 #include <cstddef>       // std::size_t
@@ -539,7 +540,8 @@ class TCPSocket {
   /**
    * @brief Listen to incoming requests. Should be called after bind.
    */
-  [[nodiscard]] Result Listen(std::int32_t backlog = 16) {
+  [[nodiscard]] Result Listen(std::int32_t backlog) {
+    backlog = std::max(backlog, 16);  // Don't be too small.
     if (listen(handle_, backlog) != 0) {
       return system::FailWithCode("Failed to listen.");
     }

src/collective/allgather.h

@@ -13,6 +13,7 @@
 #include "../common/type.h"             // for EraseType
 #include "comm.h"                       // for Comm, Channel
 #include "comm_group.h"                 // for CommGroup
+#include "topo.h"                       // for BootstrapNext, BootstrapPrev
 #include "xgboost/collective/result.h"  // for Result
 #include "xgboost/linalg.h"             // for MakeVec
 #include "xgboost/span.h"               // for Span

src/collective/broadcast.cc

@@ -1,51 +1,19 @@
 /**
- * Copyright 2023, XGBoost Contributors
+ * Copyright 2023-2024, XGBoost Contributors
  */
 #include "broadcast.h"
 
 #include <cmath>    // for ceil, log2
 #include <cstdint>  // for int32_t, int8_t
 #include <utility>  // for move
 
-#include "../common/bitfield.h"         // for TrailingZeroBits, RBitField32
+
 #include "comm.h"                       // for Comm
 #include "xgboost/collective/result.h"  // for Result
 #include "xgboost/span.h"               // for Span
+#include "topo.h"                       // for ParentRank
 
 namespace xgboost::collective::cpu_impl {
-namespace {
-std::int32_t ShiftedParentRank(std::int32_t shifted_rank, std::int32_t depth) {
-  std::uint32_t mask{std::uint32_t{0} - 1};  // Oxff...
-  RBitField32 maskbits{common::Span<std::uint32_t>{&mask, 1}};
-  RBitField32 rankbits{
-      common::Span<std::uint32_t>{reinterpret_cast<std::uint32_t*>(&shifted_rank), 1}};
-  // prepare for counting trailing zeros.
-  for (std::int32_t i = 0; i < depth + 1; ++i) {
-    if (rankbits.Check(i)) {
-      maskbits.Set(i);
-    } else {
-      maskbits.Clear(i);
-    }
-  }
-
-  CHECK_NE(mask, 0);
-  auto k = TrailingZeroBits(mask);
-  auto shifted_parent = shifted_rank - (1 << k);
-  return shifted_parent;
-}
-
-// Shift the root node to rank 0
-std::int32_t ShiftLeft(std::int32_t rank, std::int32_t world, std::int32_t root) {
-  auto shifted_rank = (rank + world - root) % world;
-  return shifted_rank;
-}
-// shift back to the original rank
-std::int32_t ShiftRight(std::int32_t rank, std::int32_t world, std::int32_t root) {
-  auto orig = (rank + root) % world;
-  return orig;
-}
-}  // namespace
-
 Result Broadcast(Comm const& comm, common::Span<std::int8_t> data, std::int32_t root) {
   // Binomial tree broadcast
   // * Wiki
@@ -56,28 +24,47 @@ Result Broadcast(Comm const& comm, common::Span<std::int8_t> data, std::int32_t
   auto rank = comm.Rank();
   auto world = comm.World();
 
-  // shift root to rank 0
-  auto shifted_rank = ShiftLeft(rank, world, root);
+  // Send data to the root to preserve the topology. Alternative is to shift the rank, but
+  // it requires a all-to-all connection.
+  //
+  // Most of the use of broadcasting in XGBoost are short messages, this should be
+  // fine. Otherwise, we can implement a linear pipeline broadcast.
+  if (root != 0) {
+    auto rc = Success() << [&] {
+      return (rank == 0) ? comm.Chan(root)->RecvAll(data) : Success();
+    } << [&] {
+      return (rank == root) ? comm.Chan(0)->SendAll(data) : Success();
+    } << [&] {
+      return comm.Block();
+    };
+    if (!rc.OK()) {
+      return Fail("Broadcast failed to send data to root.", std::move(rc));
+    }
+    root = 0;
+  }
+
   std::int32_t depth = std::ceil(std::log2(static_cast<double>(world))) - 1;
 
-  if (shifted_rank != 0) {  // not root
-    auto parent = ShiftRight(ShiftedParentRank(shifted_rank, depth), world, root);
-    auto rc = Success() << [&] { return comm.Chan(parent)->RecvAll(data); }
-                        << [&] { return comm.Chan(parent)->Block(); };
+  if (rank != 0) {  // not root
+    auto parent = ParentRank(rank, depth);
+    auto rc = Success() << [&] {
+      return comm.Chan(parent)->RecvAll(data);
+    } << [&] {
+      return comm.Chan(parent)->Block();
+    };
     if (!rc.OK()) {
-      return Fail("broadcast failed.", std::move(rc));
+      return Fail("Broadcast failed to send data to parent.", std::move(rc));
     }
   }
 
   for (std::int32_t i = depth; i >= 0; --i) {
     CHECK_GE((i + 1), 0);  // weird clang-tidy error that i might be negative
-    if (shifted_rank % (1 << (i + 1)) == 0 && shifted_rank + (1 << i) < world) {
-      auto sft_peer = shifted_rank + (1 << i);
-      auto peer = ShiftRight(sft_peer, world, root);
+    if (rank % (1 << (i + 1)) == 0 && rank + (1 << i) < world) {
+      auto peer = rank + (1 << i);
       CHECK_NE(peer, root);
       auto rc = comm.Chan(peer)->SendAll(data);
       if (!rc.OK()) {
-        return rc;
+        return Fail("Failed to seed to " + std::to_string(peer), std::move(rc));
       }
     }
   }