Method InitCopyAndMemoryCost

oneflow/core/auto_parallel/sbp_edge.cpp:324–414

Assemble copy cost

// Assemble copy cost
void SbpEdge::InitCopyAndMemoryCost(const std::string& ibn, bool use_sbp_collector,
                                    bool nccl_not_use_compute_stream) {
  std::vector<int64_t> consumer_nd_sbp_sig2memory;
  if (nccl_not_use_compute_stream) {
    in_memory_support_ = true;
    // Compute and store the memory for consumer
    const auto& consumer_operator = end_node_->op_node_->op();
    const auto& end_sbp_sig_list = end_node_->sbp_sig_list_;
    consumer_nd_sbp_sig2memory.resize(end_sbp_sig_list.size(), 0);
    const auto& lbi = consumer_operator.BnInOp2Lbi(ibn);
    const auto& consumer_hierarchy =
        *CHECK_JUST(consumer_operator.GetParallelDesc4BnInOp(ibn))->hierarchy();
    const auto& logical_blob_desc = start_node_->op_node_->LogicalBlobDesc4Lbi(lbi);
    HashMap<NdSbp, int64_t> consumer_nd_sbp2memory;
    for (int32_t sbp_sig_id = 0; sbp_sig_id < end_sbp_sig_list.size(); sbp_sig_id++) {
      const NdSbp& nd_sbp = end_sbp_sig_list[sbp_sig_id].bn_in_op2nd_sbp().at(ibn);
      auto it = consumer_nd_sbp2memory.find(nd_sbp);
      if (it == consumer_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 = consumer_nd_sbp2memory
                 .insert({nd_sbp,
                          MaxByteSize4BlobDescSbp(logical_blob_desc, nd_sbp, consumer_hierarchy)})
                 .first;
      }
      consumer_nd_sbp_sig2memory[sbp_sig_id] += it->second;
    }
  }

  // In this part, we assemble the cost from nodes to nodes.
  if (start_node_->op_node_ && end_node_->op_node_) {
    OpNode* consumer = end_node_->op_node_;

    // Add copy cost for each blob
    const LogicalBlobId& lbi = consumer->op().BnInOp2Lbi(ibn);

    // Check whether lbi is transferred by this edge
    if (use_sbp_collector && !SearchLbi(lbi)) { return; }

    OpNode* producer = start_node_->op_node_;
    const std::string& producer_lbn = *CHECK_JUST(producer->op().obn4lbi(lbi));
    const ParallelDesc& producer_parallel_desc =
        *CHECK_JUST(producer->op().GetParallelDesc4BnInOp(producer_lbn));
    const ParallelDesc& consumer_parallel_desc =
        *CHECK_JUST(consumer->op().GetParallelDesc4BnInOp(ibn));

    // Need to be careful, the logical blob description should be independent to current
    // SbpParallel. Use producer or op_node?
    const BlobDesc& logical_blob_desc = producer->LogicalBlobDesc4Lbi(lbi);
    const std::string& obn = *CHECK_JUST(producer->op().obn4lbi(lbi));
    // If we are deciding whether we need the wait time, then make require_same_sbp true.
    // B->S cause cudaEventSynchronize in current implementation.
    bool require_same_sbp = RequireSameSbp(consumer, ibn);
    int32_t consumer_sbp_size = end_node_->sbp_sig_list_.size();
    LazyMode::Guard enable_lazy_mode(true);

    // look through sbp signature in producer
    for (int32_t sbp_id_producer = 0; sbp_id_producer < start_node_->sbp_sig_list_.size();