| 727 | } |
| 728 | |
| 729 | Status Conditional::BuildIfNode(Graph* graph, |
| 730 | FunctionLibraryDefinition* library) { |
| 731 | VLOG(2) << "Build cond function for " << name(); |
| 732 | NodeDebugInfo debug_info((*merges_.begin())->def()); |
| 733 | NodeDefBuilder builder(name(), "If", library, &debug_info); |
| 734 | const string branch_name[] = {"else_branch", "then_branch"}; |
| 735 | for (auto branch : {BranchType::kElseBranch, BranchType::kThenBranch}) { |
| 736 | int branch_index = static_cast<int>(branch); |
| 737 | static std::atomic<int64> sequence_num(0LL); |
| 738 | int64 id = ++sequence_num; |
| 739 | |
| 740 | NameAttrList body_name; |
| 741 | body_name.set_name( |
| 742 | absl::StrCat("_functionalize_if_", branch_name[branch_index], "_", id)); |
| 743 | |
| 744 | VLOG(3) << "FunctionalizeControlFlow (" << branch_name[branch_index] |
| 745 | << "): " |
| 746 | << DumpGraphToFile( |
| 747 | "functionalize_cond_body_" + branch_name[branch_index], |
| 748 | *bodies_[branch_index], nullptr); |
| 749 | |
| 750 | FunctionDef body_fdef; |
| 751 | TF_RETURN_IF_ERROR(GraphToFunctionDef(*bodies_[branch_index], |
| 752 | body_name.name(), &body_fdef)); |
| 753 | TF_RETURN_IF_ERROR(library->AddFunctionDef(body_fdef)); |
| 754 | builder.Attr(branch_name[branch_index], body_name); |
| 755 | } |
| 756 | |
| 757 | VLOG(3) << "Build input type"; |
| 758 | std::vector<NodeDefBuilder::NodeOut> inputs; |
| 759 | DataTypeVector in_arg_types; |
| 760 | for (auto& kv : cond_arg_nodes_) { |
| 761 | bool inserted = false; |
| 762 | for (const Node* arg : kv.switches) { |
| 763 | const Edge* in_edge; |
| 764 | TF_RETURN_IF_ERROR(arg->input_edge(0, &in_edge)); |
| 765 | if (in_edge->IsControlEdge()) { |
| 766 | builder.ControlInput(in_edge->src()->name()); |
| 767 | } else { |
| 768 | if (!inserted) { |
| 769 | DataType dtype = arg->input_type(0); |
| 770 | inputs.emplace_back(NodeDefBuilder::NodeOut( |
| 771 | in_edge->src()->name(), in_edge->src_output(), dtype)); |
| 772 | in_arg_types.push_back(dtype); |
| 773 | inserted = true; |
| 774 | } |
| 775 | } |
| 776 | } |
| 777 | } |
| 778 | builder.Attr("Tin", in_arg_types); |
| 779 | |
| 780 | DataTypeVector out_type; |
| 781 | for (const Node* merge : merges_) { |
| 782 | DataType dtype = merge->output_type(0); |
| 783 | out_type.push_back(dtype); |
| 784 | } |
| 785 | builder.Attr("Tout", out_type); |
| 786 | VLOG(3) << "Build output type: " << DataTypeVectorString(out_type); |
nothing calls this directly
no test coverage detected