| 77 | } |
| 78 | |
| 79 | Status ApplyRewrites(OpKernelContext* ctx, |
| 80 | const std::function<RewriterConfig(void)> config_factory, |
| 81 | bool optimize_function_library, GraphDef* graph_def, |
| 82 | string* output_node) { |
| 83 | // Add an identity node as the fetch node, otherwise we might get 'placeholder |
| 84 | // is both fed and fetched' errors in some cases when using input list with |
| 85 | // placeholder dataset nodes. |
| 86 | NodeDef* node = graph_def->mutable_node()->Add(); |
| 87 | tensorflow::grappler::graph_utils::SetUniqueGraphNodeName("Sink", graph_def, |
| 88 | node); |
| 89 | node->set_op("Identity"); |
| 90 | node->add_input(*output_node); |
| 91 | (*node->mutable_attr())["T"].set_type(DT_VARIANT); |
| 92 | *output_node = node->name(); |
| 93 | |
| 94 | // Add fake sink node to graph and functions to allow rewriting the actual |
| 95 | // sink nodes. |
| 96 | // |
| 97 | // TODO(b/118820916): When MetaOptimizer adds provisions for function retvals |
| 98 | // to be optimizable, we will no longer need this. |
| 99 | for (auto& function_def : *graph_def->mutable_library()->mutable_function()) { |
| 100 | AddFakeSinks(&function_def); |
| 101 | } |
| 102 | |
| 103 | // Create metagraph. |
| 104 | MetaGraphDef meta_graph_def; |
| 105 | (*meta_graph_def.mutable_graph_def()) = *graph_def; |
| 106 | |
| 107 | // Grappler determines fetch ops from collection 'train_op'. |
| 108 | CollectionDef collection_def; |
| 109 | auto node_list = collection_def.mutable_node_list(); |
| 110 | node_list->add_value(*output_node); |
| 111 | (*meta_graph_def.mutable_collection_def())["train_op"] = collection_def; |
| 112 | |
| 113 | // Create Grappler item. |
| 114 | tensorflow::grappler::ItemConfig item_config; |
| 115 | item_config.apply_optimizations = true; |
| 116 | std::unique_ptr<tensorflow::grappler::GrapplerItem> grappler_item = |
| 117 | tensorflow::grappler::GrapplerItemFromMetaGraphDef( |
| 118 | "graph", meta_graph_def, item_config); |
| 119 | grappler_item->optimization_options().optimize_function_library = |
| 120 | optimize_function_library; |
| 121 | std::unordered_map<string, tensorflow::DeviceProperties> device_map; |
| 122 | tensorflow::grappler::VirtualCluster cluster(device_map); |
| 123 | |
| 124 | // Run data optimizer using grappler's meta optimizer. |
| 125 | tensorflow::ConfigProto config; |
| 126 | *config.mutable_graph_options()->mutable_rewrite_options() = config_factory(); |
| 127 | TF_RETURN_IF_ERROR(tensorflow::grappler::RunMetaOptimizer( |
| 128 | *grappler_item, config, ctx->device(), &cluster, graph_def)); |
| 129 | |
| 130 | // Remove fake sinks after optimizations are done. |
| 131 | // |
| 132 | // TODO(b/118820916): When MetaOptimizer adds provisions for function retvals |
| 133 | // to be optimizable, we will no longer need this. |
| 134 | for (auto& function_def : *graph_def->mutable_library()->mutable_function()) { |
| 135 | RemoveFakeSinks(&function_def); |
| 136 | } |
no test coverage detected