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 #include <iterator>
16 #include <memory>
17 #include <string>
18 #include <unordered_set>
19 #include <vector>
20
21 #include "tensorflow/lite/toco/graph_transformations/graph_transformations.h"
22 #include "tensorflow/lite/toco/model.h"
23 #include "tensorflow/lite/toco/tooling_util.h"
24 #include "tensorflow/core/platform/logging.h"
25
26 namespace toco {
27
28 namespace {
29
IsElementwiseOperator(OperatorType optype)30 bool IsElementwiseOperator(OperatorType optype) {
31 switch (optype) {
32 case OperatorType::kCast:
33 case OperatorType::kCeil:
34 case OperatorType::kExp:
35 case OperatorType::kFloor:
36 case OperatorType::kNeg:
37 case OperatorType::kRelu:
38 case OperatorType::kRelu1:
39 case OperatorType::kRelu6:
40 case OperatorType::kTanh:
41 case OperatorType::kSqrt:
42 case OperatorType::kSquare:
43 return true;
44 default:
45 return false;
46 }
47 }
48
IsMoveOperator(OperatorType optype)49 bool IsMoveOperator(OperatorType optype) {
50 switch (optype) {
51 case OperatorType::kDepthToSpace:
52 case OperatorType::kExpandDims:
53 case OperatorType::kSpaceToDepth:
54 case OperatorType::kSqueeze:
55 case OperatorType::kReshape:
56 case OperatorType::kTranspose:
57 return true;
58 default:
59 return false;
60 }
61 }
62
63 } // namespace
64
65 // Swap elementwise operators such that all value operators occur before all
66 // element move operators, e.g. negation then transpose.
Run(Model * model,std::size_t op_index,bool * modified)67 ::tensorflow::Status ReorderElementwiseUnary::Run(Model* model,
68 std::size_t op_index,
69 bool* modified) {
70 *modified = false;
71 const auto element_op_it = model->operators.begin() + op_index;
72 std::unique_ptr<Operator>& element_op = *element_op_it;
73 if (!IsElementwiseOperator(element_op->type)) {
74 return ::tensorflow::Status::OK();
75 }
76
77 const string intermediate_name = element_op->inputs[0];
78 auto it = FindOpWithOutput(*model, intermediate_name);
79 if (it == model->operators.end()) {
80 AddMessageF("No preceding operator");
81 return ::tensorflow::Status::OK();
82 }
83
84 std::unique_ptr<Operator>& move_op = *it;
85 if (!IsMoveOperator(move_op->type)) {
86 AddMessageF("Preceding operator is not a move operator");
87 return ::tensorflow::Status::OK();
88 }
89
90 if (CountOpsWithInput(*model, intermediate_name) != 1) {
91 AddMessageF("Input %s used elsewhere", intermediate_name);
92 return ::tensorflow::Status::OK();
93 }
94
95 // Check that the intermediate is discardable.
96 if (!IsDiscardableArray(*model, intermediate_name)) {
97 AddMessageF(
98 "Cannot swap elementwise as it would invalidate %s which is "
99 "an output array.",
100 intermediate_name);
101 return ::tensorflow::Status::OK();
102 }
103
104 // op->inputs may change so we need to keep a value by copy.
105 const string input_name = move_op->inputs[0];
106 const string output_name = element_op->outputs[0];
107
108 AddMessageF("Swapping around operators with %s and %s", LogName(*element_op),
109 LogName(*move_op));
110
111 // If the output array is an exit node for the graph then we need to retain
112 // the name as an output node. This makes the naming scheme a little confusing
113 // but is required in this rare case.
114 if (!IsDiscardableArray(*model, output_name)) {
115 // The output name of the sequence needs to stay static, so create a new
116 // array new use for the intermediate.
117 const auto new_intermediate_name =
118 AvailableArrayName(*model, element_op->outputs[0] + "_reorder");
119 AddMessageF("Adding new array %s to preserve output array name %s",
120 new_intermediate_name, output_name);
121
122 element_op->inputs[0] = input_name;
123 element_op->outputs[0] = new_intermediate_name;
124 model->EraseArray(intermediate_name);
125 move_op->inputs[0] = new_intermediate_name;
126 move_op->outputs[0] = output_name;
127 } else {
128 // The intermediate array is now the output array.
129 for (int i = 0; i < model->operators.size(); i++) {
130 Operator* consumer = model->operators[i].get();
131 for (int j = 0; j < consumer->inputs.size(); j++) {
132 if (consumer->inputs[j] == output_name) {
133 consumer->inputs[j] = intermediate_name;
134 }
135 }
136 }
137
138 element_op->inputs[0] = input_name;
139 move_op->inputs[0] = output_name;
140 }
141
142 // Reset both arrays as shape, type, min/max, etc can all change because of
143 // the position swap.
144 model->EraseArray(element_op->outputs[0]);
145 model->EraseArray(move_op->outputs[0]);
146
147 // Reconstruct.
148 model->GetOrCreateArray(element_op->outputs[0]);
149 model->GetOrCreateArray(move_op->outputs[0]);
150
151 // Swap the order of the operators.
152 element_op.swap(move_op);
153
154 *modified = true;
155 return ::tensorflow::Status::OK();
156 }
157
158 } // namespace toco
159