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