| 928 | } |
| 929 | |
| 930 | Status OptimizeGraph( |
| 931 | std::vector<string> ret_node_names, std::vector<string> keep_node_names, |
| 932 | FunctionLibraryDefinition* flib, const DeviceSet& device_set, |
| 933 | Device* cpu_device, const ConfigProto& config_proto, |
| 934 | const string& grappler_item_id, |
| 935 | const GrapplerItem::OptimizationOptions& optimization_options, |
| 936 | std::unique_ptr<tensorflow::Graph>* g) { |
| 937 | if (!tensorflow::grappler::MetaOptimizerEnabled(config_proto)) { |
| 938 | return Status::OK(); |
| 939 | } |
| 940 | |
| 941 | tensorflow::grappler::GrapplerItem item; |
| 942 | item.id = grappler_item_id; |
| 943 | item.optimization_options() = optimization_options; |
| 944 | |
| 945 | // Add all available devices so that inlined function can be placed. |
| 946 | for (const Device* d : device_set.devices()) { |
| 947 | Status added_device = item.AddDevice(d->name()); |
| 948 | if (!added_device.ok()) VLOG(3) << added_device.error_message(); |
| 949 | } |
| 950 | VLOG(3) << "Grappler available devices: " |
| 951 | << absl::StrJoin(item.devices(), ", "); |
| 952 | |
| 953 | // Add fetches so that the graph can be pruned. |
| 954 | item.fetch.swap(ret_node_names); |
| 955 | |
| 956 | // Add noes that can't be removed from the graph. |
| 957 | item.keep_ops = std::move(keep_node_names); |
| 958 | |
| 959 | (*g)->ToGraphDef(&item.graph); |
| 960 | |
| 961 | if (flib) { |
| 962 | *item.graph.mutable_library() = flib->ToProto(); |
| 963 | } |
| 964 | |
| 965 | tensorflow::GraphDef out_graph; |
| 966 | tensorflow::grappler::VirtualCluster cluster(&device_set); |
| 967 | // TODO(nareshmodi): Consider adding and using the more generic GraphOptions |
| 968 | // proto (which also contain the OptimizerOptions). |
| 969 | TF_RETURN_IF_ERROR(tensorflow::grappler::RunMetaOptimizer( |
| 970 | item, config_proto, cpu_device, &cluster, &out_graph)); |
| 971 | |
| 972 | std::unique_ptr<tensorflow::Graph> optimized_graph( |
| 973 | new tensorflow::Graph(OpRegistry::Global())); |
| 974 | |
| 975 | // Copy optimized functions back to the overlay lib. |
| 976 | if (flib) { |
| 977 | for (const FunctionDef& fdef : out_graph.library().function()) { |
| 978 | const string& func_name = fdef.signature().name(); |
| 979 | if (flib->Contains(func_name)) { |
| 980 | TF_RETURN_IF_ERROR(flib->ReplaceFunction(func_name, fdef)); |
| 981 | } else { |
| 982 | TF_RETURN_IF_ERROR(flib->AddFunctionDef(fdef)); |
| 983 | } |
| 984 | } |
| 985 | } |
| 986 | |
| 987 | TF_RETURN_IF_ERROR(ConvertGraphDefToGraph( |
no test coverage detected