| 261 | } |
| 262 | |
| 263 | Status FunctionalizeLoop(const FunctionLibraryDefinition* lookup_library, |
| 264 | Graph* graph, WhileLoopFrame* frame, |
| 265 | FunctionLibraryDefinition* library) { |
| 266 | VLOG(2) << "Frame " << frame->name << " before: " |
| 267 | << DumpGraphToFile("functionalize_before", *graph, library); |
| 268 | |
| 269 | // Split loop-varying Enter nodes with multiple successors. If the same |
| 270 | // Tensor is fed as input to multiple loop arguments, we may end up with a |
| 271 | // shared Enter node. We clone Enter nodes with multiple successors to |
| 272 | // maintain the invariant of a unique Enter node per argument of the final |
| 273 | // loop. |
| 274 | std::vector<WhileLoopArg> args; |
| 275 | for (const WhileLoopArg& arg : frame->args) { |
| 276 | if (arg.is_loop_invariant) { |
| 277 | args.push_back(arg); |
| 278 | } else { |
| 279 | std::vector<const Edge*> edges(arg.enter->out_edges().begin(), |
| 280 | arg.enter->out_edges().end()); |
| 281 | for (int i = 0; i < edges.size(); ++i) { |
| 282 | if (edges[i]->IsControlEdge() && edges[i]->dst()->IsSink()) { |
| 283 | continue; |
| 284 | } |
| 285 | TF_RET_CHECK(!edges[i]->IsControlEdge()) << edges[i]->src()->name(); |
| 286 | WhileLoopArg new_arg; |
| 287 | new_arg.is_loop_invariant = false; |
| 288 | if (i == 0) { |
| 289 | new_arg.enter = arg.enter; |
| 290 | } else { |
| 291 | new_arg.enter = graph->CopyNode(arg.enter); |
| 292 | frame->nodes.insert(new_arg.enter); |
| 293 | for (Edge const* e : arg.enter->in_edges()) { |
| 294 | graph->AddEdge(e->src(), e->src_output(), new_arg.enter, |
| 295 | e->IsControlEdge() ? Graph::kControlSlot : 0); |
| 296 | } |
| 297 | Node* dst = edges[i]->dst(); |
| 298 | int dst_input = edges[i]->dst_input(); |
| 299 | graph->RemoveEdge(edges[i]); |
| 300 | graph->AddEdge(new_arg.enter, 0, dst, dst_input); |
| 301 | } |
| 302 | args.push_back(new_arg); |
| 303 | } |
| 304 | } |
| 305 | } |
| 306 | frame->args = std::move(args); |
| 307 | |
| 308 | std::sort(frame->args.begin(), frame->args.end(), |
| 309 | [](const WhileLoopArg& a, const WhileLoopArg& b) { |
| 310 | return NodeCmpByNameResourcesLast()(a.enter, b.enter); |
| 311 | }); |
| 312 | |
| 313 | if (frame->loop_cond == nullptr) { |
| 314 | return errors::InvalidArgument("Loop ", frame->name, |
| 315 | " has no LoopCond node"); |
| 316 | } |
| 317 | |
| 318 | // Find the set of Switch nodes that are successors of the LoopCond. |
| 319 | std::unordered_set<Node*> switches; |
| 320 | for (const Edge* edge : frame->loop_cond->out_edges()) { |
no test coverage detected