Method InitializeMemory

oneflow/core/auto_parallel/sbp_node.cpp:739–791

Assemble memory cost

// Assemble memory cost
void SbpNode::InitializeMemory(bool is_reusable, const HashMap<LogicalBlobId, int32_t>& lbi2id,
                               const std::vector<int32_t>& id2count, bool nccl_use_compute_stream) {
  const auto& curr_operator = op_node_->op();
  // An edge should not be initialized twice
  // During each initialization, we are computing sum(memory of consumer) - sum(memory of producer)
  // This is why we need to pre-store memory of producer
  HashMap<SbpEdge*, std::vector<int64_t>> sbp_edge2nd_sbp_sig2memory;
  for (const auto& obn : curr_operator.output_bns()) {
    const LogicalBlobId& lbi = curr_operator.BnInOp2Lbi(obn);
    // Fixed memory or in the support of the reusable memory
    if (!is_reusable || id2count.at(lbi2id.at(lbi)) > 0) {
      // If not in support, memory_ would be empty.
      in_memory_support_ = true;
      memory_.resize(sbp_sig_list_.size(), 0);
      const auto& logical_blob_desc = op_node_->LogicalBlobDesc4Lbi(lbi);
      const auto& hierarchy = *CHECK_JUST(curr_operator.GetParallelDesc4BnInOp(obn))->hierarchy();
      // There are some operators with a fixed sbp for some blobs, such as conv.
      // {in: S0, kernel: B, out: S0}
      // {in: B, kernel: B, out: B}
      // The blob kernel have the same sbp for different signatures.
      // We pre-store the results for the same sbp while accessing the same blobs.
      HashMap<NdSbp, int64_t> nd_sbp2memory;
      SbpEdge* edge_contain_lbi = nullptr;
      for (const auto& edge_out : edges_out_) {
        if (edge_out->SearchLbi(lbi)) { edge_contain_lbi = edge_out; }
      }
      // There exist some lbi which does not have a consumer
      // At this moment edge_contain_lbi == nullptr
      auto& nd_sbp_sig2memory = sbp_edge2nd_sbp_sig2memory[edge_contain_lbi];
      nd_sbp_sig2memory.resize(sbp_sig_list_.size(), 0);
      for (int32_t sbp_sig_id = 0; sbp_sig_id < sbp_sig_list_.size(); sbp_sig_id++) {
        const NdSbp& nd_sbp = sbp_sig_list_[sbp_sig_id].bn_in_op2nd_sbp().at(obn);
        auto it = nd_sbp2memory.find(nd_sbp);
        if (it == nd_sbp2memory.end()) {
          // This compute the memory at rank 0, the largest one.
          // We could be faster if we just compute the average memory.
          it = nd_sbp2memory
                   .insert({nd_sbp, MaxByteSize4BlobDescSbp(logical_blob_desc, nd_sbp, hierarchy)})
                   .first;
        }
        memory_[sbp_sig_id] += it->second;
        nd_sbp_sig2memory[sbp_sig_id] += it->second;
      }
    }
  }
  // Even after the correction in the memory of edges, the relative error still have 0.73%.
  if (nccl_use_compute_stream && in_memory_support_ && is_reusable) {
    for (const auto& pair : sbp_edge2nd_sbp_sig2memory) {
      // Init memory for each out-going edge
      pair.first->InitializeMemory(lbi2id, id2count, pair.second);
    }
  }
}

// Reduce and set the wait time for op in the trunk
void SbpNode::SetTrunkWaitTime(double trunk_wait_time) {