Method RunTree

tensorflow/core/common_runtime/hierarchical_tree_broadcaster.cc:292–406

Executes a hierarchical tree broadcast. Each subdiv is a broadcast between a subset of the devices. If there is only one task, there is one subdiv comprising a broadcast between all devices belonging to the task. If there are n tasks, n>1, then there are n+1 subdivs. In the first (global) subdiv, one device from each task participates in a binary tree broadcast. Each task receives a copy of the t

// corresponds to broadcast between all devices on task i.  Thus, each task
// participates in at most 2 subdivs.
void HierarchicalTreeBroadcaster::RunTree() {
  int num_subdivs = static_cast<int>(col_params_->subdiv_rank.size());
  // TODO(b/78352018): this is easily improved when a node participates in both
  // first and second subdivision.  It would first send to its descendents in
  // the first subdiv, then wait until all pending ops are finished before
  // sending to descendents in second subdiv.  A better implementation would
  // collapse the two send blocks.
  for (int si = 0; si < num_subdivs; si++) {
    int my_rank = col_params_->subdiv_rank[si];
    // If rank is -1, this device does not participate in this subdiv.
    if (-1 == my_rank) continue;
    int source_rank = col_params_->instance.impl_details.subdiv_source_rank[si];
    if (VLOG_IS_ON(1)) {
      string subdiv_buf;
      for (int r : col_params_->instance.impl_details.subdiv_permutations[si]) {
        strings::StrAppend(&subdiv_buf, r, ",");
      }
      VLOG(1) << "Running Broadcast tree device=" << col_ctx_->device_name
              << " subdiv=" << si << " perm=" << subdiv_buf
              << " my_rank=" << my_rank << " source_rank=" << source_rank;
    }

    mutex mu;               // also guards status_ while callbacks are pending
    int pending_count = 0;  // GUARDED_BY(mu)
    condition_variable all_done;

    if (my_rank >= 0 && my_rank != source_rank) {
      // Begin by receiving the value.
      profiler::TraceMe activity(
          [&] { return strings::StrCat("ReceiveValue:", si); },
          profiler::TraceMeLevel::kInfo);
      int recv_from_rank = TreeRecvFrom(*col_params_, si);
      Notification note;
      DispatchRecv(si, recv_from_rank, my_rank, col_ctx_->output,
                   [this, &mu, &note](const Status& s) {
                     mutex_lock l(mu);
                     status_.Update(s);
                     note.Notify();
                   });
      note.WaitForNotification();
    }

    // Then forward value to all descendent devices.
    {
      profiler::TraceMe activity(
          [&] { return strings::StrCat("ForwardValue:", si); },
          profiler::TraceMeLevel::kInfo);
      if (my_rank >= 0 && status_.ok()) {
        std::vector<int> send_to_ranks;
        TreeSendTo(*col_params_, si, &send_to_ranks);
        for (int i = 0; i < send_to_ranks.size(); ++i) {
          int target_rank = send_to_ranks[i];
          {
            mutex_lock l(mu);
            ++pending_count;
          }
          DispatchSend(si, target_rank, my_rank,
                       (is_source_ ? col_ctx_->input : col_ctx_->output),