1/* Copyright 2019 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
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
42namespace mlir {
43namespace 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
49namespace {
50
51#define GEN_PASS_CLASSES
52#include "tensorflow/compiler/mlir/lite/transforms/passes.h.inc"
53
54class 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
101void 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
153void 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
173void 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
199quant::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
222quant::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.
234std::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
240std::unique_ptr<OperationPass<func::FuncOp>> CreateDefaultQuantParamsPass() {
241 return std::make_unique<DefaultQuantParamsPass>();
242}
243
244} // namespace TFL
245} // namespace mlir
246