Method PrintSBPGraphDebugInfo

oneflow/core/auto_parallel/sbp_constructor.cpp:500–573

Print the graph with SBP in order

// Print the graph with SBP in order
void SbpConstructor::PrintSBPGraphDebugInfo() {
  // sbp constructor information
  std::cout << "cost_ratio_:" << cost_ratio_ << std::endl;
  std::cout << "wait_time_:" << sbp_graph_.wait_time_ << std::endl;
  std::cout << "use_sbp_collector_" << use_sbp_collector_ << std::endl;
  std::cout << "Total auto parallel guessed memory: " << sbp_graph_.GetMemory() << std::endl;
  std::cout << "Final memory ratio: " << kMemoryRatio << std::endl;
  // test debug
  std::cout << "Get Into Print Op Graph" << std::endl;
  // Collect op_node
  std::vector<OpNode*> node_list;
  for (const auto& op_name_sbp_node : op_name2sbp_node_) {
    auto* op_node_ = op_name_sbp_node.second->op_node_;
    if (op_node_) { node_list.push_back(op_node_); }
  }

  // test debug
  std::cout << "Deciding order" << std::endl;
  // Decide the order to visit the op
  std::vector<int32_t> order;
  auto_parallel::DecideOrder(node_list, order, [&](OpNode* a, OpNode* b) {
    return a->op().op_name().compare(b->op().op_name()) > 0;
  });
  std::vector<int32_t> str_order;

  // test debug
  std::cout << "Finish deciding order" << std::endl;

  for (int32_t i = 0; i < node_list.size(); i++) {
    OpNode* op_node = node_list[order[i]];
    std::cout << op_node->op().op_name() << " (^_^):" << std::endl;
    // get corresponding sbp node
    const auto& it = op_name2sbp_node_.find(op_node->op().op_name());
    // Print debug information for sbp graph
    CHECK(it != op_name2sbp_node_.end());
    const SbpNode* sbp_node = it->second;
    std::cout << "Computation Cost: " << sbp_node->weighted_cost_[sbp_node->final_sbp_sig_id_];
    std::cout << ", Min Layer: " << sbp_node->min_layer_ << ", Max Layer: " << sbp_node->max_layer_
              << ", Tributary Layer: " << sbp_node->tributary_layer_
              << ", in trunk: " << sbp_node->on_trunk_
              << ", Remain Cost: " << sbp_node->acc_trunk_cost_ << std::endl;
    // Sort before printing
    const auto& op_input_bns = op_node->op().input_bns();
    auto CompareString = [](const std::string& a, const std::string& b) {
      return a.compare(b) > 0;
    };
    auto_parallel::DecideOrder(op_input_bns, str_order, CompareString);
    const NdSbpSignature& sbp_signature = sbp_node->FinalSbpSignature();
    // Print out SBP information for input operator
    for (int32_t j : str_order) {
      const auto& ibn = op_input_bns[j];
      const auto& producer_node = op_node->SrcNode4Ibn(ibn);
      std::cout << "Pre Op:" << producer_node.op().op_name() << ": " << ibn;
      const auto& this_sbp_parallel = sbp_signature.bn_in_op2nd_sbp().at(ibn);
      std::cout << ", " << NdSbpToString(this_sbp_parallel);
      if (RequireSameSbp(op_node, ibn)) { std::cout << ", require same SBP"; }
      std::cout << ", " << op_node->LogicalBlobDesc4Lbi(op_node->op().BnInOp2Lbi(ibn)).shape();
      std::cout << std::endl;