1 /* Copyright 2018 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 #include "tensorflow/core/grappler/utils/symbolic_shapes.h"
17 #include "tensorflow/core/util/bcast.h"
18
19 namespace tensorflow {
20 namespace grappler {
21 namespace {
22
ShapeDims(const TensorShapeProto & shape)23 BCast::Vec ShapeDims(const TensorShapeProto& shape) {
24 BCast::Vec dims;
25 dims.reserve(shape.dim_size());
26 for (int i = 0; i < shape.dim_size(); ++i)
27 dims.push_back(shape.dim(i).size());
28 return dims;
29 }
30
31 } // namespace
32
IsKnown(const TensorShapeProto::Dim & dim)33 bool IsKnown(const TensorShapeProto::Dim& dim) { return dim.size() >= 0; }
34
IsKnownSymbolically(const TensorShapeProto::Dim & dim)35 bool IsKnownSymbolically(const TensorShapeProto::Dim& dim) {
36 return dim.size() <= -2;
37 }
38
IsUnknown(const TensorShapeProto::Dim & dim)39 bool IsUnknown(const TensorShapeProto::Dim& dim) { return dim.size() == -1; }
40
ShapeIsSymbolicallyDefined(const TensorShapeProto & shape)41 bool ShapeIsSymbolicallyDefined(const TensorShapeProto& shape) {
42 return !shape.unknown_rank() &&
43 std::all_of(
44 shape.dim().begin(), shape.dim().end(),
45 [](const TensorShapeProto::Dim& dim) { return !IsUnknown(dim); });
46 }
47
ShapeIsSymbolicallyDefined(const OpInfo::TensorProperties & properties)48 bool ShapeIsSymbolicallyDefined(const OpInfo::TensorProperties& properties) {
49 return ShapeIsSymbolicallyDefined(properties.shape());
50 }
51
Rank(const TensorShapeProto & shape)52 int Rank(const TensorShapeProto& shape) {
53 if (shape.unknown_rank()) {
54 return -1;
55 }
56 return shape.dim_size();
57 }
58
NumCoefficients(const TensorShapeProto & shape)59 int64 NumCoefficients(const TensorShapeProto& shape) {
60 if (shape.unknown_rank()) {
61 return -1;
62 }
63 int64 num_coefficients = 1;
64 for (const auto& dim : shape.dim()) {
65 if (dim.size() < 0) {
66 return -1;
67 }
68 num_coefficients *= dim.size();
69 }
70 return num_coefficients;
71 }
72
ShapesSymbolicallyEqual(const TensorShapeProto & left,const TensorShapeProto & right)73 bool ShapesSymbolicallyEqual(const TensorShapeProto& left,
74 const TensorShapeProto& right) {
75 if (left.unknown_rank() || right.unknown_rank() ||
76 left.dim_size() != right.dim_size()) {
77 return false;
78 }
79 for (int i = 0; i < left.dim_size(); ++i) {
80 const auto& ldim = left.dim(i);
81 const auto& rdim = right.dim(i);
82 if (IsUnknown(ldim) || IsUnknown(rdim) || ldim.size() != rdim.size()) {
83 return false;
84 }
85 }
86 return true;
87 }
88
ShapesSymbolicallyEqual(const OpInfo::TensorProperties & left,const OpInfo::TensorProperties & right)89 bool ShapesSymbolicallyEqual(const OpInfo::TensorProperties& left,
90 const OpInfo::TensorProperties& right) {
91 return ShapesSymbolicallyEqual(left.shape(), right.shape());
92 }
93
ShapesBroadcastable(const TensorShapeProto & left,const TensorShapeProto & right)94 bool ShapesBroadcastable(const TensorShapeProto& left,
95 const TensorShapeProto& right) {
96 if (!ShapeIsSymbolicallyDefined(left) || !ShapeIsSymbolicallyDefined(right)) {
97 return false;
98 }
99 BCast bcast(ShapeDims(left), ShapeDims(right),
100 /*fewer_dims_optimization*/ false);
101 return bcast.IsValid();
102 }
103
ShapesBroadcastable(const OpInfo::TensorProperties & left,const OpInfo::TensorProperties & right)104 bool ShapesBroadcastable(const OpInfo::TensorProperties& left,
105 const OpInfo::TensorProperties& right) {
106 return ShapesBroadcastable(left.shape(), right.shape());
107 }
108
ShapeAfterBroadcast(const TensorShapeProto & left,const TensorShapeProto & right,TensorShapeProto * output_shape)109 bool ShapeAfterBroadcast(const TensorShapeProto& left,
110 const TensorShapeProto& right,
111 TensorShapeProto* output_shape) {
112 if (!ShapeIsSymbolicallyDefined(left) || !ShapeIsSymbolicallyDefined(right)) {
113 return false;
114 }
115 BCast bcast(ShapeDims(left), ShapeDims(right),
116 /*fewer_dims_optimization*/ false);
117 if (!bcast.IsValid()) {
118 return false;
119 }
120 output_shape->set_unknown_rank(false);
121 output_shape->clear_dim();
122 for (const auto& dim : bcast.output_shape()) {
123 output_shape->add_dim()->set_size(dim);
124 }
125 return true;
126 }
127
CompareSymbolicallyShapedTensorSizes(const TensorShapeProto & left,const TensorShapeProto & right)128 bool CompareSymbolicallyShapedTensorSizes(const TensorShapeProto& left,
129 const TensorShapeProto& right) {
130 // if one of the ranks is unknown, it's impossible to compare tensor sizes
131 if (left.unknown_rank() || right.unknown_rank()) {
132 return false;
133 }
134
135 // Tensor size, computed as a product of defined dimensions
136 int64 left_defined_size = 1;
137 int64 right_defined_size = 1;
138
139 // Keep how many times each unknown dimension appeared on the left and right
140 std::unordered_map<int64, int64> left_unknown_dims;
141 std::unordered_map<int64, int64> right_unknown_dims;
142
143 // Assign unique id to every unknown dimension (-1). We are going to
144 // assign positive ids, because negative values are already used by
145 // symbolic dimensions.
146 int64 unknown_dim_id = 1;
147
148 // For each shape dimension update "defined tensor size", if shape is defined,
149 // or increment a counter for unknown dim.
150 auto process_dimensions =
151 [&unknown_dim_id](const TensorShapeProto& shape, int64* defined_size,
152 std::unordered_map<int64, int64>* unknown_dims) {
153 for (int i = 0; i < shape.dim_size(); ++i) {
154 const auto& dim = shape.dim(i);
155 int64 dim_size = dim.size();
156 if (dim_size > 0) {
157 *defined_size *= dim_size;
158 } else if (IsUnknown(dim)) {
159 ++(*unknown_dims)[unknown_dim_id++];
160 } else if (IsKnownSymbolically(dim)) {
161 ++(*unknown_dims)[dim_size];
162 }
163 }
164 };
165
166 process_dimensions(left, &left_defined_size, &left_unknown_dims);
167 process_dimensions(right, &right_defined_size, &right_unknown_dims);
168
169 // Compute a union of unknown dimension ids appeared in both shapes
170 std::set<int64> unknown_dims;
171 for (const auto& el : left_unknown_dims) unknown_dims.insert(el.first);
172 for (const auto& el : right_unknown_dims) unknown_dims.insert(el.first);
173
174 // Cancel unknown dimensions that appeared in both shapes
175 for (int64 unknown_dim : unknown_dims) {
176 int64 co_occurrence = std::min(left_unknown_dims[unknown_dim],
177 right_unknown_dims[unknown_dim]);
178 left_unknown_dims[unknown_dim] -= co_occurrence;
179 right_unknown_dims[unknown_dim] -= co_occurrence;
180 }
181
182 // Count unbalanced unknown dimensions
183 int64 left_unbalanced_unknown_dims = 0;
184 int64 right_unbalanced_unknown_dims = 0;
185 for (const auto& el : left_unknown_dims)
186 left_unbalanced_unknown_dims += el.second;
187 for (const auto& el : right_unknown_dims)
188 right_unbalanced_unknown_dims += el.second;
189
190 if (left_unbalanced_unknown_dims == 0 && right_unbalanced_unknown_dims == 0) {
191 // If unknown dimensions cancelled each other, compare tensor sizes
192 // represented by defined dimensions
193 return left_defined_size < right_defined_size;
194 }
195
196 if (left_defined_size <= right_defined_size &&
197 left_unbalanced_unknown_dims == 0 && right_unbalanced_unknown_dims > 0) {
198 // If size of a 'left" tensor computed from defined dimensions less or
199 // equal, and shape on the right has unbalanced unknown dimensions, we can
200 // guarantee that shape on the left is strictly smaller (assuming that
201 // unknown dimension size is larger than 1)
202 return true;
203 }
204
205 // In every other case, assuming that unknown dimensions can be arbitrary
206 // large in size, we can't guarantee any ordering
207 return false;
208 }
209
CompareSymbolicallyShapedTensorSizes(const OpInfo::TensorProperties & left,const OpInfo::TensorProperties & right)210 bool CompareSymbolicallyShapedTensorSizes(
211 const OpInfo::TensorProperties& left,
212 const OpInfo::TensorProperties& right) {
213 return CompareSymbolicallyShapedTensorSizes(left.shape(), right.shape());
214 }
215
ComputeSizeRatio(const TensorShapeProto & numerator,const TensorShapeProto & denominator)216 int64 ComputeSizeRatio(const TensorShapeProto& numerator,
217 const TensorShapeProto& denominator) {
218 if (numerator.unknown_rank() || denominator.unknown_rank()) {
219 return -1;
220 }
221 std::multiset<int> symbolic_dims;
222 int64 num = 1;
223 for (const auto& dim : numerator.dim()) {
224 if (dim.size() == -1) {
225 return -1;
226 } else if (dim.size() < -1) {
227 symbolic_dims.insert(dim.size());
228 } else {
229 num *= dim.size();
230 }
231 }
232 int64 denom = 1;
233 for (const auto& dim : denominator.dim()) {
234 if (dim.size() == -1) {
235 return -1;
236 } else if (dim.size() < -1) {
237 auto it = symbolic_dims.find(dim.size());
238 if (it == symbolic_dims.end()) {
239 return -1;
240 }
241 symbolic_dims.erase(it);
242 } else {
243 denom *= dim.size();
244 }
245 }
246 if (denom == 0) {
247 return -1;
248 }
249 if (!symbolic_dims.empty()) {
250 return -1;
251 }
252 return num / denom;
253 }
254
255 } // end namespace grappler
256 } // end namespace tensorflow
257