1 /* Copyright 2019 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 "mlir/Dialect/StandardOps/IR/Ops.h"
17 #include "mlir/IR/AffineExpr.h"
18 #include "mlir/IR/AffineMap.h"
19 #include "mlir/IR/Attributes.h"
20 #include "mlir/IR/BuiltinTypes.h"
21 #include "mlir/IR/PatternMatch.h"
22 #include "mlir/Pass/Pass.h"
23 #include "mlir/Support/LLVM.h"
24 #include "absl/memory/memory.h"
25 #include "llvm/ADT/STLExtras.h"
26 #include "llvm/ADT/StringSwitch.h"
27 #include "mlir/Dialect/Quant/FakeQuantSupport.h" // from @llvm-project
28 #include "mlir/Dialect/Quant/QuantOps.h" // from @llvm-project
29 #include "mlir/IR/Location.h" // from @llvm-project
30 #include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
31 #include "tensorflow/compiler/mlir/lite/quantization/quantization_utils.h"
32 #include "tensorflow/compiler/mlir/lite/transforms/prepare_quantize_helper.h"
33
34 //===----------------------------------------------------------------------===//
35 // The Pass to add default quantization parameters for the activations which
36 // don't have quantization information. These default parameters are usually
37 // not from real measurement, so this pass is only for test purpose.
38
39 namespace mlir {
40 namespace TFL {
41 // Includs an auto-generated function, which can retrieve the quantization
42 // specification for an TFL operation. The signature of the function is
43 // std::unique_pointer<OpQuantSpec> TFL::GetOpQuantSpec(Operation *)
44 #include "tensorflow/compiler/mlir/lite/utils/generated_op_quant_spec_getters.inc"
45
46 namespace {
47 class DefaultQuantParamsPass
48 : public PassWrapper<DefaultQuantParamsPass, FunctionPass> {
49 public:
DefaultQuantParamsPass(double default_min,double default_max,bool is_signed)50 explicit DefaultQuantParamsPass(double default_min, double default_max,
51 bool is_signed)
52 : default_min_(default_min),
53 default_max_(default_max),
54 is_signed_(is_signed) {}
55
56 void runOnFunction() override;
57
getArgument() const58 StringRef getArgument() const final {
59 // This is the argument used to refer to the pass in
60 // the textual format (on the commandline for example).
61 return "tfl-default-quant";
62 }
getDescription() const63 StringRef getDescription() const final {
64 // This is a brief description of the pass.
65 return "Apply quantization with default quantization parameter";
66 }
67
68 private:
69 // Whether the value is used as a bias input of another op. Here we assume
70 // bias is used immediately by the user. This assumption is always correct
71 // after constant folding.
UsedAsBias(Value value)72 bool UsedAsBias(Value value) {
73 for (auto &use : value.getUses()) {
74 auto biases = TFL::GetOpQuantSpec(use.getOwner())->biases_params;
75 if (biases.find(use.getOperandNumber()) != biases.end()) return true;
76 }
77 return false;
78 }
79
80 // Uses `quant_params` to quantize `value` and inserting a pair of
81 // tfl.quantize and tfl.dequantize ops for this `value`.
82 void QuantizeValue(OpBuilder builder, Value value,
83 quant::QuantParams quant_params);
84
85 // If the value hasn't been quantized, the functions adds it to `values`.
86 void AddToWorkListIfUnquantized(Value value, std::vector<Value> *values);
87
88 // Converts the default min/max to the default quantization parameters.
89 quant::QuantParams GetDefaultQuantParams(Builder builder);
90
91 // Gets the quantization parameters for the bias of an operation by using the
92 // quantization parameters from the non-biases operands.
93 quant::QuantParams GetQuantParamsForBias(Operation *op, int bias,
94 const std::vector<int> &non_biases,
95 quant::AccumulatorScaleFunc func);
96
97 double default_min_;
98 double default_max_;
99 bool is_signed_;
100 quant::QuantParams default_quant_params_;
101 };
102 } // namespace
103
runOnFunction()104 void DefaultQuantParamsPass::runOnFunction() {
105 FuncOp func = getFunction();
106 OpBuilder builder(func);
107
108 std::vector<Value> activation_values;
109 std::vector<Value> bias_values;
110
111 // First of all, collect all the values (block arguments and op results) which
112 // are required to be quantized.
113 for (auto arg : func.getBody().begin()->getArguments()) {
114 if (UsedAsBias(arg)) {
115 AddToWorkListIfUnquantized(arg, &bias_values);
116 } else {
117 AddToWorkListIfUnquantized(arg, &activation_values);
118 }
119 }
120
121 func.walk([&](Operation *op) {
122 if (quant::IsOpNotQuantizable(op) ||
123 op->getParentOfType<TFL::CustomTfOp>()) {
124 return;
125 }
126
127 for (auto res : op->getResults()) {
128 if (UsedAsBias(res)) {
129 AddToWorkListIfUnquantized(res, &bias_values);
130 } else {
131 AddToWorkListIfUnquantized(res, &activation_values);
132 }
133 }
134 });
135
136 // Apply the default quantization parameters for these activation values.
137 quant::QuantParams default_params = GetDefaultQuantParams(builder);
138 for (Value value : activation_values) {
139 QuantizeValue(builder, value, default_params);
140 }
141
142 // Since all the non-biases operands have quantization parameters now, we
143 // should be able to propagate them to the bias operand.
144 for (Value bias : bias_values) {
145 Operation *op = *bias.user_begin();
146 auto spec = TFL::GetOpQuantSpec(op);
147 for (auto &it : spec->biases_params) {
148 quant::QuantParams bias_params = GetQuantParamsForBias(
149 op, it.first, it.second.first, it.second.second);
150 if (!bias_params) continue;
151 QuantizeValue(builder, bias, bias_params);
152 }
153 }
154 }
155
AddToWorkListIfUnquantized(Value value,std::vector<Value> * values)156 void DefaultQuantParamsPass::AddToWorkListIfUnquantized(
157 Value value, std::vector<Value> *values) {
158 // If the result isn't with float type, this result is an integer tensor and
159 // doesn't require quantization.
160 auto tensor_type = value.getType().dyn_cast<TensorType>();
161 if (!tensor_type) {
162 // There are none type values.
163 return;
164 }
165 if (!tensor_type.getElementType().isF32()) return;
166
167 // If the result is consumed by a quantize op, it has been quantized.
168 if (value.hasOneUse() &&
169 llvm::isa<TFL::QuantizeOp>(*value.getUsers().begin()))
170 return;
171
172 // Add this result to the list to apply the default value.
173 values->push_back(value);
174 }
175
QuantizeValue(OpBuilder builder,Value value,quant::QuantParams quant_params)176 void DefaultQuantParamsPass::QuantizeValue(OpBuilder builder, Value value,
177 quant::QuantParams quant_params) {
178 Type expressed_type = value.getType();
179 Type new_type = quant_params.castFromExpressedType(expressed_type);
180 // This value isn't an expressed type (float), skip.
181 if (!new_type) return;
182
183 Block &block = value.getParentRegion()->front();
184 Operation *op = value.getDefiningOp();
185 if (op) {
186 builder.setInsertionPoint(&block, ++Block::iterator(op));
187 } else {
188 builder.setInsertionPointToStart(&block);
189 }
190 TypeAttr type_attr = TypeAttr::get(new_type);
191 auto quantize = builder.create<TFL::QuantizeOp>(value.getLoc(), new_type,
192 value, type_attr);
193 auto dequantize = builder.create<TFL::DequantizeOp>(
194 value.getLoc(), expressed_type, quantize.output());
195 value.replaceAllUsesWith(dequantize);
196
197 // `quantize` is using `dequantize` now, so we should set its operand to
198 // `value`.
199 quantize.getOperation()->replaceUsesOfWith(dequantize, value);
200 }
201
GetQuantParamsForBias(Operation * op,int bias,const std::vector<int> & non_biases,quant::AccumulatorScaleFunc func)202 quant::QuantParams DefaultQuantParamsPass::GetQuantParamsForBias(
203 Operation *op, int bias, const std::vector<int> &non_biases,
204 quant::AccumulatorScaleFunc func) {
205 std::vector<quant::QuantizedType> non_bias_types;
206 non_bias_types.reserve(non_biases.size());
207 for (int non_bias : non_biases) {
208 Operation *non_bias_define = op->getOperand(non_bias).getDefiningOp();
209 if (auto dequant = llvm::dyn_cast<TFL::DequantizeOp>(non_bias_define)) {
210 auto non_bias_type = dequant.input().getType().cast<TensorType>();
211 auto non_bias_ele_type =
212 non_bias_type.getElementType().cast<quant::QuantizedType>();
213 non_bias_types.push_back(non_bias_ele_type);
214 } else {
215 // The non-bias hasn't been quantized, let's skip this bias.
216 break;
217 }
218 }
219 // The non-bias hasn't been quantized, let's skip this bias.
220 if (non_bias_types.size() != non_biases.size()) return {};
221
222 return func(non_bias_types, false);
223 }
224
GetDefaultQuantParams(Builder builder)225 quant::QuantParams DefaultQuantParamsPass::GetDefaultQuantParams(
226 Builder builder) {
227 if (!default_quant_params_) {
228 default_quant_params_ = quant::fakeQuantAttrsToType(
229 builder.getUnknownLoc(),
230 /*numBits=*/8, default_min_, default_max_, /*narrowRange=*/false,
231 builder.getF32Type(), is_signed_);
232 }
233 return default_quant_params_;
234 }
235
236 // Creates an instance of the default quant parameters pass.
CreateDefaultQuantParamsPass(double default_min,double default_max,bool is_signed)237 std::unique_ptr<OperationPass<FuncOp>> CreateDefaultQuantParamsPass(
238 double default_min, double default_max, bool is_signed) {
239 return absl::make_unique<DefaultQuantParamsPass>(default_min, default_max,
240 is_signed);
241 }
242
243 // Registers this pass with default values, only for test
__anon0881ac2d0302null244 static PassRegistration<DefaultQuantParamsPass> pass([] {
245 return CreateDefaultQuantParamsPass(/*default_min=*/-1.0,
246 /*default_max=*/1.0,
247 /*is_signed=*/false);
248 });
249
250 } // namespace TFL
251 } // namespace mlir
252