| 381 | } |
| 382 | |
| 383 | Status LoopInvariantNodeMotionOptimizer::Optimize() { |
| 384 | node_map_.reset(new NodeMap(optimized_graph_)); |
| 385 | FrameView frame_view; |
| 386 | // TODO(ezhulenev): Use GraphView when migrated from NodeMap. |
| 387 | TF_RETURN_IF_ERROR(frame_view.InferFromGraph(*optimized_graph_)); |
| 388 | |
| 389 | std::deque<int> worklist; |
| 390 | for (const NodeDef& node : optimized_graph_->node()) { |
| 391 | const std::vector<int>& frame_ids = frame_view.Frames(node); |
| 392 | |
| 393 | if (frame_ids.size() >= 3) { |
| 394 | for (unsigned int i = 1; i < frame_ids.size() - 1; ++i) { |
| 395 | frame_parent_[frame_ids[i]] = frame_ids[i - 1]; |
| 396 | frame_children_[frame_ids[i]].insert(frame_ids[i + 1]); |
| 397 | } |
| 398 | } |
| 399 | if (frame_ids.size() >= 2) { |
| 400 | frame_children_[frame_ids[0]].insert(frame_ids[1]); |
| 401 | frame_parent_[frame_ids.back()] = frame_ids[frame_ids.size() - 2]; |
| 402 | } |
| 403 | if (!frame_ids.empty()) { |
| 404 | frame_children_.insert(std::make_pair(frame_ids.back(), empty_set_)); |
| 405 | if (node.op() == "LoopCond") { |
| 406 | if (loop_cond_.count(frame_ids.back())) { |
| 407 | return errors::InvalidArgument( |
| 408 | "Loop ", frame_ids.back(), |
| 409 | " has more than one LoopCond node: ", node.name(), " and ", |
| 410 | loop_cond_[frame_ids.back()]->name()); |
| 411 | } |
| 412 | loop_cond_[frame_ids.back()] = &node; |
| 413 | } |
| 414 | if (IsEnter(node) && node.attr().at("is_constant").b()) { |
| 415 | invariant_enters_[frame_ids.back()].push_back( |
| 416 | const_cast<NodeDef*>(&node)); |
| 417 | } |
| 418 | } |
| 419 | } |
| 420 | |
| 421 | for (auto it = frame_children_.begin(); it != frame_children_.end(); ++it) { |
| 422 | if (it->second.empty()) { |
| 423 | worklist.push_back(it->first); |
| 424 | } |
| 425 | } |
| 426 | |
| 427 | while (!worklist.empty()) { |
| 428 | int frame_id = worklist.front(); |
| 429 | new_enter_id_ = 0; |
| 430 | worklist.pop_front(); |
| 431 | auto parent_it = frame_parent_.find(frame_id); |
| 432 | if (parent_it != frame_parent_.end()) { |
| 433 | int parent_id = parent_it->second; |
| 434 | frame_children_[parent_id].erase(frame_id); |
| 435 | if (frame_children_[parent_id].empty()) { |
| 436 | worklist.push_back(parent_id); |
| 437 | } |
| 438 | } |
| 439 | |
| 440 | if (invariant_enters_[frame_id].empty()) { |
nothing calls this directly
no test coverage detected