| 653 | } |
| 654 | |
| 655 | Status GraphExecutionState::InitBaseGraph(std::unique_ptr<Graph>&& new_graph) { |
| 656 | // Save stateful placements before placing. |
| 657 | RestoreStatefulNodes(new_graph.get()); |
| 658 | |
| 659 | GraphOptimizationPassOptions optimization_options; |
| 660 | optimization_options.session_handle = session_handle_; |
| 661 | optimization_options.session_options = session_options_; |
| 662 | optimization_options.graph = &new_graph; |
| 663 | optimization_options.flib_def = flib_def_.get(); |
| 664 | optimization_options.device_set = device_set_; |
| 665 | |
| 666 | const OptimizerOptions& session_optimizer_options = |
| 667 | session_options_->config.graph_options().optimizer_options(); |
| 668 | if (session_optimizer_options.do_op_fusion()) { |
| 669 | OptimizerOptions opts; |
| 670 | opts.set_opt_level(OptimizerOptions::L0); |
| 671 | opts.set_do_op_fusion( |
| 672 | session_optimizer_options.do_op_fusion()); |
| 673 | GraphOptimizer optimizer(opts); |
| 674 | VLOG(2) << "RUN Graph Optimization Passes before PRE_PLACEMENT"; |
| 675 | optimizer.Optimize(nullptr, nullptr, |
| 676 | /*device=*/nullptr, &new_graph, /*shape_map=*/nullptr); |
| 677 | } |
| 678 | |
| 679 | TF_RETURN_IF_ERROR(OptimizationPassRegistry::Global()->RunGrouping( |
| 680 | OptimizationPassRegistry::PRE_PLACEMENT, optimization_options)); |
| 681 | |
| 682 | // Assigned per gpu device to a session when use multi-stream in SessionGroup mode. |
| 683 | // The device name is listed in config.per_session_devices. |
| 684 | Device* default_local_device = nullptr; |
| 685 | DeviceSet devices; |
| 686 | if (session_options_->config.per_session_devices_size() > 0) { |
| 687 | const auto& dname = session_options_->config.per_session_devices(0); |
| 688 | for (auto& d : device_set_->devices()) { |
| 689 | if (d->name() == dname) { |
| 690 | devices.AddDevice(d); |
| 691 | default_local_device = d; |
| 692 | LOG(INFO) << "Assign " << d->name() << " to a session in session group."; |
| 693 | break; |
| 694 | } |
| 695 | } |
| 696 | if (session_options_->config.per_session_devices_size() > 1) { |
| 697 | const auto& dname1 = session_options_->config.per_session_devices(1); |
| 698 | for (auto& d : device_set_->devices()) { |
| 699 | if (d->name() == dname1) { |
| 700 | devices.AddDevice(d); |
| 701 | break; |
| 702 | } |
| 703 | } |
| 704 | } |
| 705 | } |
| 706 | |
| 707 | Placer placer(new_graph.get(), "", flib_def_.get(), |
| 708 | session_options_->config.per_session_devices_size() > 0 ? |
| 709 | &devices : device_set_, |
| 710 | default_local_device, |
| 711 | session_options_ == nullptr || |
| 712 | session_options_->config.allow_soft_placement(), |