1/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations under the License.
14==============================================================================*/
15#ifndef TENSORFLOW_LITE_TOCO_TOOLING_UTIL_H_
16#define TENSORFLOW_LITE_TOCO_TOOLING_UTIL_H_
17
18#include <algorithm>
19#include <cmath>
20#include <functional>
21#include <iostream>
22#include <limits>
23#include <memory>
24#include <string>
25#include <vector>
26
27#include "absl/strings/string_view.h"
28#include "tensorflow/core/lib/core/errors.h"
29#include "tensorflow/core/lib/core/status.h"
30#include "tensorflow/core/platform/logging.h"
31#include "tensorflow/lite/kernels/internal/types.h"
32#include "tensorflow/lite/toco/model.h"
33#include "tensorflow/lite/toco/model_flags.pb.h"
34#include "tensorflow/lite/toco/runtime/types.h"
35#include "tensorflow/lite/toco/toco_flags.pb.h"
36#include "tensorflow/lite/toco/types.pb.h"
37
38// TODO(aselle): Replace with using a container specific hash override instead.
39namespace std {
40template <>
41struct hash<toco::OperatorType> {
42 size_t operator()(const toco::OperatorType& op) const {
43 return std::hash<size_t>()(static_cast<size_t>(op));
44 }
45};
46} // namespace std
47
48namespace toco {
49
50constexpr int kLogLevelModelChanged = 1;
51constexpr int kLogLevelModelUnchanged = 2;
52
53absl::string_view FindLongestCommonPrefix(absl::string_view a,
54 absl::string_view b);
55std::string LogName(const Operator& op);
56
57std::string ArrayDataTypeName(ArrayDataType data_type);
58
59// Returns true if the given array is specified as a model input array.
60bool IsInputArray(const Model& model, const std::string& array_name);
61// Returns true if the given array is specified as a model output array.
62bool IsOutputArray(const Model& model, const std::string& array_name);
63
64bool IsArrayConsumed(const Model& model, const std::string& name);
65int CountTrueOutputs(const Model& model, const Operator& op);
66
67int CountOpsWithInput(const Model& model, const std::string& array_name);
68bool DeleteArrayIfUnused(const std::string& array_name, Model* model);
69
70// Deletes the op and any of its input and output arrays if they are unused
71// after the op has been deleted.
72void DeleteOpAndArrays(Model* model, const Operator* op);
73
74std::vector<std::unique_ptr<Operator>>::const_iterator FindOpWithOutput(
75 const Model& model, const std::string& array_name);
76Operator* GetOpWithOutput(const Model& model, const std::string& array_name);
77
78std::vector<std::unique_ptr<Operator>>::iterator FindOpWithOutput(
79 Model& model, const std::string& array_name);
80
81std::vector<std::unique_ptr<Operator>>::const_iterator FindOpWithInput(
82 const Model& model, const std::string& array_name);
83
84std::vector<std::unique_ptr<Operator>>::iterator FindOpWithInput(
85 Model& model, const std::string& array_name);
86
87Operator* GetOpWithInput(const Model& model, const std::string& array_name);
88Operator* GetFirstOpWithInput(const Model& model,
89 const std::string& array_name);
90
91// Replaces all uses of the |old_array_name| with the |new_array_name|.
92void ReplaceArrayUsage(Model* model, const std::string& old_array_name,
93 const std::string& new_array_name);
94
95std::vector<std::unique_ptr<Operator>>::const_iterator FindOp(
96 const Model& model, const Operator* op);
97std::vector<std::unique_ptr<Operator>>::iterator FindOp(Model& model,
98 const Operator* op);
99
100const char* OperatorTypeName(OperatorType type);
101std::string HelpfulOperatorTypeName(const Operator& op);
102
103// Whether the operator can be fused with an activation function. Note that this
104// will return false by default for new operators; fusing support is opt-in.
105bool OperatorSupportsFusedActivation(OperatorType type);
106
107void DumpGraphvizVideoFrame(const Model& model);
108void LogDump(int log_level, const std::string& message, const Model& model);
109void LogSummary(int log_level, const std::string& message, const Model& model);
110
111// TODO(b/36075966): Clean up when dims superseded by array shape.
112void ExtendShape(Shape* shape, int new_shape_size);
113
114// TODO(b/36075966): Clean up when dims superseded by array shape.
115void UnextendShape(Shape* shape, int new_shape_size);
116
117// Checks that all dimensions of 'shape' are at least 1. Note that scalars,
118// lacking dimensions, satisfy this condition and are considered non-empty.
119bool IsNonEmpty(const Shape& shape);
120
121// Given two shapes with potentially different dimensionality and dimension
122// arrays d0 and d1. Without loss of generality, assume that shape0 may have
123// higher dimensionality (length(d0) >= length(d1)). Then shape0 and shape1
124// "agree up to broadcasting" if:
125// - When walking the d0 and d1 from back to front with indices i0, i1,
126// d0[i0] == d1[i1] or d0[i0] == 1 or d1[i1] == 1, for each dimension until
127// i1 == 0 (inclusive).
128bool ShapesAgreeUpToBroadcasting(const Shape& shape0, const Shape& shape1);
129
130// A stricter constraint than ShapesAgreeUpToBroadcasting().
131//
132// Given two shapes with potentially different dimensionality and dimension
133// arrays d0 and d1. Without loss of generality, assume that shape0 may have
134// higher dimensionality (length(d0) >= length(d1)). Then shape0 and shape1
135// "agree up to extending" if:
136// - When walking the d0 and d1 from back to front with indices i0, i1,
137// d0[i0] == d1[i1] for each dimension until i1 == 0 (inclusive).
138// - For the remaining indices [0..i0), d0[i0] == 1.
139bool ShapesAgreeUpToExtending(const Shape& shape0, const Shape& shape1);
140
141inline ::tflite::RuntimeShape ToRuntimeShape(const Shape& shape) {
142 return ::tflite::RuntimeShape(shape.dimensions_count(), shape.dims().data());
143}
144
145bool IsArrayFullyConnectedWeights(const Model& model, const std::string& name);
146
147// If there is a wildcard dimension (-1), this may return a negative value.
148int RequiredBufferSizeForShape(const Shape& shape);
149
150bool IsConstantParameterArray(const Model& model, const std::string& name);
151
152// Compares two constant parameter arrays for exact equality.
153bool CompareConstantArrays(const Array& lhs_array, const Array& rhs_array);
154
155void CheckNoMissingArray(const Model& model);
156void CheckInvariants(const Model& model);
157
158void CheckModelCounts(const Model& model);
159
160void FixOperatorOrdering(Model* model);
161void FixNoMissingArray(Model* model);
162void FixNoOrphanedArray(Model* model);
163
164// Fixes input/output arrays that may have issues during export or inference.
165void FixEdgeArrays(Model* model);
166
167// Finds and deduplicates large constant arrays in the model.
168// After constant propagation runs it's possible to end up with several of the
169// same large array (whether they be zeros or otherwise).
170//
171// |min_size| is used to adjust the minimum size in bytes of an array before
172// it's considered for deduping. As deduping can make the graphs more difficult
173// to read this helps prevent small arrays from spidering out.
174void DedupeConstantArrays(Model* model, size_t min_size);
175
176// Copies the contents of an array into another.
177// Expects that the shape and data type match.
178template <ArrayDataType A>
179void CopyArrayBuffer(const Array& source_array, Array* target_array) {
180 int source_buffer_size = RequiredBufferSizeForShape(source_array.shape());
181 int target_buffer_size = RequiredBufferSizeForShape(target_array->shape());
182 CHECK_EQ(source_buffer_size, target_buffer_size)
183 << "Buffer sizes must match in element count";
184 CHECK(source_array.data_type == target_array->data_type)
185 << "Data types must match";
186 if (source_array.buffer) {
187 const auto& source_buffer = source_array.GetBuffer<A>();
188 auto& target_buffer = target_array->GetMutableBuffer<A>();
189 target_buffer.data = source_buffer.data;
190 }
191}
192
193// Inserts a no-op reshape operator between the source array and the target
194// array. This effectively just copies the data.
195void InsertCopyOperator(Model* model, const std::string& source_array_name,
196 const std::string& target_array_name);
197
198// Clones an array with all data and parameters.
199void CloneArray(Model* model, const std::string& source_array_name,
200 const std::string& target_array_name);
201
202void ResolveModelFlags(const ModelFlags& model_flags, Model* model);
203
204template <typename T>
205T ConvertOperator(Operator* o, OperatorType type) {
206 if (o != nullptr && o->type == type) {
207 return static_cast<T>(o);
208 }
209
210 return nullptr;
211}
212
213void CheckIsReadyForQuantization(const Model& model);
214
215bool ReshapeIsEquivalentToTranspose(const Model& model,
216 const TensorFlowReshapeOperator* op,
217 bool allow_extra_unary_dims);
218
219inline int Offset(const Shape& shape, const std::vector<int>& indices) {
220 DCHECK_EQ(shape.dimensions_count(), indices.size());
221 const int dims_count = shape.dimensions_count();
222 int offset = 0;
223 for (int i = 0; i < dims_count; i++) {
224 const int index = indices[i];
225 DCHECK(index >= 0 && index < shape.dims(i));
226 offset *= shape.dims(i);
227 offset += index;
228 }
229 return offset;
230}
231
232inline std::vector<int> ReverseOffset(const Shape& shape, int index) {
233 DCHECK_GE(index, 0);
234 DCHECK_LT(index, RequiredBufferSizeForShape(shape));
235 const int dims_count = shape.dimensions_count();
236 std::vector<int> indices(dims_count);
237 int residual = index;
238 for (int i = dims_count - 1; i >= 0; i--) {
239 indices[i] = residual % shape.dims(i);
240 residual /= shape.dims(i);
241 }
242 return indices;
243}
244
245int ElementSize(ArrayDataType data_type);
246
247void DropMinMax(Model* model, const std::string& array_name);
248
249bool IsAllocatableTransientArray(const Model& model,
250 const std::string& array_name);
251
252void CreateOrCheckRnnStateArray(const std::string& name, int size,
253 int state_num_dims, Model* model);
254
255std::string AvailableArrayName(const Model& model, const std::string& name);
256
257// Formats a shape as a string: [ dims(0), dims(1), ..., dims(num_dims-1) ].
258std::string ShapeToString(const Shape& shape);
259
260void PrintArrayShape(Model* model, const std::string& name);
261
262void MakeArrayDims(int num_dims, int batch, int height, int width, int depth,
263 std::vector<int>* out_dims);
264
265// Defines a constant int32 array with the provided values formatted for use
266// as op parameters.
267std::string CreateInt32Array(Model* model, const std::string& param_name,
268 const std::vector<int>& value);
269
270bool EstimateArithmeticOpsCount(const Model& model, const Operator& op,
271 int64_t* result);
272bool EstimateArithmeticOpsCount(const Model& model, int64_t* result);
273std::string FormattedNumber(int64_t x);
274
275int AxesCount(AxesOrder axes_order);
276
277// Returns the permutation of the dimensions based on the input axes order and
278// output axes order.
279void GetShuffleShape(AxesOrder input_axes_order, AxesOrder output_axes_order,
280 std::vector<int>* shuffle);
281
282// Extend shuffle is designed to match ExtendShape, which pads the shape with
283// unit dimensions at the beginning.
284void ExtendShuffle(const std::vector<int>& input_shuffle, int newdim,
285 std::vector<int>* extended_shuffle);
286
287void ShuffleDims(const Shape& input_shape, AxesOrder input_axes_order,
288 AxesOrder output_axes_order, Shape* output_shape);
289void ShuffleArray(const Shape& input_shape, AxesOrder input_axes_order,
290 AxesOrder output_axes_order, const Shape& output_shape,
291 const float* input_data, float* output_data);
292void ShuffleArray(const Shape& input_shape, AxesOrder input_axes_order,
293 AxesOrder output_axes_order, const Shape& output_shape,
294 const uint8* input_data, uint8* output_data);
295
296// Returns true if it may be OK for any graph transformation to ever discard
297// that array. The idea is that we can't ever discard arrays that are either
298// an input or an output of the whole graph, or that appear in RNN back-edges,
299// as that would undercut explicit flags that the user might pass.
300bool IsDiscardableArray(const Model& model, const std::string& array_name);
301
302void CheckFinalDataTypesSatisfied(const Model& model);
303
304ArrayDataType ConvertIODataTypeToArrayDataType(IODataType type);
305
306// The process of building models varies according to the import format.
307//
308// (a) In some cases, such as model-proto format, the model should be fully
309// specified. In these cases, no extra action should be taken by this function.
310// (b) In other cases, such as TF graphdef format, the desired types of RNN
311// arrays are not specified directly in the model, neither can they be inferred.
312// However, we can set the types of RNN destination arrays to float. This breaks
313// any cycles such as when resolution of the type of an RNN source array depends
314// on the type of its destination array.
315//
316// This function is applied after the main import, after resolution of flags and
317// after application of ArraysExtraInfo. It only defaults destination RNN arrays
318// to float. If the model is subsequently quantized, it is assumed that the
319// model contains sufficient information for that to be completed. If it is
320// already quantized, then case (a) should hold.
321void FinishBuildingRNNStates(Model* model);
322
323void UseArraysExtraInfo(Model* model, bool quantize_output);
324
325// Calculates the number of elements in tensor given a shape. Shape elements
326// are assumed to be of type T, while the result total is of type U. If U
327// doesn't have enough range to represent the sum of elements, an error is
328// returned.
329template <typename T, typename U>
330tensorflow::Status NumElements(const std::vector<T>& shape, U* num_elements) {
331 static_assert(
332 std::numeric_limits<T>::max() <= std::numeric_limits<uint64_t>::max(),
333 "vector type exceed capabilities of NumElements");
334
335 *num_elements = 1;
336 for (const T& dim : shape) {
337 if (dim < 0) {
338 // TensorFlow's shapes sometimes include -1 to represent an "unknown"
339 // size but TOCO isn't able to create arrays of unknown sizes and will
340 // crash in RequiredBufferSizeForShape().
341 return tensorflow::errors::InvalidArgument(
342 "Tensor shape should not include negative values");
343 }
344 if (*num_elements != 0 &&
345 static_cast<uint64_t>(dim) >
346 std::numeric_limits<U>::max() / *num_elements) {
347 *num_elements = 0;
348 return tensorflow::errors::InvalidArgument("Tensor shape is too large");
349 }
350 *num_elements *= dim;
351 }
352 return ::tensorflow::OkStatus();
353}
354
355// A model file may have shuffled FC weights.
356// When that happens, we want to de-shuffle them immediately on import,
357// so that the rest of toco doesn't need to know about shuffled weights.
358void UndoWeightsShuffling(Model* model);
359
360// Copies minmax, quantization_params, and narrow_range.
361void CopyMinMaxAndQuantizationRelatedFields(const Array& src, Array* dst);
362
363// Delete Array if it's discardable and not referenced as input or output array
364// by any other op than the specified op.
365bool DeleteArrayIfUnusedOutsideOfOp(const std::string& array_name,
366 const Operator* op, Model* model);
367
368} // namespace toco
369
370#endif // TENSORFLOW_LITE_TOCO_TOOLING_UTIL_H_
371