Function ForEachOpGraphNecessaryCtrlEdge

oneflow/core/graph/task_graph.cpp:446–483

}

void ForEachOpGraphNecessaryCtrlEdge(
    const OpGraph* op_graph, const std::function<void(const OpNode*, const OpNode*)>& Handler) {
  auto IsOpGraphDataReachable = op_graph->CreatePredicatorIsReachable();
  op_graph->ForEachNode([&](OpNode* dst) {
    for (const auto& ctrl_in_op_name : dst->op().op_conf().ctrl_in_op_name()) {
      const OpNode* src = op_graph->OpNode4OpName(ctrl_in_op_name);
      CHECK(!IsOpGraphDataReachable(dst, src));
      // src has ctrl to dst, but src has no data path to dst.
      if (!IsOpGraphDataReachable(src, dst)) {
        CHECK_EQ(dst->parallel_desc().parallel_num(), src->parallel_desc().parallel_num());
        const Shape* src_time_shape = CHECK_JUST(src->op().GetOpTimeShape()).get();
        const Shape* dst_time_shape = CHECK_JUST(dst->op().GetInputBlobFastestTimeShape()).get();
        if (dst_time_shape == nullptr) {
          dst_time_shape = CHECK_JUST(dst->op().GetOpTimeShape()).get();
        }
        if (src_time_shape->elem_cnt() != dst_time_shape->elem_cnt()) {
          // NOTE(chengcheng): acc / pack op node can be merged and add ctrl edge.
          CHECK(src->op().op_conf().has_user_conf());
          const std::string& op_type_name = src->op().op_conf().user_conf().op_type_name();
          CHECK(op_type_name == "acc" || op_type_name == "pack");
          const Shape* src_input_time_shape =
              CHECK_JUST(src->op().GetInputBlobFastestTimeShape()).get();
          CHECK_EQ(src_input_time_shape->elem_cnt(), dst_time_shape->elem_cnt());
        } else {
          CHECK_EQ(src_time_shape->elem_cnt(), dst_time_shape->elem_cnt());
        }
        if (!src->parallel_desc().EqualsIgnoringHierarchy(dst->parallel_desc())) {
          LOG(WARNING) << " Warning, there is a ctrl edge connected across placement from: "
                       << src->op().op_name() << " ["
                       << src->parallel_desc().parallel_conf().DebugString()
                       << "] to: " << dst->op().op_name() << " ["
                       << dst->parallel_desc().parallel_conf().DebugString() << "]";
        }
        Handler(src, dst);
      }
    }
  });
}

void GetHostInputLbis4OpNode(const OpNode* op_node,
                             std::vector<LogicalBlobId>* host_mem_input_lbis) {