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_CORE_GRAPPLER_COSTS_GRAPH_PROPERTIES_H_ 17 #define TENSORFLOW_CORE_GRAPPLER_COSTS_GRAPH_PROPERTIES_H_ 18 19 #include <unordered_map> 20 #include <vector> 21 #include "tensorflow/core/framework/shape_inference.h" 22 #include "tensorflow/core/grappler/clusters/cluster.h" 23 #include "tensorflow/core/grappler/costs/op_performance_data.pb.h" 24 #include "tensorflow/core/grappler/grappler_item.h" 25 26 namespace tensorflow { 27 28 namespace grappler { 29 30 // Optional attributes that tell about node output information. 31 // We use these side information, if provided, for static shape inference 32 // and VirtualScheduler scheduling. 33 34 // Switch op attribute as a vector of int that tells which branch the 35 // Switch output is taken on every round of execution. 36 // Used for scheduling ops after Switch correctly (e.g., While loop). 37 ABSL_CONST_INIT const char kOutputSlots[] = "_output_slot_vector"; 38 39 // Example: 40 // Assume a node has two outputs and iterated for three times. Then it has: 41 // _execution_count = 3 42 // _output_sizes_vector = [2, 2, 2] 43 // _output_dtype_vector.size = 6 44 // _output_shape_vector.size = 6 45 46 // If all the iterations have same output shapes, then 47 // _execution_count = 3 48 // _same_output_for_iterations = true 49 // _output_sizes_vector = [2] 50 // _output_dtype_vector.size = 2 51 // _output_shape_vector.size = 2 52 53 // How many times this node has been executed. 54 ABSL_CONST_INIT const char kExecutionCount[] = "_execution_count"; 55 56 // Records the output sizes for each round of execution. 57 ABSL_CONST_INIT const char kOutputSizes[] = "_output_sizes_vector"; 58 59 // The node has been scheduled multiple times with outputs that have the same 60 // shape. 61 ABSL_CONST_INIT const char kOutputSame[] = "_same_output_for_iterations"; 62 63 // Outputs DataType vector. 64 ABSL_CONST_INIT const char kOutputTypes[] = "_output_dtype_vector"; 65 66 // Outputs TensorShapeProto vector. 67 ABSL_CONST_INIT const char kOutputShapes[] = "_output_shape_vector"; 68 69 class SymbolicShapeRefiner; 70 class TopoQueue; 71 72 // Infer OpInfo::TensorProperties for graph nodes inputs/outputs. 73 // 74 // Typical use case, is to infer tensor properties from a graph, before doing 75 // optimization pass. Nodes modified during optimization pass have to be 76 // invalidated, to prevent further incorrect optimizations based on wrong shape 77 // and data type properties. 78 class GraphProperties { 79 public: 80 // The item must outlive the properties GraphProperties(const GrapplerItem & item)81 explicit GraphProperties(const GrapplerItem& item) : item_(item) {} 82 83 // Infer the shapes through abstract interpretation. Feed information can be 84 // incorrect so it should be discarded to ensure correctness of the analysis. 85 // However, it can help infer shapes in the fanout of fed nodes (even though 86 // the correctness of these shapes can't be guaranteed), so in some cases 87 // (such as simulation or scheduling) it makes sense of keep these shapes. 88 // aggressive_shape_inference option executes nodes on the host to identify 89 // output values when possible and does other aggressive strategies. 90 // Similar to assuming_valid_feeds, this may cause incorrectness in graph 91 // analyses, but is useful for simulation or scheduling. 92 Status InferStatically(bool assume_valid_feeds, 93 bool aggressive_shape_inference); InferStatically(bool assume_valid_feeds)94 Status InferStatically(bool assume_valid_feeds) { 95 return InferStatically(assume_valid_feeds, 96 /*aggressive_shape_inference=*/false); 97 } 98 // Infer the shape by running the graph on the specified cluster and recording 99 // the shapes of the processed tensors. 100 Status InferDynamically(Cluster* cluster); 101 // Extract the properties from a cost graph. For testing only since there is 102 // no way to ensure that the cost graph match the item. 103 Status InferFromCostGraph(const CostGraphDef& cost_graph); 104 105 // Stores `item_.graph` with the inferred output shapes to `output_graph_def`. 106 Status AnnotateOutputShapes(GraphDef* output_graph_def) const; 107 108 // Return the properties of node inputs/outputs, including data types and 109 // shapes. Note that the dimensions in the shapes can be negative. We use the 110 // -1 value to denote that we don't know anything about a dimension. We use 111 // values strictly less than -1 to encode symbolic dimensions: although we 112 // don't know the actual value of the symbolic dimension, we know that all the 113 // dimensions denoted by the same negative value are the equal. 114 bool HasInputProperties(const string& node_name) const; 115 bool HasOutputProperties(const string& node_name) const; 116 const std::vector<OpInfo::TensorProperties>& GetInputProperties( 117 const string& node_name) const; 118 const std::vector<OpInfo::TensorProperties>& GetOutputProperties( 119 const string& node_name) const; 120 // Invalidate input/output properties for nodes modified during graph 121 // optimization pass, to prevent potential optimizations, based on incorrect 122 // shape information. 123 void ClearInputProperties(const string& node_name); 124 void ClearOutputProperties(const string& node_name); 125 // Returns true if we have *any* properties. has_properties()126 bool has_properties() const { 127 return input_properties_.size() > 0 || output_properties_.size() > 0; 128 } 129 130 private: 131 // Relaxes shapes <shapes_and_types>, determined from an EnqueueV2 node, into 132 // <*queue_shapes_and_types>. 133 static Status RelaxEnqueueShapesAndMergeTypes( 134 SymbolicShapeRefiner* shape_refiner, const NodeDef* qnode, 135 const std::vector<shape_inference::ShapeAndType>& shapes_and_types, 136 std::vector<shape_inference::ShapeAndType>* queue_shapes_and_types); 137 138 // Update the shapes of the enqueue node, port them over to the corresponding 139 // queue, and schedule the reprocessing of the queue if needed. 140 static Status UpdateEnqueue( 141 const NodeDef* enqueue_node, 142 const std::unordered_map<const NodeDef*, const NodeDef*>& 143 resource_handles, 144 SymbolicShapeRefiner* shape_refiner, bool* new_shapes); 145 146 // Update the shapes and types of the Queue node, if not set by Enqueue node. 147 static Status UpdateQueue(const NodeDef* queue_node, 148 SymbolicShapeRefiner* shape_refiner, 149 bool* new_shapes); 150 151 // Update the output shapes of a Merge node, and enqueue its fanout in 152 // new_shapes if needed. 153 Status UpdateMerge(SymbolicShapeRefiner* shape_refiner, const NodeDef* node, 154 bool* new_shapes) const; 155 // Process the Enter node, and enqueue its fanout in new_shapes if needed. 156 static Status UpdateEnter(SymbolicShapeRefiner* shape_refiner, 157 const NodeDef* node, bool* new_shapes); 158 // Update the shapes for node 'n'. If output shapes for n have changed, 159 // enqueue its fanout in 'new_shapes'. 160 Status UpdateShapes(SymbolicShapeRefiner* shape_refiner, 161 const std::unordered_map<const NodeDef*, const NodeDef*>& 162 resource_handles, 163 const NodeDef* n, bool* new_shapes) const; 164 // Propagate the shapes for the nodes enqueued in new_shapes and their 165 // transitive fanout until a fixed point is reached. 166 Status PropagateShapes( 167 SymbolicShapeRefiner* shape_refiner, TopoQueue* new_shapes, 168 const std::unordered_map<const NodeDef*, const NodeDef*>& 169 resource_handles, 170 int num_loops) const; 171 172 // Data members 173 const GrapplerItem& item_; 174 std::unordered_map<string, std::vector<OpInfo::TensorProperties>> 175 input_properties_; 176 std::unordered_map<string, std::vector<OpInfo::TensorProperties>> 177 output_properties_; 178 const std::vector<OpInfo::TensorProperties> missing_properties_; 179 }; 180 181 } // end namespace grappler 182 } // end namespace tensorflow 183 184 #endif // TENSORFLOW_CORE_GRAPPLER_COSTS_GRAPH_PROPERTIES_H_ 185