1 /* Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #ifndef TENSORFLOW_GRAPPLER_OPTIMIZERS_GRAPH_REWRITER_H_ 17 #define TENSORFLOW_GRAPPLER_OPTIMIZERS_GRAPH_REWRITER_H_ 18 19 #include <unordered_map> 20 #include <unordered_set> 21 #include "tensorflow/core/grappler/grappler_item.h" 22 23 namespace tensorflow { 24 namespace grappler { 25 26 // Tools and utilities to simplify common graph rewrites. 27 class GraphRewriter { 28 public: 29 GraphRewriter(const GrapplerItem& item); 30 31 // Forward the inputs of original_node as needed to skip over the nodes that 32 // are to be deleted. In other words, if I is an input of 'original_node', and 33 // I doesn't belong to one of the nodes in 'nodes_to_delete', I will be an 34 // input to 'new_node'. On the other hand, if I belong to a node that will be 35 // deleted, I will be replaced with the inputs J of the deleted node (unless J 36 // belong to nodes that will be deleted, in which case we'll look for 37 // preserved inputs further down the graph). 38 void ForwardInputs(const NodeDef& original_node, 39 const std::unordered_set<const NodeDef*>& nodes_to_delete, 40 NodeDef* new_node); 41 42 // Returns true if at least one of the edges in the direct fanout of 'node' is 43 // a control dependency edge. 44 bool DrivesControlDependency(const NodeDef& node) const; 45 46 // Returns true if at least one of the incident edges is a control dependency 47 // edge. 48 bool IsDrivenByControlDependency(const NodeDef& node) const; 49 50 // Returns true if at least one of the nodes in the direct fanin or the direct 51 // fanout (excluding control dependencies) of 'node' is a function. 52 bool IsConnectedToFunction(const NodeDef& node) const; 53 54 // Returns true if the node is driven by at least one node placed on another 55 // device. 56 bool IsDrivenByAnotherDevice(const NodeDef& node) const; 57 58 // Returns true if the node has input from a stateful op. 59 bool ReceivesRefValue(const NodeDef& node) const; 60 61 // Returns true if the node is driven by a Switch node. 62 bool IsDrivenBySwitch(const NodeDef& node) const; 63 64 // Returns true if the node feeds a Merge node. 65 bool FeedsMerge(const NodeDef& node) const; 66 67 // Returns true if removal of this degree would increase edge count, i.e. if 68 // in-degree * out-degree > in-degree + out-degree or if the condition could 69 // not be verified. 70 bool RemovalIncreasesEdgeCount(const NodeDef& node) const; 71 72 private: 73 void RecordConnectivity(const NodeDef& node, 74 const std::unordered_set<string>& function_names); 75 void ForwardInputsInternal( 76 const NodeDef& original_node, 77 const std::unordered_set<const NodeDef*>& nodes_to_delete, 78 bool add_as_control, NodeDef* new_node); 79 80 struct NodeInfo { 81 int out_degree = 0; 82 const NodeDef* def; 83 84 // These are filled in when the NodeInfo is built, but not that they 85 // may be empty - if the op could not be loaded from the registry. 86 DataTypeVector outputs; 87 }; 88 89 std::unordered_map<string, std::unique_ptr<NodeInfo>> nodes_; 90 std::unordered_map<string, const NodeDef*> optimized_nodes_; 91 std::unordered_set<const NodeDef*> control_dependency_drivers_; 92 std::unordered_set<const NodeDef*> function_neighbors_; 93 std::unordered_set<const NodeDef*> cross_device_receivers_; 94 std::unordered_set<const NodeDef*> ref_receivers_; 95 std::unordered_set<const NodeDef*> switch_receivers_; 96 std::unordered_set<const NodeDef*> merge_feeders_; 97 }; 98 99 } // end namespace grappler 100 } // end namespace tensorflow 101 102 #endif // TENSORFLOW_GRAPPLER_OPTIMIZERS_GRAPH_REWRITER_H_ 103
