wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 6 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/compiler/mlir/lite/quantization/common/BUILD
+41
View File
@@ -0,0 +1,41 @@
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
# By default, these targets should only be used within the quantization library.
default_visibility = [
"//learning/brain/mlir/quantization:__subpackages__",
"//tensorflow/compiler/mlir/lite:__subpackages__",
"//tensorflow/compiler/mlir/quantization:__subpackages__",
"//tensorflow/compiler/mlir/stablehlo:__subpackages__",
],
licenses = ["notice"],
)
cc_library(
name = "test_base",
testonly = 1,
srcs = [],
hdrs = ["test_base.h"],
compatible_with = get_compatible_with_portable(),
deps = [
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/lite/quantization/ir:QuantOps",
"//tensorflow/compiler/mlir/quantization/common:func",
"//tensorflow/compiler/mlir/quantization/common/ir:QuantOps",
"//tensorflow/compiler/mlir/quantization/stablehlo/cc:context",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/core:test",
"@com_google_absl//absl/strings:string_view",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
"@llvm-project//mlir:QuantOps",
"@llvm-project//mlir:Support",
"@stablehlo//:stablehlo_ops",
],
)
tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/BUILD
+150
View File
@@ -0,0 +1,150 @@
load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library")
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
# By default, these targets should only be used within the quantization library.
# copybara:uncomment_begin(google-only)
# default_visibility = [
# "//learning/brain/mlir/quantization:__subpackages__",
# "//third_party/odml/litert:__subpackages__",
# "//tensorflow:__subpackages__",
# ],
# copybara:uncomment_end_and_comment_begin
default_visibility = [
"//visibility:public",
],
# copybara:comment_end
licenses = ["notice"],
)
cc_library(
name = "tfl_quantization_driver",
srcs = [
"tfl_quantization_driver.cc",
],
hdrs = [
"tfl_quantization_driver.h",
],
deps = [
":quantization_config",
":quantization_lib",
"//tensorflow/compiler/mlir/lite:tensorflow_lite_ops",
"//tensorflow/compiler/mlir/lite/quantization/ir:QuantOps",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:QuantOps",
"@llvm-project//mlir:Support",
],
)
cc_library(
name = "quantization_lib",
srcs = [
"quantization_driver.cc",
"quantization_interface.cc.inc",
"quantization_utils.cc",
],
hdrs = [
"quantization_driver.h",
"quantization_interface.h.inc",
"quantization_traits.h",
"quantization_utils.h",
],
deps = [
":quantization_config",
":quantization_interfaces_inc_gen",
"//tensorflow/compiler/mlir/lite/quantization/ir:QuantOps",
"//tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy:portable_tensor_utils",
"//tensorflow/compiler/mlir/quantization/common/ir:QuantOps",
"//tensorflow/compiler/mlir/tools/optimize:quantization_utils",
"//tensorflow/core:lib_proto_parsing",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:QuantOps",
"@llvm-project//mlir:Support",
],
)
tf_cc_test(
name = "quantization_driver_test",
srcs = ["quantization_driver_test.cc"],
deps = [
":quantization_lib",
"//tensorflow/compiler/mlir/lite/quantization/common:test_base",
"//tensorflow/compiler/mlir/lite/quantization/ir:QuantOps",
"//tensorflow/compiler/mlir/quantization/common:attrs_and_constraints",
"//tensorflow/compiler/mlir/quantization/common:func",
"//tensorflow/compiler/mlir/tensorflow",
"@com_google_absl//absl/strings:string_view",
"@com_google_googletest//:gtest_main",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:QuantOps",
"@llvm-project//mlir:Support",
],
)
td_library(
name = "quantization_td_files",
srcs = [
"quantization.td",
],
compatible_with = get_compatible_with_portable(),
deps = [
"//tensorflow/compiler/mlir/lite/quantization/ir:QuantizationOpsTdFiles",
"@llvm-project//mlir:OpBaseTdFiles",
],
)
gentbl_cc_library(
name = "quantization_interfaces_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {
"quantization_interface.h.inc": ["-gen-op-interface-decls"],
"quantization_interface.cc.inc": ["-gen-op-interface-defs"],
},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "quantization.td",
deps = [
":quantization_td_files",
],
)
cc_library(
name = "quantization_config",
srcs = [
"quantization_config.cc",
],
hdrs = [
"quantization_config.h",
],
deps = [
"//tensorflow/compiler/mlir/lite/tools/optimize:reduced_precision_metadata",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
],
)
exports_files([
"quantization_traits.h",
"quantization_config.h",
"quantization_utils.h",
])
tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization.td
+248
View File
@@ -0,0 +1,248 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
// This is the quantization definition file for TensorFlow.
#ifdef TF_Quantization
#else
#define TF_Quantization
include "mlir/IR/OpBase.td"
include "mlir/Dialect/Quant/IR/QuantBase.td"
//===----------------------------------------------------------------------===//
// QuantizedType definitions.
//===----------------------------------------------------------------------===//
// The base class of a quantized type. Signed quantized types may be expressed
// as signless integers (i.e. up to op interpretation), but we include an
// explicit signedness check to differentiate the signed/unsigned constraints
// predicates from one another at the TD level.
class QuantizedType<string n, list<int> params, bit signed>
: Type<And<[CPred<"llvm::isa<mlir::quant::QuantizedType>($_self)">,
CPred<"llvm::cast<mlir::quant::QuantizedType>($_self)" #
".getStorageTypeIntegralWidth() == " # !head(params)>,
Or<[CPred<"llvm::cast<mlir::quant::QuantizedType>($_self)" #
".getStorageType().isSignlessInteger()">,
CPred<"llvm::cast<mlir::quant::QuantizedType>($_self)" #
".getStorageType().isSignedInteger() == " # signed>]>]>,
"Q" # !if (signed, "I", "UI") # !head(params) # " type"> {
string name = n;
string asTraitArgsStr =
!interleave(params, ", ") # !if(signed, ", true", ", false");
}
// Uniform quantized types. Two integers "smantissa" and "sexp" are used to
// express the Mantissa and Exponent components of the floating-point scale so
// the scale of the quantized type is "smantissa * 10 ^ sexp".
class UInt8UniformQuantizedType<int zero_pt, int smantissa, int sexp>
: QuantizedType<"Uniform",
[8, zero_pt, smantissa, sexp, 0, 255], 0>;
class Int8UniformQuantizedType<int zero_pt, int smantissa, int sexp>
: QuantizedType<"Uniform",
[8, zero_pt, smantissa, sexp, -128, 127], 1>;
// General uniform quantized types. The definitions can be used to specify
// operand's tensor types.
def QI2 : QuantizedType<"Uniform", [2], 1>;
def QI4 : QuantizedType<"Uniform", [4], 1>;
def QUI4 : QuantizedType<"Uniform", [4], 0>;
def QUI8 : QuantizedType<"Uniform", [8], 0>;
def QI8 : QuantizedType<"Uniform", [8], 1>;
def QUI16 : QuantizedType<"Uniform", [16], 0>;
def QI16 : QuantizedType<"Uniform", [16], 1>;
def QUI32 : QuantizedType<"Uniform", [32], 0>;
def QI32 : QuantizedType<"Uniform", [32], 1>;
//===----------------------------------------------------------------------===//
// TFL native op traits (for quantization).
//
// Ops in this link should have those traits specified:
// https://www.tensorflow.org/lite/performance/quantization_spec
//===----------------------------------------------------------------------===//
def FixedOutputRangeInterface : OpInterface<
"FixedOutputRangeInterface"> {
let cppNamespace = "TFL";
let description = [{
Interface for defining the fixed output range.
}];
let methods = [
InterfaceMethod<
[{Returns the fixed output range.}],
"UniformQuantizedType", "GetFixedOutputRange",
(ins "bool":$sign, "int":$bit_width)
>,
];
}
def AffineQuantizedOpInterface : OpInterface<
"AffineQuantizedOpInterface"> {
let cppNamespace = "TFL";
let description = [{
Interface for affine quantized ops (conv2d, fully_connected, etc.)
}];
let methods = [
InterfaceMethod<
[{Returns the affine operand index.}],
"int", "GetAffineOperandIndex",
(ins), [{}], [{return 1;}]>,
InterfaceMethod<
[{Returns whether narrow range is required for the affine operand.}],
"bool", "RequiredNarrowRangeAffineOperand",
(ins), [{}], [{return true;}]>,
InterfaceMethod<
[{Returns quantization dim for the affine operand.}],
"int", "GetQuantizationDimIndex",
(ins)>,
InterfaceMethod<
[{Returns the dimension index of the output channels.}],
"int", "GetChannelDimIndex", (ins)
>,
];
}
def SameOperandsAndResultsScale : OpInterface<"SameScalesOpInterface"> {
let cppNamespace = "TFL";
let description = [{
Interface for ops potentially have same operands and results scales.
}];
let methods = [
InterfaceMethod<
[{Returns whether same operands and results scales are required.}],
"bool", "RequiredSameOperandsAndResultsScale",
(ins "bool":$sign, "int":$bit_width), [{}], [{return true;}]
>,
InterfaceMethod<
[{Returns whether operands and results must have the same quantized axis.}],
"bool", "RequiredSameQuantizedAxes",
(ins), [{}], [{return true;}]
>,
];
let verify = [{
return mlir::TFL::VerifySameScales($_op);
}];
}
def RequiresQuantizedBiasInterface : OpInterface<"RequiresQuantizedBiasInterface"> {
let description = [{
Interface for ops that require the converter to quantize their bias based on
the scales of the other operands.
}];
let methods = [
InterfaceMethod<
[{Returns the bias operand index.}],
"int", "GetBiasOperandIndex",
(ins), [{}], [{return -1;}]>,
InterfaceMethod<
[{Returns the non-bias operand indices.}],
"std::vector<int>", "GetNonBiasOperandIndices",
(ins), [{}], [{return {};}]>,
];
}
def DynamicRangeQuantizedOpInterface : OpInterface<
"DynamicRangeQuantizedOpInterface"> {
let cppNamespace = "TFL";
let description = [{
Interface for ops dynamic range quantization is supported.
If the op has the kernel support for dynamic range quantization, Q/DQ op
pairs connected to the op are rewritten by its quantized alternatives where
a new op uses Q ops for its operands instead of DQ op. Otherwise, it is
left as is for weight-only which means the weight is dequantized at runtime.
For example, if the kernel does not support dynamic range quantization the
graph will be converted into the following IR:
%q_w = "tfl.pseudo_qconst"() {
qtype = tensor<64x3x3x3x!quant.uniform<i8<-127:127>:f32, 1.000000e+00>>
%w = "tfl.dequantize"(%q_w) :
(tensor<64x3x3x3x!quant.uniform<i8<-127:127>:f32, 1.000000e+00>>) ->
tensor<64x3x3x3xf32>
%conv = "tfl.conv_2d"(%input_act, %w, %bias)
but if it is supported, it will be rewritten as:
%q_w = "tfl.pseudo_qconst"() {
qtype = tensor<64x3x3x3x!quant.uniform<i8<-127:127>:f32, 1.000000e+00>>
%conv = "tfl.conv_2d"(%input_act, %q_w, %bias)
Note that this is part of reaching feature parity with the old quantizer for
dynamic range quantization except:
- Only use_updated_hybrid_scheme=True is supported which means the ops with
the asymmetrically quantizing input support is enabled to use this feature
during MLIR graph rewriting passes while it is configurable in the old
quantizer. So when those ops are matched during graph rewriting passes,
MLIR quantizer will always ignore the pre-set value of the attribute, if
there's any, and set it to True. The reason behind this decision is that
generally activations of these ops show better accuracy with asymmetric
input quantization so we want to deprecate symmetric activation quantization
for those ops eventually.
- Unlike to the old quantizer, per-channel quantization is supported for
weight-only TransposeConvOp.
}];
let methods = [
InterfaceMethod<
[{Returns the quantizable operand indices of the op.}],
"std::vector<int>", "GetQuantizableOperandIndices",
(ins), [{}], [{return {};}]>,
InterfaceMethod<
[{Returns whether the op has the kernel support for dynamic range
quantization.}],
"bool", "GetDynamicRangeQuantKernelSupport",
(ins), [{}], [{return false;}]>,
InterfaceMethod<
[{Returns whether the op requires asymmetric quantize input attribute
setting.}],
"bool", "RequireAsymmetricQuantizeInputsAttr",
(ins), [{}], [{return false;}]>,
];
}
// Specify this trait if the op has a fixed output value range.
class FixedResultScale<QuantizedType qt> : NativeOpTrait<!strconcat(
"TFL::FixedResult", qt.name, "Scale<", qt.asTraitArgsStr, ">::Impl")>;
// Specify this trait if the bias-th input of the op is a bias input, which
// needs a scale based on the scales of op1 and op2.
class AccumulatorUniformScale<int bias, int op1, int op2> : NativeOpTrait<
!strconcat("TFL::AccumulatorUniformScale<",
!interleave([bias, op1, op2], ", "),
">::Impl")>;
// Specify the operand index of the coefficient operand for an affine op
// and also the quantization dimension if per-axis quantization is support.
// If the quantization dimension is -1, per-axis quantization isn't supported.
class AffineOpCoefficient<int dim, int index> : NativeOpTrait<
!strconcat("TFL::AffineOpCoefficient<",
!interleave([dim, index], ", "),
">::Impl")>;
// Specify this trait if the op does have quantizable output. Quantizers will
// apply quantization on this op.
def QuantizableResult : NativeOpTrait<"TFL::QuantizableResult">;
#endif // TF_Quantization
tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_config.cc
+183
View File
@@ -0,0 +1,183 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_config.h"
#include <ios>
#include <optional>
#include <sstream>
#include <string>
#include <vector>
#include "absl/strings/numbers.h"
#include "absl/strings/str_split.h"
#include "absl/strings/string_view.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h"
#include "tensorflow/core/framework/types.pb.h"
// Returns whether the given dtype is a quantization type in TensorFlow.
static bool IsQuantizationType(tensorflow::DataType dtype) {
switch (dtype) {
case tensorflow::DT_QINT8:
case tensorflow::DT_QUINT8:
case tensorflow::DT_QINT16:
case tensorflow::DT_QUINT16:
case tensorflow::DT_QINT32:
return true;
default:
return false;
}
}
namespace mlir {
namespace TFL {
namespace {
bool GetBooleanSpecs(const std::string& bool_val) {
bool result;
std::stringstream iss(bool_val);
iss >> std::boolalpha >> result;
return result;
}
} // namespace
void ParseCustomOpSpecs(const absl::string_view node_names,
const CustomOpUpdateOptions& update_option,
CustomOpMap& custom_op_map) {
if (node_names.empty()) return;
const std::vector<std::string> custom_nodes = absl::StrSplit(node_names, ',');
for (const std::string& cur_node : custom_nodes) {
const std::vector<std::string> node_infos = absl::StrSplit(cur_node, '=');
const std::string& node_name = node_infos[0];
const std::string& node_specification = node_infos[1];
CustomOpInfo new_node_info;
switch (update_option) {
case CustomOpUpdateOptions::kInputIndices: {
const std::vector<std::string> indices =
absl::StrSplit(node_specification, '-');
for (const std::string& cur_index : indices) {
custom_op_map[node_name].quantizable_input_indices.push_back(
std::stoi(cur_index));
}
break;
}
case CustomOpUpdateOptions::kWeightOnly:
custom_op_map[node_name].is_weight_only =
GetBooleanSpecs(node_specification);
break;
case CustomOpUpdateOptions::kNoSideEffect:
custom_op_map[node_name].no_side_effect =
GetBooleanSpecs(node_specification);
break;
}
}
}
bool ParseInputNodeQuantSpecs(const absl::string_view node_names,
const absl::string_view min_values,
const absl::string_view max_values,
const absl::string_view inference_type,
QuantizationSpecs* quant_specs) {
const std::vector<std::string> input_nodes = absl::StrSplit(node_names, ',');
std::vector<std::optional<double>> node_mins;
if (!min_values.empty()) {
std::vector<std::string> node_mins_str = absl::StrSplit(min_values, ',');
for (const std::string& node_mins_str : node_mins_str) {
double value;
if (!absl::SimpleAtod(node_mins_str, &value)) {
llvm::errs() << "Unexpected mins: " << node_mins_str << "\n";
return true;
}
node_mins.push_back(value);
}
}
std::vector<std::optional<double>> node_maxs;
if (!max_values.empty()) {
const std::vector<std::string> node_maxs_str =
absl::StrSplit(max_values, ',');
for (const std::string& node_maxs_str : node_maxs_str) {
double value;
if (!absl::SimpleAtod(node_maxs_str, &value)) {
llvm::errs() << "Unexpected mins: " << node_maxs_str << "\n";
return true;
}
node_maxs.push_back(value);
}
}
tensorflow::DataType final_type = tensorflow::DT_FLOAT;
if (!inference_type.empty() && !DataType_Parse(inference_type, &final_type)) {
return true;
}
return GetInputNodeQuantSpecs(input_nodes, node_mins, node_maxs, final_type,
quant_specs);
}
bool GetInputNodeQuantSpecs(const std::vector<std::string>& node_names,
const std::vector<std::optional<double>>& node_mins,
const std::vector<std::optional<double>>& node_maxs,
const tensorflow::DataType inference_type,
QuantizationSpecs* quant_specs) {
quant_specs->inference_type = inference_type;
// If min/max are not specified, just return;
if (node_mins.empty() || node_maxs.empty()) return false;
// Otherwise make sure min/max has the same size as inputs.
if (IsQuantizationType(inference_type)) {
// min/max should have same size as inputs, or shouldn't be specified.
if (node_names.size() != node_mins.size() ||
node_names.size() != node_maxs.size()) {
return true;
}
for (int i = 0; i < node_names.size(); ++i) {
quant_specs->input_ranges.push_back({node_mins[i], node_maxs[i]});
}
return false;
}
if (!node_mins.empty()) {
llvm::dbgs() << "Ignored input_min_values.";
}
if (!node_maxs.empty()) {
llvm::dbgs() << "Ignored input_max_values.";
}
return false;
}
std::string GetQDQQuantModeString(const QDQConversionMode mode) {
switch (mode) {
case QDQConversionMode::kQDQStatic:
return "Static";
case QDQConversionMode::kQDQDynamic:
return "Dynamic";
case QDQConversionMode::kQDQStrict:
return "Strict";
default:
return "NoQDQ";
}
}
QDQConversionMode GetQDQQuantModeFromString(const std::string& mode_str) {
if (mode_str == "Static") return QDQConversionMode::kQDQStatic;
if (mode_str == "Dynamic") return QDQConversionMode::kQDQDynamic;
if (mode_str == "Strict") return QDQConversionMode::kQDQStrict;
return QDQConversionMode::kQDQNone;
}
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_config.h
+255
View File
@@ -0,0 +1,255 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
// This header file defines node specs for quantization and the methods to parse
// command line flags to these specs.
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_COMMON_QUANTIZATION_LIB_QUANTIZATION_CONFIG_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_COMMON_QUANTIZATION_LIB_QUANTIZATION_CONFIG_H_
#include <cstdint>
#include <optional>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
#include "absl/container/flat_hash_set.h"
#include "absl/strings/string_view.h"
#include "tensorflow/compiler/mlir/lite/tools/optimize/reduced_precision_metadata.h"
#include "tensorflow/core/framework/types.pb.h"
namespace mlir {
namespace TFL {
// Stores information about how to quantize a user-specified custom operation.
struct CustomOpInfo {
std::vector<std::int32_t> quantizable_input_indices;
bool is_weight_only = false;
bool no_side_effect = true;
};
using CustomOpMap = std::unordered_map<std::string, CustomOpInfo>;
enum CustomOpUpdateOptions { kInputIndices, kWeightOnly, kNoSideEffect };
enum class QDQConversionMode { kQDQNone, kQDQStatic, kQDQDynamic, kQDQStrict };
struct QuantizationSpecs {
// Which function this node quant specifications belong to.
std::string target_func = "main";
// Whether to trigger quantization passses for post-training quantization.
// If true, the model input doesn't require user specified input ranges.
bool post_training_quantization = false;
// Whether to allow dynamic range quantization. This is the easiest
// quantization mode which doesn't require QAT or sample inputs.
// This option only targets `DT_HALF` and `DT_QINT8` inference type.
bool weight_quantization = false;
// Whether to use the MLIR dynamic range quantizer instead of TOCO.
bool enable_mlir_dynamic_range_quantizer = false;
// Whether to allow weight-only quantization. This scheme quantizes
// weights but will dequantize them back at runtime which is useful for
// memory bound case without kernel support available in lower precisions.
// Used in MLIR dynamic range quantizer.
bool weight_only_quantization = false;
// The minimum number of elements in a weights array required to apply
// quantization. This is especially useful not to quantize small tensors as
// it is hard to get performance benefits from them with quantization. Used
// in MLIR dynamic range quantizer with int8 weight data type.
int64_t minimum_elements_for_weights = 1024;
// Whether to calculate scales in float to keep quantized values the same with
// old TOCO quantizer.
bool legacy_float_scale = false;
// Whether to perform per-tensor quantization. Currently, this option is only
// valid when the quantization parameters need to be created by scanning the
// constant content (post-training quantization or QAT without weight
// FakeQuant).
bool disable_per_channel = false;
// Whether to disable per-channel weight quantization and enable legacy per
// tensor quantization. The legacy quantization for Dense layers is
// inconsistent with Conv 1x1 which always performs per channel quantization.
bool disable_per_channel_for_dense_layers = false;
// Whether to use fixed output ranges of the activation ops (tanh, sigmoid,
// etc.) and not infer weight constants.
// If this option is set, quantization emulation ops should be placed after
// the ops in the input graph. This flag should be set to false for
// post-training quantization.
bool disable_infer_tensor_range = false;
// Whether to use the unfrozen variable quantization in MLIR. Typically,
// variables are frozen for passing passes, but some variables aren't frozen.
// If it is true, QuantizeVariables pass will be added after the
// PrepareQuantizePass.
bool enable_mlir_variable_quantization = false;
// The node type when the model is exported. Currently this is limited to
// DT_FLOAT, DT_HALF, DT_QINT8, and DT_QUINT8. When DT_HALF is used, the
// `weight_quantization` flag needs to set to true. When DT_QUINT8 is used,
// the `weight_quantization` flag needs to set to false.
tensorflow::DataType inference_type = tensorflow::DT_FLOAT;
// The input and output data type during inference. This flag is only used
// when `inference_type` is different from DT_FLOAT. This flag can only be set
// to DT_FLOAT or as same as `inference_type`. If this flag is different
// from `inference_type`, adaptor ops are inserted as heading and tailing ops
// in the result model.
tensorflow::DataType inference_input_type = tensorflow::DT_FLOAT;
// Input node ranges. These ranges are stored as the same order of function
// arguments. They are only used when `weight_quantization` is set to false,
// and the model is required to have quantization parameters, either from
// quantization aware training or calibration, for the remaining tensors.
std::vector<std::pair<std::optional<double>, std::optional<double>>>
input_ranges;
// Whether to disable setting the quantization parameters of the input nodes
// using input ranges.
bool disable_set_input_nodes_quantization_params = false;
// The default ranges can be used when a tensor doesn't have quantization
// parameters and couldn't be quantized. Used only for latency tests.
std::pair<std::optional<double>, std::optional<double>> default_ranges;
// A serialized "QuantizationInfo" object to specify value ranges for some of
// the tensors with known names.
std::string serialized_quant_stats = "";
// A bitmask to encode support for reduced precision inference in the model.
tflite::optimize::ReducedPrecisionSupport support_mask =
tflite::optimize::ReducedPrecisionSupport::None;
// Whether to run the passes to propagate the quantization parameters and
// graph rewrites. Returns false if the inference_type is DT_FLOAT or
// `weight_quantization` flag is set.
bool RunPropagationAndRewriteQuantizationPasses() const {
return inference_type != tensorflow::DT_FLOAT && !weight_quantization;
}
// TODO: b/202075505 - make implicit weight type clearer
// Whether run the passes and graph rewrites for dynamic range quantization.
bool RunAndRewriteDynamicRangeQuantizationPasses() const {
bool dynamic_range_quantize =
(inference_type != tensorflow::DT_FLOAT) && weight_quantization &&
!post_training_quantization && !disable_infer_tensor_range &&
enable_mlir_dynamic_range_quantizer;
return dynamic_range_quantize;
}
// Returns whether this inference type represents a signed storage type.
bool IsSignedInferenceType() const {
switch (inference_type) {
case tensorflow::DT_QUINT8:
case tensorflow::DT_QUINT16:
return false;
default:
return true;
}
}
// Gets the width of this quantization type. Returns 0 if it isn't a
// quantization type.
int64_t GetQuantizationTypeWidth() const {
switch (inference_type) {
case tensorflow::DT_INT8:
case tensorflow::DT_UINT8:
case tensorflow::DT_QINT8:
case tensorflow::DT_QUINT8:
return 8;
case tensorflow::DT_INT16:
case tensorflow::DT_UINT16:
case tensorflow::DT_QINT16:
case tensorflow::DT_QUINT16:
return 16;
case tensorflow::DT_INT32:
case tensorflow::DT_QINT32:
return 32;
default:
return 0;
}
}
// Whether to add the NumericVerify ops to verify numbers before and after
// quantization.
bool verify_numeric = false;
// Whether to add verification for layer by layer, or on whole model. When
// disabled (per-layer) float and quantized ops will be run from same input
// (output of previous quantized layer). When enabled, float and quantized ops
// will run with respective float and quantized output of previous ops.
bool whole_model_verify = false;
// Whether to use fake quant attributes to calculate quantization parameters.
bool use_fake_quant_num_bits = false;
// Names of ops to block from quantization. Used in QuantizePass.
// For dynamic range quantization, ops in blocklist are quantized in weight-
// only manner.
absl::flat_hash_set<std::string> ops_blocklist;
// Names of locations to block from quantization. Used in QuantizePass.
absl::flat_hash_set<std::string> nodes_blocklist;
// Map from custom op code to custom op quantization information.
// For dynamic range quantization, among the custom ops in the graph those
// specified in this map are subject to quantization.
CustomOpMap custom_map;
// If other than kQDQNone, the model is a floating point graph with QDQ ops
// to be eliminated and fused into quantized kernels.
QDQConversionMode qdq_conversion_mode = QDQConversionMode::kQDQNone;
// When set, adheres to the QDQ annotations added by the framework when
// possible rather than quantizing any op that is possible to quantize.
bool strict_qdq_mode = false;
};
// Parses the command line flag strings to the CustomOpMap specification.
void ParseCustomOpSpecs(absl::string_view node_names,
const CustomOpUpdateOptions& update_option,
CustomOpMap& custom_op_map);
// Parses the command line flag strings to the quantization specification for
// input arrays of a graph. The array names are not stored in the spec, and will
// be matched by position. Returns true if failed.
bool ParseInputNodeQuantSpecs(absl::string_view node_names,
absl::string_view min_values,
absl::string_view max_values,
absl::string_view inference_type,
QuantizationSpecs* quant_specs);
// Gets the quantization specification for input arrays. The array names are not
// stored in the spec, and will be matched by position. The min/max will be
// ignored if the inference_type isn't a quantized type. Returns true if failed.
bool GetInputNodeQuantSpecs(const std::vector<std::string>& node_names,
const std::vector<std::optional<double>>& node_mins,
const std::vector<std::optional<double>>& node_maxs,
tensorflow::DataType inference_type,
QuantizationSpecs* quant_specs);
// Returns a human-readable string of the QDQQuantMode enum class
std::string GetQDQQuantModeString(QDQConversionMode mode);
// Returns the QDQQuantMode enum class from a human-readable string
QDQConversionMode GetQDQQuantModeFromString(const std::string& mode_str);
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_COMMON_QUANTIZATION_LIB_QUANTIZATION_CONFIG_H_
tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_driver.cc
+959
View File
@@ -0,0 +1,959 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_driver.h"
#include <cmath>
#include <cstdint>
#include <iterator>
#include <limits>
#include <memory>
#include <utility>
#include <vector>
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/SmallVector.h"
#include "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/Quant/IR/QuantTypes.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinTypeInterfaces.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Matchers.h" // from @llvm-project
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_traits.h"
#include "tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_utils.h"
#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
namespace mlir {
namespace TFL {
namespace {
constexpr int32_t kBiasMax = std::numeric_limits<int32_t>::max() / 2;
// Uses the type of `value` to set the initial state of the index-th result if
// `as_result` is true or index-th operand if `as_result` is false. The state
// is immutable if the type is a quantized type. Returns the index of this
// new state in the state vector.
void InitializeStateForValue(
Operation* op, const int index, const Value value, const bool as_result,
std::vector<QuantState>& states,
DenseMap<Value, QuantizationDriver::QuantStateIndex>& value_to_state,
DenseMap<QuantizationDriver::OpWithOperandIndex,
QuantizationDriver::QuantStateIndex>& operand_states,
DenseMap<QuantizationDriver::OpWithResultIndex,
QuantizationDriver::QuantStateIndex>& result_states) {
const auto [cached, inserted] = value_to_state.try_emplace(value, 0);
if (!inserted) {
if (as_result) {
result_states[{op, index}] = cached->second;
} else {
operand_states[{op, index}] = cached->second;
}
return;
}
const QuantizedType quantized_type =
QuantizedType::getQuantizedElementType(value.getType());
const bool immutable = quantized_type != nullptr;
const QuantizationDriver::QuantStateIndex next_state_index = states.size();
states.push_back({quantized_type, immutable});
if (as_result) {
result_states[{op, index}] = next_state_index;
} else {
operand_states[{op, index}] = next_state_index;
}
cached->second = next_state_index;
}
bool HasPerAxisQuantizedOperand(Operation* op) {
for (int i = 0; i < op->getNumOperands(); ++i) {
if (auto dq_op = dyn_cast_or_null<quantfork::DequantizeCastOp>(
op->getOperand(i).getDefiningOp())) {
auto type =
mlir::cast<TensorType>(dq_op.getArg().getType()).getElementType();
if (auto per_axis_qtype =
mlir::dyn_cast_or_null<quant::UniformQuantizedPerAxisType>(
QuantizedType::getQuantizedElementType(type))) {
return true;
}
}
}
return false;
}
} // namespace
void QuantizationDriver::InitializeArgState(const BlockArgument arg,
const Value arg_value) {
const auto [cached, inserted] = value_to_state_.try_emplace(arg_value, 0);
if (!inserted) {
arg_states_[arg] = cached->second;
return;
}
const QuantizedType quantized_type =
QuantizedType::getQuantizedElementType(arg_value.getType());
const bool immutable = quantized_type != nullptr;
const QuantizationDriver::QuantStateIndex next_state_index = states_.size();
states_.push_back({quantized_type, immutable});
arg_states_[arg] = next_state_index;
cached->second = next_state_index;
}
void QuantizationDriver::InitializeOperandState(Operation* op, const int index,
const Value value) {
InitializeStateForValue(op, index, value, /*as_result=*/false, states_,
value_to_state_, operand_states_, result_states_);
}
void QuantizationDriver::InitializeResultState(Operation* op, const int index,
const Value value) {
InitializeStateForValue(op, index, value, /*as_result=*/true, states_,
value_to_state_, operand_states_, result_states_);
}
std::unique_ptr<OpQuantSpec> QuantizationDriver::GetQuantSpec(Operation* op) {
return op_quant_spec_getter_(op);
}
std::unique_ptr<OpQuantScaleSpec> QuantizationDriver::GetQuantScaleSpec(
Operation* op) {
return op_quant_scale_spec_getter_(op);
}
bool QuantizationDriver::IsQuantized(Operation* op) {
for (int i = 0; i < op->getNumResults(); ++i) {
if (GetResultQuantState(op, i).IsEmpty()) return false;
}
return true;
}
bool QuantizationDriver::SetConstantResultParams(Operation* op) {
DenseFPElementsAttr attr;
const Value result = op->getResult(0);
if (!matchPattern(result, m_Constant(&attr))) {
return false;
}
// TODO: b/323478683 - Make storage_type_width and narrow_range configurable.
Type final_type;
const auto it = optimized_weights_.find(op);
const bool is_weight = it != optimized_weights_.end();
const bool is_weight_with_per_channel_support =
is_weight && it->second != -1 && is_signed_;
if (is_weight_with_per_channel_support && !disable_per_channel_) {
// When `disable_per_channel_` is false, per-channel symmetric quantization
// parameters are created from the weights when the ops support per-channel
// quantization. Otherwise, uses per-tensor asymmetric quantization with
// narrow range.
// per-axis quantization weight, with symmetric min/max enforced.
final_type = GetUniformQuantizedPerAxisTypeForWeight(
attr, it->second, /*symmetric=*/true, /*num_bits=*/8, is_signed_,
/*narrow_range=*/true, legacy_float_scale_);
} else {
// per-tensor quantization weight
final_type = GetUniformQuantizedTypeForWeight(
attr, /*symmetric=*/is_weight && is_signed_,
/*num_bits=*/8, is_signed_,
/*narrow_range=*/is_weight, legacy_float_scale_);
}
if (const auto quant_type = mlir::dyn_cast_or_null<QuantizedType>(final_type);
quant_type != nullptr) {
return SetResultParams(op, /*result_index=*/0, quant_type);
}
return false;
}
bool QuantizationDriver::SetResultParams(Operation* op, const int result_index,
const QuantizedType quantized_type) {
QuantState& state = GetResultQuantState(op, result_index);
if (state.params == quantized_type) {
return false;
}
if (!state.IsEmpty()) {
RequantizeStates& rescales = GetResultRequantizeStates(op, result_index);
RequantizeState& rescale = rescales.emplace_back();
rescale.pos = RequantizeState::ON_INPUT;
rescale.params = quantized_type;
return true;
}
state.params = quantized_type;
AddUserToList(op, result_index);
return true;
}
QuantizedType QuantizationDriver::GetBiasParams(
Operation* op, const int bias_index,
const ArrayRef<int> non_bias_operand_indices,
const AccumulatorScaleFunc func) {
QuantState& bias_state = GetOperandQuantState(op, bias_index);
if (!bias_state.IsEmpty()) {
return bias_state.params;
}
std::vector<QuantizedType> op_types{};
op_types.reserve(non_bias_operand_indices.size());
int adjusted_quant_dim = -1;
if (op->getNumOperands() > bias_index) {
// Some kernels allow 1D bias, broadcasting it inside the kernel. In this
// case, the `quantizedDimension=0` when quantizing per-channel.
// However, for some kernels which require bias to be already broadcasted
// to match the accumulation shape, the very last index should be used.
Operation* bias_op = op->getOperand(bias_index).getDefiningOp();
if (bias_op != nullptr) {
Type bias_type = bias_op->getResult(0).getType();
if (bias_type != builder_.getNoneType()) {
const int bias_rank = mlir::dyn_cast<ShapedType>(bias_type).getRank();
adjusted_quant_dim = bias_rank > 1 ? bias_rank - 1 : 0;
}
}
}
for (const int non_bias_operand_index : non_bias_operand_indices) {
const QuantState& non_bias_state =
GetOperandQuantState(op, non_bias_operand_index);
op_types.push_back(non_bias_state.params);
}
return func(op_types, adjusted_quant_dim, legacy_float_scale_);
}
bool QuantizationDriver::SetOperandParams(Operation* op,
const int operand_index,
const QuantizedType quantized_type,
const bool override) {
QuantState& state = GetOperandQuantState(op, operand_index);
if (state.params == quantized_type) {
return false;
}
if (!state.IsEmpty() && !override) {
RequantizeStates& rescales = GetOperandRequantizeStates(op, operand_index);
for (RequantizeState& rescale : rescales) {
if (rescale.params == quantized_type) {
rescale.users.emplace_back(op, operand_index);
return true;
}
}
RequantizeState& rescale = rescales.emplace_back();
rescale.pos = RequantizeState::ON_OUTPUT;
rescale.params = quantized_type;
rescale.users.emplace_back(op, operand_index);
return true;
}
state.params = quantized_type;
AddOperandToList(op, operand_index);
return true;
}
void QuantizationDriver::QuantizeOpResult(Operation* op, const int result_index,
const QuantizedType quantized_type) {
builder_.setInsertionPointAfter(op);
const Value original_result = op->getResult(result_index);
QuantizeValue(original_result, quantized_type, op->getLoc());
}
void QuantizationDriver::QuantizeArg(BlockArgument arg,
const QuantizedType quantized_type) {
builder_.setInsertionPointToStart(arg.getOwner());
QuantizeValue(arg, quantized_type, builder_.getUnknownLoc());
}
void QuantizationDriver::QuantizeValue(Value value,
QuantizedType quantized_type,
const Location loc) {
const Type expressed_type = value.getType();
const Type new_value_type =
quantized_type.castFromExpressedType(expressed_type);
// Skip if `value` or `value`'s element type doesn't match the expressed type
// of `quantized_type`.
if (new_value_type == nullptr) return;
auto quantize =
quantfork::QuantizeCastOp::create(builder_, loc, new_value_type, value);
auto dequantize = quantfork::DequantizeCastOp::create(
builder_, loc, expressed_type, quantize.getResult());
// This attribute is set to distinguish the quantize ops being added by the
// quantization pass. These ops can be removed without losing original
// program accuracy.
// TODO: b/323478683 - Make the attribute being part of op definition.
quantize->setAttr(kVolatileOpAttrName, builder_.getUnitAttr());
// `original_result` has a use to `quantize`, so this will replace that use
// by the result of `dequantize`. Remember to reset that use afterwards
value.replaceAllUsesWith(dequantize);
quantize.getOperation()->replaceUsesOfWith(dequantize, value);
}
void QuantizationDriver::RequantizeOpResult(Operation* op,
const int result_index,
RequantizeStates& states) {
if (states.empty()) return;
builder_.setInsertionPointAfter(op);
Value value = op->getResult(result_index);
RequantizeState::RequantizePosition pos = states.front().pos;
if (pos == RequantizeState::NO_REQUANTIZE) {
return;
}
for (const RequantizeState& state : states) {
// Check that all requantization positions are the same for each state.
// Unsure if this check is required.
if (state.pos != pos) {
return;
}
}
if (pos == RequantizeState::ON_OUTPUT) {
Operation* user = value.getUses().begin().getUser();
if (isa<quantfork::QuantizeCastOp>(user)) {
// The requantize op is inserted between `quantize` and `dequantize` ops.
value = user->getResult(0);
builder_.setInsertionPointAfter(user);
}
}
RequantizeValue(value, states, op->getLoc());
}
void QuantizationDriver::RequantizeArg(const BlockArgument arg,
RequantizeStates& states) {
Value value = arg;
builder_.setInsertionPointToStart(arg.getOwner());
if (value.hasOneUse()) {
Operation* user = value.use_begin().getUser();
if (auto q = dyn_cast<quantfork::QuantizeCastOp>(user)) {
value = q.getResult();
builder_.setInsertionPoint(arg.getOwner(), ++Block::iterator(user));
}
}
RequantizeValue(value, states, builder_.getUnknownLoc());
}
void QuantizationDriver::RequantizeValue(Value value, RequantizeStates& states,
const Location loc) {
if (states.empty() || states.front().pos == RequantizeState::NO_REQUANTIZE) {
return;
}
if (states.front().pos == RequantizeState::ON_INPUT) {
RequantizeState& state = states.front();
const Type expressed_type = value.getType();
// The value needs to be requantized. A Quantize op will be created to use
// it as the operand and replace its uses.
const Type new_type = state.params.castFromExpressedType(expressed_type);
if (!new_type) return;
auto requantize_op =
quantfork::QuantizeCastOp::create(builder_, loc, new_type, value);
value.replaceAllUsesWith(requantize_op);
requantize_op.getOperation()->replaceUsesOfWith(requantize_op, value);
// This requantization was defined as required for the result value, so
// there should be only one requant state.
return;
}
// If this is an operand that requires requantization, then the value should
// only have one `DequantizeCastOp` user which produces the operand value.
if (!value.hasOneUse()) {
return;
}
auto dequant_op = dyn_cast_or_null<quantfork::DequantizeCastOp>(
value.use_begin().getUser());
if (!dequant_op) {
return;
}
// It is possible that the dequant value is used by a op that doesn't require
// requant, so only overwrite the first if that is not the case.
const int num_uses = std::distance(dequant_op.getResult().use_begin(),
dequant_op.getResult().use_end());
// Whether to replace quantization params of the first dequantize op
// after the quantized value is produced.
// If there is a use other than the requantize states, then we can't clobber.
bool clobber_first = num_uses <= states.size();
for (RequantizeState& state : states) {
Type expressed_type = QuantizedType::castToExpressedType(value.getType());
if (!expressed_type) continue;
// The value needs to be requantized. A Quantize op will be created to use
// it as the operand and replace its uses.
const Type new_type = state.params.castFromExpressedType(expressed_type);
// This value isn't an expressed type (float), skip.
if (!new_type) continue;
auto requantize_op =
quantfork::QuantizeCastOp::create(builder_, loc, new_type, value);
if (clobber_first) {
dequant_op.setOperand(requantize_op.getResult());
// All ops requiring this value already use the result of dequant.
clobber_first = false;
} else {
auto new_dequant_op = quantfork::DequantizeCastOp::create(
builder_, loc, dequant_op.getResult().getType(),
requantize_op.getResult());
for (auto [op, operand_idx] : state.users) {
op->setOperand(operand_idx, new_dequant_op.getResult());
}
}
}
}
// A heuristic to get quantization parameters satisfies the same scale
// constraints:
// - If there are immutable states,
// - use the single input, or,
// - use the single output, or,
// - use the first one in the collection,
// - use the single input if it is ready, or,
// - use the single output if it is ready, or,
// - use the first ready one in the collection.
QuantizedType QuantizationDriver::GetQuantParamsForSameScaleConstraint(
Operation* op) {
// Two vector to collect Non-empty operands and results states.
std::vector<QuantState*> mutable_states, immutable_states;
for (int i = 0; i < op->getNumOperands(); ++i) {
QuantState& state = GetOperandQuantState(op, i);
if (state.immutable) {
immutable_states.push_back(&state);
} else if (!state.IsEmpty()) {
mutable_states.push_back(&state);
}
}
const int immutable_operands_num = immutable_states.size();
const int mutable_operands_num = mutable_states.size();
// Use the operand's state if it is immutable and it is the only one
// operand.
if (op->getNumOperands() == 1 && immutable_operands_num == 1) {
return immutable_states.front()->params;
}
for (int i = 0; i < op->getNumResults(); ++i) {
QuantState& state = GetResultQuantState(op, i);
if (state.immutable) {
immutable_states.push_back(&state);
} else if (!state.IsEmpty()) {
mutable_states.push_back(&state);
}
}
const int immutable_results_num =
immutable_states.size() - immutable_operands_num;
const int mutable_results_num = mutable_states.size() - mutable_operands_num;
// Use the result's state if it is immutable and it is the only one result.
if (op->getNumResults() == 1 && immutable_results_num == 1) {
return immutable_states.back()->params;
}
// Use the first immutable state to quantize the rest operands and results.
if (!immutable_states.empty()) return immutable_states.front()->params;
// If there are no immutable states, use the operand's state if it is the
// only one operand and has parameters propagated.
if (op->getNumOperands() == 1 && mutable_operands_num == 1) {
return mutable_states.front()->params;
}
// If there are no immutable states, use the result's state if it is the
// only one result and has parameters propagated.
if (op->getNumResults() == 1 && mutable_results_num == 1) {
return mutable_states.back()->params;
}
// Use the first propagated state to quantize the rest operands and results.
if (!mutable_states.empty()) return mutable_states.front()->params;
// None operands/results have parameters propagated, skip this node for now.
return {};
}
void QuantizationDriver::PreprocessConstantOps() {
fn_.walk([&](arith::ConstantOp cst) {
// Non-float tensors are neither weights nor require quantization.
const auto type = mlir::dyn_cast<ShapedType>(cst.getType());
if (!type || !mlir::isa<FloatType>(type.getElementType())) return;
// Skip if the value is NaN or INF.
// Otherwise the illegal scale/zp will be calculated.
auto float_attr = mlir::dyn_cast<DenseFPElementsAttr>(cst.getValueAttr());
if (float_attr && (float_attr.getValues<APFloat>().empty() ||
!float_attr.getValues<APFloat>()[0].isFinite())) {
return;
}
const Value value = cst.getResult();
builder_.setInsertionPoint(cst);
// The following loop will change the value uses, thus we cache all the uses
// needs to be changed.
SmallVector<std::pair<Operation*, int>> uses;
for (OpOperand& use : value.getUses()) {
uses.push_back({use.getOwner(), use.getOperandNumber()});
}
for (const auto [user, operand_num] : uses) {
const std::unique_ptr<OpQuantSpec> spec = GetQuantSpec(user);
const std::unique_ptr<OpQuantScaleSpec> scale_spec =
GetQuantScaleSpec(user);
const BiasParamsMap biases = spec->biases_params;
// The quantization parameters of a `weight` shouldn't be determined by
// other values. So any constants which are not bias, an operand of an
// op with same scale requirements, and haven't been quantized are
// weights.
if (!biases.contains(operand_num) &&
!scale_spec->has_same_scale_requirement &&
!dyn_cast<quantfork::QuantizeCastOp>(user)) {
// Needs to scan the content of weights to get the quantization
// parameters if there are no quantization parameters (FakeQuant ops).
// For this case, the weight will not be duplicated.
weights_.insert(cst);
if (spec->coeff_op_quant_dim.find(operand_num) !=
spec->coeff_op_quant_dim.end()) {
optimized_weights_.insert(
{cst, spec->coeff_op_quant_dim[operand_num]});
}
} else {
// This is a bias or an operand of an op with same scale requirements,
// so the quantization parameter are propagated from or determined by
// other values. Duplicate this constant in case it is shared by
// different users.
if (uses.size() > 1) {
auto new_constant_op =
arith::ConstantOp::create(builder_, cst.getLoc(), cst.getValue());
user->setOperand(operand_num, new_constant_op);
}
}
}
});
}
void QuantizationDriver::SetupAllStates() {
for (BlockArgument arg : fn_.getArguments()) {
args_.push_back(arg);
Value value = arg;
// If the argument is quantized, it should only has one user.
if (arg.hasOneUse()) {
Operation* user = value.use_begin().getUser();
if (auto q = dyn_cast<quantfork::QuantizeCastOp>(user)) {
value = q.getResult();
}
}
InitializeArgState(arg, value);
}
fn_.walk([&](Operation* op) {
std::unique_ptr<OpQuantScaleSpec> scale_spec = GetQuantScaleSpec(op);
if (!IsOpQuantizable(op) && !scale_spec->has_same_scale_requirement) {
return;
}
work_list_.push_back(op);
for (int i = 0; i < op->getNumOperands(); ++i) {
Value operand = op->getOperand(i);
if (Operation* inst = operand.getDefiningOp()) {
// If the operand comes from a `quantfork::DequantizeCastOp`, we use
// the quantized input of this `quantfork::DequantizeCastOp` to set the
// state.
if (auto dq = dyn_cast<quantfork::DequantizeCastOp>(inst)) {
operand = dq.getArg();
}
}
InitializeOperandState(op, i, operand);
}
for (int i = 0; i < op->getNumResults(); ++i) {
Value result = op->getResult(i);
// If the result has been quantized, it should only be used by a
// `quantfork::QuantizeCastOp`. For this case, we uses the quantized
// result to create the state and mark it immutable.
if (result.hasOneUse()) {
Operation* user = result.use_begin().getUser();
if (auto q = dyn_cast<quantfork::QuantizeCastOp>(user)) {
result = q.getResult();
}
}
InitializeResultState(op, i, result);
}
});
}
arith::ConstantOp QuantizationDriver::DuplicateConstantOpIfNeeded(
arith::ConstantOp op, Operation* target_op, const int operand_index) {
if (op.getResult().hasOneUse()) {
return op;
}
OpBuilder builder(op->getContext());
builder.setInsertionPointAfter(op);
arith::ConstantOp new_op = cast<arith::ConstantOp>(builder.clone(*op));
target_op->getOpOperand(operand_index).set(new_op.getResult());
InitializeOperandState(target_op, operand_index, new_op.getResult());
InitializeResultState(new_op, 0, new_op.getResult());
return new_op;
}
bool QuantizationDriver::ShouldCheckBiasScale(
Operation* op, const int bias_index, ArrayRef<int> input_indices,
const QuantizedType quantized_type, int& input_index, int& filter_index) {
// For now, restrict scale adjustment to ops with affine quantized weights,
// and having weights and biases as constants. This currently only applies to
// FC and Conv* ops. Restriction for the weight can be relaxed if there are
// needs for adjusting scale of variable weights.
auto affine_op = dyn_cast<AffineQuantizedOpInterface>(op);
auto bias_op = op->getOperand(bias_index).getDefiningOp<arith::ConstantOp>();
if (!affine_op || !bias_op || input_indices.size() != 2) return false;
if (!mlir::isa<DenseFPElementsAttr>(bias_op.getValue())) return false;
filter_index = affine_op.GetAffineOperandIndex();
if (!op->getOperand(filter_index).getDefiningOp<arith::ConstantOp>()) {
return false;
}
if (filter_index == input_indices[0]) {
input_index = input_indices[1];
} else if (filter_index == input_indices[1]) {
input_index = input_indices[0];
} else {
return false;
}
const QuantState& input_state = GetOperandQuantState(op, input_index);
const QuantState& filter_state = GetOperandQuantState(op, filter_index);
// If quantization parameter for the filter is fixed, should return it as-is.
// Only checks ops with 8-bit input and weights, and 32-bit biases.
return input_state.params.getStorageTypeIntegralWidth() == 8 &&
filter_state.params.getStorageTypeIntegralWidth() == 8 &&
quantized_type.getStorageTypeIntegralWidth() == 32;
}
bool QuantizationDriver::SetBiasParamsWithAdjustments(
Operation* op, const int bias_index, ArrayRef<int> input_indices,
const QuantizedType params) {
bool changed = false;
int input_index;
int filter_index;
if (!ShouldCheckBiasScale(op, bias_index, input_indices, params, input_index,
filter_index)) {
return SetOperandParams(op, bias_index, params);
}
QuantState input_state = GetOperandQuantState(op, input_index);
QuantState filter_state = GetOperandQuantState(op, filter_index);
auto bias_op = op->getOperand(bias_index).getDefiningOp<arith::ConstantOp>();
const double input_scale =
mlir::cast<UniformQuantizedType>(input_state.params).getScale();
auto bias_values = mlir::cast<DenseFPElementsAttr>(bias_op.getValue());
// Restrict maximum absolute value of bias within INT_MAX / 2, to make some
// room for accumulator.
if (auto bias_quantized_type = mlir::dyn_cast<UniformQuantizedType>(params);
bias_quantized_type != nullptr) {
double bias_half_range = 0.0f;
for (auto bias : bias_values.getValues<APFloat>()) {
if (bias_half_range < std::abs(bias.convertToFloat())) {
bias_half_range = std::abs(bias.convertToFloat());
}
}
if (bias_half_range / bias_quantized_type.getScale() < kBiasMax) {
return SetOperandParams(op, bias_index, params);
}
const double new_bias_scale =
static_cast<double>(bias_half_range) / kBiasMax;
changed |= SetOperandParams(
op, bias_index,
UniformQuantizedType::getChecked(
bias_op->getLoc(), params.getFlags(), params.getStorageType(),
params.getExpressedType(), new_bias_scale, 0,
params.getStorageTypeMin(), params.getStorageTypeMax()));
arith::ConstantOp filter_op = DuplicateConstantOpIfNeeded(
op->getOperand(filter_index).getDefiningOp<arith::ConstantOp>(), op,
filter_index);
if (!filter_op) {
return SetOperandParams(op, bias_index, params);
}
const auto filter_quantized_type =
mlir::cast<UniformQuantizedType>(filter_state.params);
changed |= SetOperandParams(
op, filter_index,
UniformQuantizedType::getChecked(
filter_op->getLoc(), filter_quantized_type.getFlags(),
filter_quantized_type.getStorageType(),
filter_quantized_type.getExpressedType(),
new_bias_scale / input_scale, 0,
filter_quantized_type.getStorageTypeMin(),
filter_quantized_type.getStorageTypeMax()),
/*override=*/true);
} else if (auto bias_quantized_type =
mlir::dyn_cast<quant::UniformQuantizedPerAxisType>(params);
bias_quantized_type != nullptr) {
const auto filter_quantized_type =
mlir::cast<quant::UniformQuantizedPerAxisType>(filter_state.params);
std::vector<double> new_bias_scales = bias_quantized_type.getScales().vec();
std::vector<double> new_filter_scales =
filter_quantized_type.getScales().vec();
bool needs_adjustment = false;
for (int i = 0; i < bias_quantized_type.getScales().size(); ++i) {
const float abs_bias = std::abs(bias_values.getValues<float>()[i]);
if (abs_bias / new_bias_scales[i] > kBiasMax) {
new_bias_scales[i] = static_cast<double>(abs_bias) / kBiasMax;
new_filter_scales[i] = new_bias_scales[i] / input_scale;
needs_adjustment = true;
}
}
if (!needs_adjustment) {
return SetOperandParams(op, bias_index, params);
}
changed |= SetOperandParams(
op, bias_index,
quant::UniformQuantizedPerAxisType::getChecked(
bias_op->getLoc(), params.getFlags(), params.getStorageType(),
params.getExpressedType(), new_bias_scales,
bias_quantized_type.getZeroPoints(),
bias_quantized_type.getQuantizedDimension(),
params.getStorageTypeMin(), params.getStorageTypeMax()));
arith::ConstantOp filter_op = DuplicateConstantOpIfNeeded(
op->getOperand(filter_index).getDefiningOp<arith::ConstantOp>(), op,
filter_index);
changed |= SetOperandParams(
op, filter_index,
quant::UniformQuantizedPerAxisType::getChecked(
filter_op->getLoc(), filter_quantized_type.getFlags(),
filter_quantized_type.getStorageType(),
filter_quantized_type.getExpressedType(), new_filter_scales,
filter_quantized_type.getZeroPoints(),
filter_quantized_type.getQuantizedDimension(),
filter_quantized_type.getStorageTypeMin(),
filter_quantized_type.getStorageTypeMax()),
/*override=*/true);
}
return changed;
}
// This method scans the operations in the function to setup the initial
// states for quantization parameter propagation.
// TODO: b/323478683 - This algorithm assumes there are only one pair of
// `quantfork::QuantizeCastOp` and `quantfork::DequantizeCastOp` ops between two
// quantizable ops. A sanity check should be applied.
void QuantizationDriver::Initialize() {
// Duplicate the bias constant, so the states can be setup correctly.
// TODO: b/323478683 - Function definition should also be duplicated if there
// are multiple call sites.
PreprocessConstantOps();
// Setup all the internal states.
SetupAllStates();
}
// Propagates the quantization parameters to the operands, results, and biases.
// TODO: b/323478683 - Do not use while loop to handle this logic.
bool QuantizationDriver::PropagateParamsAndReturnIfChanged() {
// TODO: b/323478683 - Use a typed indicator instead of a bool value.
bool changed = false;
while (!work_list_.empty()) {
Operation* op = work_list_.back();
work_list_.pop_back();
// This op has been quantized, so we should not consider it again.
if (quantized_.contains(op)) continue;
quantized_.insert(op);
if (auto constant_op = dyn_cast<arith::ConstantOp>(op); constant_op) {
// If the workflow requires inferring ranges from the content
// (post-training quantization) and it is weight (filter) and hasn't
// been quantized, we infer the quantization parameters from the content.
if (infer_tensor_range_ && IsWeight(constant_op) && !IsQuantized(op)) {
// The quantization parameters are determined by the content of the
// constant.
changed |= SetConstantResultParams(op);
}
continue;
}
std::unique_ptr<OpQuantScaleSpec> scale_spec = GetQuantScaleSpec(op);
if (scale_spec->has_same_scale_requirement) {
const QuantizedType params = GetQuantParamsForSameScaleConstraint(op);
// The quantization parameters haven't been propagated to any operands
// or results. Skip this node for now.
if (!params) {
quantized_.erase(op);
continue;
}
// If this is a QDQ conversion only, the op could have a same-scale
// requirement for the floating point kernel but allow per-axis
// quantization for the quantized kernel. If the quantized dimension
// changes, the following logic no longer works as the same `params`
// shouldn't be used for both input and output quantization params.
// E.g. During TransposeOp's quantization propagation in
// PrepareQuantize, if the quantization is per-axis and the
// QuantizedDimension is transposed, then the output q-dq params must
// reflect the new QuantizedDimension. So, check and skip the
// propagation if any of the operands has a per-axis quantized type param
// and `RequiredSameQuantizedAxes` set to false.
// Currently, these lines of code are only applicable to TFL_TransposeOp
// and TFL_ReshapeOp. And the output q-dq propagation for this Op is
// performed in `PropagateTransposedPerAxisQuantDim` and
// `PropagateReshapedPerAxisQuantDim` respectively.
if (is_qdq_conversion_ &&
!scale_spec->required_same_quantized_axes_func()) {
if (HasPerAxisQuantizedOperand(op)) continue;
}
// Use the final state to set all the operands' parameters.
for (int i = 0; i < op->getNumOperands(); ++i) {
if (auto type =
mlir::dyn_cast<ShapedType>(op->getOperand(i).getType())) {
// Without this check, it will accidentally propagate the quantization
// information by the shared non-float tensors.
if (mlir::isa<FloatType>(type.getElementType()))
changed |= SetOperandParams(op, i, params);
}
}
// Use the final state to set all the results' parameters.
for (int i = 0; i < op->getNumResults(); ++i)
if (auto type = mlir::dyn_cast<ShapedType>(op->getResult(i).getType());
type != nullptr) {
// Without this check, it will accidentally propagate the quantization
// information by the shared non-float-tensors.
if (mlir::isa<FloatType>(type.getElementType()))
changed |= SetResultParams(op, i, params);
}
}
// If the model already contains immutable QDQs, require upstream to
// explicitly fix output range instead.
if (scale_spec->has_fixed_output_range && infer_tensor_range_ &&
!is_qdq_conversion_) {
// Infer ranges from the activation ops. This is usually required for
// the post-training quantization workflow.
// TODO: b/323478683 - Different result can have different fixed range.
const QuantizedType params =
scale_spec->fixed_output_range_func(is_signed_, bit_width_);
for (auto i = 0; i < op->getNumResults(); ++i) {
// The range is null if the result has been quantized.
if (params) {
changed |= SetResultParams(op, i, params);
}
}
}
const std::unique_ptr<OpQuantSpec> spec = GetQuantSpec(op);
for (const auto& [bias_operand_idx, non_bias_params] :
spec->biases_params) {
const auto& [non_bias_operand_indices, accumulator_scale_func] =
non_bias_params;
const QuantizedType params =
GetBiasParams(op, bias_operand_idx, non_bias_operand_indices,
accumulator_scale_func);
if (!params) {
quantized_.erase(op);
continue;
}
changed |= SetBiasParamsWithAdjustments(op, bias_operand_idx,
non_bias_operand_indices, params);
}
}
return changed;
}
// Finalizes the arguments and result states in the function.
void QuantizationDriver::Finalize() {
for (BlockArgument arg : args_) {
const QuantState& state = GetArgQuantState(arg);
RequantizeStates& requantizes = GetArgRequantizeStates(arg);
if (state.IsEmpty() || (state.immutable && requantizes.empty())) {
continue;
}
if (!state.immutable) {
QuantizeArg(arg, state.params);
}
if (!requantizes.empty()) {
RequantizeArg(arg, requantizes);
}
}
for (const auto& [op_with_result_idx, quant_state_idx] : result_states_) {
const auto [op, result_idx] = op_with_result_idx;
const QuantState& state = GetResultQuantState(op, result_idx);
RequantizeStates& requantizes = GetResultRequantizeStates(op, result_idx);
if (state.IsEmpty() || (state.immutable && requantizes.empty())) {
continue;
}
if (!state.immutable) {
QuantizeOpResult(op, result_idx, state.params);
}
if (!requantizes.empty()) {
RequantizeOpResult(op, result_idx, requantizes);
}
}
}
// Runs quantization in following steps:
// 1. Scans the operations in the function to setup the initial
// states for quantization parameter propagation.
// 2. Propagates the quantization parameters to the operands, results, and
// biases.
// 3. Finalizes the arguments and result states in the function.
void QuantizationDriver::Run() {
Initialize();
if (PropagateParamsAndReturnIfChanged()) {
Finalize();
}
}
void ApplyQuantizationParamsPropagation(
const func::FuncOp func, const bool is_signed, const int bit_width,
const bool disable_per_channel,
const OpQuantSpecGetter op_quant_spec_getter,
const bool infer_tensor_ranges, const bool legacy_float_scale,
const bool is_qdq_conversion) {
ApplyQuantizationParamsPropagation(
func, is_signed, bit_width, disable_per_channel, op_quant_spec_getter,
GetDefaultQuantScaleSpec, infer_tensor_ranges, legacy_float_scale,
is_qdq_conversion);
}
void ApplyQuantizationParamsPropagation(
const func::FuncOp func, const bool is_signed, const int bit_width,
const bool disable_per_channel,
const OpQuantSpecGetter op_quant_spec_getter,
const OpQuantScaleSpecGetter op_quant_scale_spec_getter,
const bool infer_tensor_ranges, const bool legacy_float_scale,
const bool is_qdq_conversion) {
QuantizationDriver(func, is_signed, bit_width, disable_per_channel,
op_quant_spec_getter, op_quant_scale_spec_getter,
infer_tensor_ranges, legacy_float_scale, is_qdq_conversion)
.Run();
}
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_driver.h
+387
View File
@@ -0,0 +1,387 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_COMMON_QUANTIZATION_LIB_QUANTIZATION_DRIVER_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_COMMON_QUANTIZATION_LIB_QUANTIZATION_DRIVER_H_
#include <memory>
#include <utility>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
#include "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/Quant/IR/QuantTypes.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_utils.h"
namespace mlir {
namespace TFL {
// The state for each op result during the quantization parameters propagation.
struct QuantState {
// Quantization parameters propagated to an op result.
QuantizedType params;
// A flag indicates this state (the params) shouldn't be changed after it is
// initialized. This flag will be set to true if the quantization parameters
// are from the quantization-aware training.
const bool immutable;
bool IsEmpty() const { return params == nullptr; }
};
// The state for rescaling the propagated quantization parameters. This can be
// on the input side to satisfy the constraint of previous operation, or on the
// output side to satisfy the constraint of the next operation.
struct RequantizeState {
// Sometimes, we have to "requantize" the quantization result to satisfy all
// the constraints. The "requantize" can happen either on the input or output
// of the quantization result.
enum RequantizePosition {
NO_REQUANTIZE,
ON_INPUT,
ON_OUTPUT
} pos = NO_REQUANTIZE;
// Quantization parameters will be used to add the requantize ops.
QuantizedType params;
// Avoid clobbering all uses of the value, limit to just these ops.
SmallVector<std::pair<Operation*, int>> users;
};
using RequantizeStates = SmallVector<RequantizeState>;
// This is a worklist-driven driver for propagating quantization parameters
// across operations.
//
// The initial quantization parameters are extracted from the quantized type
// between adjacent `quantfork::QuantizeCastOp` and
// `quantfork::DequantizeCastOp`s. All these initial parameters are marked as
// immutable because they are from quantization-aware training.
//
// The algorithm traverses each op and sets the quantization parameters of its
// operands and results, according to its quantization specification, and then
// adds the operands and results to the worklist. If there are any conflicts
// (for example, there are quantization parameters propagated from the previous
// iteration), this process stops if the existing parameters are the immutable,
// or adding `requantize` op to resolve the conflicts.
//
// After the algorithm is converged, pairs of `quantfork::QuantizeCastOp` and
// `quantfork::DequantizeCastOp` are inserted to the right position to
// materialize the propagation and requantize results.
//
class QuantizationDriver {
public:
// Type alias of int used to access `states_`.
using QuantStateIndex = int;
// (op, operand index) pair.
using OpWithOperandIndex = std::pair<Operation*, int>;
// (op, result index) pair.
using OpWithResultIndex = std::pair<Operation*, int>;
explicit QuantizationDriver(func::FuncOp func_op, const bool is_signed,
const int bit_width,
const bool disable_per_channel,
OpQuantSpecGetter op_quant_spec_getter,
OpQuantScaleSpecGetter op_quant_scale_spec_getter,
const bool infer_tensor_range,
const bool legacy_float_scale = false,
const bool is_qdq_conversion = false)
: fn_(func_op),
builder_(func_op.getBody()),
is_signed_(is_signed),
bit_width_(bit_width),
disable_per_channel_(disable_per_channel),
op_quant_spec_getter_(op_quant_spec_getter),
op_quant_scale_spec_getter_(op_quant_scale_spec_getter),
infer_tensor_range_(infer_tensor_range),
legacy_float_scale_(legacy_float_scale),
is_qdq_conversion_(is_qdq_conversion) {}
// The entry point of the quantization parameters propagation.
void Run();
// Sets up the states for all the op results in the function.
void Initialize();
// Propagates the quantization parameters across all the ops.
bool PropagateParamsAndReturnIfChanged();
// Inserts the Quantize and Dequantize ops according to the propagation
// result.
void Finalize();
SmallVector<BlockArgument, 4> GetArgs() { return args_; }
llvm::DenseMap<std::pair<mlir::Operation*, int>, int> GetResultStates() {
return result_states_;
}
DenseMap<OpWithResultIndex, QuantStateIndex> result_states_;
// Returns the state of the block argument.
QuantState& GetArgQuantState(BlockArgument arg) {
return states_[arg_states_[arg]];
}
// Returns the state of the index-th result of the op.
QuantState& GetResultQuantState(Operation* op, const int index) {
return states_[result_states_[{op, index}]];
}
private:
// Duplicates the constant op if it has multiple uses, and replaces
// target_op->operand[operand_index] with the newly created op. This also
// replaces corresponsing quantization states.
arith::ConstantOp DuplicateConstantOpIfNeeded(arith::ConstantOp op,
Operation* target_op,
int operand_index);
// Adjusts bias scale that is derived from other scales (fc, conv ops) to
// prevent overflow of quantized bias values. This also changes quantization
// state of other inputs when needed.
bool SetBiasParamsWithAdjustments(Operation* op, int bias_index,
ArrayRef<int> input_indices,
QuantizedType params);
// Checks preconditions to adjust bias scale.
bool ShouldCheckBiasScale(Operation* op, int bias_index,
ArrayRef<int> input_indices,
QuantizedType quantized_type, int& input_index,
int& filter_index);
// Preprocesses the constants by doing the following:
// - Duplicates constants if it is used by multiple ops. For example, if a
// constant is used by multiple ops as a bias, duplicate constants and
// let each op assign its own quantization parameter for bias.
// - Adds all the non-bias constants (weights) to a set for looking up
// later.
// - Adds all per-channel weights to a set for looking up later.
void PreprocessConstantOps();
// Sets up all the data structures for quantization propagation.
void SetupAllStates();
// Returns Whether the constant is a weight, which shouldn't be shared by
// different ops.
bool IsWeight(Operation* cst) { return llvm::is_contained(weights_, cst); }
// Returns all the related quantization constraints of the op.
std::unique_ptr<OpQuantSpec> GetQuantSpec(Operation* op);
std::unique_ptr<OpQuantScaleSpec> GetQuantScaleSpec(Operation* op);
// Returns whether quantization parameters have been propagated to the results
// of this op.
bool IsQuantized(Operation* op);
// Adds all the users of index-th result of op to the work list.
void AddUserToList(Operation* op, const int index) {
for (Operation* user : op->getResult(index).getUsers()) {
work_list_.push_back(user);
}
}
// Adds the defining op of index-th operand of op to the work list.
void AddOperandToList(Operation* op, const int index) {
if (Operation* operand_op = op->getOperand(index).getDefiningOp();
operand_op != nullptr) {
work_list_.push_back(operand_op);
}
}
// Returns the quantization params for the bias input from the non-bias
// operands which have their indexes in the `non_biases` vector. The returned
// parameters are calculated by `func`.
QuantizedType GetBiasParams(Operation* op, int bias_index,
ArrayRef<int> non_bias_operand_indices,
AccumulatorScaleFunc func);
// Sets the quantization parameters of the result to `quantized_type`. If
// any quantization parameters have been propagated, a requantize will
// happen on the input of propagated quantization. Returns `true` if internal
// state has been modified.
bool SetResultParams(Operation* op, int result_index,
QuantizedType quantized_type);
// Sets the quantization parameters of the operand to `quantized_type`. If any
// quantization parameters have been propagated, a `requantize` will happen on
// the output of propagated quantization. When `override` is set, quantization
// state of the value is replaced instead of adding requantization. Returns
// `true` if internal state has been modified.
bool SetOperandParams(Operation* op, int operand_index,
QuantizedType quantized_type, bool override = false);
// Sets the quantization parameters of the constant result according to its
// content.
bool SetConstantResultParams(Operation* op);
// Inserts the Quantize and Dequantize ops after `op`'s `index`-th result. The
// quantized element type for the result is `quantized_type`.
void QuantizeOpResult(Operation* op, int result_index,
QuantizedType quantized_type);
// Inserts the Quantize and Dequantize ops after `arg`. The quantized element
// type for `arg` is `quantized_type`.
void QuantizeArg(BlockArgument arg, QuantizedType quantized_type);
// Inserts the Quantize and Dequantize ops (i.e. QDQ) after `value`. The
// quantized element type for `value` is `quantized_type`.
void QuantizeValue(Value value, QuantizedType quantized_type, Location loc);
// Inserts the Quantize ops for requantizing the index-th result of the op.
void RequantizeOpResult(Operation* op, int result_index,
RequantizeStates& states);
// Inserts the Quantize ops for requantizing a block argument.
void RequantizeArg(BlockArgument arg, RequantizeStates& states);
// Inserts the Quantize and Dequantize ops to quantize the value and returns
// the Quantize op.
void RequantizeValue(Value value, RequantizeStates& states, Location loc);
// Returns the quantization parameter satisfies the same scale
// constraints for the op. Returns an empty option if this quantization
// parameter doesn't exist.
QuantizedType GetQuantParamsForSameScaleConstraint(Operation* op);
// Returns the state of the index-th operand of the op.
QuantState& GetOperandQuantState(Operation* op, const int index) {
return states_[operand_states_[{op, index}]];
}
// Returns the states of the index-th operand of the op.
RequantizeStates& GetOperandRequantizeStates(Operation* op, const int index) {
return rescale_states_[operand_states_[{op, index}]];
}
// Returns the states of the index-th result of the op.
RequantizeStates& GetResultRequantizeStates(Operation* op, const int index) {
return rescale_states_[result_states_[{op, index}]];
}
// Returns the states of the arg.
RequantizeStates& GetArgRequantizeStates(BlockArgument arg) {
return rescale_states_[arg_states_[arg]];
}
// Sets the state of an argument. If this value is cached, uses the cached
// result without creating new entry in the state vector. Otherwise, allocate
// a new entry in the state vector.
void InitializeArgState(BlockArgument arg, Value arg_value);
// Sets the state of the index-th operand of the op. If this operand is
// cached, uses the cached result without creating new entry in the state
// vector. Otherwise, allocate a new entry in the state vector.
void InitializeOperandState(Operation* op, int index, Value value);
// Sets the state of the index-th result of the op. If this result is cached,
// uses the cached result without creating new entry in the state vector.
// Otherwise, allocate a new entry in the state vector.
void InitializeResultState(Operation* op, int index, Value value);
func::FuncOp fn_;
OpBuilder builder_;
const bool is_signed_;
const int bit_width_;
const bool disable_per_channel_;
// We should distinguish weights and bias constants. Biases are specified by
// the quantization spec or are the operands of ops with same scale spec. The
// rest are weights.
DenseSet<Operation*> weights_;
// The weights require narrow_range quantization. This map collects all the
// weight operands defined by the op quant spec. The value of each entry is
// the quantization dimension. If it is positive, per-channel quantization is
// required.
DenseMap<Operation*, int> optimized_weights_;
// All the ops needs to propagate the quantization parameters to.
std::vector<Operation*> work_list_;
absl::flat_hash_set<Operation*> quantized_;
// The vector contains all the quantization parameters propagated from the
// defining operations of the value, or from the quantization aware training.
std::vector<QuantState> states_;
// The map contains all the quantization parameters which are required to
// satisfy the same operands and results constraint. The keys of this map are
// the values from `operand_states_` and `result_state_`.
absl::flat_hash_map<QuantStateIndex, RequantizeStates> rescale_states_;
// Maps of indexes to the propagation state vector from the ops operands,
// results and arguments.
DenseMap<OpWithOperandIndex, QuantStateIndex> operand_states_;
DenseMap<BlockArgument, QuantStateIndex> arg_states_;
DenseMap<Value, QuantStateIndex> value_to_state_;
// This vector is to preserve the arguments order, so the newly inserted
// quantized ops for the arguments are deterministically ordered.
SmallVector<BlockArgument, 4> args_;
OpQuantSpecGetter op_quant_spec_getter_;
OpQuantScaleSpecGetter op_quant_scale_spec_getter_;
// Infer output ranges for activation ops and constants. This is usually
// required for post-training quantization.
const bool infer_tensor_range_;
// Calculate scales in float instead of double, so that the scales and
// quantized values are exactly the same with the TOCO quantizer.
const bool legacy_float_scale_;
// If true, the model is a floating point graph with QDQ ops to be eliminated
// and fused into quantized kernels.
const bool is_qdq_conversion_;
};
// Propagates quantization parameters across ops in this function and satisfies
// the quantization specification of the ops. This methods assumes the initial
// quantization parameters are stored as adjacent quantize and dequantize ops
// and the propagation results are materialized by inserting pairs of quantize
// and dequantize ops to this function. Set `disable_per_channel` to true to not
// use per channel quantization even the op supports it.
// Setting `infer_tensor_range` to true, to infer quantization parameters from
// the activation ops and weight constants. This is only used for post-training
// quantization.
void ApplyQuantizationParamsPropagation(func::FuncOp func, bool is_signed,
int bit_width, bool disable_per_channel,
OpQuantSpecGetter op_quant_spec_getter,
bool infer_tensor_ranges,
bool legacy_float_scale,
bool is_qdq_conversion);
void ApplyQuantizationParamsPropagation(
func::FuncOp func, bool is_signed, int bit_width, bool disable_per_channel,
OpQuantSpecGetter op_quant_spec_getter,
OpQuantScaleSpecGetter op_quant_scale_spec_getter, bool infer_tensor_ranges,
bool legacy_float_scale, bool is_qdq_conversion);
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_COMMON_QUANTIZATION_LIB_QUANTIZATION_DRIVER_H_
tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_driver_test.cc
+169
View File
@@ -0,0 +1,169 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_driver.h"
#include <cstdint>
#include <functional>
#include <memory>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/strings/string_view.h"
#include "llvm/ADT/DenseMap.h"
#include "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/Quant/IR/QuantTypes.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_utils.h"
#include "tensorflow/compiler/mlir/lite/quantization/common/test_base.h"
#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
#include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.h" // IWYU pragma: keep
#include "tensorflow/compiler/mlir/quantization/common/func.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
namespace mlir::TFL {
namespace {
using ApplyQuantizationParamsPropagationTest =
mlir::quant::QuantizationTestBase;
using ::testing::IsEmpty;
using ::testing::Not;
constexpr absl::string_view kModuleTFLite = R"mlir(
module {
func.func @main(%arg0: tensor<1x4x4x3xf32>) -> tensor<1x4x4x3xf32> attributes {_from_xla_call_module} {
%cst_0 = arith.constant dense<1.0> : tensor<3x1x1x3xf32>
%cst_1 = arith.constant dense<2.0> : tensor<3xf32>
%0 = "tf.XlaCallModule"(%arg0, %cst_0, %cst_1) <{Sout = [#tf_type.shape<1x4x4x3>], module = "", version = 9 : i64}> {_entry_function = @composite_fn_1, _stablehlo_version = "1.0.0", _original_entry_function = "composite_fn_1", _tfl_quant_trait = "fully_quantizable"} : (tensor<1x4x4x3xf32>, tensor<3x1x1x3xf32>, tensor<3xf32>) -> tensor<1x4x4x3xf32>
%1 = "tf.XlaCallModule"(%0, %cst_0, %cst_1) <{Sout = [#tf_type.shape<1x4x4x3>], module = "", version = 9 : i64}> {_entry_function = @composite_fn_2, _stablehlo_version = "1.0.0", _original_entry_function = "composite_fn_2", _tfl_quant_trait = "fully_quantizable"} : (tensor<1x4x4x3xf32>, tensor<3x1x1x3xf32>, tensor<3xf32>) -> tensor<1x4x4x3xf32>
return %1 : tensor<1x4x4x3xf32>
}
func.func private @composite_fn_1(%arg0: tensor<1x4x4x3xf32>, %arg1: tensor<3x1x1x3xf32>, %arg2: tensor<3xf32>) -> tensor<1x4x4x3xf32> attributes {tf_quant.composite_function} {
%0 = "tfl.conv_2d"(%arg0, %arg1, %arg2) {dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "RELU", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32} : (tensor<1x4x4x3xf32>, tensor<3x1x1x3xf32>, tensor<3xf32>) -> tensor<1x4x4x3xf32>
return %0 : tensor<1x4x4x3xf32>
}
func.func private @composite_fn_2(%arg0: tensor<1x4x4x3xf32>, %arg1: tensor<3x1x1x3xf32>, %arg2: tensor<3xf32>) -> tensor<1x4x4x3xf32> attributes {tf_quant.composite_function} {
%0 = "tfl.conv_2d"(%arg0, %arg1, %arg2) {dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "RELU", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32} : (tensor<1x4x4x3xf32>, tensor<3x1x1x3xf32>, tensor<3xf32>) -> tensor<1x4x4x3xf32>
return %0 : tensor<1x4x4x3xf32>
}
}
)mlir";
// TOOD: b/323478683 - Directly use types rather than creating a `unique_ptr`.
std::unique_ptr<OpQuantSpec> GetOpQuantSpec(
const mlir::Operation* op,
bool disable_per_channel_for_dense_layers = false) {
auto spec = std::make_unique<OpQuantSpec>();
spec->coeff_op_quant_dim[1] = 3;
spec->biases_params[2] = {{0, 1}, GetUniformQuantizedTypeForBias};
for (const auto& [key, value] : spec->coeff_op_quant_dim) {
spec->quantizable_operands.insert(key);
}
return spec;
}
TEST_F(ApplyQuantizationParamsPropagationTest,
ConstsUsedMultipleTimesAreDuplicated) {
const OwningOpRef<ModuleOp> module_op_ref =
mlir::quant::QuantizationTestBase::ParseModuleOpString(kModuleTFLite);
func::FuncOp main_fn = mlir::quant::FindMainFuncOp(*module_op_ref);
auto op_quant_spec_getter = [&](mlir::Operation* op) {
return GetOpQuantSpec(op, /*disable_per_channel_for_dense_layers=*/false);
};
QuantizationDriver quantization_driver(
main_fn, /*is_signed=*/true, /*bit_width=*/8,
/*disable_per_channel=*/false, op_quant_spec_getter,
GetDefaultQuantScaleSpec,
/*infer_tensor_range=*/true, /*legacy_float_scale=*/false,
/*is_qdq_conversion=*/false);
quantization_driver.Initialize();
int64_t num_constant_op = 0;
main_fn.walk([&](arith::ConstantOp cst) { ++num_constant_op; });
EXPECT_EQ(num_constant_op, 4);
}
TEST_F(ApplyQuantizationParamsPropagationTest,
PropagateParamsCreatesQuantState) {
const OwningOpRef<ModuleOp> module_op_ref =
ParseModuleOpString(kModuleTFLite);
func::FuncOp main_fn = mlir::quant::FindMainFuncOp(*module_op_ref);
auto op_quant_spec_getter = [&](mlir::Operation* op) {
return GetOpQuantSpec(op, /*disable_per_channel_for_dense_layers=*/false);
};
QuantizationDriver quantization_driver(
main_fn, /*is_signed=*/true, /*bit_width=*/8,
/*disable_per_channel=*/false, op_quant_spec_getter,
GetDefaultQuantScaleSpec,
/*infer_tensor_range=*/true, /*legacy_float_scale=*/false,
/*is_qdq_conversion=*/false);
quantization_driver.Initialize();
ASSERT_TRUE(quantization_driver.PropagateParamsAndReturnIfChanged());
EXPECT_THAT(quantization_driver.GetArgs(), Not(IsEmpty()));
for (const auto& arg : quantization_driver.GetArgs()) {
const QuantState& state = quantization_driver.GetArgQuantState(arg);
EXPECT_TRUE(isa<quant::QuantizedType>(state.params));
}
for (const auto& result : quantization_driver.GetResultStates()) {
mlir::Operation* op = result.first.first;
const int res_index = result.first.second;
const QuantState state =
quantization_driver.GetResultQuantState(op, res_index);
EXPECT_TRUE(isa<quant::QuantizedType>(state.params));
}
}
TEST_F(ApplyQuantizationParamsPropagationTest, FinalizeInsertsQDQOps) {
const OwningOpRef<ModuleOp> module_op_ref =
ParseModuleOpString(kModuleTFLite);
func::FuncOp main_fn = mlir::quant::FindMainFuncOp(*module_op_ref);
auto op_quant_spec_getter = [&](mlir::Operation* op) {
return GetOpQuantSpec(op, /*disable_per_channel_for_dense_layers=*/false);
};
ApplyQuantizationParamsPropagation(
main_fn, /*is_signed=*/true, /*bit_width=*/8,
/*disable_per_channel=*/false, op_quant_spec_getter,
/*infer_tensor_ranges=*/true, /*legacy_float_scale=*/false,
/*is_qdq_conversion=*/false);
mlir::Operation* xla_call_module_op =
quant::FindOperationOfType<TF::XlaCallModuleOp>(main_fn);
mlir::Operation* filter_dcast_op =
xla_call_module_op->getOperand(1).getDefiningOp();
mlir::Operation* filter_qcast_op =
filter_dcast_op->getOperand(0).getDefiningOp();
ASSERT_NE(filter_qcast_op, nullptr);
EXPECT_TRUE(isa<quantfork::QuantizeCastOp>(filter_qcast_op));
EXPECT_TRUE(isa<quantfork::DequantizeCastOp>(filter_dcast_op));
EXPECT_TRUE(isa<quant::UniformQuantizedPerAxisType>(
mlir::cast<TensorType>(filter_qcast_op->getResult(0).getType())
.getElementType()));
}
} // namespace
} // namespace mlir::TFL
tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_traits.h
+152
View File
@@ -0,0 +1,152 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
// This file defines the op traits used in the MLIR TensorFlow Lite dialect.
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_COMMON_QUANTIZATION_LIB_QUANTIZATION_TRAITS_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_COMMON_QUANTIZATION_LIB_QUANTIZATION_TRAITS_H_
#include <cmath>
#include <cstdint>
#include <vector>
#include "mlir/Dialect/Quant/IR/QuantTypes.h" // from @llvm-project
#include "mlir/IR/BuiltinTypeInterfaces.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
using QuantizedType = mlir::quant::QuantizedType;
using UniformQuantizedType = mlir::quant::UniformQuantizedType;
namespace mlir {
namespace TFL {
// Verifies that the op satisfies the same operands and results scales
// constraints. Note that this constraint can only be applied on some
// storage types of the op.
LogicalResult VerifySameScales(Operation* op);
} // namespace TFL
// This includes the interface class definition. It couldn't be in a namespace
// because the table gen doesn't emit the namespace when it is used.
#include "tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_interface.h.inc"
namespace OpTrait {
namespace TFL {
// The base class that all the quantization related OpTrait implements.
template <typename ConcreteType, template <typename> class TraitType>
struct QuantizationSpecTraitBase : public TraitBase<ConcreteType, TraitType> {
static bool IsBias(int index) { return false; }
static bool IsQuantizable() { return true; }
};
// This class provides the API for ops that has a fixed output value range.
// This is used as a trait like this:
//
// class SoftmaxOp
// : public Op<SoftmaxOp,
// OpTrait::TFL::FixedResultUniformScale<
// 8, -128, 390625, -8, 0, 255, false>::Impl> {
//
// TODO(fengliuai): create a better way to express floating point scale in the
// template argument list.
template <unsigned BitWidth, int ZeroPoint, int ScaleMantissa, int ScaleExp,
int64_t StorageTypeMin, int64_t StorageTypeMax, bool Sign>
class FixedResultUniformScale {
public:
template <typename ConcreteType>
class Impl
: public QuantizationSpecTraitBase<
ConcreteType, FixedResultUniformScale<
BitWidth, ZeroPoint, ScaleMantissa, ScaleExp,
StorageTypeMin, StorageTypeMax, Sign>::Impl> {
public:
QuantizedType GetResultQuantizedType(int index) {
auto op = this->getOperation();
const auto result_type =
op->getResult(index).getType().template cast<ShapedType>();
if (!result_type.getElementType().template isa<FloatType>()) return {};
Builder builder(op->getContext());
const IntegerType storage_type = builder.getIntegerType(BitWidth);
const double scale = static_cast<double>(ScaleMantissa) *
std::pow(10.0, static_cast<double>(ScaleExp));
return UniformQuantizedType::getChecked(
Sign, storage_type, result_type.getElementType(), scale, ZeroPoint,
StorageTypeMin, StorageTypeMax, builder.getUnknownLoc());
}
};
};
// This class provides the API for ops that has input as bias. This is used
// as a trait like this:
//
// class Conv2DOp
// : public Op<Conv2DOp, OpTrait::TFL::AccumulatorScale<2, 0, 1>::Impl>
//
// TODO(fengliuai): supports a configurable accumulator bit width.
template <int Bias, int... Operands>
class AccumulatorUniformScale {
public:
template <typename ConcreteType>
class Impl
: public QuantizationSpecTraitBase<
ConcreteType, AccumulatorUniformScale<Bias, Operands...>::Impl> {
public:
// Whether the index-th operand is a bias.
static bool IsBias(int index) { return index == Bias; }
// Returns the indexes of all the non-bias operands.
static std::vector<int> GetAllNonBiasOperands() {
return std::vector<int>({Operands...});
}
};
};
// The trait to specify the operand index of the coefficient for an affine op
// and also the quantization dimension if per-axis quantization is support.
// If the quantization dimension is -1, per-axis quantization isn't supported.
//
// class Conv2DOp
// : public Op<Conv2DOp, OpTrait::TFL::AffineOpCoefficient<0>::Impl>
//
template <int QuantDim, int OperandIndex = 1>
class AffineOpCoefficient {
public:
template <typename ConcreteType>
class Impl
: public TraitBase<ConcreteType,
AffineOpCoefficient<QuantDim, OperandIndex>::Impl> {
public:
static int GetCoefficientOperandIndex() { return OperandIndex; }
static int GetQuantizationDim() { return QuantDim; }
};
};
// This class provides the API for ops that can be quantized.
// This is as a trait like this:
//
// class LessOp : public Op<LessOp, OpTrait::TFL::QuantizableResult> {
//
template <typename ConcreteType>
class QuantizableResult
: public QuantizationSpecTraitBase<ConcreteType, QuantizableResult> {};
} // namespace TFL
} // namespace OpTrait
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_COMMON_QUANTIZATION_LIB_QUANTIZATION_TRAITS_H_
tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_utils.cc
+1075
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_utils.h
+973
View File
@@ -0,0 +1,973 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
// This header file defines common utils used by TFLite transformation
// passes to work with op attributes.
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_COMMON_QUANTIZATION_LIB_QUANTIZATION_UTILS_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_COMMON_QUANTIZATION_LIB_QUANTIZATION_UTILS_H_
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <functional>
#include <memory>
#include <string>
#include <type_traits>
#include <utility>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/strings/string_view.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Casting.h"
#include "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/Quant/IR/QuantTypes.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinTypeInterfaces.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/IRMapping.h" // from @llvm-project
#include "mlir/IR/Location.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Matchers.h" // from @llvm-project
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/OperationSupport.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_config.h"
#include "tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_traits.h"
#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
#include "tensorflow/compiler/mlir/quantization/common/ir/FakeQuantSupport.h"
#include "tensorflow/core/framework/types.pb.h"
namespace mlir {
namespace TFL {
// A unit attribute can be attached to the quantize/dequantize ops which are
// added by the quantization passes. These ops can be removed erased without
// losing accuracy.
inline constexpr char kVolatileOpAttrName[] = "volatile";
// Following attributes are used to mark ops that are not quantizable during
// debug model generation process for whole-model verify mode. If these
// attributes are attached, the upstream float/quantized ops know which ops to
// connect to, and it also prevents these ops from being copied again.
inline constexpr char kDebugModeOpFloatAttrName[] = "debug_float";
inline constexpr char kDebugModeOpQuantAttrName[] = "debug_quant";
// Used to annotate custom ops if they are quantizable.
inline constexpr char kQuantTraitAttrName[] = "_tfl_quant_trait";
enum QuantizationTrait { FullyQuantizable = 0, NotQuantizable = 1 };
inline constexpr absl::string_view QuantTraitValues[] = {"fully_quantizable",
"not_quantizable"};
inline constexpr char kOutputQuantized[] = "_output_quantized";
inline constexpr double kNearZeroTolerance = 1.0e-6;
using QuantParams = QuantizedType;
using QuantSpec = QuantizationSpecs;
using SignedInteger = std::pair<unsigned, unsigned>; // bitwidth and sign
using QuantParamsForResults = llvm::SmallVector<QuantizedType, 4>;
using AccumulatorScaleFunc =
std::function<QuantizedType(const std::vector<QuantizedType>&, int, bool)>;
using BiasParamsMap =
absl::flat_hash_map<int, std::pair<std::vector<int>, AccumulatorScaleFunc>>;
// UniformQuantizedType GetFixedOutputRange(bool sign, int bit_width)
using GetFixedOutputRangeFunc = std::function<UniformQuantizedType(bool, int)>;
// bool RequiredSameOperandsAndResultsScale(bool sign, int $bit_width)
using RequiredSameOperandsAndResultsScaleFunc = std::function<bool(bool, int)>;
// bool RequiredSameQuantizedAxes()
using RequiredSameQuantizedAxesFunc = std::function<bool()>;
using CustomMap = CustomOpMap;
using Operation = ::mlir::Operation;
// Quantization spec of an op, driving the quantization algorithm.
struct OpQuantSpec {
// Maps the operand index of a bias input to its quantization specifications,
// including the non-bias operand indexes and the method retrieving
// quantization parameters from list of parameters of the non-bias operands.
// This map is empty if the op doesn't have a bias operand.
BiasParamsMap biases_params;
// Quantization parameters for value restricted outputs. This is the
// "hard-coded" parameters and should be used unconditionally for the
// quantized op. This vector is empty if the op doesn't have value restricted
// outputs.
llvm::DenseMap<SignedInteger, QuantParamsForResults> restricted_output_params;
// Coefficient operand index and whether supporting per-channel quantization.
// For QAT, this information is carried by the FakeQuant*/Quantize/Dequantize
// ops, but post-training quantization, the quantization parameters need to be
// inferred from the tensor content and op property. A "-1" value indicates
// the operand doesn't support per-channel quantization.
llvm::DenseMap<int, int> coeff_op_quant_dim;
// Indices of quantizable operands. Biases are not included in this field,
// the indices of biases can be found in the `biases_params`.
absl::flat_hash_set<int> quantizable_operands;
};
// A function signature for getting the particular OpQuantSpec for the provided
// op.
using OpQuantSpecGetter =
std::function<std::unique_ptr<OpQuantSpec>(mlir::Operation*)>;
// Quantization scale spec of an op. The information defined in the MLIR
// interfaces FixedOutputRangeInterface and SameOperandsAndResultsScale should
// be checked first if present.
// TODO: b/323478683: Consider deprecating this.
struct OpQuantScaleSpec {
// Whether this op has a fixed range requirement (e.g. sigmoid)
bool has_fixed_output_range = false;
// Whether this op should have same operand and result scales (e.g. concat)
bool has_same_scale_requirement = false;
// Whether this op should have same operand and result type (e.g. gather)
bool has_same_operand_and_result_type_requirement = false;
// Returns the fixed output range, when has_fixed_output_range is set.
GetFixedOutputRangeFunc fixed_output_range_func;
// Returns whether same operands and results scales are required.
RequiredSameOperandsAndResultsScaleFunc required_same_scale_func =
[](bool sign, int bit_width) { return true; };
// Returns whether operands and results must have the same quantized axis.
RequiredSameQuantizedAxesFunc required_same_quantized_axes_func = []() {
return true;
};
};
// A function signature for getting the particular OpQuantScaleSpec for the
// provided op.
using OpQuantScaleSpecGetter =
std::function<std::unique_ptr<OpQuantScaleSpec>(mlir::Operation*)>;
// Used in TFL Numeric Verify
struct NumericVerifySpec {
// Whether to enable numeric verification
bool verify_numeric = false;
// Tolerance level from the quantized value for verification. If the tolerance
// is very small(<0.1), only the stats of the diff is displayed.
float error_tolerance = 5.0f;
// Whether to verify numerical correctness layer by layer or by whole model
bool whole_model_verify = false;
// Whether to enable log for failures
bool log_if_failed_flag = false;
};
// Used in TFL Quantize Pass
struct QuantPassSpec {
// Variables to control TFL Numeric Verify
NumericVerifySpec numeric_verify_spec;
// Variables related to quantization
QuantSpec quant_spec;
};
// Re-calculates scales again in float instead of simply downcasting existing
// scales.
quant::QuantizedType DownCastScale(quant::QuantizedType type,
const SmallVectorImpl<double>& mins,
const SmallVectorImpl<double>& maxs,
Location loc);
quant::QuantizedType DownCastScale(quant::QuantizedType type, double min,
double max, Location loc);
bool IsOpQuantizable(mlir::Operation* op);
bool QuantizableOpSupportsFloatOutputType(mlir::Operation* op);
// Specialized version of location to string for flatbuffer exported locations.
inline std::string GetTensorNameFromLoc(Location loc) {
if (auto name_loc = llvm::dyn_cast<NameLoc>(loc)) {
return name_loc.getName().str();
}
return "";
}
template <typename QuantizeOpT, typename DequantizeOpT>
struct ConvertStatsToQDQs : public OpRewritePattern<quantfork::StatisticsOp> {
ConvertStatsToQDQs(int num_bits, bool narrow_range, bool is_signed,
bool legacy_float_scale, MLIRContext* context)
: OpRewritePattern<quantfork::StatisticsOp>(context),
num_bits(num_bits),
narrow_range(narrow_range),
is_signed(is_signed),
legacy_float_scale(legacy_float_scale) {}
LogicalResult matchAndRewrite(quantfork::StatisticsOp op,
PatternRewriter& rewriter) const override {
Type expressed = llvm::cast<ShapedType>(op.getType()).getElementType();
quant::QuantizedType quant_type;
SmallVector<double, 4> mins, maxs;
if (op.getAxisStats().has_value()) {
// Per axis quantization (or per channel quantization)
int stats_num = op.getAxisStats()->getNumElements();
if (stats_num == 0 || stats_num % 2 != 0) return failure();
auto stats = llvm::dyn_cast<DenseFPElementsAttr>(*op.getAxisStats());
if (!stats) return failure();
for (auto it = stats.begin(), e = stats.end(); it != e; ++it) {
double rmin = FloatAttr::getValueAsDouble(*it++);
double rmax = FloatAttr::getValueAsDouble(*it);
// The default nudging implementation of mlir quant library might cause
// clamping during inference if the calibration range isn't wide enough.
// So here we adjust the range to include 0.0.
rmin = std::min(rmin, 0.0);
rmax = std::max(rmax, 0.0);
if (num_bits == 16) {
// TODO: b/266536261 - Since the kernel implementation assumes that
// 16x8 integer quantization is symmetric, this MLIR quantizer
// supports only symmetric quantization.
rmax = std::max(std::abs(rmin), std::abs(rmax));
rmin = -rmax;
}
TensorRangeSanityCheck(op, rmin, rmax);
mins.push_back(rmin);
maxs.push_back(rmax);
}
quant_type = quantfork::fakeQuantAttrsToType(
op.getLoc(), num_bits, *op.getAxis(), mins, maxs, narrow_range,
expressed, is_signed);
if (legacy_float_scale) {
quant_type =
mlir::TFL::DownCastScale(quant_type, mins, maxs, op->getLoc());
}
} else if (auto stats =
llvm::dyn_cast<DenseFPElementsAttr>(op.getLayerStats())) {
// Per tensor quantization
auto statValues = stats.getValues<APFloat>();
double rmin = FloatAttr::getValueAsDouble(statValues[0]);
double rmax = FloatAttr::getValueAsDouble(statValues[1]);
// The default nudging implementation of mlir quant library might cause
// clamping during inference if the calibration range isn't wide enough.
// So here we adjust the range to include 0.0.
rmin = std::min(rmin, 0.0);
rmax = std::max(rmax, 0.0);
if (num_bits == 16) {
// TODO: b/266536261 - Since the kernel implementation assumes that
// 16x8 integer quantization is symmetric, this MLIR quantizer supports
// only symmetric quantization.
rmax = std::max(std::abs(rmin), std::abs(rmax));
rmin = -rmax;
}
TensorRangeSanityCheck(op, rmin, rmax);
quant_type =
quantfork::fakeQuantAttrsToType(op.getLoc(), num_bits, rmin, rmax,
narrow_range, expressed, is_signed);
if (legacy_float_scale) {
quant_type =
mlir::TFL::DownCastScale(quant_type, rmin, rmax, op->getLoc());
}
} else {
return failure();
}
rewriter.setInsertionPointAfter(op.getOperation());
Type result_type = quant_type.castFromExpressedType(op.getType());
auto q =
QuantizeOpT::create(rewriter, op.getLoc(), result_type, op.getArg());
q->setAttr(kVolatileOpAttrName, rewriter.getUnitAttr());
auto dq = DequantizeOpT::create(rewriter, op.getLoc(), op.getType(), q);
op.getResult().replaceAllUsesWith(dq);
q.getOperation()->replaceUsesOfWith(dq, op.getArg());
op.erase();
return success();
}
private:
int num_bits;
bool narrow_range;
bool is_signed;
bool legacy_float_scale;
// Emits an op warning message if the calibrated range is larger than 10.0 and
// the storage type is less than or equal to 8 bits.
void TensorRangeSanityCheck(quantfork::StatisticsOp op, double& min,
double& max) const {
double range = std::fabs(max - min);
if (num_bits <= 8 && range >= 10.0) {
op.emitWarning()
<< "Tensor range is too wide to be quantized. Use tf.clip_by_value "
"or tf.relu6 to narrow the tensor range. Range: "
<< range << ", bit width: " << num_bits;
}
if (std::abs(max - min) < kNearZeroTolerance) {
op.emitWarning() << "Tensor range (" << min << ", " << max
<< ") is too narrow and it might cause overflow. "
"Expanding range symmetrically by "
<< kNearZeroTolerance;
min -= kNearZeroTolerance;
max += kNearZeroTolerance;
}
}
};
template <typename VerifierT>
bool UsedBy(mlir::Operation* op) {
for (mlir::Operation* user : op->getUsers()) {
if (llvm::isa_and_nonnull<VerifierT>(user)) return true;
}
return false;
}
template <typename VerifierT>
void CreateVerifier(mlir::Operation* quantizing_op,
mlir::Operation* quantized_op, PatternRewriter& rewriter,
int result_idx, const QuantPassSpec& quant_params) {
rewriter.setInsertionPointAfter(quantized_op);
FloatAttr tolerance = rewriter.getF32FloatAttr(
quant_params.numeric_verify_spec.error_tolerance);
BoolAttr log =
rewriter.getBoolAttr(quant_params.numeric_verify_spec.log_if_failed_flag);
// Verify the quantized value by sending the result to the verifier.
VerifierT::create(rewriter, quantizing_op->getLoc(),
quantized_op->getResult(result_idx).getType(),
quantized_op->getResult(result_idx),
quantizing_op->getResult(result_idx), tolerance, log);
}
template <>
inline bool UsedBy<void>(mlir::Operation* op) {
return false;
}
// This specialization is not going to be called, but needed for compilation.
template <>
inline void CreateVerifier<void>(mlir::Operation* quantizing_op,
mlir::Operation* quantized_op,
PatternRewriter& rewriter, int result_idx,
const QuantPassSpec& quant_params) {}
// A base rewrite pattern which matches any N-in-M-out operations with
// quantization parameters propagated to at least one of its operands. The
// quantization parameters are annotated by the QuantizeOp/DequantizeOp pairs.
// Each matched pattern are rewritten by its quantized alternatives.
//
// The concrete pattern, extends from this base pattern, can specify whether it
// allows dynamic range quantized operands and results for the operations in the
// current context. These "DynamicRangeQuantized" operands and results don't
// have quantization parameters propagated to, so will be in float in the
// quantized results. The concrete pattern should define the following two
// functions:
//
// bool AllowDynamicRangeQuantizedOperand(Operation *) const
// bool AllowDynamicRangeQuantizedResult(Operation *) const
//
// Full integer quantization disallows "DynamicRangeQuantized" operands or
// results. Dynamic range quantization allows "DynamicRangeQuantized" operands
// and results.
template <typename ConcreteT, typename QuantizeOpT, typename DequantizeOpT,
typename VerifierT, typename RootOpT = DequantizeOpT>
class QuantizationPattern : public RewritePattern {
public:
using BaseType = QuantizationPattern<ConcreteT, QuantizeOpT, DequantizeOpT,
VerifierT, RootOpT>;
explicit QuantizationPattern(MLIRContext* context,
const QuantPassSpec& quant_params)
// Set the score to a large number so it is always preferred.
: RewritePattern(RootOpT::getOperationName(), 300, context),
quant_params_(quant_params) {}
LogicalResult matchAndRewrite(mlir::Operation* op,
PatternRewriter& rewriter) const override {
llvm::SmallVector<mlir::Operation*, 4> quantizing_ops;
// Collect all the ops to quantize, as the user / producer of the root op.
if constexpr (std::is_same_v<RootOpT, DequantizeOpT>) {
if (op->getNumResults() != 1) {
return failure();
}
auto users = op->getResult(0).getUsers();
quantizing_ops.append(users.begin(), users.end());
} else if constexpr (std::is_same_v<RootOpT, QuantizeOpT>) {
if (op->getNumOperands() != 1) {
return failure();
}
Value quantize_operand = op->getOperand(0);
if (QuantizedType::getQuantizedElementType(quantize_operand.getType())) {
// The input of this QuantizeOp has already been quantized, i.e.
// rescale.
return failure();
}
DenseFPElementsAttr attr;
if (matchPattern(quantize_operand, m_Constant(&attr))) {
// Const-> QuantizeOp pattern will be handled separately.
return failure();
}
if (mlir::Operation* quantizing_op = quantize_operand.getDefiningOp()) {
quantizing_ops.push_back(quantizing_op);
}
}
tensorflow::DataType inference_type =
quant_params_.quant_spec.inference_type;
bool weight_only_quantization =
quant_params_.quant_spec.weight_only_quantization;
bool enable_verify = quant_params_.numeric_verify_spec.verify_numeric;
bool enable_whole_model_verify =
quant_params_.numeric_verify_spec.whole_model_verify;
absl::flat_hash_set<std::string> ops_blocklist =
quant_params_.quant_spec.ops_blocklist;
absl::flat_hash_set<std::string> nodes_blocklist =
quant_params_.quant_spec.nodes_blocklist;
CustomMap custom_map = quant_params_.quant_spec.custom_map;
// Rewrite the floating-point ops to the quantized version, by fusing
// preceding dequantize ops and succeding quantize ops.
for (mlir::Operation* quantizing_op : quantizing_ops) {
// If it is requantize op, we shouldn't rewrite this op.
if (llvm::isa<QuantizeOpT, DequantizeOpT>(quantizing_op)) {
return failure();
}
// If the op is terminator, not quantizable or any ops from the mlir quant
// ops dialect, we shouldn't rewrite. In case of whole-model verify debug
// mode, not-quantizable ops should be duplicated to keep parallel
// float/quant model execution.
if (quantizing_op->hasTrait<OpTrait::IsTerminator>()) {
return failure();
}
if (!IsOpQuantizable(quantizing_op) &&
!static_cast<const ConcreteT*>(this)->IsQuantizableCustomOp(
quantizing_op, custom_map)) {
if (!(enable_verify && enable_whole_model_verify)) {
return failure();
}
if (quantizing_op->hasAttr(kDebugModeOpQuantAttrName) ||
quantizing_op->hasAttr(kDebugModeOpFloatAttrName)) {
return failure();
}
rewriter.setInsertionPoint(quantizing_op);
mlir::Operation* float_op = rewriter.clone(*quantizing_op);
quantizing_op->setAttr(kDebugModeOpQuantAttrName,
rewriter.getUnitAttr());
float_op->setAttr(kDebugModeOpFloatAttrName, rewriter.getUnitAttr());
RewireFloatModelBackbone(quantizing_op, float_op);
return success();
}
// Blocklist op is checked in advance for non-dynamic range quantization
// case.
if (!quant_params_.quant_spec.weight_quantization &&
(ops_blocklist.find(quantizing_op->getName().getStringRef().str()) !=
ops_blocklist.end())) {
return failure();
}
if (!nodes_blocklist.empty()) {
if (auto name_loc = llvm::dyn_cast<NameLoc>(quantizing_op->getLoc())) {
std::string sloc = name_loc.getName().str();
if (!sloc.empty() &&
(nodes_blocklist.find(sloc) != nodes_blocklist.end())) {
return failure();
}
}
}
// An op with float inputs and outputs are expected when it's used by a
// NumericVerify op. Skip this op.
if (enable_verify && UsedBy<VerifierT>(quantizing_op)) {
continue;
}
bool is_operand_or_result_modified = false;
// Collect all the quantized inputs and "clone" the matched op by these
// inputs.
SmallVector<Value, 4> inputs;
inputs.reserve(quantizing_op->getNumOperands());
for (auto operand : quantizing_op->getOperands()) {
Type operand_type = operand.getType();
if (isa<NoneType>(operand_type)) {
inputs.push_back(operand);
continue;
}
auto ele_type =
llvm::cast<TensorType>(operand.getType()).getElementType();
if (static_cast<const ConcreteT*>(this)
->AllowDynamicRangeQuantizedOperand(quantizing_op,
custom_map)) {
auto dq_op = dyn_cast_or_null<DequantizeOpT>(operand.getDefiningOp());
if (dq_op && inference_type == tensorflow::DT_QINT8 &&
!static_cast<const ConcreteT*>(this)->IsWeightOnlyOp(
quantizing_op, ops_blocklist, weight_only_quantization,
custom_map)) {
// Dynamic range quantization is applied by having QuantizeOp as an
// input. Only int8 weight is supported for now.
inputs.push_back(dq_op.getOperand());
is_operand_or_result_modified = true;
} else {
// Otherwise, it's the case where the operand is activations or the
// quantizing_op is non-supported/weight-only.
inputs.push_back(operand);
}
} else {
if (auto dq_op =
dyn_cast_or_null<DequantizeOpT>(operand.getDefiningOp())) {
is_operand_or_result_modified = true;
inputs.push_back(dq_op.getOperand());
} else if (!ele_type.isF32()) {
// If the operand is an integer tensor, then it doesn't require the
// DequantizeOp in the pattern.
inputs.push_back(operand);
} else {
return failure();
}
}
}
mlir::Operation* quantized_op;
if (QuantizableOpSupportsFloatOutputType(quantizing_op)) {
rewriter.setInsertionPointAfter(quantizing_op);
OperationState new_state(
quantizing_op->getLoc(), quantizing_op->getName().getStringRef(),
inputs, quantizing_op->getResultTypes(), quantizing_op->getAttrs());
for (const auto& indexed_regions :
llvm::enumerate(quantizing_op->getRegions())) {
Region* target_region = new_state.addRegion();
IRMapping mapping;
indexed_regions.value().cloneInto(target_region, mapping);
}
quantized_op = rewriter.create(new_state);
rewriter.replaceOp(quantizing_op, quantized_op);
} else {
// Collect all the quantized outputs and replace them by the results of
// the new quantized op.
llvm::SmallDenseMap<Value, int> outputs_replaced;
SmallVector<Type, 4> output_types;
output_types.reserve(quantizing_op->getNumResults());
for (const auto& enumerated_result :
llvm::enumerate(quantizing_op->getResults())) {
Value result = enumerated_result.value();
Type result_type = result.getType();
// Add this to the test coverage once we create test ops with none
// type results.
if (isa<NoneType>(result_type)) {
outputs_replaced.insert({result, enumerated_result.index()});
output_types.push_back(result_type);
continue;
}
Type result_ele_type =
llvm::cast<TensorType>(result.getType()).getElementType();
// If the user is the QuantizeOp, it must be the only user.
if (result.hasOneUse() &&
llvm::isa<QuantizeOpT>(*result.user_begin())) {
auto user = llvm::cast<QuantizeOpT>(*result.user_begin());
outputs_replaced.insert(
{user.getResult(), enumerated_result.index()});
output_types.push_back(user.getType());
is_operand_or_result_modified = true;
} else if (!result_ele_type.isF32()) {
// If the result is an integer tensor, then it doesn't require the
// D op in the pattern.
outputs_replaced.insert({result, enumerated_result.index()});
output_types.push_back(result.getType());
} else if (static_cast<const ConcreteT*>(this)
->AllowDynamicRangeQuantizedResult(quantizing_op,
custom_map)) {
outputs_replaced.insert({result, enumerated_result.index()});
output_types.push_back(result.getType());
} else {
return failure();
}
}
// For float16 quantization if none of the operand or result is
// modified, replacing the op. See b/335025403.
if (inference_type == tensorflow::DT_HALF &&
!is_operand_or_result_modified) {
return failure();
}
rewriter.setInsertionPointAfter(quantizing_op);
OperationState new_state(
quantizing_op->getLoc(), quantizing_op->getName().getStringRef(),
inputs, output_types, quantizing_op->getAttrs());
for (int i = 0; i < quantizing_op->getNumRegions(); ++i) {
new_state.addRegion();
}
quantized_op = rewriter.create(new_state);
if (quantizing_op->getNumRegions() != 0) {
for (const auto& indexed_regions :
llvm::enumerate(quantizing_op->getRegions())) {
Region& target_region =
quantized_op->getRegion(indexed_regions.index());
IRMapping mapping;
indexed_regions.value().cloneInto(&target_region, mapping);
}
}
for (auto output : outputs_replaced) {
output.getFirst().replaceAllUsesWith(
quantized_op->getResult(output.getSecond()));
}
}
// To verify the numericals, the original floating-point ops are
// preserved in the graph. The result of these floating-point ops are sent
// to a numeric verifier op as the reference.
if (enable_verify && !std::is_same_v<VerifierT, void>) {
// For constant operands, the floating-point constant is duplicated in
// case it is quantized.
for (int i = 0, e = quantized_op->getNumOperands(); i < e; ++i) {
auto def = quantized_op->getOperand(i).getDefiningOp();
if (auto q = llvm::dyn_cast_or_null<QuantizeOpT>(def)) {
DenseFPElementsAttr attr;
if (!matchPattern(q.getOperand(), m_Constant(&attr))) {
continue;
}
auto cst = arith::ConstantOp::create(rewriter,
quantized_op->getLoc(), attr);
quantizing_op->setOperand(i, cst.getResult());
}
}
for (int i = 0, e = quantized_op->getNumResults(); i < e; ++i) {
if (!isa<FloatType>(
cast<ShapedType>(quantizing_op->getResult(i).getType())
.getElementType())) {
continue;
}
CreateVerifier<VerifierT>(quantizing_op, quantized_op, rewriter, i,
quant_params_);
if (enable_whole_model_verify) {
RewireFloatModelBackbone(quantized_op, quantizing_op);
}
}
}
}
return success();
}
private:
// Reconnects float ops in the whole-model verify mode. Works for both
// Quantizable ops and Unquantizable ops
void RewireFloatModelBackbone(mlir::Operation* quantized_op,
mlir::Operation* float_op) const {
for (int i = 0, e = quantized_op->getNumResults(); i < e; ++i) {
if (!llvm::cast<ShapedType>(float_op->getResult(i).getType())
.getElementType()
.isF32()) {
continue;
}
// Find the Quantize/Dequantize users of the new op results, and replace
// the usage. Then all the floating-point ops are connected, forming a
// separate float "backbone" model that the quantized model can be
// compared against in parallel.
// N.B. the return op will use this floating-point result.
Value result;
if (!IsOpQuantizable(float_op)) {
// For not quantizable ops, search for dequantize attached to the
// quantized op of the output.
if (mlir::Operation* quantize_op = dyn_cast_or_null<QuantizeOpT>(
*quantized_op->getResult(i).getUsers().begin())) {
result = quantize_op->getResult(0);
} else {
quantized_op->emitError()
<< "Output[" << i
<< "] is expected to have only one user [QUANTIZE]";
return;
}
} else {
result = quantized_op->getResult(i);
}
for (auto user : result.getUsers()) {
// Skip the Requantize op and set the user to the following dequantize
// op. This happens when the quantizer tries to match the scale conflict
// with QuantizeOp - QuantizeOp(requant) - DequantizeOp triples. The
// correct float op should be the user of the last DequantizeOp.
if (llvm::isa<QuantizeOpT>(user)) {
user = *user->getResult(0).getUsers().begin();
}
if (auto dequantize = llvm::dyn_cast<DequantizeOpT>(user)) {
// Replace all uses, except not quantizable ops that are being used in
// the float backbone.
dequantize.getResult().replaceUsesWithIf(
float_op->getResult(i), [&](OpOperand& use) {
return !use.getOwner()->hasAttr(kDebugModeOpQuantAttrName);
});
}
}
}
}
QuantPassSpec quant_params_;
};
// A pattern that removes debug attributes that are annotated to ops during
// the debug model creation.
class RemoveDebugAttrPattern : public RewritePattern {
public:
explicit RemoveDebugAttrPattern(MLIRContext* context)
: RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}
LogicalResult matchAndRewrite(mlir::Operation* op,
PatternRewriter& rewriter) const override;
};
// Converts quantized tensor type with signed integer type to quantized tensor
// type with unsigned integer type.
Type ConvertSignedQuantizedToUnsigned(Type signed_tensor_type, Location loc);
// Converts quantize ops with unsigned quantized types to these with signed
// quantized types and preserves the scales.
template <typename QuantizeOpT>
struct ConvertUnsignedToSigned : public OpRewritePattern<QuantizeOpT> {
using BaseType = ConvertUnsignedToSigned<QuantizeOpT>;
using QType = quant::QuantizedType;
explicit ConvertUnsignedToSigned(MLIRContext* context)
: OpRewritePattern<QuantizeOpT>(context, 1) {}
LogicalResult matchAndRewrite(QuantizeOpT op,
PatternRewriter& rewriter) const override {
Type output_type = op.getResult().getType();
auto qtype = QType::getQuantizedElementType(output_type);
if (!qtype || qtype.isSigned()) return failure();
int num_bits = qtype.getStorageTypeIntegralWidth();
if (num_bits == 8) {
// If storage is 8-bit, trained num bits may be less than 8 so check here.
num_bits =
static_cast<int>(std::ceil(std::log2(qtype.getStorageTypeMax())));
}
// This is a positive value, and will be applied on zero points and fixed
// point ranges.
int64_t offset =
QType::getDefaultMinimumForInteger(/*isSigned=*/false, num_bits) -
QType::getDefaultMinimumForInteger(/*isSigned=*/true, num_bits);
auto flags = quant::QuantizationFlags::Signed;
QType new_qtype;
if (auto uqtype = llvm::dyn_cast<quant::UniformQuantizedType>(qtype)) {
new_qtype = quant::UniformQuantizedType::getChecked(
op.getLoc(), flags, qtype.getStorageType(), qtype.getExpressedType(),
uqtype.getScale(), uqtype.getZeroPoint() - offset,
uqtype.getStorageTypeMin() - offset,
uqtype.getStorageTypeMax() - offset);
} else if (auto aqtype =
llvm::dyn_cast<quant::UniformQuantizedPerAxisType>(qtype)) {
auto zero_points = aqtype.getZeroPoints();
llvm::SmallVector<int64_t, 4> new_zero_points(zero_points.begin(),
zero_points.end());
for (int i = 0, e = new_zero_points.size(); i < e; ++i) {
new_zero_points[i] -= offset;
}
new_qtype = quant::UniformQuantizedPerAxisType::getChecked(
op.getLoc(), flags, qtype.getStorageType(), qtype.getExpressedType(),
aqtype.getScales(), new_zero_points, aqtype.getQuantizedDimension(),
aqtype.getStorageTypeMin() - offset,
aqtype.getStorageTypeMax() - offset);
} else {
return failure();
}
if (!new_qtype) return failure();
Type new_output_type = new_qtype.castFromExpressedType(
QType::castToExpressedType(output_type));
rewriter.replaceOpWithNewOp<QuantizeOpT>(op, new_output_type, op.getArg());
return success();
}
};
// Fold Extra Requantize ops if the preceding ops has free scale requirement.
template <typename RequantizeOpT>
struct FoldTrivalRequantizeOp : public OpRewritePattern<RequantizeOpT> {
explicit FoldTrivalRequantizeOp(MLIRContext* context)
: OpRewritePattern<RequantizeOpT>(context, 1) {}
LogicalResult matchAndRewrite(RequantizeOpT op,
PatternRewriter& rewriter) const override {
Value pre_quantized = op->getOperand(0);
auto pre_quantized_type =
quant::QuantizedType::getQuantizedElementType(pre_quantized.getType());
if (!pre_quantized_type) return failure();
mlir::Operation* def = pre_quantized.getDefiningOp();
if (!def) return failure();
if (llvm::isa<FixedOutputRangeInterface, SameScalesOpInterface>(def) ||
!def->hasTrait<OpTrait::TFL::QuantizableResult>()) {
return failure();
}
// This op should not clobber def, if more than one requant of this value.
if (!pre_quantized.hasOneUse()) {
return failure();
}
op.emitWarning("Remove trivial `rescale` op. Please fix the source graph.");
llvm::SmallVector<Type, 4> new_output_types;
for (auto result : def->getResults()) {
if (result.hasOneUse() && *result.getUsers().begin() == op) {
new_output_types.push_back(op.getResult().getType());
} else {
new_output_types.push_back(result.getType());
}
}
// Remove this rescale op.
rewriter.replaceOp(op, {pre_quantized});
// Replace the output scale of the preceding op.
rewriter.setInsertionPointAfter(def);
OperationState new_state(def->getLoc(), def->getName().getStringRef(),
def->getOperands(), new_output_types,
def->getAttrs());
Operation* new_op = rewriter.create(new_state);
rewriter.replaceOp(def, new_op->getResults());
return success();
}
};
// Given a quantized type `input`, magnifying its scales by the factor stored in
// `factor`. If `input` isn't a quantized type or the `factor` doesn't match the
// dimension size of `input` or isn't floating-point, nullptr will be returned.
TypeAttr RescaleQuantizedType(Type input, Attribute factor);
// Converts the min/max/num_bits/narrow_range information to a
// QuantizedType, and then returns the attribute containing the QuantizedType.
// The `min` and `max` arguments can be FloatAttr or DenseFPElementsAttr and
// returns UniformQuantizedType or UniformQuantizedPerAxisType respectively.
// `narrow_range` is set to true for weights and `is_signed` is set to true
// if it is using signed int symmetric quantization.
//
// Note that this method may broadcast min and max to match the dimension length
// of `input_type`, if the `quant_dim` is valid. On the other hand, the
// symmetry of min and max is not adjusted by this method. The QAT workflow
// should set min/max correctly (and use `narrow_range`=true, `is_signed`=true)
// if symmetric quantization is required.
TypeAttr GetQuantizedTypeAttr(Builder builder, Type input_type, Attribute min,
Attribute max, int quant_dim,
IntegerAttr num_bits, BoolAttr narrow_range,
bool is_signed, bool legacy_float_scale = false,
bool use_fake_quant_num_bits = false);
// Casts the `target` type to a quantized type by using the quantization
// parameters from the type in the `source` type attribute.
// Examples:
// f32 -> !quant.uniform<i8:f32, 1.0>
// tensor<4xf32> -> tensor<4x!quant.uniform<i8:f32, 1.0>>
// The result is wrapped by a type attribute. Returns nullptr if the cast
// isn't valid.
//
// `axis` is to specify the quantization dimension in the `target` and only
// used if the element type of `source` is a per-channel quantized type. During
// the casting, the quantization dimension of the result type needs to be set
// this new `axis` value.
TypeAttr CastQuantizedTypeAttrFromExpressedType(Builder builder,
TypeAttr source, Type target,
int axis);
// Quantizes the elements in the attribute `real_value` by the quantization
// parameters in `tensor_type`. Returns empty Attribute if the
// `tensor_type` is not a QuantizedType or the quantization fails.
ElementsAttr Quantize(Attribute real_value, Type tensor_type);
// Quantizes the elements in "legacy mode", where it calls TOCO's methods to
// to quantize values with float scale.
ElementsAttr QuantizeLegacy(Attribute real_value, Type tensor_type);
// Returns the quantized type for an element attribute. The quantization
// parameters in this type is based on the min and max element of the
// attribute. When the elements in the `attr` are not in floating-point, or
// the value range isn't straddling zero, an empty type is returned. The min/max
// are adjusted to be symmetric if `symmetric` flag is set to True. And
// `symmetric` can only be set to true when it is signed and narrow_range.
Type GetUniformQuantizedTypeForWeight(ElementsAttr attr, bool symmetric,
unsigned num_bits, bool is_signed,
bool narrow_range,
bool legacy_float_scale = false,
bool use_fake_quant_num_bits = false);
// Returns the per channel quantized type for an element attribute.
// `quant_dim` defines the quantization axis. The channel min/max are adjusted
// to be symmetric if `symmetric` flag is set to True. And `symmetric` can only
// be set to true when it is signed and narrow_range.
Type GetUniformQuantizedPerAxisTypeForWeight(
ElementsAttr attr, int quant_dim, bool symmetric, unsigned num_bits,
bool is_signed, bool narrow_range, bool legacy_float_scale = false,
bool use_fake_quant_num_bits = false);
// Returns the quantized type of a bias input, given the quantized types of
// other operands which are multiply-accumulated (the bias is added to the
// accumulated value).
quant::QuantizedType GetUniformQuantizedTypeForBias(
const std::vector<quant::QuantizedType>& op_types, int adjusted_quant_dim,
bool legacy_float_scale = false);
// Gets quantization scale specs (e.g. fixed output range, same result and
// operand scales) from the default quantization interfaces. The op should
// outlive returned spec for its interface methods to be properly referenced.
std::unique_ptr<OpQuantScaleSpec> GetDefaultQuantScaleSpec(Operation* op);
// The function might contain more stats ops than required, and it will
// introduce requantize if the calibration stats have conflicts. This method
// tries to remove all the redundant stats ops.
bool RemoveRedundantStatsOps(mlir::func::FuncOp func,
OpQuantSpecGetter op_quant_spec_getter,
OpQuantScaleSpecGetter op_quant_scale_spec_getter =
GetDefaultQuantScaleSpec);
// Given quantization parameters for int8, compute the quantization parameters
// for uint if it is required, and wrap the result in an UniformQuantizedType.
quant::UniformQuantizedType GetFixedOutputRange(bool is_signed, int bit_width,
Type tensor_type, double scale,
int64_t zero_point,
int64_t storage_min,
int64_t storage_max);
quant::UniformQuantizedType GetFixedOutputRange(bool is_signed, int bit_width,
Type tensor_type, double scale,
int64_t zero_point);
// Extracts min and max values from the DenseFPElementsAttr, and stores them
// into `mins` and `maxs`. When mins and maxs are extracted per-channel,
// `dim_size` is number of channels and `slice_size` is the size of slice per
// each channel. When `symmetric` is true, the range is expanded to [-M, M].
void ExtractMinMaxFromAttr(DenseFPElementsAttr values, int dim_size,
int slice_size, bool symmetric,
SmallVectorImpl<double>& mins,
SmallVectorImpl<double>& maxs);
// Returns the quantized type for the
// input_type/min/max/storage_type_width/narrow_range.
Type GetQuantizedType(Builder builder, Type input_type, ArrayRef<double> min,
ArrayRef<double> max, int quant_dim,
int storage_type_width, bool narrow_range, bool is_signed,
bool legacy_float_scale = false,
bool use_fake_quant_num_bits = false);
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_COMMON_QUANTIZATION_LIB_QUANTIZATION_UTILS_H_
tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/tfl_quantization_driver.cc
+1143
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/tfl_quantization_driver.h
+391
View File
@@ -0,0 +1,391 @@
/* Copyright 2025 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_COMMON_QUANTIZATION_LIB_TFL_QUANTIZATION_DRIVER_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_COMMON_QUANTIZATION_LIB_TFL_QUANTIZATION_DRIVER_H_
#include <memory>
#include <utility>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
#include "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/Quant/IR/QuantTypes.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_config.h"
#include "tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_utils.h"
namespace mlir {
namespace TFL {
// TODO(b/413355305): Remove temp namespace after TFL's 2 quantization_drivers
// are merged.
namespace temp {
// The state for each op result during the quantization parameters propagation.
struct QuantState {
// Quantization parameters propagated to an op result.
QuantizedType params;
// A flag indicates this state (the params) shouldn't be changed after it is
// initialized. This flag will be set to true if the quantization parameters
// are from the quantization-aware training.
const bool immutable;
bool IsEmpty() const { return params == nullptr; }
};
// The state for rescaling the propagated quantization parameters. This can be
// on the input side to satisfy the constraint of previous operation, or on the
// output side to satisfy the constraint of the next operation.
struct RequantizeState {
// Sometimes, we have to "requantize" the quantization result to satisfy all
// the constraints. The "requantize" can happen either on the input or output
// of the quantization result.
enum RequantizePosition {
NO_REQUANTIZE,
ON_INPUT,
ON_OUTPUT
} pos = NO_REQUANTIZE;
// Quantization parameters will be used to add the requantize ops.
QuantizedType params;
// Avoid clobbering all uses of the value, limit to just these ops.
SmallVector<std::pair<Operation*, int>> users;
};
using RequantizeStates = SmallVector<RequantizeState>;
// This is a worklist-driven driver for propagating quantization parameters
// across operations.
//
// The initial quantization parameters are extracted from the quantized type
// between adjacent `quantfork::QuantizeCastOp` and
// `quantfork::DequantizeCastOp`s. All these initial parameters are marked as
// immutable because they are from quantization-aware training.
//
// The algorithm traverses each op and sets the quantization parameters of its
// operands and results, according to its quantization specification, and then
// adds the operands and results to the worklist. If there are any conflicts
// (for example, there are quantization parameters propagated from the previous
// iteration), this process stops if the existing parameters are the immutable,
// or adding `requantize` op to resolve the conflicts.
//
// After the algorithm is converged, pairs of `quantfork::QuantizeCastOp` and
// `quantfork::DequantizeCastOp` are inserted to the right position to
// materialize the propagation and requantize results.
//
class QuantizationDriver {
public:
// Type alias of int used to access `states_`.
using QuantStateIndex = int;
// (op, operand index) pair.
using OpWithOperandIndex = std::pair<Operation*, int>;
// (op, result index) pair.
using OpWithResultIndex = std::pair<Operation*, int>;
explicit QuantizationDriver(func::FuncOp func_op, const bool is_signed,
const int bit_width,
const bool disable_per_channel,
OpQuantSpecGetter op_quant_spec_getter,
OpQuantScaleSpecGetter op_quant_scale_spec_getter,
const bool infer_tensor_range,
const QDQConversionMode qdq_conversion_mode,
const bool legacy_float_scale = false)
: fn_(func_op),
builder_(func_op.getBody()),
is_signed_(is_signed),
bit_width_(bit_width),
disable_per_channel_(disable_per_channel),
op_quant_spec_getter_(op_quant_spec_getter),
op_quant_scale_spec_getter_(op_quant_scale_spec_getter),
infer_tensor_range_(infer_tensor_range),
legacy_float_scale_(legacy_float_scale),
qdq_conversion_mode_(qdq_conversion_mode) {}
// The entry point of the quantization parameters propagation.
void Run();
// Sets up the states for all the op results in the function.
void Initialize();
// Propagates the quantization parameters across all the ops.
bool PropagateParamsAndReturnIfChanged();
// Inserts the Quantize and Dequantize ops according to the propagation
// result.
void Finalize();
SmallVector<BlockArgument, 4> GetArgs() { return args_; }
llvm::DenseMap<std::pair<mlir::Operation*, int>, int> GetResultStates() {
return result_states_;
}
DenseMap<OpWithResultIndex, QuantStateIndex> result_states_;
// Returns the state of the block argument.
QuantState& GetArgQuantState(BlockArgument arg) {
return states_[arg_states_[arg]];
}
// Returns the state of the index-th result of the op.
QuantState& GetResultQuantState(Operation* op, const int index) {
return states_[result_states_[{op, index}]];
}
private:
// Duplicates the constant op if it has multiple uses, and replaces
// target_op->operand[operand_index] with the newly created op. This also
// replaces corresponsing quantization states.
arith::ConstantOp DuplicateConstantOpIfNeeded(arith::ConstantOp op,
Operation* target_op,
int operand_index);
// Adjusts bias scale that is derived from other scales (fc, conv ops) to
// prevent overflow of quantized bias values. This also changes quantization
// state of other inputs when needed.
bool SetBiasParamsWithAdjustments(Operation* op, int bias_index,
ArrayRef<int> input_indices,
QuantizedType params);
// Checks preconditions to adjust bias scale.
bool ShouldCheckBiasScale(Operation* op, int bias_index,
ArrayRef<int> input_indices,
QuantizedType quantized_type, int& input_index,
int& filter_index);
// Preprocesses the constants by doing the following:
// - Duplicates constants if it is used by multiple ops. For example, if a
// constant is used by multiple ops as a bias, duplicate constants and
// let each op assign its own quantization parameter for bias.
// - Adds all the non-bias constants (weights) to a set for looking up
// later.
// - Adds all per-channel weights to a set for looking up later.
void PreprocessConstantOps();
// Sets up all the data structures for quantization propagation.
void SetupAllStates();
// Returns Whether the constant is a weight, which shouldn't be shared by
// different ops.
bool IsWeight(Operation* cst) { return llvm::is_contained(weights_, cst); }
// Returns all the related quantization constraints of the op.
std::unique_ptr<OpQuantSpec> GetQuantSpec(Operation* op);
std::unique_ptr<OpQuantScaleSpec> GetQuantScaleSpec(Operation* op);
// Returns whether quantization parameters have been propagated to the results
// of this op.
bool IsQuantized(Operation* op);
// Adds all the users of index-th result of op to the work list.
void AddUserToList(Operation* op, const int index) {
for (Operation* user : op->getResult(index).getUsers()) {
work_list_.push_back(user);
}
}
// Adds the defining op of index-th operand of op to the work list.
void AddOperandToList(Operation* op, const int index) {
if (Operation* operand_op = op->getOperand(index).getDefiningOp();
operand_op != nullptr) {
work_list_.push_back(operand_op);
}
}
// Returns the quantization params for the bias input from the non-bias
// operands which have their indexes in the `non_biases` vector. The returned
// parameters are calculated by `func`.
QuantizedType GetBiasParams(Operation* op, int bias_index,
ArrayRef<int> non_bias_operand_indices,
AccumulatorScaleFunc func);
// Sets the quantization parameters of the result to `quantized_type`. If
// any quantization parameters have been propagated, a requantize will
// happen on the input of propagated quantization. Returns `true` if internal
// state has been modified.
bool SetResultParams(Operation* op, int result_index,
QuantizedType quantized_type);
// Sets the quantization parameters of the operand to `quantized_type`. If any
// quantization parameters have been propagated, a `requantize` will happen on
// the output of propagated quantization. When `override` is set, quantization
// state of the value is replaced instead of adding requantization. Returns
// `true` if internal state has been modified.
bool SetOperandParams(Operation* op, int operand_index,
QuantizedType quantized_type, bool override = false);
// Sets the quantization parameters of the constant result according to its
// content.
bool SetConstantResultParams(Operation* op);
// Inserts the Quantize and Dequantize ops after `op`'s `index`-th result. The
// quantized element type for the result is `quantized_type`.
void QuantizeOpResult(Operation* op, int result_index,
QuantizedType quantized_type);
// Inserts the Quantize and Dequantize ops after `arg`. The quantized element
// type for `arg` is `quantized_type`.
void QuantizeArg(BlockArgument arg, QuantizedType quantized_type);
// Inserts the Quantize and Dequantize ops (i.e. QDQ) after `value`. The
// quantized element type for `value` is `quantized_type`.
void QuantizeValue(Value value, QuantizedType quantized_type, Location loc);
// Inserts the Quantize ops for requantizing the index-th result of the op.
void RequantizeOpResult(Operation* op, int result_index,
RequantizeStates& states);
// Inserts the Quantize ops for requantizing a block argument.
void RequantizeArg(BlockArgument arg, RequantizeStates& states);
// Inserts the Quantize and Dequantize ops to quantize the value and returns
// the Quantize op.
void RequantizeValue(Value value, RequantizeStates& states, Location loc);
// Returns the quantization parameter satisfies the same scale
// constraints for the op. Returns an empty option if this quantization
// parameter doesn't exist.
QuantizedType GetQuantParamsForSameScaleConstraint(Operation* op);
// Returns the state of the index-th operand of the op.
QuantState& GetOperandQuantState(Operation* op, const int index) {
return states_[operand_states_[{op, index}]];
}
// Returns the states of the index-th operand of the op.
RequantizeStates& GetOperandRequantizeStates(Operation* op, const int index) {
return rescale_states_[operand_states_[{op, index}]];
}
// Returns the states of the index-th result of the op.
RequantizeStates& GetResultRequantizeStates(Operation* op, const int index) {
return rescale_states_[result_states_[{op, index}]];
}
// Returns the states of the arg.
RequantizeStates& GetArgRequantizeStates(BlockArgument arg) {
return rescale_states_[arg_states_[arg]];
}
// Sets the state of an argument. If this value is cached, uses the cached
// result without creating new entry in the state vector. Otherwise, allocate
// a new entry in the state vector.
void InitializeArgState(BlockArgument arg, Value arg_value);
// Sets the state of the index-th operand of the op. If this operand is
// cached, uses the cached result without creating new entry in the state
// vector. Otherwise, allocate a new entry in the state vector.
void InitializeOperandState(Operation* op, int index, Value value);
// Sets the state of the index-th result of the op. If this result is cached,
// uses the cached result without creating new entry in the state vector.
// Otherwise, allocate a new entry in the state vector.
void InitializeResultState(Operation* op, int index, Value value);
func::FuncOp fn_;
OpBuilder builder_;
const bool is_signed_;
const int bit_width_;
const bool disable_per_channel_;
// We should distinguish weights and bias constants. Biases are specified by
// the quantization spec or are the operands of ops with same scale spec. The
// rest are weights.
DenseSet<Operation*> weights_;
// The weights require narrow_range quantization. This map collects all the
// weight operands defined by the op quant spec. The value of each entry is
// the quantization dimension. If it is positive, per-channel quantization is
// required.
DenseMap<Operation*, int> optimized_weights_;
// All the ops needs to propagate the quantization parameters to.
std::vector<Operation*> work_list_;
absl::flat_hash_set<Operation*> quantized_;
// The vector contains all the quantization parameters propagated from the
// defining operations of the value, or from the quantization aware training.
std::vector<QuantState> states_;
// The map contains all the quantization parameters which are required to
// satisfy the same operands and results constraint. The keys of this map are
// the values from `operand_states_` and `result_state_`.
absl::flat_hash_map<QuantStateIndex, RequantizeStates> rescale_states_;
// Maps of indexes to the propagation state vector from the ops operands,
// results and arguments.
DenseMap<OpWithOperandIndex, QuantStateIndex> operand_states_;
DenseMap<BlockArgument, QuantStateIndex> arg_states_;
DenseMap<Value, QuantStateIndex> value_to_state_;
// This vector is to preserve the arguments order, so the newly inserted
// quantized ops for the arguments are deterministically ordered.
SmallVector<BlockArgument, 4> args_;
OpQuantSpecGetter op_quant_spec_getter_;
OpQuantScaleSpecGetter op_quant_scale_spec_getter_;
// Infer output ranges for activation ops and constants. This is usually
// required for post-training quantization.
const bool infer_tensor_range_;
// Calculate scales in float instead of double, so that the scales and
// quantized values are exactly the same with the TOCO quantizer.
const bool legacy_float_scale_;
// The type of qdq conversion.
const QDQConversionMode qdq_conversion_mode_;
};
// Propagates quantization parameters across ops in this function and satisfies
// the quantization specification of the ops. This methods assumes the initial
// quantization parameters are stored as adjacent quantize and dequantize ops
// and the propagation results are materialized by inserting pairs of quantize
// and dequantize ops to this function. Set `disable_per_channel` to true to not
// use per channel quantization even the op supports it.
// Setting `infer_tensor_range` to true, to infer quantization parameters from
// the activation ops and weight constants. This is only used for post-training
// quantization.
void ApplyQuantizationParamsPropagation(func::FuncOp func, bool is_signed,
int bit_width, bool disable_per_channel,
OpQuantSpecGetter op_quant_spec_getter,
bool infer_tensor_ranges,
bool legacy_float_scale,
QDQConversionMode qdq_conversion_mode);
void ApplyQuantizationParamsPropagation(
func::FuncOp func, bool is_signed, int bit_width, bool disable_per_channel,
OpQuantSpecGetter op_quant_spec_getter,
OpQuantScaleSpecGetter op_quant_scale_spec_getter, bool infer_tensor_ranges,
bool legacy_float_scale, QDQConversionMode qdq_conversion_mode);
} // namespace temp
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_COMMON_QUANTIZATION_LIB_TFL_QUANTIZATION_DRIVER_H_
tensorflow/compiler/mlir/lite/quantization/common/test_base.h
+88
View File
@@ -0,0 +1,88 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_COMMON_TEST_BASE_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_COMMON_TEST_BASE_H_
#include <memory>
#include <gtest/gtest.h>
#include "absl/strings/string_view.h"
#include "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/Quant/IR/Quant.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "mlir/Parser/Parser.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "stablehlo/dialect/StablehloOps.h" // from @stablehlo
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
#include "tensorflow/compiler/mlir/quantization/common/func.h"
#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/context.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/core/platform/test.h"
namespace mlir::quant {
using ::testing::Test;
class QuantizationTestBase : public Test {
protected:
QuantizationTestBase()
: ctx_(stablehlo::CreateMlirContextForQuantization()),
builder_(ctx_.get()) {
ctx_->loadDialect<
arith::ArithDialect, mlir::stablehlo::StablehloDialect,
func::FuncDialect, TF::TensorFlowDialect, TFL::TensorFlowLiteDialect,
tf_saved_model::TensorFlowSavedModelDialect,
tf_executor::TensorFlowExecutorDialect, quant::QuantDialect,
quantfork::QuantizationForkDialect, ir::TFQuantDialect>();
}
// Parses `module_op_str` to create a `ModuleOp`.
OwningOpRef<ModuleOp> ParseModuleOpString(
const absl::string_view module_op_str) {
return parseSourceString<ModuleOp>(module_op_str, ctx_.get());
}
// Convenience function that returns the first operation of type `OpT` from
// the `@main` function in `module_op`. Useful when testing with a text
// representation of a `ModuleOp` containing a single function `@main`.
// Returns `failure` iff there is no `@main` or no such operation is found in
// `@main`.
template <typename OpT>
FailureOr<OpT> FindFirstOpFromMainFunc(ModuleOp module_op) {
func::FuncOp main_func_op = FindMainFuncOp(module_op);
if (main_func_op == nullptr) return failure();
auto ops = main_func_op.getOps<OpT>();
if (ops.empty()) return failure();
return *ops.begin();
}
std::unique_ptr<MLIRContext> ctx_;
OpBuilder builder_;
};
} // namespace mlir::quant
#endif // TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_COMMON_TEST_BASE_H_