tree.

Author: trivialfis

src/collective/allreduce.h

@@ -1,5 +1,5 @@
 /**
- * Copyright 2023-2024, XGBoost Contributors
+ * Copyright 2023-2026, XGBoost Contributors
  */
 #pragma once
 #include <cstdint>      // for int8_t
@@ -147,25 +147,24 @@ AllreduceV(Comm const& comm, std::vector<T>* data, Fn redop) {
     };
   };
 
-  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) {
+  for (std::int32_t level = 0; (std::int32_t{1} << level) < world; ++level) {
     if (!continue_reduce) {
       continue;
     }
-    if (shifted_rank % (step * 2) == step) {
-      auto parent = shifted_rank - step;
+    if (rank > 0 && binomial_tree::ParentLevel(rank) == level) {
+      auto parent = binomial_tree::Parent(rank);
       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;
+    if (binomial_tree::HasChild(rank, level, world)) {
+      auto child = binomial_tree::Child(rank, level);
       auto rc = recv(child, &incoming);
       if (!rc.OK()) {
         return Fail("AllreduceV failed to receive data from child.", std::move(rc));

src/collective/broadcast.cc

@@ -3,13 +3,11 @@
  */
 #include "broadcast.h"
 
-#include <cmath>    // for ceil, log2
 #include <cstdint>  // for int32_t, int8_t
 #include <utility>  // for move
 #include <vector>   // for vector
 
-#include "../common/bitfield.h"         // for TrailingZeroBits, RBitField32
-#include "comm.h"                       // for Comm
+#include "comm.h"                       // for Comm, binomial_tree
 #include "xgboost/collective/result.h"  // for Result
 #include "xgboost/span.h"               // for Span
 
@@ -19,11 +17,9 @@ namespace {
 Result BroadcastTree(Comm const& comm, common::Span<std::int8_t> data) {
   auto rank = comm.Rank();
   auto world = comm.World();
-  std::int32_t depth = std::ceil(std::log2(static_cast<double>(world))) - 1;
 
   if (rank != 0) {
-    auto k = TrailingZeroBits(static_cast<std::uint32_t>(rank));
-    auto parent = rank - (1 << k);
+    auto parent = binomial_tree::Parent(rank);
     auto rc = Success() << [&] {
       return comm.Chan(parent)->RecvAll(data);
     } << [&] {
@@ -34,9 +30,9 @@ Result BroadcastTree(Comm const& comm, common::Span<std::int8_t> data) {
     }
   }
 
-  for (std::int32_t i = depth; i >= 0; --i) {
-    if (rank % (1 << (i + 1)) == 0 && rank + (1 << i) < world) {
-      auto child = rank + (1 << i);
+  for (std::int32_t level = binomial_tree::Depth(world); level >= 0; --level) {
+    if (binomial_tree::HasChild(rank, level, world)) {
+      auto child = binomial_tree::Child(rank, level);
       auto rc = comm.Chan(child)->SendAll(data);
       if (!rc.OK()) {
         return rc;
@@ -48,27 +44,25 @@ Result BroadcastTree(Comm const& comm, common::Span<std::int8_t> data) {
 }
 
 // Compute the path from `src` to rank 0 through the binomial tree (excluding 0).
-std::vector<std::int32_t> TreePathToRoot(std::int32_t src) {
+std::vector<std::int32_t> TreePathToRoot(std::int32_t node) {
   std::vector<std::int32_t> path;
-  auto cursor = src;
+  auto cursor = node;
   while (cursor > 0) {
     path.push_back(cursor);
-    auto k = TrailingZeroBits(static_cast<std::uint32_t>(cursor));
-    cursor -= (1 << k);
+    cursor = binomial_tree::Parent(cursor);
   }
   return path;
 }
 
-// Relay data from `root` up to rank 0 through the binomial tree.
-// Only nodes on the path from root to 0 participate; all others skip.
-Result RelayToRoot(Comm const& comm, common::Span<std::int8_t> data, std::int32_t root) {
+// Relay data from `node` up to rank 0 through the binomial tree.
+// Only nodes on the path from `node` to 0 participate; all others skip.
+Result RelayToRoot(Comm const& comm, common::Span<std::int8_t> data, std::int32_t node) {
   auto rank = comm.Rank();
-  auto path = TreePathToRoot(root);
+  auto path = TreePathToRoot(node);
 
   for (auto node : path) {
     CHECK_GT(node, 0);
-    auto k = TrailingZeroBits(static_cast<std::uint32_t>(node));
-    auto parent = node - (std::int32_t{1} << k);
+    auto parent = binomial_tree::Parent(node);
 
     if (rank == node) {
       auto rc = Success() << [&] {

src/collective/comm.cc

@@ -217,6 +217,9 @@ Result ConnectTrackerImpl(proto::PeerInfo info, std::chrono::seconds timeout, st
     workers[prev_rank] = std::move(prev);
   }
 
+  // For each peer pair, the lower-ranked worker initiates (connect) and the higher-ranked
+  // worker accepts. Ring sockets are already in place, so only tree-only peers need new
+  // connections.
   for (auto r : my_peers) {
     if (r > comm.Rank() && !workers[r]) {
       rc = ConnectPeer(peers[r], comm.Rank(), timeout, retry, &workers[r]);
@@ -226,6 +229,9 @@ Result ConnectTrackerImpl(proto::PeerInfo info, std::chrono::seconds timeout, st
     }
   }
 
+  // Accept connections from lower-ranked tree peers that weren't already covered by the
+  // ring. The exact arrival order is unspecified, so we accept n_accept times and use the
+  // rank sent by each peer to place the socket in the right slot.
   std::int32_t n_accept = 0;
   for (auto r : my_peers) {
     if (r < comm.Rank() && !workers[r]) {

src/collective/comm.h

@@ -25,7 +25,7 @@ namespace xgboost::collective {
 inline constexpr std::int64_t DefaultTimeoutSec() { return 60 * 30; }  // 30min
 inline constexpr std::int32_t DefaultRetry() { return 3; }
 
-// indexing into the ring
+// Indexing into the ring
 inline std::int32_t BootstrapNext(std::int32_t r, std::int32_t world) {
   auto nrank = (r + world + 1) % world;
   return nrank;
@@ -36,6 +36,41 @@ inline std::int32_t BootstrapPrev(std::int32_t r, std::int32_t world) {
   return nrank;
 }
 
+/**
+ * @brief Helpers for the binomial tree rooted at rank 0.
+ *
+ * References
+ * - https://people.mpi-inf.mpg.de/~mehlhorn/ftp/NewToolbox/collective.pdf
+ * - https://en.wikipedia.org/wiki/Broadcast_(parallel_pattern)
+ */
+namespace binomial_tree {
+inline std::int32_t ParentLevel(std::int32_t rank) {
+  CHECK_GT(rank, 0);
+  return static_cast<std::int32_t>(TrailingZeroBits(static_cast<std::uint32_t>(rank)));
+}
+
+inline std::int32_t Parent(std::int32_t rank) {
+  return rank - (std::int32_t{1} << ParentLevel(rank));
+}
+
+inline std::int32_t Child(std::int32_t rank, std::int32_t level) {
+  return rank + (std::int32_t{1} << level);
+}
+
+inline bool HasChild(std::int32_t rank, std::int32_t level, std::int32_t world) {
+  return rank % (std::int32_t{1} << (level + 1)) == 0 && Child(rank, level) < world;
+}
+
+inline std::int32_t Depth(std::int32_t world) {
+  if (world <= 1) return -1;
+  std::int32_t depth = 0;
+  while ((std::int32_t{1} << (depth + 1)) < world) {
+    ++depth;
+  }
+  return depth;
+}
+}  // namespace binomial_tree
+
 /**
  * @brief Compute the sparse peer set for a given rank: ring neighbors union binomial tree
  *        neighbors (rooted at rank 0).
@@ -49,17 +84,14 @@ inline std::vector<std::int32_t> SparsePeers(std::int32_t rank, std::int32_t wor
   peers.insert(BootstrapNext(rank, world));
   peers.insert(BootstrapPrev(rank, world));
 
+  // Connect tree parents and children
   if (rank > 0) {
-    auto k = TrailingZeroBits(static_cast<std::uint32_t>(rank));
-    peers.insert(rank - (1 << k));
+    peers.insert(binomial_tree::Parent(rank));
   }
 
-  for (std::int32_t k = 0; (1 << k) < world; ++k) {
-    if (rank % (1 << (k + 1)) == 0) {
-      auto child = rank + (1 << k);
-      if (child < world) {
-        peers.insert(child);
-      }
+  for (std::int32_t level = 0; (std::int32_t{1} << level) < world; ++level) {
+    if (binomial_tree::HasChild(rank, level, world)) {
+      peers.insert(binomial_tree::Child(rank, level));
     }
   }