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