| 400 | } |
| 401 | |
| 402 | Status Conditional::BuildArgumentNodes() { |
| 403 | VLOG(1) << "Build function arguments"; |
| 404 | struct Hash { |
| 405 | size_t operator()(const std::pair<Node*, int>& item) const { |
| 406 | return Hash64Combine(hash<Node*>()(item.first), |
| 407 | std::hash<int>()(item.second)); |
| 408 | } |
| 409 | }; |
| 410 | |
| 411 | std::unordered_map<std::pair<Node*, int>, int, Hash> input_index; |
| 412 | for (Node* switch_node : switches_) { |
| 413 | const Edge* e; |
| 414 | TF_RETURN_IF_ERROR(switch_node->input_edge(0, &e)); |
| 415 | std::pair<Node*, int> key = std::make_pair(e->src(), e->src_output()); |
| 416 | if (input_index.find(key) == input_index.end()) { |
| 417 | input_index[key] = cond_arg_nodes_.size(); |
| 418 | cond_arg_nodes_.emplace_back(key.first, key.second); |
| 419 | } |
| 420 | cond_arg_nodes_.at(input_index.at(key)).switches.push_back(switch_node); |
| 421 | } |
| 422 | VLOG(5) << "CondArg nodes created: " << DebugString(cond_arg_nodes_); |
| 423 | |
| 424 | int arg_count = 0; |
| 425 | for (CondArgNode& cond_arg_node : cond_arg_nodes_) { |
| 426 | DataType dtype = cond_arg_node.src->output_type(cond_arg_node.src_output); |
| 427 | for (auto branch : {BranchType::kElseBranch, BranchType::kThenBranch}) { |
| 428 | int branch_index = static_cast<int>(branch); |
| 429 | TF_RETURN_IF_ERROR( |
| 430 | NodeBuilder(absl::StrCat("_Arg", arg_count), |
| 431 | FunctionLibraryDefinition::kArgOp) |
| 432 | .Attr("T", dtype) |
| 433 | .Attr("index", arg_count) |
| 434 | .Finalize(bodies_[branch_index].get(), |
| 435 | &cond_arg_node.branch_copy[branch_index])); |
| 436 | } |
| 437 | for (Node* node : cond_arg_node.switches) { |
| 438 | for (const Edge* e : node->out_edges()) { |
| 439 | if (e->IsControlEdge()) continue; |
| 440 | int branch_index = e->src_output(); |
| 441 | Node* src_copy = cond_arg_node.branch_copy[branch_index]; |
| 442 | Node* dst_copy = node_maps_[branch_index][e->dst()->id()]; |
| 443 | |
| 444 | // The graph may contain dead switch nodes, |
| 445 | if (dst_copy == nullptr) continue; |
| 446 | |
| 447 | TF_RET_CHECK(dst_copy != nullptr) |
| 448 | << "Unable to find copied node for " << e->dst()->DebugString() |
| 449 | << " on branch " << Branch_Name(BranchType(branch_index)); |
| 450 | // If the input goes directly to a merge then the merge has |
| 451 | // been replaced by a retval so the dst input is 0 instead of |
| 452 | // dst_input. |
| 453 | int dst_input = IsMerge(e->dst()) ? 0 : e->dst_input(); |
| 454 | bodies_[branch_index]->AddEdge(src_copy, 0, dst_copy, dst_input); |
| 455 | } |
| 456 | } |
| 457 | ++arg_count; |
| 458 | } |
| 459 |
nothing calls this directly
no test coverage detected