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/BUILD
+136
View File
@@ -0,0 +1,136 @@
load("//tensorflow:tensorflow.bzl", "tf_cc_test", "tf_native_cc_binary")
load("//tensorflow/core/platform:build_config.bzl", "tf_proto_library")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
# copybara:uncomment_begin(google-only)
# default_visibility = [
# ":friends",
# "//tensorflow:__pkg__",
# ],
# copybara:uncomment_end_and_comment_begin
default_visibility = [
"//visibility:public",
],
# copybara:comment_end
licenses = ["notice"],
)
package_group(
name = "friends",
packages = [
"//learning/brain/mlir/quantization/...",
"//tensorflow/compiler/mlir/...",
"//tensorflow/lite/tools/optimize/...",
],
)
tf_proto_library(
name = "quantization_info_proto",
srcs = [
"quantization_info.proto",
],
)
cc_library(
name = "quantization_passes",
srcs = [
"import_quant_stats_pass.cc",
],
hdrs = [
"quantization_passes.h",
],
deps = [
":quantization_info_proto_cc",
"//tensorflow/compiler/mlir/lite/quantization/ir:QuantOps",
"//tensorflow/compiler/mlir/tensorflow:import_utils",
"//tensorflow/compiler/mlir/tensorflow:location_utils",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:QuantOps",
"@llvm-project//mlir:Support",
],
alwayslink = 1,
)
tf_native_cc_binary(
name = "op_quant_spec_getters_gen",
srcs = [
"tools/op_quant_spec_getters_gen.cc",
],
deps = [
"@llvm-project//llvm:Support",
"@llvm-project//llvm:TableGen",
"@llvm-project//mlir:TableGen",
],
)
tf_native_cc_binary(
name = "tflite_op_coverage_spec_getters_gen",
srcs = [
"tools/tflite_op_coverage_spec_getters_gen.cc",
],
deps = [
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//llvm:TableGen",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:TableGen",
"@tsl//tsl/platform:regexp",
"@xla//xla/tsl/platform:logging",
],
)
cc_library(
name = "numerical_utils",
srcs = ["numerical_utils.cc"],
hdrs = ["numerical_utils.h"],
deps = [
"@com_google_absl//absl/types:optional",
],
)
cc_library(
name = "device_target",
srcs = ["device_target.cc"],
hdrs = ["device_target.h"],
deps = [
":numerical_utils",
"//tensorflow/compiler/mlir/lite/quantization/ir:QuantOps",
"@com_google_absl//absl/types:optional",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:QuantOps",
"@llvm-project//mlir:Support",
],
)
cc_library(
name = "quantization_context",
srcs = ["quantization_context.cc"],
hdrs = ["quantization_context.h"],
deps = [
":device_target",
"//tensorflow/compiler/mlir/lite/quantization/common/quantization_lib",
"//tensorflow/compiler/mlir/lite/quantization/ir:QuantOps",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:QuantOps",
"@llvm-project//mlir:Support",
],
)
tf_cc_test(
name = "numerical_utils_test",
srcs = ["numerical_utils_test.cc"],
deps = [
":numerical_utils",
"@com_google_absl//absl/types:optional",
"@com_google_googletest//:gtest_main",
],
)
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_
tensorflow/compiler/mlir/lite/quantization/device_target.cc
+184
View File
@@ -0,0 +1,184 @@
/* Copyright 2020 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/device_target.h"
#include <cassert>
#include <cmath>
#include <functional>
#include <optional>
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"
#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/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
#include "tensorflow/compiler/mlir/lite/quantization/numerical_utils.h"
namespace mlir {
namespace quant {
constexpr int k8Bits = 8;
constexpr int k32Bits = 32;
constexpr unsigned kSigned = QuantizationFlags::Signed;
DeviceTarget::DeviceTarget(MLIRContext* ctx) : ctx_(ctx) {
f32_ = Float32Type::get(ctx_);
i8_ = IntegerType::get(ctx_, k8Bits);
i8_min_ = QuantizedType::getDefaultMinimumForInteger(kSigned, k8Bits);
i8_max_ = QuantizedType::getDefaultMaximumForInteger(kSigned, k8Bits);
i32_ = IntegerType::get(ctx_, k32Bits);
i32_min_ = QuantizedType::getDefaultMinimumForInteger(kSigned, k32Bits);
i32_max_ = QuantizedType::getDefaultMaximumForInteger(kSigned, k32Bits);
any_ = AnyQuantizedType();
qi8_ = AnyQuantizedType::get(kSigned, i8_, f32_, i8_min_, i8_max_);
qi8n_ = AnyQuantizedType::get(kSigned, i8_, f32_, i8_min_ + 1, i8_max_);
qi32_ = AnyQuantizedType::get(kSigned, i32_, f32_, i32_min_, i32_max_);
assert(qi8n_ == qi8n_);
}
std::optional<KernelSpec> DeviceTarget::GetKernelSpec(
llvm::StringRef kernel, const KernelSpecs::Signature& signature) const {
auto kernel_specs_it = specs_.find(kernel);
if (kernel_specs_it == specs_.end()) return std::nullopt;
return kernel_specs_it->getValue().Find(signature);
}
ScaleDecomposeFn DeviceTarget::GetDecomposeFn(
quantfork::QuantizeRegionOp op) const {
auto kernel_specs_it = specs_.find(op.getLogicalKernel());
if (kernel_specs_it == specs_.end()) return ScaleDecomposeFn(nullptr);
return kernel_specs_it->second.GetDecomposeFn();
}
void DeviceTarget::AppendToSignature(Type spec,
KernelSpecs::Signature* signature) {
if (auto quant = llvm::dyn_cast_or_null<UniformQuantizedType>(spec)) {
signature->push_back(AnyQuantizedType::get(
quant.getFlags(), quant.getStorageType(), quant.getExpressedType(),
quant.getStorageTypeMin(), quant.getStorageTypeMax()));
} else if (auto any = llvm::dyn_cast_or_null<AnyQuantizedType>(spec)) {
signature->push_back(any);
} else { // float
signature->push_back(AnyQuantizedType());
}
}
LogicalResult DeviceTarget::RegisterKernel(
llvm::StringRef kernel, const KernelSpecs::Signature& signature,
const ScaleFn& fn, const ScaleDecomposeFn& dfn) {
return specs_[kernel].Add(signature, {ScaleConstraintType::CustomScale, fn});
}
namespace ph = std::placeholders;
LogicalResult DeviceTarget::RegisterKernel(
llvm::StringRef kernel, const KernelSpecs::Signature& signature,
const ScaleConstraintType constraint) {
if (failed(specs_[kernel].Add(signature, {constraint, {}}))) return failure();
switch (constraint) {
case ScaleConstraintType::OutputInputSameScale:
specs_[kernel].WithImpl(std::bind(&DeviceTarget::DecomposeSameScale,
ph::_1, ph::_2, ph::_3, ph::_4));
return success();
default:
return failure();
}
}
LogicalResult DeviceTarget::DecomposeMultiplyAccumulateScale(
Operation* op, QuantizedMultipliers* input_multipliers,
QuantizedMultipliers* output_multipliers, QuantizedRanges* output_ranges) {
auto rop = llvm::dyn_cast<quantfork::QuantizeRegionOp>(op);
if (!rop) return failure();
llvm::SmallVector<Type, 4> input_specs, out_specs;
for (auto spec : rop.getInputSpecs()) {
input_specs.push_back(llvm::cast<TypeAttr>(spec).getValue());
}
for (auto spec : rop.getOutputSpecs()) {
out_specs.push_back(llvm::cast<TypeAttr>(spec).getValue());
}
auto in_spec = llvm::dyn_cast<UniformQuantizedType>(input_specs[0]);
// TODO(fengliuai): handles the PerAxis QuantizedType.
auto w_spec = llvm::dyn_cast<UniformQuantizedType>(input_specs[1]);
auto b_spec = llvm::dyn_cast<UniformQuantizedType>(input_specs[2]);
auto o_spec = llvm::dyn_cast<UniformQuantizedType>(out_specs[0]);
if (!in_spec || !w_spec || !b_spec || !o_spec) return failure();
double scale_product = in_spec.getScale() * w_spec.getScale();
if (fabs(scale_product - b_spec.getScale()) >= 1e-6) return failure();
// input multipliers
input_multipliers->append(3, kUnitQuantizedMultiplier);
// output multipliers
double real_multiplier = scale_product / o_spec.getScale();
output_multipliers->push_back(QuantizeMultiplier(real_multiplier));
// output ranges
auto min = rop->getAttrOfType<FloatAttr>("min");
auto max = rop->getAttrOfType<FloatAttr>("max");
output_ranges->push_back(CalculateQuantizedRange(
o_spec.getScale(), o_spec.getZeroPoint(),
(min ? std::optional<double>(min.getValueAsDouble()) : std::nullopt),
(max ? std::optional<double>(max.getValueAsDouble()) : std::nullopt),
o_spec.getStorageTypeMin(), o_spec.getStorageTypeMax()));
return success();
}
LogicalResult DeviceTarget::DecomposeSameScale(
Operation* op, QuantizedMultipliers* input_multipliers,
QuantizedMultipliers* output_multipliers, QuantizedRanges* output_ranges) {
auto rop = llvm::dyn_cast<quantfork::QuantizeRegionOp>(op);
if (!rop) return failure();
// input multipliers
for (int i = 0; i < op->getNumOperands(); ++i) {
input_multipliers->push_back(kUnitQuantizedMultiplier);
}
// output multipliers
for (int i = 0; i < op->getNumResults(); ++i) {
output_multipliers->push_back(kUnitQuantizedMultiplier);
}
auto o_spec = llvm::dyn_cast<UniformQuantizedType>(
llvm::cast<TypeAttr>(rop.getOutputSpecs()[0]).getValue());
if (!o_spec) return failure();
// output ranges
auto min = rop->getAttrOfType<FloatAttr>("min");
auto max = rop->getAttrOfType<FloatAttr>("max");
output_ranges->push_back(CalculateQuantizedRange(
o_spec.getScale(), o_spec.getZeroPoint(),
(min ? std::optional<double>(min.getValueAsDouble()) : std::nullopt),
(max ? std::optional<double>(max.getValueAsDouble()) : std::nullopt),
o_spec.getStorageTypeMin(), o_spec.getStorageTypeMax()));
return success();
}
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/lite/quantization/device_target.h
+193
View File
@@ -0,0 +1,193 @@
/* Copyright 2020 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_DEVICE_TARGET_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_DEVICE_TARGET_H_
#include <cstdint>
#include <functional>
#include <optional>
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/Hashing.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringMap.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/ErrorHandling.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/Dialect/Quant/IR/QuantTypes.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
#include "tensorflow/compiler/mlir/lite/quantization/numerical_utils.h"
namespace mlir {
namespace quant {
class QuantizeContext;
using AdjacentOperations = llvm::SmallVectorImpl<Operation*>;
using QuantizedMultipliers = llvm::SmallVector<QuantizedMultiplier, 4>;
using QuantizedRanges = llvm::SmallVector<QuantizedRange, 4>;
using ScaleFn = std::function<LogicalResult(QuantizeContext*, Operation*,
AdjacentOperations*, bool*)>;
using ScaleDecomposeFn =
std::function<LogicalResult(Operation*, QuantizedMultipliers*,
QuantizedMultipliers*, QuantizedRanges*)>;
static const QuantizedMultiplier kUnitQuantizedMultiplier{1, 0};
enum class ScaleConstraintType {
OutputInputSameScale,
OutputInputFreeScale,
CustomScale,
};
// Each kernel signature has its own specification for scales.
struct KernelSpec {
// Scale constraint
ScaleConstraintType type;
// Custom function to derive the scales. Only available when the scale
// constraint is `CustomScale`.
ScaleFn scale_fn;
};
class KernelSpecs {
public:
using Signature = llvm::SmallVector<quant::AnyQuantizedType, 4>;
// Returns the kernel specification for the kernel signature.
std::optional<KernelSpec> Find(const Signature& signature) const {
auto spec_it = all_signatures_.find(signature);
if (spec_it != all_signatures_.end()) {
return spec_it->second;
} else {
return std::nullopt;
}
}
ScaleDecomposeFn GetDecomposeFn() const { return decompose_fn_; }
// Adds the kernel signature with the kernel specification.
LogicalResult Add(const Signature& signature, const KernelSpec& spec) {
if (all_signatures_.insert({signature, spec}).second) return success();
return failure();
}
KernelSpecs& WithSignature(const KernelSpecs::Signature& signature,
const ScaleFn& fn) {
(void)Add(signature, {ScaleConstraintType::CustomScale, fn});
return *this;
}
KernelSpecs& WithImpl(const ScaleDecomposeFn& dfn) {
decompose_fn_ = dfn;
return *this;
}
private:
// The signature is pattern match based.
struct SignatureInfo : public llvm::DenseMapInfo<Signature> {
static Signature getEmptyKey() { return {}; }
static unsigned getHashValue(Signature val) {
return llvm::hash_combine_range(val.begin(), val.end());
}
static bool isEqual(Signature LHS, Signature RHS) {
if (LHS.size() != RHS.size()) return false;
for (auto arg : llvm::zip(LHS, RHS)) {
if (std::get<0>(arg) != std::get<1>(arg)) return false;
}
return true;
}
};
// Maps the signature to the kernel spec. Note that the matching is
// pattern match based.
llvm::DenseMap<Signature, KernelSpec, SignatureInfo> all_signatures_;
// A method to compute the effective multipliers. This is independent on the
// bits of the ports, thus all the signature shares the same here.
ScaleDecomposeFn decompose_fn_;
};
class DeviceTarget {
public:
explicit DeviceTarget(MLIRContext* ctx);
// Retrieves the kernel spec for the quant region op.
std::optional<KernelSpec> GetKernelSpec(
llvm::StringRef kernel, const KernelSpecs::Signature& signature) const;
// Retrieves the scale decomposition function for the quant region op.
ScaleDecomposeFn GetDecomposeFn(quantfork::QuantizeRegionOp op) const;
// converts specification to signature:
// - UniformedQuantizedType -> AnyQuantizedType
// - AnyQuantizedType (int) -> AnyQuantizedType
// - Float -> {}
static void AppendToSignature(Type spec, KernelSpecs::Signature* signature);
protected:
// Adds the kernel spec with the custom scale function for the kernel.
LogicalResult RegisterKernel(llvm::StringRef kernel,
const KernelSpecs::Signature& signature,
const ScaleFn& fn, const ScaleDecomposeFn& dfn);
// Adds the kernel spec with the scale constraint type for the kernel.
LogicalResult RegisterKernel(llvm::StringRef kernel,
const KernelSpecs::Signature& signature,
ScaleConstraintType constraint);
// Adds the kernel with the name. Retrun an existing one if it has been
// added before.
KernelSpecs& RegisterKernel(llvm::StringRef kernel) { return specs_[kernel]; }
// For "mulmat->add" type of kernels, convert the scales of all the ports to
// multipliers.
static LogicalResult DecomposeMultiplyAccumulateScale(
Operation* op, QuantizedMultipliers* input_multipliers,
QuantizedMultipliers* output_multipliers, QuantizedRanges* output_ranges);
// For "reshape" type of kernels.
static LogicalResult DecomposeSameScale(
Operation* op, QuantizedMultipliers* input_multipliers,
QuantizedMultipliers* output_multipliers, QuantizedRanges* output_ranges);
// A set of parameters are required to build the signatures.
FloatType f32_;
IntegerType i8_, i32_;
int64_t i8_min_, i8_max_, i32_min_, i32_max_;
quant::AnyQuantizedType any_, qi8_, qi8n_, qi32_;
private:
// Maps the kernel names to all the available kernels.
llvm::StringMap<KernelSpecs> specs_;
// Points to the global MLIRContext.
MLIRContext* ctx_;
};
} // namespace quant
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_DEVICE_TARGET_H_
tensorflow/compiler/mlir/lite/quantization/import_quant_stats_pass.cc
+260
View File
@@ -0,0 +1,260 @@
/* 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 <cstdint>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/strings/str_split.h"
#include "llvm/ADT/APFloat.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringMap.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Regex.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/Quant/IR/Quant.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributeInterfaces.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/DialectRegistry.h" // from @llvm-project
#include "mlir/IR/Location.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/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassRegistry.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
#include "tensorflow/compiler/mlir/lite/quantization/quantization_info.pb.h"
#include "tensorflow/compiler/mlir/lite/quantization/quantization_passes.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/import_utils.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/location_utils.h"
// NOLINTNEXTLINE
static llvm::cl::opt<std::string> quantize_stats(
"quant-test-stats", llvm::cl::value_desc("string"),
llvm::cl::desc("serialized quant info string. Only used in tests"),
llvm::cl::init(""));
//===----------------------------------------------------------------------===//
// The Pass to import quantization stats to the ops in a function. This requires
// a custom method to retrieve the unique name of the operation.
namespace mlir {
namespace quant {
using QuantParamsEntry = QuantizationInfo::QuantParams;
namespace {
class ImportQuantStatsPass
: public PassWrapper<ImportQuantStatsPass, OperationPass<func::FuncOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ImportQuantStatsPass)
explicit ImportQuantStatsPass(OperationToName op_to_name)
: op_to_name_(op_to_name) {}
StringRef getArgument() const final {
// This is the argument used to refer to the pass in
// the textual format (on the commandline for example).
return "quant-import-stats";
}
StringRef getDescription() const final {
// This is a brief description of the pass.
return "Import quantization stats to the model";
}
void runOnOperation() override;
void getDependentDialects(DialectRegistry &registry) const override {
registry.insert<quant::QuantDialect,
quantfork::QuantizationForkDialect>();
}
// Parses the serialized quant stats protobuf and initialize the internal
// data structure. This method must be called after the pass is created.
bool ParseQuantStats(const std::string &stats_str);
private:
void ImportAsStatsOps(OpBuilder b, Operation *op, int index,
const QuantParamsEntry &info);
void InsertStatsOpAtResult(OpBuilder b, Value res, ElementsAttr layer_stats,
ElementsAttr axis_stats, IntegerAttr axis);
// If the index is out of range, this method returns false. Otherwise it
// returns true if the value is a float tensor.
bool IsQuantizableResult(Operation *op, int index) {
if (index < 0 || index >= static_cast<int>(op->getNumResults()))
return false;
Value res = op->getResult(index);
return isa<ShapedType>(res.getType()) &&
isa<FloatType>(cast<ShapedType>(res.getType()).getElementType());
}
// A method to retrieve the name for the given op.
OperationToName op_to_name_;
// We split the normal names and regex names, since the former can use hash
// map to lookup and the latter needs to iterate all the regex to find the
// match.
// The `int` in the following two containers are to specify the result index
// of the given op. -1 indicates all the floating-point results.
llvm::StringMap<std::pair<int, const QuantParamsEntry>> name_to_info_;
llvm::StringMap<std::pair<int, const QuantParamsEntry>> regex_to_info_;
};
} // namespace
bool ImportQuantStatsPass::ParseQuantStats(const std::string &stats_str) {
QuantizationInfo quant_stats;
if (!tensorflow::LoadProtoFromBuffer(stats_str, &quant_stats).ok()) {
return true;
}
for (const auto &entry : quant_stats.entries()) {
if (!entry.name().empty()) {
std::vector<std::string> name_and_port =
absl::StrSplit(entry.name(), ':');
int port = name_and_port.size() == 2 ? std::stoi(name_and_port[1]) : -1;
name_to_info_.insert({name_and_port[0], {port, entry}});
} else if (!entry.name_regex().empty()) {
std::vector<std::string> name_and_port =
absl::StrSplit(entry.name_regex(), ':');
int port = name_and_port.size() == 2 ? std::stoi(name_and_port[1]) : -1;
regex_to_info_.insert({name_and_port[0], {port, entry}});
}
}
return false;
}
void ImportQuantStatsPass::InsertStatsOpAtResult(OpBuilder b, Value res,
ElementsAttr layer_stats,
ElementsAttr axis_stats,
IntegerAttr axis) {
auto stats_op = quantfork::StatisticsOp::create(
b, b.getUnknownLoc(), res, layer_stats, axis_stats, axis);
res.replaceAllUsesWith(stats_op);
stats_op.getOperation()->replaceUsesOfWith(stats_op, res);
}
void ImportQuantStatsPass::ImportAsStatsOps(OpBuilder b, Operation *op,
int index,
const QuantParamsEntry &info) {
if (info.params_size() == 0) return;
SmallVector<APFloat, 4> min_maxs;
min_maxs.reserve(info.params_size() * 2);
for (const auto &param : info.params()) {
llvm::APFloat min(param.min_max().min());
llvm::APFloat max(param.min_max().max());
min_maxs.push_back(min);
min_maxs.push_back(max);
}
// The layer stats contain only the first min/max pairs.
ElementsAttr layer_stats = DenseFPElementsAttr::get(
RankedTensorType::get({2}, b.getF32Type()), {min_maxs[0], min_maxs[1]});
ElementsAttr axis_stats;
IntegerAttr axis;
if (info.params_size() > 1) {
SmallVector<int64_t, 4> axis_stats_shape{info.params_size(), 2};
axis_stats = DenseFPElementsAttr::get(
RankedTensorType::get(axis_stats_shape, b.getF32Type()), min_maxs);
axis = b.getI64IntegerAttr(info.meta().quantize_axis());
}
b.setInsertionPointAfter(op);
if (IsQuantizableResult(op, index)) {
InsertStatsOpAtResult(b, op->getResult(index), layer_stats, axis_stats,
axis);
} else {
for (int i = 0, e = op->getNumResults(); i < e; ++i) {
if (IsQuantizableResult(op, i)) {
InsertStatsOpAtResult(b, op->getResult(i), layer_stats, axis_stats,
axis);
}
}
}
}
void ImportQuantStatsPass::runOnOperation() {
func::FuncOp func = getOperation();
OpBuilder builder(func);
func.walk([&](Operation *op) {
if (op->hasTrait<OpTrait::IsTerminator>()) return;
auto op_name = op_to_name_(op);
// Check the named info collection first.
auto it = name_to_info_.find(op_name);
if (it != name_to_info_.end()) {
ImportAsStatsOps(builder, op, it->second.first, it->second.second);
return;
}
// Iterate all the regex names and matches the first one.
for (auto &regex : regex_to_info_) {
if (llvm::Regex(regex.first()).match(op_name)) {
ImportAsStatsOps(builder, op, regex.second.first, regex.second.second);
break;
}
}
});
}
// Creates an instance of the default quant parameters pass.
std::unique_ptr<OperationPass<func::FuncOp>> CreateImportQuantStatsPass(
OperationToName op_to_name, const std::string &stats_str) {
auto pass = std::make_unique<ImportQuantStatsPass>(op_to_name);
if (pass->ParseQuantStats(stats_str)) return nullptr;
return pass;
}
// Creates an instance pass to import quantization stats to the operations in
// the function. A custom method to get the name from the op is used because
// different dialect ops might have different ways to assign the name.
std::unique_ptr<OperationPass<func::FuncOp>>
CreateImportQuantStatsPassForTFControlDialect(const std::string &stats_str) {
auto get_name_func = [](Operation *op) {
Location loc = tensorflow::GetLocationWithoutOpType(op->getLoc());
if (auto name = llvm::dyn_cast<NameLoc>(loc)) {
return name.getName().strref();
} else if (auto fused_name = llvm::dyn_cast<FusedLoc>(loc)) {
for (auto sub_loc : fused_name.getLocations()) {
if (auto named_sub_loc = llvm::dyn_cast<NameLoc>(sub_loc)) {
return named_sub_loc.getName().strref();
}
}
}
return llvm::StringRef("");
};
return CreateImportQuantStatsPass(get_name_func, stats_str);
}
// Registers this pass with default values, only for test
static PassRegistration<ImportQuantStatsPass> pass([] {
return CreateImportQuantStatsPassForTFControlDialect(quantize_stats);
});
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/lite/quantization/ir/.clang-format
+2
View File
@@ -0,0 +1,2 @@
BasedOnStyle: LLVM
AlwaysBreakTemplateDeclarations: Yes
tensorflow/compiler/mlir/lite/quantization/ir/.clang-tidy
+76
View File
@@ -0,0 +1,76 @@
Checks: >
# LINT.IfChange(checks)
-*,clang-diagnostic-*,llvm-*,misc-*,
-misc-unused-parameters,
-misc-non-private-member-variables-in-classes,
-misc-no-recursion,
bugprone-argument-comment,
bugprone-assert-side-effect,
bugprone-branch-clone,
bugprone-copy-constructor-init,
bugprone-dangling-handle,
bugprone-dynamic-static-initializers,
bugprone-macro-parentheses,
bugprone-macro-repeated-side-effects,
bugprone-misplaced-widening-cast,
bugprone-move-forwarding-reference,
bugprone-multiple-statement-macro,
bugprone-suspicious-semicolon,
bugprone-swapped-arguments,
bugprone-terminating-continue,
bugprone-unused-raii,
bugprone-unused-return-value,
misc-redundant-expression,
misc-static-assert,
misc-unused-parameters,
misc-unused-using-decls,
modernize-use-bool-literals,
modernize-loop-convert,
modernize-make-unique,
modernize-raw-string-literal,
modernize-use-equals-default,
modernize-use-default-member-init,
modernize-use-emplace,
modernize-use-nullptr,
modernize-use-override,
modernize-use-using,
performance-for-range-copy,
performance-implicit-conversion-in-loop,
performance-inefficient-algorithm,
performance-inefficient-vector-operation,
performance-move-const-arg,
performance-no-automatic-move,
performance-trivially-destructible,
performance-unnecessary-copy-initialization,
performance-unnecessary-value-param,
readability-avoid-const-params-in-decls,
readability-const-return-type,
readability-container-size-empty,
readability-identifier-naming,
readability-inconsistent-declaration-parameter-name,
readability-misleading-indentation,
readability-redundant-control-flow,
readability-simplify-boolean-expr,
readability-simplify-subscript-expr,
readability-use-anyofallof
# LINT.ThenChange(METADATA:checks)
CheckOptions:
# LINT.IfChange(check-options)
- key: readability-identifier-naming.ClassCase
value: CamelCase
- key: readability-identifier-naming.EnumCase
value: CamelCase
- key: readability-identifier-naming.FunctionCase
value: camelBack
- key: readability-identifier-naming.MemberCase
value: camelBack
- key: readability-identifier-naming.ParameterCase
value: camelBack
- key: readability-identifier-naming.UnionCase
value: CamelCase
- key: readability-identifier-naming.VariableCase
value: camelBack
- key: readability-identifier-naming.IgnoreMainLikeFunctions
value: 1
# LINT.ThenChange(METADATA:check-options)
tensorflow/compiler/mlir/lite/quantization/ir/BUILD
+88
View File
@@ -0,0 +1,88 @@
load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library")
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"],
default_visibility = ["//visibility:public"],
licenses = ["notice"],
)
td_library(
name = "QuantizationOpsTdFiles",
srcs = [
"QuantOps.td",
"QuantOpsBase.td",
],
compatible_with = get_compatible_with_portable(),
deps = [
"@llvm-project//mlir:InferTypeOpInterfaceTdFiles",
"@llvm-project//mlir:OpBaseTdFiles",
"@llvm-project//mlir:QuantizationOpsTdFiles",
"@llvm-project//mlir:SideEffectInterfacesTdFiles",
],
)
gentbl_cc_library(
name = "QuantOpsIncGen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {
"QuantOps.h.inc": ["-gen-op-decls"],
"QuantOps.cc.inc": ["-gen-op-defs"],
"QuantOpsDialect.h.inc": [
"-gen-dialect-decls",
"-dialect=quantfork",
],
"QuantOpsDialect.cc.inc": [
"-gen-dialect-defs",
"-dialect=quantfork",
],
},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "QuantOps.td",
deps = [":QuantizationOpsTdFiles"],
)
gentbl_cc_library(
name = "QuantPassIncGen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {"Passes.h.inc": [
"-gen-pass-decls",
"-name=quantfork",
]},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "Passes.td",
deps = ["@llvm-project//mlir:PassBaseTdFiles"],
)
cc_library(
name = "QuantOps",
srcs = [
"ConvertConst.cc",
"ConvertSimQuant.cc",
"QuantOps.cc",
"QuantizeUtils.cc",
],
hdrs = [
"Passes.h",
"QuantOps.h",
"QuantizeUtils.h",
],
compatible_with = get_compatible_with_portable(),
deps = [
":QuantOpsIncGen",
":QuantPassIncGen",
"//tensorflow/compiler/mlir/quantization/common/ir:QuantOps",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:BytecodeOpInterface",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:InferTypeOpInterface",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:QuantOps",
"@llvm-project//mlir:SideEffectInterfaces",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
],
)
tensorflow/compiler/mlir/lite/quantization/ir/ConvertConst.cc
+122
View File
@@ -0,0 +1,122 @@
/* Copyright 2022 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 <memory>
#include <utility>
#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/BuiltinAttributeInterfaces.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/Matchers.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/quantization/ir/Passes.h"
#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantizeUtils.h"
using namespace mlir;
using namespace mlir::quantfork;
using mlir::quant::QuantizedType;
namespace {
#define GEN_PASS_DEF_QUANTCONVERTCONST
#include "tensorflow/compiler/mlir/lite/quantization/ir/Passes.h.inc"
struct ConvertConstPass : public impl::QuantConvertConstBase<ConvertConstPass> {
void runOnOperation() override;
};
struct QuantizedConstRewrite : public OpRewritePattern<QuantizeCastOp> {
using OpRewritePattern<QuantizeCastOp>::OpRewritePattern;
LogicalResult matchAndRewrite(QuantizeCastOp qbarrier,
PatternRewriter &rewriter) const override;
};
} // namespace
/// Matches a [constant] -> [qbarrier] where the qbarrier results type is
/// quantized and the operand type is quantizable.
LogicalResult QuantizedConstRewrite::matchAndRewrite(
QuantizeCastOp qbarrier, PatternRewriter &rewriter) const {
Attribute value;
// Is the operand a constant?
if (!matchPattern(qbarrier.getArg(), m_Constant(&value))) {
return failure();
}
// Does the qbarrier convert to a quantized type. This will not be true
// if a quantized type has not yet been chosen or if the cast to an equivalent
// storage type is not supported.
Type qbarrierResultType = qbarrier.getResult().getType();
QuantizedType quantizedElementType =
QuantizedType::getQuantizedElementType(qbarrierResultType);
if (!quantizedElementType) {
return failure();
}
if (!QuantizedType::castToStorageType(qbarrierResultType)) {
return failure();
}
// Is the operand type compatible with the expressed type of the quantized
// type? This will not be true if the qbarrier is superfluous (converts
// from and to a quantized type).
if (!quantizedElementType.isCompatibleExpressedType(
qbarrier.getArg().getType())) {
return failure();
}
// Is the constant value a type expressed in a way that we support?
if (!mlir::isa<FloatAttr, DenseElementsAttr, SparseElementsAttr>(value)) {
return failure();
}
Type newConstValueType;
auto newConstValue =
quantizeAttr(value, quantizedElementType, newConstValueType);
if (!newConstValue) {
return failure();
}
// When creating the new const op, use a fused location that combines the
// original const and the qbarrier that led to the quantization.
auto fusedLoc = rewriter.getFusedLoc(
{qbarrier.getArg().getDefiningOp()->getLoc(), qbarrier.getLoc()});
auto newConstOp = arith::ConstantOp::create(
rewriter, fusedLoc, newConstValueType, cast<TypedAttr>(newConstValue));
rewriter.replaceOpWithNewOp<StorageCastOp>(qbarrier, qbarrier.getType(),
newConstOp);
return success();
}
void ConvertConstPass::runOnOperation() {
RewritePatternSet patterns(&getContext());
auto func = getOperation();
auto *context = &getContext();
patterns.add<QuantizedConstRewrite>(context);
(void)applyPatternsGreedily(func, std::move(patterns));
}
std::unique_ptr<OperationPass<func::FuncOp>>
mlir::quantfork::createConvertConstPass() {
return std::make_unique<ConvertConstPass>();
}
tensorflow/compiler/mlir/lite/quantization/ir/ConvertSimQuant.cc
+158
View File
@@ -0,0 +1,158 @@
/* Copyright 2022 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 <cassert>
#include <memory>
#include <utility>
#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/MLIRContext.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/quantization/ir/Passes.h"
#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
#include "tensorflow/compiler/mlir/quantization/common/ir/FakeQuantSupport.h"
#include "tensorflow/compiler/mlir/quantization/common/ir/UniformSupport.h"
namespace mlir::quantfork {
#define GEN_PASS_DEF_QUANTCONVERTSIMULATEDQUANT
#include "tensorflow/compiler/mlir/lite/quantization/ir/Passes.h.inc"
struct ConvertSimulatedQuantPass
: public impl::QuantConvertSimulatedQuantBase<ConvertSimulatedQuantPass> {
void runOnOperation() override;
};
/// Base class rewrites ConstFakeQuant into a qbarrier/dbarrier pair.
template <typename ConcreteRewriteClass, typename FakeQuantOp>
class FakeQuantRewrite : public OpRewritePattern<FakeQuantOp> {
public:
using OpRewritePattern<FakeQuantOp>::OpRewritePattern;
FakeQuantRewrite(MLIRContext *ctx, bool *hadFailure)
: OpRewritePattern<FakeQuantOp>(ctx), hadFailure(hadFailure) {}
LogicalResult matchAndRewrite(FakeQuantOp op,
PatternRewriter &rewriter) const override {
// TODO: If this pattern comes up more frequently, consider adding core
// support for failable rewrites.
if (failableRewrite(op, rewriter)) {
*hadFailure = true;
return failure();
}
return success();
}
private:
bool *hadFailure;
bool failableRewrite(FakeQuantOp op, PatternRewriter &rewriter) const {
auto converter =
mlir::quant::ir::ExpressedToQuantizedConverter::forInputType(
op.getType());
if (!converter) {
return (op.emitError("unsupported quantized type conversion"), true);
}
quant::QuantizedType elementType =
static_cast<const ConcreteRewriteClass *>(this)
->convertFakeQuantAttrsToType(op, converter.expressed_type);
if (!elementType) {
// Note that the fakeQuantAttrsToType will have emitted the error.
return true;
}
Type quantizedType = converter.convert(elementType);
assert(quantizedType &&
"Converter accepted a type that it did not convert");
// TODO: Map to a qbarrier with an attribute like [Forced] to signal that
// this is a forced/hard-coded constraint.
auto qbarrier = QuantizeCastOp::create(rewriter, op.getLoc(), quantizedType,
op.getInputs());
rewriter.replaceOpWithNewOp<DequantizeCastOp>(op, converter.input_type,
qbarrier.getResult());
return false;
}
};
class ConstFakeQuantRewrite
: public FakeQuantRewrite<ConstFakeQuantRewrite, ConstFakeQuant> {
public:
using BaseRewrite = FakeQuantRewrite<ConstFakeQuantRewrite, ConstFakeQuant>;
ConstFakeQuantRewrite(MLIRContext *ctx, bool *hadFailure)
: BaseRewrite(ctx, hadFailure) {}
quant::QuantizedType convertFakeQuantAttrsToType(ConstFakeQuant fqOp,
Type expressedType) const {
return fakeQuantAttrsToType(
fqOp.getLoc(), fqOp.getNumBits(), fqOp.getMin().convertToFloat(),
fqOp.getMax().convertToFloat(), fqOp.getNarrowRange(), expressedType,
fqOp.getIsSigned());
}
};
class ConstFakeQuantPerAxisRewrite
: public FakeQuantRewrite<ConstFakeQuantPerAxisRewrite,
ConstFakeQuantPerAxis> {
public:
using BaseRewrite =
FakeQuantRewrite<ConstFakeQuantPerAxisRewrite, ConstFakeQuantPerAxis>;
ConstFakeQuantPerAxisRewrite(MLIRContext *ctx, bool *hadFailure)
: BaseRewrite(ctx, hadFailure) {}
quant::QuantizedType convertFakeQuantAttrsToType(ConstFakeQuantPerAxis fqOp,
Type expressedType) const {
SmallVector<double, 4> min, max;
min.reserve(fqOp.getMin().size());
max.reserve(fqOp.getMax().size());
for (auto m : fqOp.getMin())
min.push_back(cast<FloatAttr>(m).getValueAsDouble());
for (auto m : fqOp.getMax())
max.push_back(cast<FloatAttr>(m).getValueAsDouble());
return fakeQuantAttrsToType(fqOp.getLoc(), fqOp.getNumBits(),
fqOp.getAxis(), min, max, fqOp.getNarrowRange(),
expressedType, fqOp.getIsSigned());
}
};
void ConvertSimulatedQuantPass::runOnOperation() {
bool hadFailure = false;
auto func = getOperation();
RewritePatternSet patterns(func.getContext());
auto *ctx = func.getContext();
patterns.add<ConstFakeQuantRewrite, ConstFakeQuantPerAxisRewrite>(
ctx, &hadFailure);
(void)applyPatternsGreedily(func, std::move(patterns));
if (hadFailure) signalPassFailure();
}
std::unique_ptr<OperationPass<func::FuncOp>> createConvertSimulatedQuantPass() {
return std::make_unique<ConvertSimulatedQuantPass>();
}
} // namespace mlir::quantfork
tensorflow/compiler/mlir/lite/quantization/ir/Passes.h
+57
View File
@@ -0,0 +1,57 @@
/* Copyright 2022 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 all of the passes owned by the quantization dialect. As
// things mature, it is expected that passes specific to certain frontend or
// backend dialects will move to those dialects directly. For now, they are
// incubated here.
//
//===----------------------------------------------------------------------===//
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_IR_PASSES_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_IR_PASSES_H_
#include "mlir/Pass/Pass.h" // from @llvm-project
namespace mlir {
namespace func {
class FuncOp;
} // namespace func
namespace quantfork {
/// Creates a pass that converts quantization simulation operations (i.e.
/// FakeQuant and those like it) to casts into/out of supported QuantizedTypes.
std::unique_ptr<OperationPass<func::FuncOp>> createConvertSimulatedQuantPass();
/// Creates a pass that converts constants followed by a qbarrier to a
/// constant whose value is quantized. This is typically one of the last
/// passes done when lowering to express actual quantized arithmetic in a
/// low level representation. Because it modifies the constant, it is
/// destructive and cannot be undone.
std::unique_ptr<OperationPass<func::FuncOp>> createConvertConstPass();
//===----------------------------------------------------------------------===//
// Registration
//===----------------------------------------------------------------------===//
/// Generate the code for registering passes.
#define GEN_PASS_REGISTRATION
#include "tensorflow/compiler/mlir/lite/quantization/ir/Passes.h.inc"
} // namespace quantfork
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_IR_PASSES_H_
tensorflow/compiler/mlir/lite/quantization/ir/Passes.td
+34
View File
@@ -0,0 +1,34 @@
/* Copyright 2022 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 QUANT_FORK_PASSES
#define QUANT_FORK_PASSES
include "mlir/Pass/PassBase.td"
def QuantConvertConst : Pass<"quant-convert-const", "func::FuncOp"> {
let summary = "Converts constants followed by qbarrier to actual quantized "
"values";
let constructor = "mlir::quantfork::createConvertConstPass()";
}
def QuantConvertSimulatedQuant
: Pass<"quant-convert-simulated-quantization", "func::FuncOp"> {
let summary = "Converts training-time simulated quantization ops to "
"corresponding quantize/dequantize casts";
let constructor = "mlir::quantfork::createConvertSimulatedQuantPass()";
}
#endif // QUANT_FORK_PASSES
tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.cc
+143
View File
@@ -0,0 +1,143 @@
/* Copyright 2022 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/ir/QuantOps.h"
#include <cstdint>
#include <functional>
#include <iterator>
#include <numeric>
#include "llvm/ADT/STLExtras.h"
#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/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
using namespace mlir;
using namespace mlir::quantfork;
using mlir::quant::QuantizedType;
#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOpsDialect.cc.inc"
void QuantizationForkDialect::initialize() {
addOperations<
#define GET_OP_LIST
#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.cc.inc"
>();
}
OpFoldResult StorageCastOp::fold(FoldAdaptor) {
// Matches x -> [scast -> scast] -> y, replacing the second scast with the
// value of x if the casts invert each other.
auto srcScastOp = getArg().getDefiningOp<StorageCastOp>();
if (!srcScastOp || srcScastOp.getArg().getType() != getType())
return OpFoldResult();
return srcScastOp.getArg();
}
/// The quantization specification should match the expressed type.
static bool isValidQuantizationSpec(Attribute quantSpec, Type expressed) {
if (auto typeAttr = mlir::dyn_cast<TypeAttr>(quantSpec)) {
Type spec = typeAttr.getValue();
if (mlir::isa<TensorType, VectorType>(spec)) return false;
// The spec should be either a quantized type which is compatible to the
// expressed type, or a primitive type which is as same as the
// (element type of) the expressed type.
if (auto quantizedType = mlir::dyn_cast<QuantizedType>(spec))
return quantizedType.isCompatibleExpressedType(expressed);
if (auto tensorType = mlir::dyn_cast<TensorType>(expressed))
return spec == tensorType.getElementType();
if (auto vectorType = mlir::dyn_cast<VectorType>(expressed))
return spec == vectorType.getElementType();
}
return false;
}
LogicalResult QuantizeRegionOp::verify() {
// There are specifications for both inputs and outputs.
if (getNumOperands() != getInputSpecs().size() ||
getNumResults() != getOutputSpecs().size())
return emitOpError(
"has unmatched operands/results number and spec attributes number");
// Verify that quantization specifications are valid.
for (auto input : llvm::zip(getOperandTypes(), getInputSpecs())) {
Type inputType = std::get<0>(input);
Attribute inputSpec = std::get<1>(input);
if (!isValidQuantizationSpec(inputSpec, inputType)) {
return emitOpError() << "has incompatible specification " << inputSpec
<< " and input type " << inputType;
}
}
for (auto result : llvm::zip(getResultTypes(), getOutputSpecs())) {
Type outputType = std::get<0>(result);
Attribute outputSpec = std::get<1>(result);
if (!isValidQuantizationSpec(outputSpec, outputType)) {
return emitOpError() << "has incompatible specification " << outputSpec
<< " and output type " << outputType;
}
}
return success();
}
LogicalResult StatisticsOp::verify() {
auto tensorArg = mlir::dyn_cast<TensorType>(getArg().getType());
if (!tensorArg) return emitOpError("arg needs to be tensor type.");
// Verify layerStats attribute.
{
auto layerStatsType = getLayerStats().getShapedType();
if (!mlir::isa<FloatType>(layerStatsType.getElementType())) {
return emitOpError("layerStats must have a floating point element type");
}
if (layerStatsType.getRank() != 1 || layerStatsType.getDimSize(0) != 2) {
return emitOpError("layerStats must have shape [2]");
}
}
// Verify axisStats (optional) attribute.
if (getAxisStats()) {
if (!getAxis()) return emitOpError("axis must be specified for axisStats");
auto shape = tensorArg.getShape();
auto argSliceSize =
std::accumulate(std::next(shape.begin(), *getAxis()), shape.end(), 1,
std::multiplies<int64_t>());
auto axisStatsType = getAxisStats()->getShapedType();
if (!mlir::isa<FloatType>(axisStatsType.getElementType())) {
return emitOpError("axisStats must have a floating point element type");
}
if (axisStatsType.getRank() != 2 || axisStatsType.getDimSize(1) != 2 ||
axisStatsType.getDimSize(0) != argSliceSize) {
return emitOpError(
"axisStats must have shape [N,2] "
"where N = the slice size defined by the axis dim");
}
}
return success();
}
#define GET_OP_CLASSES
#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.cc.inc"
tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h
+34
View File
@@ -0,0 +1,34 @@
/* Copyright 2022 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_IR_QUANTOPS_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_IR_QUANTOPS_H_
#include "llvm/Support/MathExtras.h"
#include "mlir/Bytecode/BytecodeOpInterface.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Dialect.h" // from @llvm-project
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
#include "mlir/Interfaces/InferTypeOpInterface.h" // from @llvm-project
#include "mlir/Interfaces/SideEffectInterfaces.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOpsDialect.h.inc"
#define GET_OP_CLASSES
#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h.inc"
#endif // TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_IR_QUANTOPS_H_
tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.td
+294
View File
@@ -0,0 +1,294 @@
/* Copyright 2022 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 operation definition file for Quantization.
//
//===----------------------------------------------------------------------===//
#ifndef QUANT_FORK_OPS
#define QUANT_FORK_OPS
include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOpsBase.td"
include "mlir/Dialect/Quant/IR/QuantBase.td"
include "mlir/Interfaces/InferTypeOpInterface.td"
include "mlir/Interfaces/SideEffectInterfaces.td"
class quant_TypedPrimitiveOrContainer<Type etype> :
Type<Or<[etype.predicate,
TensorOf<[etype]>.predicate,
VectorOfNonZeroRankOf<[etype]>.predicate]>,
"primitive/tensor/vector of " # etype.summary>;
// A primitive type that can represent a real value. This is either a
// floating point value or a quantized type.
def quant_RealPrimitiveType :
Type<Or<[AnyFloat.predicate, quant_QuantizedType.predicate]>,
"real valued primitive (float or quantized type)">;
// A primitive type that can represent a storage value. This is either an
// integer or quantized type.
def quant_StoragePrimitiveType :
Type<Or<[AnySignlessInteger.predicate, quant_QuantizedType.predicate]>,
"quantized storage primitive (integer or quantized type)">;
// A primitive or container of RealPrimitiveType.
def quant_RealValueType :
quant_TypedPrimitiveOrContainer<quant_RealPrimitiveType>;
// A primitive or container of StoragePrimitiveType.
def quant_StorageValueType :
quant_TypedPrimitiveOrContainer<quant_StoragePrimitiveType>;
// Either a real valued or storage primitive or container type.
def quant_RealOrStorageValueType :
Type<Or<[quant_RealValueType.predicate, quant_StorageValueType.predicate]>,
"real valued or storage primitive or container type">;
//===----------------------------------------------------------------------===//
// Base classes
//===----------------------------------------------------------------------===//
class quantfork_Op<string mnemonic, list<Trait> traits> :
Op<QuantizationFork_Dialect, mnemonic, traits>;
//===----------------------------------------------------------------------===//
// Quantization casts
//===----------------------------------------------------------------------===//
// A QuantizeCast (qcast) represents a potential type shift from a quantizable
// type to a quantized type.
//
// At runtime, a qcast will apply the transformation expressed by its
// operand and result type. For flexibility during transformation, it is also
// possible to have a qcast that performs no transformation (both its
// operand and result type are quantizable).
//
// A qcast will typically originate from either:
// a) An expressed or implied constraint in the source dialect which signals
// that a certain level of quantization is possible or required.
// b) An inference made by a quantization algorithm indicating that a
// quantized representation may be acceptable.
//
// Especially early in transformation, it is common to have pairs of
// qcast/dcast at points where a transition to a quantized type is
// required. In addition, it is also common to have an identity qcast
// (where the operand and result type are not quantized) at all points where
// it is legal to use a quantized representation (but is not known to be
// acceptable).
def quantfork_QuantizeCastOp : quantfork_Op<"qcast", [Pure]> {
let arguments = (ins quant_RealValueType:$arg);
let results = (outs quant_RealValueType);
}
// A DequantizeCast op (dcast) represents the inverse of a qcast,
// converting back from a quantized to quantizable (expressed) type.
//
// Like qcasts, a dcast is allowed to have both its operand and result
// as non quantized types. This facilitates transformations and marks edges
// where the computation must be carried out in the expressed type.
//
// Especially early in transformation, it is common to have dcasts on
// all operands to ops that must operate with the expressed type (typically
// math ops prior to lowering to target-specific, quantized kernels).
def quantfork_DequantizeCastOp : quantfork_Op<"dcast", [Pure]> {
let arguments = (ins quant_RealValueType:$arg);
let results = (outs quant_RealValueType);
}
// A StorageCast (scast) represents a cast from or to a type based on the
// storage type and a type based on a corresponding quantized type.
//
// This op exists to ensure type coherency for between parts of the computation
// which are operating directly on an underlying storage type and those which
// operate on quantized values.
//
// Examples from storage to quantized type:
// i8 -> !quant<"uniform[i8:f32]{1.0}">
// tensor<4xi8> -> tensor<4x!quant<"uniform[i8:f32]{1.0}">>
// vector<4xi8> -> vector<4x!quant<"uniform[i8:f32]{1.0}">>
def quantfork_StorageCastOp : quantfork_Op<"scast", [Pure]> {
let arguments = (ins quant_RealOrStorageValueType:$arg);
let results = (outs quant_RealOrStorageValueType);
let hasFolder = 1;
}
// A QuantizeRegion (region) represents a quantization unit which wraps
// high-precision ops with quantization specifications for all the inputs
// and outputs. Some quantization specifications can be undetermined and
// derived from other ports by the target specification of the kernel.
def quantfork_QuantizeRegionOp : quantfork_Op<"region", [
Pure,
IsolatedFromAbove,
SingleBlockImplicitTerminator<"ReturnOp">]> {
let summary = [{
The `region` operation wraps high-precision ops as a logical low-precision
quantized kernel.
}];
let arguments = (ins Variadic<AnyType>:$inputs,
TypeArrayAttr:$input_specs,
TypeArrayAttr:$output_specs,
StrAttr:$logical_kernel);
let results = (outs Variadic<AnyType>:$outputs);
let regions = (region SizedRegion<1>:$body);
let hasVerifier = 1;
}
def quantfork_ReturnOp : quantfork_Op<"return", [Terminator]> {
let summary = [{
The `return` operation terminates a quantize region and returns values.
}];
let arguments = (ins Variadic<AnyTensor>:$results);
}
//===----------------------------------------------------------------------===//
// Training integration and instrumentation ops
//===----------------------------------------------------------------------===//
def quantfork_ConstFakeQuant : quantfork_Op<"const_fake_quant",
[SameOperandsAndResultType, Pure]> {
let summary = [{
Simulates the effect of uniform quantization with const range.
}];
let description = [{
Given a const min, max, num_bits and narrow_range attribute, applies the
same uniform quantization simulation as is done by the TensorFlow
fake_quant_with_min_max_args op. See the fakeQuantAttrsToType() utility
method and the quant-convert-simulated-quantization pass for further details.
}];
let arguments = (ins
F32Tensor:$inputs,
F32Attr:$min,
F32Attr:$max,
// The bitwidth of the quantization; between 2 and 16, inclusive.
I64Attr:$num_bits,
// Quantization range starts from 0 or 1; starts from 1 if true.
DefaultValuedOptionalAttr<BoolAttr, "false">:$narrow_range,
// The sign of the quantization.
DefaultValuedOptionalAttr<BoolAttr, "false">:$is_signed
);
let results = (outs
F32Tensor:$outputs
);
}
def quantfork_ConstFakeQuantPerAxis : quantfork_Op<"const_fake_quant_per_axis",
[SameOperandsAndResultType, Pure]> {
let summary = [{
Simulates the effect of per axis uniform quantization with const range.
}];
let description = [{
Given a const min, max, num_bits and narrow_range attribute, applies the
same per axis uniform quantization simulation as is done by the TensorFlow
fake_quant_with_min_max_vars_per_channel op. See the fakeQuantAttrsToType()
utility method and the quant-convert-simulated-quantization pass for further
details.
}];
let arguments = (ins
F32Tensor:$inputs,
F32ArrayAttr:$min,
F32ArrayAttr:$max,
// The quantized dimension of the inputs tensor.
I64Attr:$axis,
// The bitwidth of the quantization; between 2 and 16, inclusive.
I64Attr:$num_bits,
// Quantization range starts from 0 or 1; starts from 1 if true.
DefaultValuedOptionalAttr<BoolAttr, "false">:$narrow_range,
// The sign of the quantization.
DefaultValuedOptionalAttr<BoolAttr, "false">:$is_signed
);
let results = (outs
F32Tensor:$outputs
);
}
def quantfork_StatisticsRefOp : quantfork_Op<"stats_ref", [SameOperandsAndResultType]> {
let summary = "Indicates that statistics are resolved by reference.";
let description = [{
This op acts as an identity that, when encountered at runtime, should result
in statistics being collected about about the value of its operand/result.
Such statistics will be stored with the provided key, allowing this node
to later be converted to a 'stats' op if statistics with that key have been
encountered.
}];
let arguments = (ins
quant_RealValueType:$arg,
StrAttr:$statsKey
);
let results = (outs quant_RealValueType);
}
def quantfork_StatisticsOp : quantfork_Op<"stats", [SameOperandsAndResultType]> {
let summary = "Identity op which associates statistics with the value.";
let description = [{
Associates statistics about the runtime ranges of values observed for
evaluations of this node.
Statistics about the entire type are reported in the 'layerStats' attribute
and those for each axis, in the (optional) `axisStats` attribute. The
interpretation of each is determined by the last dimension of its shape.
Currently, only dim=2 is supported, which is interpreted as [min, max].
`layerStats` must be a rank 1 tensor: [2]
`axisStats` must be a rank 2 tensor: [N, 2], where N=the slice size
splitted by the `axis` dimension. For example:
```
<?x?x3x2>, axis=3 => N=2
<?x?x3x2>, axis=2 => N=6
```
}];
let arguments = (ins
quant_RealValueType:$arg,
ElementsAttr:$layerStats,
OptionalAttr<ElementsAttr>:$axisStats,
OptionalAttr<I64Attr>:$axis);
let results = (outs quant_RealValueType);
let hasVerifier = 1;
}
def quantfork_CoupledRefOp : quantfork_Op<"coupled_ref", [SameOperandsAndResultType]> {
let summary = [{
Indicates that one point of the computation is coupled to another.
}];
let description = [{
Ordinarily, relationships between ops for the purposes of determining
compatible quantized types is explicit based on the use-def chain. However,
in some situations, a use may be separated from its def by arbitrary
external connections. In such a case, during analysis, all coupled_ref
nodes in a module which share a coupledKey will be considered to be
directly connected as via an identity op for the purpose of type inference.
}];
let arguments = (ins
quant_RealValueType:$arg,
StrAttr:$coupledKey);
let results = (outs quant_RealValueType);
}
#endif // QUANT_FORK_OPS
tensorflow/compiler/mlir/lite/quantization/ir/QuantOpsBase.td
+32
View File
@@ -0,0 +1,32 @@
/* Copyright 2022 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.
==============================================================================*/
//
//===----------------------------------------------------------------------===//
//
// Predicates for types in the Quantization dialect.
//
//===----------------------------------------------------------------------===//
#ifndef QUANT_FORK_BASE
#define QUANT_FORK_BASE
include "mlir/IR/OpBase.td"
def QuantizationFork_Dialect : Dialect {
let name = "quantfork";
let cppNamespace = "::mlir::quantfork";
}
#endif // QUANT_FORK_BASE
tensorflow/compiler/mlir/lite/quantization/ir/QuantizeUtils.cc
+147
View File
@@ -0,0 +1,147 @@
/* Copyright 2022 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/ir/QuantizeUtils.h"
#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/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/common/ir/UniformSupport.h"
using namespace mlir;
using namespace mlir::quantfork;
/// Converts a possible primitive, real expressed value attribute to a
/// corresponding storage attribute (typically FloatAttr -> IntegerAttr).
/// quantizedElementType is the QuantizedType that describes the expressed
/// origValue.
/// Returns a converter Attribute or nullptr if conversion is not possible.
static Attribute convertPrimitiveValueAttr(
Attribute origRealValue, quant::QuantizedType quantizedElementType,
const mlir::quant::ir::UniformQuantizedValueConverter &converter,
Type &outConvertedType) {
if (mlir::isa<FloatAttr>(origRealValue)) {
FloatAttr floatAttr = mlir::cast<FloatAttr>(origRealValue);
outConvertedType = quantizedElementType.getStorageType();
return IntegerAttr::get(quantizedElementType.getStorageType(),
converter.quantizeFloatToInt(floatAttr.getValue()));
}
return nullptr;
}
/// Converts a real expressed DenseFPElementsAttr to a corresponding
/// DenseElementsAttr (typically DenseIntElementsAttr) containing quantized
/// storage values assuming the given quantizedElementType and converter.
static DenseElementsAttr convertDenseFPElementsAttr(
DenseFPElementsAttr realFPElementsAttr,
quant::QuantizedType quantizedElementType,
const mlir::quant::ir::UniformQuantizedValueConverter &converter) {
return realFPElementsAttr.mapValues(
quantizedElementType.getStorageType(),
[&converter](const APFloat &realVal) {
return converter.quantizeFloatToInt(realVal);
});
}
/// Converts a real expressed SplatElementsAttr to a corresponding
/// SplatElementsAttr containing quantized storage values assuming the given
/// quantizedElementType and converter.
static SparseElementsAttr convertSparseElementsAttr(
SparseElementsAttr realSparseAttr,
quant::QuantizedType quantizedElementType,
const mlir::quant::ir::UniformQuantizedValueConverter &converter) {
DenseElementsAttr realDenseAttr = realSparseAttr.getValues();
if (!mlir::isa<DenseFPElementsAttr>(realDenseAttr)) {
return nullptr;
}
DenseElementsAttr quantDenseAttr =
convertDenseFPElementsAttr(mlir::cast<DenseFPElementsAttr>(realDenseAttr),
quantizedElementType, converter);
if (!quantDenseAttr) {
return nullptr;
}
// Cast from an expressed-type-based type to storage-type-based type,
// preserving the sparse shape (i.e. tensor<4xf32> -> tensor<4xi8>).
ShapedType newSparseType = mlir::dyn_cast_or_null<ShapedType>(
quantizedElementType.castExpressedToStorageType(
realSparseAttr.getType()));
if (!newSparseType) {
return nullptr;
}
return SparseElementsAttr::get(newSparseType, realSparseAttr.getIndices(),
quantDenseAttr);
}
/// Converts a real expressed Attribute to a corresponding Attribute containing
/// quantized storage values assuming the given uniform quantizedElementType and
/// converter.
Attribute mlir::quantfork::quantizeAttrUniform(
Attribute realValue, quant::UniformQuantizedType quantizedElementType,
const mlir::quant::ir::UniformQuantizedValueConverter &converter,
Type &outConvertedType) {
// Fork to handle different variants of constants supported.
if (mlir::isa<DenseFPElementsAttr>(realValue)) {
// Dense tensor or vector constant.
auto converted =
convertDenseFPElementsAttr(mlir::cast<DenseFPElementsAttr>(realValue),
quantizedElementType, converter);
outConvertedType = converted.getType();
return converted;
}
if (mlir::isa<SparseElementsAttr>(realValue)) {
// Sparse tensor or vector constant.
auto converted =
convertSparseElementsAttr(mlir::cast<SparseElementsAttr>(realValue),
quantizedElementType, converter);
outConvertedType = converted.getType();
return converted;
}
// Nothing else matched: try to convert a primitive.
return convertPrimitiveValueAttr(realValue, quantizedElementType, converter,
outConvertedType);
}
/// Convert an attribute from a type based on
/// quantizedElementType.getExpressedType() to one based on
/// quantizedElementType.getStorageType().
/// Returns nullptr if the conversion is not supported.
/// On success, stores the converted type in outConvertedType.
Attribute mlir::quantfork::quantizeAttr(
Attribute realValue, quant::QuantizedType quantizedElementType,
Type &outConvertedType) {
if (auto uniformQuantized =
mlir::dyn_cast<quant::UniformQuantizedType>(quantizedElementType)) {
mlir::quant::ir::UniformQuantizedValueConverter converter(uniformQuantized);
return quantizeAttrUniform(realValue, uniformQuantized, converter,
outConvertedType);
}
if (auto uniformQuantizedPerAxis =
mlir::dyn_cast<quant::UniformQuantizedPerAxisType>(
quantizedElementType)) {
mlir::quant::ir::UniformQuantizedPerAxisValueConverter converter(
uniformQuantizedPerAxis);
auto converted = converter.convert(realValue);
// TODO: why we need this outConvertedType? remove it?
if (converted) {
outConvertedType = converted.getType();
}
return converted;
}
return nullptr;
}
tensorflow/compiler/mlir/lite/quantization/ir/QuantizeUtils.h
+73
View File
@@ -0,0 +1,73 @@
/* Copyright 2022 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_IR_QUANTIZEUTILS_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_IR_QUANTIZEUTILS_H_
namespace mlir {
class Attribute;
class Type;
namespace quant {
class QuantizedType;
class UniformQuantizedType;
namespace ir {
class UniformQuantizedValueConverter;
} // namespace ir
} // namespace quant
namespace quantfork {
/// Converts an attribute from a type based on
/// quantizedElementType.getExpressedType() to one based on
/// quantizedElementType.getStorageType(), where quantizedElementType is as from
/// QuantizedType::getQuantizedElementType().
/// Returns nullptr if the conversion is not supported. On success, stores the
/// converted type in outConvertedType.
///
/// Examples:
/// 1. realValue is a primitive value attribute:
/// (realValue: FloatAttr, quantizedElementType: UniformQuantizedType[i8:f32])
/// -> (IntegerAttr, outConvertedType: i8)
/// 2. realValue is an elements attribute:
/// (realValue: DenseElementsAttr[tensor<2x2xf32>],
/// quantizedElementType: UniformQuantizedType[i8:f32])
/// -> (DenseElementsAttr[tensor<2x2xi8>], outConvertedType: tensor<2x2xi8>)
Attribute quantizeAttr(Attribute realValue,
quant::QuantizedType quantizedElementType,
Type &outConvertedType);
/// Converts an attribute from a type based on
/// quantizedElementType.getExpressedType() to one based on
/// quantizedElementType.getStorageType(), where quantizedElementType is as from
/// QuantizedType::getQuantizedElementType() and casted to an
/// UniformQuantizedType. Returns nullptr if the conversion is not supported. On
/// success, stores the converted type in outConvertedType.
///
/// Examples:
/// 1. realValue is a primitive value attribute:
/// (realValue: FloatAttr, quantizedElementType: UniformQuantizedType[i8:f32])
/// -> (IntegerAttr, outConvertedType: i8)
/// 2. realValue is an elements attribute:
/// (realValue: DenseElementsAttr[tensor<2x2xf32>],
/// quantizedElementType: UniformQuantizedType[i8:f32])
/// -> (DenseElementsAttr[tensor<2x2xi8>], outConvertedType: tensor<2x2xi8>)
Attribute quantizeAttrUniform(
Attribute realValue, quant::UniformQuantizedType quantizedElementType,
const mlir::quant::ir::UniformQuantizedValueConverter &converter,
Type &outConvertedType);
} // namespace quantfork
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_IR_QUANTIZEUTILS_H_
tensorflow/compiler/mlir/lite/quantization/lite/BUILD
+231
View File
@@ -0,0 +1,231 @@
load("//tensorflow:tensorflow.bzl", "tf_cc_binary", "tf_cc_test")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
default_visibility = [
":friends",
"//tensorflow:__pkg__",
],
licenses = ["notice"],
)
exports_files(glob([
"testdata/*.bin",
]))
package_group(
name = "friends",
packages = [
"//tensorflow/compiler/mlir/lite/...",
"//tensorflow/lite/...",
],
)
cc_library(
name = "quantize_model",
srcs = [
"quantize_model.cc",
],
hdrs = [
"quantize_model.h",
],
deps = [
"//tensorflow/compiler/mlir/lite:common",
"//tensorflow/compiler/mlir/lite:flatbuffer_translate_lib",
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/lite:tf_tfl_passes",
"//tensorflow/compiler/mlir/lite/debug",
"//tensorflow/compiler/mlir/lite/debug:debug_options_proto_cc",
"//tensorflow/compiler/mlir/lite/quantization/common/quantization_lib:quantization_config",
"//tensorflow/compiler/mlir/lite/schema:schema_fbs",
"//tensorflow/compiler/mlir/tensorflow:error_util",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
],
)
cc_library(
name = "quantize_weights",
srcs = [
"quantize_weights.cc",
],
hdrs = [
"quantize_weights.h",
],
deps = [
"//tensorflow/compiler/mlir/lite:common",
"//tensorflow/compiler/mlir/lite:flatbuffer_translate_lib",
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/lite:tf_tfl_passes",
"//tensorflow/compiler/mlir/lite/quantization/common/quantization_lib:quantization_config",
"//tensorflow/compiler/mlir/lite/schema:schema_fbs",
"//tensorflow/compiler/mlir/tensorflow:error_util",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@flatbuffers//:runtime_cc",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
],
)
cc_library(
name = "tfl_to_std",
srcs = [
"tfl_to_std.cc",
],
hdrs = [
"tfl_to_std.h",
"//tensorflow/compiler/mlir/lite/quantization/common/quantization_lib:quantization_utils.h",
],
deps = [
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/lite/quantization/common/quantization_lib",
"//tensorflow/compiler/mlir/lite/quantization/common/quantization_lib:quantization_config",
"//tensorflow/compiler/mlir/lite/quantization/ir:QuantOps",
"//tensorflow/compiler/mlir/quantization/common/ir:QuantOps",
"//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/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",
],
)
# Binary to apply quantization on the annotated files.
tf_cc_binary(
name = "tfl_quantizer",
srcs = [
"tfl_quantizer.cc",
],
deps = [
":quantize_model",
"//tensorflow/compiler/mlir/lite/schema:schema_fbs",
"@com_google_absl//absl/status",
"@llvm-project//llvm:Support",
],
)
tf_cc_test(
name = "quantize_model_test",
srcs = ["quantize_model_test.cc"],
args = [
"--test_model_file=$(location //tensorflow/compiler/mlir/lite/quantization/lite:testdata/single_conv_weights_min_0_max_plus_10.bin)",
],
data = [
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/add_with_const_input.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/argmax.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/broadcast_to.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/concat.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/fc.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/fc_qat.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/gather_nd.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/lstm_calibrated.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/lstm_calibrated2.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/lstm_quantized.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/lstm_quantized2.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/maximum.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/minimum.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/mixed.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/mixed16x8.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/multi_input_add_reshape.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/pack.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/single_avg_pool_min_minus_5_max_plus_5.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/single_conv_no_bias.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/single_conv_weights_min_0_max_plus_10.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/single_conv_weights_min_minus_127_max_plus_127.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/single_softmax_min_minus_5_max_plus_5.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/split.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/svdf_calibrated.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/svdf_quantized.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/transpose.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/unidirectional_sequence_lstm_calibrated.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/unidirectional_sequence_lstm_quantized.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/unpack.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/where.bin",
],
tags = [
"tflite_not_portable_android",
"tflite_not_portable_ios",
],
deps = [
":quantize_model",
":test_util",
"//tensorflow/compiler/mlir/lite/core:absl_error_model_builder",
"//tensorflow/compiler/mlir/lite/schema:schema_fbs",
"//tensorflow/compiler/mlir/lite/schema:schema_utils",
"//tensorflow/core:framework_internal",
"//tensorflow/core:lib",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/status",
"@com_google_googletest//:gtest",
"@flatbuffers",
"@xla//xla/tsl/lib/core:status_test_util",
],
)
tf_cc_test(
name = "quantize_weights_test",
srcs = ["quantize_weights_test.cc"],
args = [
"--test_model_file=$(location //tensorflow/compiler/mlir/lite/quantization/lite:testdata/single_conv_weights_min_0_max_plus_10.bin)",
],
data = [
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/custom_op.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/quantized_with_gather.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/single_conv_weights_min_0_max_plus_10.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/weight_shared_between_convs.bin",
],
tags = [
# TODO(b/327796566): re-enable after the bug is fixed
"manual",
"noasan",
"noguitar",
"notap",
"tflite_not_portable_android",
"tflite_not_portable_ios",
],
deps = [
":quantize_weights",
":test_util",
"//tensorflow/compiler/mlir/lite/core:absl_error_model_builder",
"//tensorflow/compiler/mlir/lite/schema:schema_fbs",
"//tensorflow/compiler/mlir/lite/schema:schema_utils",
"//tensorflow/core:framework_internal",
"//tensorflow/core:lib",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_googletest//:gtest",
"@flatbuffers",
"@xla//xla/tsl/platform:logging",
],
)
cc_library(
name = "test_util",
testonly = 1,
srcs = ["test_util.cc"],
hdrs = ["test_util.h"],
visibility = ["//visibility:public"],
deps = [
"//tensorflow/compiler/mlir/lite/core/api:error_reporter",
"@com_google_googletest//:gtest",
],
)
tensorflow/compiler/mlir/lite/quantization/lite/quantize_model.cc
+147
View File
@@ -0,0 +1,147 @@
/* 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/lite/quantize_model.h"
#include <optional>
#include <string>
#include <unordered_set>
#include "absl/container/flat_hash_set.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/Twine.h"
#include "llvm/Support/Debug.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Location.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/common/tfl_pass_config.h"
#include "tensorflow/compiler/mlir/lite/debug/debug.h"
#include "tensorflow/compiler/mlir/lite/debug/debug_options.pb.h"
#include "tensorflow/compiler/mlir/lite/flatbuffer_export.h"
#include "tensorflow/compiler/mlir/lite/flatbuffer_import.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_config.h"
#include "tensorflow/compiler/mlir/lite/schema/schema_generated.h"
#include "tensorflow/compiler/mlir/lite/tf_tfl_passes.h"
#include "tensorflow/compiler/mlir/lite/transforms/passes.h"
#include "tensorflow/compiler/mlir/lite/utils/convert_type.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
#include "tensorflow/core/framework/types.pb.h"
namespace mlir {
namespace lite {
std::string TfLiteToMlir(const absl::string_view tflite_op_name) {
llvm::StringRef op_name(tflite_op_name.data(), tflite_op_name.size());
return llvm::Twine("tfl.", op_name.lower()).str();
}
// TODO(fengliuai): check the result for `fully_quantize` flag.
absl::Status QuantizeModel(
const absl::string_view model_buffer, const tflite::TensorType &input_type,
const tflite::TensorType &output_type,
const tflite::TensorType &inference_type,
const std::unordered_set<std::string> &operator_names,
bool disable_per_channel, bool fully_quantize, std::string &output_buffer,
bool verify_numeric, bool whole_model_verify, bool legacy_float_scale,
const absl::flat_hash_set<std::string> &denylisted_ops,
const absl::flat_hash_set<std::string> &denylisted_nodes,
const bool enable_variable_quantization,
bool disable_per_channel_for_dense_layers,
const std::optional<const tensorflow::converter::DebugOptions>
&debug_options) {
// Translate TFLite names to mlir op names.
absl::flat_hash_set<std::string> denylisted_mlir_op_names;
for (const auto& entry : denylisted_ops) {
denylisted_mlir_op_names.insert(TfLiteToMlir(entry));
}
DialectRegistry registry;
registry.insert<mlir::TFL::TensorFlowLiteDialect>();
MLIRContext context(registry);
StatusScopedDiagnosticHandler statusHandler(&context,
/*propagate=*/true);
OwningOpRef<mlir::ModuleOp> module = tflite::FlatBufferToMlir(
model_buffer, &context, UnknownLoc::get(&context));
if (!module) {
return absl::InternalError("Couldn't import flatbuffer to MLIR.");
}
// Apply quantization passes.
PassManager pm((*module)->getName(), OpPassManager::Nesting::Implicit);
if (debug_options.has_value()) {
// Add debugging instrumentation
tensorflow::InitPassManager(pm, debug_options.value());
}
TFL::QuantizationSpecs quant_specs;
quant_specs.inference_type = tflite::TflTypeToTfType(inference_type);
quant_specs.post_training_quantization = true;
quant_specs.disable_per_channel = disable_per_channel;
quant_specs.disable_per_channel_for_dense_layers =
disable_per_channel_for_dense_layers;
quant_specs.verify_numeric = verify_numeric;
quant_specs.whole_model_verify = whole_model_verify;
quant_specs.legacy_float_scale = legacy_float_scale;
quant_specs.ops_blocklist = denylisted_mlir_op_names;
quant_specs.nodes_blocklist = denylisted_nodes;
quant_specs.enable_mlir_variable_quantization = enable_variable_quantization;
llvm::dbgs() << "fully_quantize: " << fully_quantize
<< ", inference_type: " << quant_specs.inference_type
<< ", input_inference_type: "
<< tflite::EnumNameTensorType(input_type)
<< ", output_inference_type: "
<< tflite::EnumNameTensorType(output_type) << "\n";
mlir::Builder mlir_builder(&context);
mlir::Type input_mlir_type =
tflite::ConvertElementType(input_type, mlir_builder);
mlir::Type output_mlir_type =
tflite::ConvertElementType(output_type, mlir_builder);
if (fully_quantize) {
input_mlir_type = tflite::ConvertElementType(inference_type, mlir_builder);
output_mlir_type = input_mlir_type;
}
tensorflow::AddQuantizationPasses(mlir::TFL::PassConfig(quant_specs), pm);
pm.addPass(TFL::CreateModifyIONodesPass(input_mlir_type, output_mlir_type));
// If the first or final ops are not quantized, remove QDQ.
pm.addPass(TFL::CreatePostQuantizeRemoveQDQPass());
if (failed(pm.run(module.get()))) {
const std::string err(statusHandler.ConsumeStatus().message());
return absl::InternalError(err);
}
// Export the results.
tflite::FlatbufferExportOptions options;
options.converter_flags.set_force_select_tf_ops(false);
options.converter_flags.set_enable_select_tf_ops(true);
options.converter_flags.set_allow_custom_ops(true);
if (!tflite::MlirToFlatBufferTranslateFunction(module.get(), options,
&output_buffer)) {
return absl::InternalError("Failed to export MLIR to flatbuffer.");
}
return absl::OkStatus();
}
} // namespace lite
} // namespace mlir
tensorflow/compiler/mlir/lite/quantization/lite/quantize_model.h
+66
View File
@@ -0,0 +1,66 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_LITE_QUANTIZE_MODEL_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_LITE_QUANTIZE_MODEL_H_
#include <optional>
#include <string>
#include <unordered_set>
#include "absl/container/flat_hash_set.h"
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "tensorflow/compiler/mlir/lite/debug/debug_options.pb.h"
#include "tensorflow/compiler/mlir/lite/schema/schema_generated.h"
namespace mlir {
namespace lite {
// Quantizes the input model represented as `model_buffer` and writes the result
// to the `output_buffer`. Both `model_buffer` and `output_buffer` should be a
// valid FlatBuffer format for Model supported by TFLite.
//
// The `input_type`, `output_type` and `inference_type` can be float32 / qint8 /
// int8 / int16.
//
// Returns a partially quantized model if `fully_quantize` is false. Returns a
// non-OK status if the quantization fails.
//
// When `verify_numeric` is true, the model will have it's original float ops
// and NumericVerify ops to compare output values from the quantized and float
// ops.
//
// When `legacy_float_scale` is true, the quantizer will use float scale instead
// of double, and call TOCO's quantization routines to maintain bit-exactness of
// the values with the TOCO quantizer.
absl::Status QuantizeModel(
absl::string_view model_buffer, const tflite::TensorType &input_type,
const tflite::TensorType &output_type,
const tflite::TensorType &inference_type,
const std::unordered_set<std::string> &operator_names,
bool disable_per_channel, bool fully_quantize, std::string &output_buffer,
bool verify_numeric = false, bool whole_model_verify = false,
bool legacy_float_scale = true,
const absl::flat_hash_set<std::string> &denylisted_ops = {},
const absl::flat_hash_set<std::string> &denylisted_nodes = {},
bool enable_variable_quantization = false,
bool disable_per_channel_for_dense_layers = false,
const std::optional<const tensorflow::converter::DebugOptions>
&debug_options = std::nullopt);
} // namespace lite
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_LITE_QUANTIZE_MODEL_H_
tensorflow/compiler/mlir/lite/quantization/lite/quantize_model_test.cc
+1787
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/mlir/lite/quantization/lite/quantize_weights.cc
+237
View File
@@ -0,0 +1,237 @@
/* Copyright 2022 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/lite/quantize_weights.h"
#include <cstdint>
#include <memory>
#include <string>
#include <utility>
#include "absl/container/flat_hash_set.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "flatbuffers/buffer.h" // from @flatbuffers
#include "flatbuffers/flatbuffer_builder.h" // from @flatbuffers
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/Twine.h"
#include "llvm/Support/Debug.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Location.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/common/tfl_pass_config.h"
#include "tensorflow/compiler/mlir/lite/flatbuffer_export.h"
#include "tensorflow/compiler/mlir/lite/flatbuffer_import.h"
#include "tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_config.h"
#include "tensorflow/compiler/mlir/lite/schema/schema_generated.h"
#include "tensorflow/compiler/mlir/lite/tf_tfl_passes.h"
#include "tensorflow/compiler/mlir/lite/transforms/passes.h"
#include "tensorflow/compiler/mlir/lite/utils/convert_type.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
#include "tensorflow/core/framework/types.pb.h"
namespace mlir {
namespace lite {
namespace {
using llvm::StringRef;
// Convert op represented in TFLite builtin_code to its corresponding MLIR
// OperationName.
void TfLiteBuiltinOpToMlir(const BuiltinOperatorSet& tflite_builtin_codes,
absl::flat_hash_set<std::string>& mlir_op_names) {
for (const auto& entry : tflite_builtin_codes) {
StringRef tflite_op_name = EnumNameBuiltinOperator(entry);
std::string mlir_name = llvm::Twine("tfl.", tflite_op_name.lower()).str();
mlir_op_names.insert(std::move(mlir_name));
}
}
std::string TfLiteToMlir(absl::string_view tflite_op_name) {
StringRef op_name(tflite_op_name.data(), tflite_op_name.size());
return op_name.lower();
}
std::unique_ptr<tflite::ModelT> CreateMutableModelFromFile(
const tflite::Model* input_model) {
auto copied_model = std::make_unique<tflite::ModelT>();
input_model->UnPackTo(copied_model.get(), nullptr);
return copied_model;
}
} // namespace
// TODO(b/214314076): Support MLIR model as an input for the C++ dynamic range
// quantization API
absl::Status QuantizeWeights(
flatbuffers::FlatBufferBuilder* builder, const tflite::Model* input_model,
const tflite::TensorType& inference_type,
const absl::flat_hash_set<std::string>& denylisted_ops,
const CustomOpMap& custom_op_map, int64_t minimum_elements_for_weights,
bool disable_per_channel, bool weight_only_quantization,
bool legacy_float_scale) {
// Translate TFLite names to mlir op names.
absl::flat_hash_set<std::string> denylisted_mlir_op_names;
for (auto& entry : denylisted_ops) {
denylisted_mlir_op_names.insert(TfLiteToMlir(entry));
}
DialectRegistry registry;
MLIRContext context(registry);
StatusScopedDiagnosticHandler statusHandler(&context,
/*propagate=*/true);
// Import input_model to a MLIR module
flatbuffers::FlatBufferBuilder input_builder;
flatbuffers::Offset<tflite::Model> input_model_location = tflite::Model::Pack(
input_builder, CreateMutableModelFromFile(input_model).get());
tflite::FinishModelBuffer(input_builder, input_model_location);
std::string serialized_model(
reinterpret_cast<const char*>(input_builder.GetBufferPointer()),
input_builder.GetSize());
OwningOpRef<mlir::ModuleOp> module = tflite::FlatBufferToMlir(
serialized_model, &context, UnknownLoc::get(&context));
// Apply quantization passes.
PassManager pm((*module)->getName(), OpPassManager::Nesting::Implicit);
TFL::QuantizationSpecs quant_specs;
quant_specs.inference_type = tflite::TflTypeToTfType(inference_type);
quant_specs.weight_quantization = true;
quant_specs.weight_only_quantization = weight_only_quantization;
quant_specs.minimum_elements_for_weights = minimum_elements_for_weights;
quant_specs.disable_per_channel = disable_per_channel;
quant_specs.legacy_float_scale = legacy_float_scale;
quant_specs.ops_blocklist = denylisted_mlir_op_names;
for (const auto& entry : custom_op_map) {
quant_specs.custom_map[entry.first].quantizable_input_indices =
entry.second.quantizable_input_indices;
quant_specs.custom_map[entry.first].is_weight_only =
entry.second.is_weight_only;
quant_specs.custom_map[entry.first].no_side_effect =
entry.second.no_side_effect;
}
if (quant_specs.inference_type == tensorflow::DT_INT8)
quant_specs.inference_type = tensorflow::DT_QINT8;
if (!(quant_specs.inference_type == tensorflow::DT_HALF ||
quant_specs.inference_type == tensorflow::DT_QINT8)) {
LOG(ERROR) << "Couldn't apply dynamic range quantization since unsupported "
"inference_type is passed.";
return absl::InvalidArgumentError(
"Quantize weights transformation failed.");
}
llvm::dbgs() << "weight_quantization: " << true
<< ", weight_only_quantization: "
<< quant_specs.weight_only_quantization << ", mlir_quantizer: "
<< quant_specs.enable_mlir_dynamic_range_quantizer
<< ", inference_type: " << quant_specs.inference_type << "\n";
Builder mlir_builder(&context);
tensorflow::AddDynamicRangeQuantizationPasses(
mlir::TFL::PassConfig(quant_specs), pm);
if (failed(pm.run(module.get()))) {
absl::string_view err = statusHandler.ConsumeStatus().message();
LOG(ERROR) << "Failed to quantize: " << err;
return absl::InvalidArgumentError(
"Quantize weights transformation failed.");
}
// Export the results to the builder
std::string result;
tflite::FlatbufferExportOptions options;
options.converter_flags.set_force_select_tf_ops(false);
options.converter_flags.set_enable_select_tf_ops(true);
options.converter_flags.set_allow_custom_ops(true);
if (!tflite::MlirToFlatBufferTranslateFunction(module.get(), options,
&result)) {
LOG(ERROR) << "Failed to export MLIR to flatbuffer.";
return absl::InvalidArgumentError(
"Quantize weights transformation failed.");
}
builder->PushFlatBuffer(reinterpret_cast<const uint8_t*>(result.data()),
result.size());
return absl::OkStatus();
}
absl::Status QuantizeWeights(flatbuffers::FlatBufferBuilder* builder,
const tflite::Model* input_model,
int64_t weights_min_num_elements,
bool use_hybrid_evaluation) {
return QuantizeWeights(
builder, input_model,
/*inference_type=*/tflite::TensorType_INT8,
/*denylisted_ops=*/{},
/*custom_op_map=*/{},
/*minimum_elements_for_weights=*/weights_min_num_elements,
/*disable_per_channel=*/false,
/*weight_only_quantization=*/!use_hybrid_evaluation,
/*legacy_float_scale=*/true);
}
// In MLIR use_updated_hybrid_scheme = true means per-channel operation.
absl::Status QuantizeWeights(flatbuffers::FlatBufferBuilder* builder,
const tflite::Model* input_model,
BufferType quant_type,
bool use_updated_hybrid_scheme) {
tflite::TensorType inference_type;
switch (quant_type) {
case BufferType::QUANTIZED_FLOAT16:
inference_type = tflite::TensorType_FLOAT16;
break;
default:
inference_type = tflite::TensorType_INT8;
}
return QuantizeWeights(builder, input_model, inference_type,
/*denylisted_ops=*/{},
/*custom_op_map=*/{},
/*minimum_elements_for_weights=*/1024,
/*disable_per_channel=*/!use_updated_hybrid_scheme,
/*weight_only_quantization=*/false,
/*legacy_float_scale=*/true);
}
absl::Status QuantizeWeights(flatbuffers::FlatBufferBuilder* builder,
const tflite::Model* input_model,
int64_t weights_min_num_elements,
const CustomOpMap& custom_op_map,
bool use_updated_hybrid_scheme,
const BuiltinOperatorSet& op_denylist) {
const tflite::TensorType inference_type = tflite::TensorType_INT8;
absl::flat_hash_set<std::string> mlir_op_denylist;
TfLiteBuiltinOpToMlir(op_denylist, mlir_op_denylist);
return QuantizeWeights(
builder, input_model, inference_type,
/*denylisted_ops=*/mlir_op_denylist,
/*custom_op_map=*/custom_op_map,
/*minimum_elements_for_weights=*/weights_min_num_elements,
/*disable_per_channel=*/!use_updated_hybrid_scheme,
/*weight_only_quantization=*/false,
/*legacy_float_scale=*/true);
}
} // namespace lite
} // namespace mlir
tensorflow/compiler/mlir/lite/quantization/lite/quantize_weights.h
+91
View File
@@ -0,0 +1,91 @@
/* Copyright 2022 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_LITE_QUANTIZE_WEIGHTS_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_LITE_QUANTIZE_WEIGHTS_H_
#include <cstdint>
#include <memory>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>
#include "absl/container/flat_hash_set.h"
#include "absl/status/status.h"
#include "flatbuffers/flatbuffer_builder.h" // from @flatbuffers
#include "tensorflow/compiler/mlir/lite/schema/schema_generated.h"
namespace mlir {
namespace lite {
// Supported resulting types from quantization process.
enum class BufferType { QUANTIZED_INT8, QUANTIZED_FLOAT16 };
// Stores information about how to quantize a user-specified custom operation.
// CustomOpInfo contains info of its corresponding CustomOp registered in the
// CustomOpMap. 'quantizable_input_indices' is used to determine which indices
// of the CustomOp are quantizable. 'is_weight_only' is used specify whether the
// custom op is quantized only for storage and dequantized at runtime.
// 'no_side_effect' is used to determine whether the op can be pruned if
// considered as trivially dead.
struct CustomOpInfo {
std::vector<std::int32_t> quantizable_input_indices;
bool is_weight_only = false;
bool no_side_effect = true;
};
using BuiltinOperatorSet = absl::flat_hash_set<tflite::BuiltinOperator>;
// Map from custom op code to custom op quantization information.
using CustomOpMap = std::unordered_map<std::string, CustomOpInfo>;
// Applies dynamic range quantization for the given model wehre the input_model
// type is flatbuffer but is converted to MLIR during quantization process and
// then converted back to flatbuffer for return. Note that this is part of
// reaching feature parity with the old quantizer for dynamic range
// quantization, specifically for
// third_party/tensorflow/lite/tools/optimize/quantize_weights.h.
// TODO(b/202468183): Selective quantization + quant debugger support for
// dynamic range quantization for verify_numeric and whole_model_verify flags.
absl::Status QuantizeWeights(
flatbuffers::FlatBufferBuilder* builder, const tflite::Model* input_model,
const tflite::TensorType& inference_type,
const absl::flat_hash_set<std::string>& denylisted_ops,
const CustomOpMap& custom_op_map,
int64_t minimum_elements_for_weights = 1024,
bool disable_per_channel = false, bool weight_only_quantization = false,
bool legacy_float_scale = false);
// Overloading methods to support old quantizer versions API
absl::Status QuantizeWeights(flatbuffers::FlatBufferBuilder* builder,
const tflite::Model* input_model,
int64_t weights_min_num_elements,
bool use_hybrid_evaluation = true);
absl::Status QuantizeWeights(flatbuffers::FlatBufferBuilder* builder,
const tflite::Model* input_model,
BufferType quant_type = BufferType::QUANTIZED_INT8,
bool use_updated_hybrid_scheme = true);
absl::Status QuantizeWeights(flatbuffers::FlatBufferBuilder* builder,
const tflite::Model* input_model,
int64_t weights_min_num_elements,
const CustomOpMap& custom_op_map,
bool use_updated_hybrid_scheme = true,
const BuiltinOperatorSet& op_denylist = {});
} // namespace lite
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_LITE_QUANTIZE_WEIGHTS_H_
tensorflow/compiler/mlir/lite/quantization/lite/quantize_weights_test.cc
+783
View File
@@ -0,0 +1,783 @@
/* Copyright 2022 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/lite/quantize_weights.h"
#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <cstdlib>
#include <iostream>
#include <memory>
#include <string>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "flatbuffers/buffer.h" // from @flatbuffers
#include "flatbuffers/flatbuffer_builder.h" // from @flatbuffers
#include "flatbuffers/vector.h" // from @flatbuffers
#include "tensorflow/compiler/mlir/lite/core/absl_error_model_builder.h"
#include "tensorflow/compiler/mlir/lite/quantization/lite/test_util.h"
#include "tensorflow/compiler/mlir/lite/schema/schema_generated.h"
#include "tensorflow/compiler/mlir/lite/schema/schema_utils.h"
#include "xla/tsl/platform/logging.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/util/command_line_flags.h"
// Note: branched from tensorflow/lite/tools/optimize/quantize_weights_test.cc
namespace {
std::string* g_test_model_dir = nullptr;
} // namespace
namespace tflite {
namespace optimize {
namespace {
using mlir::lite::BufferType;
using mlir::lite::CustomOpMap;
using mlir::lite::QuantizeWeights;
using mlir::TFL::FlatBufferModelAbslError;
constexpr bool kUseUpdatedHybridSchemeDefault = true;
std::unique_ptr<ModelT> CreateMutableModelFromFile(const Model* input_model) {
auto copied_model = std::make_unique<ModelT>();
input_model->UnPackTo(copied_model.get(), nullptr);
return copied_model;
}
std::unique_ptr<FlatBufferModelAbslError> ReadTestModel() {
auto model_path = tensorflow::io::JoinPath(
*g_test_model_dir, ::mlir::lite::internal::kConvModelWith0Plus10Weights);
return FlatBufferModelAbslError::BuildFromFile(model_path.c_str());
}
std::unique_ptr<FlatBufferModelAbslError> ReadSharedWeightsTestModel() {
auto model_path = tensorflow::io::JoinPath(
*g_test_model_dir, ::mlir::lite::internal::kModelWithSharedWeights);
return FlatBufferModelAbslError::BuildFromFile(model_path.c_str());
}
std::unique_ptr<FlatBufferModelAbslError> ReadGatherTestModel() {
auto model_path = tensorflow::io::JoinPath(
*g_test_model_dir, ::mlir::lite::internal::kQuantizedWithGather);
return FlatBufferModelAbslError::BuildFromFile(model_path.c_str());
}
std::unique_ptr<FlatBufferModelAbslError> ReadCustomOpTestModel() {
auto model_path = tensorflow::io::JoinPath(
*g_test_model_dir, ::mlir::lite::internal::kModelWithCustomOp);
return FlatBufferModelAbslError::BuildFromFile(model_path.c_str());
}
template <typename T>
std::vector<T> GetAsVector(const flatbuffers::Vector<T>* vec) {
return std::vector<T>(vec->begin(), vec->end());
}
class QuantizeWeightsTest : public testing::Test {
protected:
QuantizeWeightsTest() = default;
void LoadBasicModel() {
input_model_ = ReadTestModel();
model_ = input_model_->GetModel();
}
void LoadSharedWeightsModel() {
input_model_ = ReadSharedWeightsTestModel();
model_ = input_model_->GetModel();
}
void LoadGatherTestModel() {
input_model_ = ReadGatherTestModel();
model_ = input_model_->GetModel();
}
void LoadCustomOpTestModel() {
input_model_ = ReadCustomOpTestModel();
model_ = input_model_->GetModel();
}
std::unique_ptr<FlatBufferModelAbslError> input_model_;
const Model* model_;
bool IsModelInputOrOutput(const Model* model, uint32_t tensor_idx) {
for (size_t subgraph_idx = 0; subgraph_idx < model_->subgraphs()->size();
++subgraph_idx) {
const auto subgraph = model->subgraphs()->Get(subgraph_idx);
for (size_t i = 0; i < subgraph->inputs()->size(); ++i) {
if (subgraph->inputs()->Get(i) == tensor_idx) {
return true;
}
}
for (size_t i = 0; i < subgraph->outputs()->size(); ++i) {
if (subgraph->outputs()->Get(i) == tensor_idx) {
return true;
}
}
}
return false;
}
// Returns the producer op code of the specified tensor_idx.
bool GetProducerOpCode(const Model* model, uint32_t subgraph_idx,
uint32_t tensor_idx, BuiltinOperator* op_code) {
const auto subgraph = model->subgraphs()->Get(subgraph_idx);
for (size_t op_idx = 0; op_idx < subgraph->operators()->size(); ++op_idx) {
const auto op = subgraph->operators()->Get(op_idx);
for (size_t i = 0; i < op->outputs()->size(); ++i) {
if (op->outputs()->Get(i) == tensor_idx) {
const uint32_t op_code_idx = op->opcode_index();
*op_code = GetBuiltinCode(model->operator_codes()->Get(op_code_idx));
return true;
}
}
}
return false;
}
};
// Returns true if everything between the two graphs
// are identical except for name field.
// Used when comparing graphs after optimization passes,
// as tensor names may have changed via MLIR constant folding process.
bool ExpectEqualTensor(const Tensor* tensor, const Tensor* expected_tensor) {
return (tensor->is_variable() == expected_tensor->is_variable()) &&
(GetAsVector(tensor->shape()) ==
GetAsVector(expected_tensor->shape()));
}
// Finds the match of the quantized tensor from the possible tensors. Each
// possible tensors can be used only once. It checks shape and name if the
// tensor is quantized and also checks buffer contents and tensor type if not
// quantized. For the quantized case, tensor type and quantizaction params are
// expected to be checked in the test body with the match.
const Tensor* FindMatchingExpectedTensor(
const Model* quantized_model, const Model* expected_model,
const Tensor* quantized_tensor,
const flatbuffers::Vector<flatbuffers::Offset<Tensor>>* possible_tensors,
std::vector<int>& used_tensors, bool quantized = false) {
std::unique_ptr<ModelT> quant_model =
CreateMutableModelFromFile(quantized_model);
std::unique_ptr<ModelT> float_model =
CreateMutableModelFromFile(expected_model);
for (int i = 0; i < possible_tensors->size(); i++) {
// Skip if the tensor is already used for match.
auto it = std::find(used_tensors.begin(), used_tensors.end(), i);
if (it != used_tensors.end()) continue;
const Tensor* float_tensor = possible_tensors->Get(i);
LOG(INFO) << quantized_tensor->name()->str() << " "
<< float_tensor->name()->str();
if (ExpectEqualTensor(quantized_tensor, float_tensor)) {
if (quantized && quantized_tensor->name()->str().find("weights")) {
// If tensor is quantized, data type and buffer contents can be
// different between float and quantized tensors. So do those tests
// separately in the test body without checking them here.
used_tensors.push_back(i);
return float_tensor;
} else {
// Otherwise, do additional checks for data type and buffer contents.
const std::vector<uint8_t> quantized_buffer =
quant_model->buffers[quantized_tensor->buffer()].get()->data;
const std::vector<uint8_t> float_buffer =
float_model->buffers[float_tensor->buffer()].get()->data;
if ((quantized_buffer == float_buffer) &&
(quantized_tensor->type() == float_tensor->type())) {
used_tensors.push_back(i);
return float_tensor;
}
}
}
}
return nullptr;
}
TEST_F(QuantizeWeightsTest, QuantizationSucceeds) {
LoadBasicModel();
flatbuffers::FlatBufferBuilder builder;
EXPECT_OK(QuantizeWeights(&builder, model_, 0));
const uint8_t* buffer = builder.GetBufferPointer();
const Model* output_model = GetModel(buffer);
ASSERT_TRUE(output_model);
}
TEST_F(QuantizeWeightsTest, QuantizationFails) {
LoadBasicModel();
flatbuffers::FlatBufferBuilder builder;
EXPECT_EQ(QuantizeWeights(&builder, model_, TensorType_UINT8, {}, {}, 1024),
absl::InternalError("Quantize weights transformation failed."));
}
TEST_F(QuantizeWeightsTest, WeightsMinNumElements) {
LoadBasicModel();
// Make weights_min_size sufficiently large such that no quantization should
// happen, i.e. the original model is the same size as the old one.
flatbuffers::FlatBufferBuilder builder;
const uint64_t kWeightsMinNumElements = 1000000;
EXPECT_EQ(QuantizeWeights(&builder, model_, kWeightsMinNumElements),
absl::InternalError("Quantize weights transformation failed."));
const uint8_t* buffer = builder.GetBufferPointer();
const Model* output_model = GetModel(buffer);
ASSERT_TRUE(output_model);
for (size_t subgraph_idx = 0; subgraph_idx < model_->subgraphs()->size();
subgraph_idx++) {
const auto quantized_graph = output_model->subgraphs()->Get(subgraph_idx);
const auto float_graph = model_->subgraphs()->Get(subgraph_idx);
ASSERT_EQ(quantized_graph->tensors()->size(),
float_graph->tensors()->size());
std::vector<int> used_tensors;
for (size_t i = 0; i < quantized_graph->tensors()->size(); i++) {
const auto quant_tensor = quantized_graph->tensors()->Get(i);
const auto float_tensor = FindMatchingExpectedTensor(
/*quantized_model=*/output_model, /*expected_model=*/model_,
/*quantized_tensor=*/quant_tensor,
/*possible_tensors=*/float_graph->tensors(),
/*used_tensors=*/used_tensors);
EXPECT_NE(float_tensor, nullptr);
}
EXPECT_EQ(used_tensors.size(), quantized_graph->tensors()->size());
}
}
TEST_F(QuantizeWeightsTest, HybridConv) {
LoadBasicModel();
flatbuffers::FlatBufferBuilder builder;
EXPECT_OK(QuantizeWeights(&builder, model_, 0));
const uint8_t* buffer = builder.GetBufferPointer();
const Model* output_model = GetModel(buffer);
ASSERT_TRUE(output_model);
// Nothing should change.
ASSERT_EQ(output_model->subgraphs()->size(), model_->subgraphs()->size());
for (size_t subgraph_idx = 0; subgraph_idx < model_->subgraphs()->size();
subgraph_idx++) {
const auto quantized_graph = output_model->subgraphs()->Get(subgraph_idx);
const auto float_graph = model_->subgraphs()->Get(subgraph_idx);
ASSERT_EQ(quantized_graph->tensors()->size(),
float_graph->tensors()->size());
// Make sure the graph only has one Conv operation.
ASSERT_EQ(quantized_graph->operators()->size(), 1);
const auto op = quantized_graph->operators()->Get(0);
const uint32_t op_code_idx = op->opcode_index();
ASSERT_EQ(GetBuiltinCode(output_model->operator_codes()->Get(op_code_idx)),
BuiltinOperator_CONV_2D);
std::vector<int> used_tensors;
for (size_t i = 0; i < quantized_graph->tensors()->size(); i++) {
const auto quant_tensor = quantized_graph->tensors()->Get(i);
const auto float_tensor = FindMatchingExpectedTensor(
/*quantized_model=*/output_model, /*expected_model=*/model_,
/*quantized_tensor=*/quant_tensor,
/*possible_tensors=*/float_graph->tensors(),
/*used_tensors=*/used_tensors, /*quantized=*/true);
EXPECT_NE(float_tensor, nullptr);
// If the tensor is a weight, it should have type INT8, otherwise it
// should stay with type FLOAT32.
// If the tensor is a bias, it should have type FLOAT32.
//
// Check with float_tensor name since quantized tensor
// may be renamed.
if (float_tensor->name()->str() == "conv_bias") {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
} else if (IsModelInputOrOutput(output_model, i)) {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
} else if (quant_tensor->buffer() != 0) {
EXPECT_EQ(quant_tensor->type(), TensorType_INT8)
<< quant_tensor->name()->str();
auto shape = GetAsVector(quant_tensor->shape());
if (kUseUpdatedHybridSchemeDefault) {
EXPECT_EQ(quant_tensor->quantization()->scale()->size(), shape[0]);
} else {
EXPECT_EQ(quant_tensor->quantization()->scale()->size(), 1);
}
} else {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
}
}
EXPECT_EQ(used_tensors.size(), quantized_graph->tensors()->size());
}
}
TEST_F(QuantizeWeightsTest, DequantizeConv) {
LoadBasicModel();
flatbuffers::FlatBufferBuilder builder;
EXPECT_OK(QuantizeWeights(&builder, model_, 0,
/*use_hybrid_evaluation=*/false));
const uint8_t* buffer = builder.GetBufferPointer();
const Model* output_model = GetModel(buffer);
ASSERT_TRUE(output_model);
ASSERT_EQ(output_model->subgraphs()->size(), model_->subgraphs()->size());
for (size_t subgraph_idx = 0; subgraph_idx < model_->subgraphs()->size();
++subgraph_idx) {
const auto quantized_graph = output_model->subgraphs()->Get(subgraph_idx);
const auto float_graph = model_->subgraphs()->Get(subgraph_idx);
// The output graph should have an extra tensor from the added dequantize
// op.
ASSERT_EQ(quantized_graph->tensors()->size(),
float_graph->tensors()->size() + 1);
// Check that a dequantize op exists.
int32_t dequant_input_idx = -1;
int32_t dequant_output_idx = -1;
for (size_t i = 0; i < quantized_graph->operators()->size(); ++i) {
const auto op = quantized_graph->operators()->Get(i);
const uint32_t op_code_idx = op->opcode_index();
if (GetBuiltinCode(output_model->operator_codes()->Get(op_code_idx)) ==
BuiltinOperator_DEQUANTIZE) {
dequant_input_idx = op->inputs()->Get(0);
dequant_output_idx = op->outputs()->Get(0);
}
}
ASSERT_GT(dequant_input_idx, -1);
ASSERT_GT(dequant_output_idx, -1);
std::vector<int> used_tensors;
for (size_t i = 0; i < quantized_graph->tensors()->size(); ++i) {
const auto quant_tensor = quantized_graph->tensors()->Get(i);
const auto float_tensor = FindMatchingExpectedTensor(
/*quantized_model=*/output_model, /*expected_model=*/model_,
/*quantized_tensor=*/quant_tensor,
/*possible_tensors=*/float_graph->tensors(),
/*used_tensors=*/used_tensors, /*quantized=*/true);
// If the tensor is a weight, it should have type INT8.
// If the tensor is a bias, it should have type FLOAT32.
// If the tensor is an input or output it should have type FLOAT32.
// The input to dequantize should be INT8, and all other tensors should be
// FLOAT32.
if (i == dequant_input_idx) {
EXPECT_EQ(quant_tensor->type(), TensorType_INT8);
} else if (i == dequant_output_idx) {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
} else if (IsModelInputOrOutput(output_model, i)) {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
} else if (float_tensor != nullptr &&
float_tensor->name()->str() == "conv_bias") {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
} else if (quant_tensor->buffer() != 0) {
// If it's a non-bias constant tensor, it must be the weight.
EXPECT_EQ(quant_tensor->type(), TensorType_INT8);
} else {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
}
}
}
}
TEST_F(QuantizeWeightsTest, DequantizeConvFloat16) {
LoadBasicModel();
flatbuffers::FlatBufferBuilder builder;
EXPECT_OK(QuantizeWeights(&builder, model_, BufferType::QUANTIZED_FLOAT16));
const uint8_t* buffer = builder.GetBufferPointer();
const Model* output_model = GetModel(buffer);
ASSERT_TRUE(output_model);
ASSERT_EQ(output_model->subgraphs()->size(), model_->subgraphs()->size());
for (size_t subgraph_idx = 0; subgraph_idx < model_->subgraphs()->size();
++subgraph_idx) {
const auto quantized_graph = output_model->subgraphs()->Get(subgraph_idx);
const auto float_graph = model_->subgraphs()->Get(subgraph_idx);
// The output graph should have two extra tensors from the added dequantize
// op.
ASSERT_EQ(quantized_graph->tensors()->size(),
float_graph->tensors()->size() + 2);
// Check that a dequantize op exists.
int32_t dequant_input_idx = -1;
int32_t dequant_output_idx = -1;
for (size_t i = 0; i < quantized_graph->operators()->size(); ++i) {
const auto op = quantized_graph->operators()->Get(i);
const uint32_t op_code_idx = op->opcode_index();
if (GetBuiltinCode(output_model->operator_codes()->Get(op_code_idx)) ==
BuiltinOperator_DEQUANTIZE) {
dequant_input_idx = op->inputs()->Get(0);
dequant_output_idx = op->outputs()->Get(0);
}
}
ASSERT_GT(dequant_input_idx, -1);
ASSERT_GT(dequant_output_idx, -1);
for (size_t i = 0; i < quantized_graph->tensors()->size(); ++i) {
const auto quant_tensor = quantized_graph->tensors()->Get(i);
// If the tensor is a weight, it should have type FLOAT16.
// If the tensor is a bias, it should have type FLOAT16.
// If the tensor is an input or output it should have type FLOAT32.
// The input to dequantize should be FLOAT16, and all other tensors should
// be FLOAT32.
if (i == dequant_input_idx) {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT16);
} else if (i == dequant_output_idx) {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
} else if (IsModelInputOrOutput(output_model, i)) {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
} else if (quant_tensor->name()->str() == "conv_bias") {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT16);
} else if ((!CreateMutableModelFromFile(output_model)
->buffers[quant_tensor->buffer()]
.get()
->data.empty())) {
// If it's a non-bias constant tensor, it must be the weight.
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT16);
} else {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
}
}
}
}
TEST_F(QuantizeWeightsTest, SharedWeights_Hybrid) {
LoadSharedWeightsModel();
flatbuffers::FlatBufferBuilder builder;
EXPECT_OK(QuantizeWeights(&builder, model_, 0));
const uint8_t* buffer = builder.GetBufferPointer();
const Model* output_model = GetModel(buffer);
ASSERT_TRUE(output_model);
ASSERT_EQ(output_model->subgraphs()->size(), model_->subgraphs()->size());
uint32_t num_conv_ops = 0;
for (size_t subgraph_idx = 0; subgraph_idx < model_->subgraphs()->size();
++subgraph_idx) {
const auto quantized_graph = output_model->subgraphs()->Get(subgraph_idx);
for (size_t i = 0; i < quantized_graph->operators()->size(); ++i) {
const auto op = quantized_graph->operators()->Get(i);
const uint32_t op_code_idx = op->opcode_index();
const auto op_code =
GetBuiltinCode(output_model->operator_codes()->Get(op_code_idx));
if (op_code == BuiltinOperator_CONV_2D) {
num_conv_ops++;
// Ensure that each convolution's weights tensor is now INT8.
const auto weights_tensor =
quantized_graph->tensors()->Get(op->inputs()->Get(1));
EXPECT_EQ(weights_tensor->type(), TensorType_INT8);
}
}
}
// Ensure that there were exactly two convolutions in the model.
EXPECT_EQ(num_conv_ops, 2);
}
TEST_F(QuantizeWeightsTest, SharedWeights_Dequantize) {
LoadSharedWeightsModel();
flatbuffers::FlatBufferBuilder builder;
EXPECT_OK(
QuantizeWeights(&builder, model_, 0, /*use_hybrid_evaluation=*/false));
const uint8_t* buffer = builder.GetBufferPointer();
const Model* output_model = GetModel(buffer);
ASSERT_TRUE(output_model);
ASSERT_EQ(output_model->subgraphs()->size(), model_->subgraphs()->size());
uint32_t num_conv_ops = 0;
for (size_t subgraph_idx = 0; subgraph_idx < model_->subgraphs()->size();
++subgraph_idx) {
const auto quantized_graph = output_model->subgraphs()->Get(subgraph_idx);
for (size_t i = 0; i < quantized_graph->operators()->size(); ++i) {
const auto op = quantized_graph->operators()->Get(i);
const uint32_t op_code_idx = op->opcode_index();
const auto op_code =
GetBuiltinCode(output_model->operator_codes()->Get(op_code_idx));
if (op_code == BuiltinOperator_CONV_2D) {
num_conv_ops++;
// Ensure that each convolution's weights tensor is still FLOAT
// (the output of the dequantize).
uint32_t weights_tensor_index = op->inputs()->Get(1);
const auto weights_tensor =
quantized_graph->tensors()->Get(weights_tensor_index);
EXPECT_EQ(weights_tensor->type(), TensorType_FLOAT32);
// Check that it comes from a dequantize operation.
BuiltinOperator producer_op_code;
ASSERT_TRUE(GetProducerOpCode(output_model, subgraph_idx,
weights_tensor_index, &producer_op_code));
EXPECT_EQ(producer_op_code, BuiltinOperator_DEQUANTIZE);
}
}
}
// Ensure that there were exactly two convolutions in the model.
EXPECT_EQ(num_conv_ops, 2);
}
TEST_F(QuantizeWeightsTest, VerifyGatherQuantization) {
LoadGatherTestModel();
flatbuffers::FlatBufferBuilder builder;
EXPECT_OK(QuantizeWeights(&builder, model_, 0));
const uint8_t* buffer = builder.GetBufferPointer();
const Model* output_model = GetModel(buffer);
ASSERT_TRUE(output_model);
ASSERT_EQ(output_model->subgraphs()->size(), model_->subgraphs()->size());
for (size_t subgraph_idx = 0; subgraph_idx < model_->subgraphs()->size();
++subgraph_idx) {
const auto quantized_graph = output_model->subgraphs()->Get(subgraph_idx);
for (size_t i = 0; i < quantized_graph->operators()->size(); ++i) {
const auto op = quantized_graph->operators()->Get(i);
const uint32_t op_code_idx = op->opcode_index();
const auto op_code =
GetBuiltinCode(output_model->operator_codes()->Get(op_code_idx));
if (op_code == BuiltinOperator_GATHER) {
uint32_t input_tensor_index = op->inputs()->Get(0);
const auto weights_tensor =
quantized_graph->tensors()->Get(input_tensor_index);
EXPECT_EQ(weights_tensor->type(), TensorType_INT8);
}
}
}
}
TEST_F(QuantizeWeightsTest, VerifyCustomOpQuantizationDequantize) {
LoadCustomOpTestModel();
// The custom op is not hybrid, and the second input is a constant that can
// be quantized.
CustomOpMap custom_op_map;
custom_op_map["CustomTestOp"] = {
{1}, // quantizable_input_indices
true, // is_weight_only
};
flatbuffers::FlatBufferBuilder builder;
EXPECT_OK(QuantizeWeights(&builder, model_, 0, custom_op_map));
const uint8_t* buffer = builder.GetBufferPointer();
const Model* output_model = GetModel(buffer);
ASSERT_TRUE(output_model);
ASSERT_EQ(output_model->subgraphs()->size(), model_->subgraphs()->size());
const auto quantized_graph = output_model->subgraphs()->Get(0);
// A dequantize op should be added.
ASSERT_EQ(quantized_graph->operators()->size(),
model_->subgraphs()->Get(0)->operators()->size() + 1);
int num_custom_ops_found = 0;
for (size_t i = 0; i < quantized_graph->operators()->size(); ++i) {
const auto op = quantized_graph->operators()->Get(i);
const uint32_t op_code_idx = op->opcode_index();
const auto op_code =
GetBuiltinCode(output_model->operator_codes()->Get(op_code_idx));
if (op_code == BuiltinOperator_CUSTOM) {
uint32_t weights_tensor_index = op->inputs()->Get(1);
const auto weights_tensor =
quantized_graph->tensors()->Get(weights_tensor_index);
EXPECT_EQ(weights_tensor->type(), TensorType_FLOAT32);
// Check that it comes from a dequantize operation.
BuiltinOperator producer_op_code;
ASSERT_TRUE(GetProducerOpCode(output_model, 0, weights_tensor_index,
&producer_op_code));
EXPECT_EQ(producer_op_code, BuiltinOperator_DEQUANTIZE);
num_custom_ops_found++;
}
}
EXPECT_EQ(num_custom_ops_found, 1);
}
TEST_F(QuantizeWeightsTest, VerifyCustomOpQuantizationHybrid) {
LoadCustomOpTestModel();
// The custom op is dynamic range quantizable, and the second input is a
// constant that can be quantized.
CustomOpMap custom_op_map;
custom_op_map["CustomTestOp"] = {
{1}, // quantizable_input_indices
false, // is_weight_only
};
flatbuffers::FlatBufferBuilder builder;
ASSERT_OK(QuantizeWeights(&builder, model_, 0, custom_op_map));
const uint8_t* buffer = builder.GetBufferPointer();
const Model* output_model = GetModel(buffer);
ASSERT_TRUE(output_model);
ASSERT_EQ(output_model->subgraphs()->size(), model_->subgraphs()->size());
const auto quantized_graph = output_model->subgraphs()->Get(0);
ASSERT_EQ(quantized_graph->operators()->size(),
model_->subgraphs()->Get(0)->operators()->size());
int num_custom_ops_found = 0;
for (size_t i = 0; i < quantized_graph->operators()->size(); ++i) {
const auto op = quantized_graph->operators()->Get(i);
const uint32_t op_code_idx = op->opcode_index();
const auto op_code =
GetBuiltinCode(output_model->operator_codes()->Get(op_code_idx));
if (op_code == BuiltinOperator_CUSTOM) {
uint32_t weights_tensor_index = op->inputs()->Get(1);
const auto weights_tensor =
quantized_graph->tensors()->Get(weights_tensor_index);
EXPECT_EQ(weights_tensor->type(), TensorType_INT8);
num_custom_ops_found++;
}
}
EXPECT_EQ(num_custom_ops_found, 1);
}
TEST_F(QuantizeWeightsTest, VerifyUpdatedHybridSchemeFalseQuantizationHybrid) {
LoadBasicModel();
flatbuffers::FlatBufferBuilder builder;
const CustomOpMap custom_op_map;
EXPECT_OK(QuantizeWeights(&builder, model_, 0, custom_op_map, false));
const uint8_t* buffer = builder.GetBufferPointer();
const Model* output_model = GetModel(buffer);
ASSERT_TRUE(output_model);
// Nothing should change.
ASSERT_EQ(output_model->subgraphs()->size(), model_->subgraphs()->size());
for (size_t subgraph_idx = 0; subgraph_idx < model_->subgraphs()->size();
subgraph_idx++) {
const auto quantized_graph = output_model->subgraphs()->Get(subgraph_idx);
const auto float_graph = model_->subgraphs()->Get(subgraph_idx);
ASSERT_EQ(quantized_graph->tensors()->size(),
float_graph->tensors()->size());
// Make sure the graph only has one Conv operation.
ASSERT_EQ(quantized_graph->operators()->size(), 1);
const auto op = quantized_graph->operators()->Get(0);
const uint32_t op_code_idx = op->opcode_index();
ASSERT_EQ(GetBuiltinCode(output_model->operator_codes()->Get(op_code_idx)),
BuiltinOperator_CONV_2D);
std::vector<int> used_tensors;
for (size_t i = 0; i < quantized_graph->tensors()->size(); i++) {
const auto quant_tensor = quantized_graph->tensors()->Get(i);
const auto float_tensor = FindMatchingExpectedTensor(
/*quantized_model=*/output_model, /*expected_model=*/model_,
/*quantized_tensor=*/quant_tensor,
/*possible_tensors=*/float_graph->tensors(),
/*used_tensors=*/used_tensors, /*quantized=*/true);
EXPECT_NE(float_tensor, nullptr);
// If the tensor is a weight, it should have type INT8, otherwise it
// should stay with type FLOAT32.
// If the tensor is a bias, it should have type FLOAT32.
if (float_tensor->name()->str() == "conv_bias") {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
} else if (IsModelInputOrOutput(output_model, i)) {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
} else if ((!CreateMutableModelFromFile(output_model)
->buffers[quant_tensor->buffer()]
.get()
->data.empty())) {
EXPECT_EQ(quant_tensor->type(), TensorType_INT8)
<< quant_tensor->name()->str();
auto shape = GetAsVector(quant_tensor->shape());
EXPECT_EQ(quant_tensor->quantization()->scale()->size(), 1);
} else {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
}
}
EXPECT_EQ(used_tensors.size(), quantized_graph->tensors()->size());
}
}
TEST_F(QuantizeWeightsTest, DequantizeConvBlocklisted) {
LoadBasicModel();
flatbuffers::FlatBufferBuilder builder;
const CustomOpMap custom_op_map;
EXPECT_OK(QuantizeWeights(&builder, model_, 0, custom_op_map,
/*use_updated_hybrid_scheme=*/true,
{BuiltinOperator_CONV_2D}));
const uint8_t* buffer = builder.GetBufferPointer();
const Model* output_model = GetModel(buffer);
ASSERT_TRUE(output_model);
ASSERT_EQ(output_model->subgraphs()->size(), model_->subgraphs()->size());
for (size_t subgraph_idx = 0; subgraph_idx < model_->subgraphs()->size();
++subgraph_idx) {
const auto quantized_graph = output_model->subgraphs()->Get(subgraph_idx);
const auto float_graph = model_->subgraphs()->Get(subgraph_idx);
// The output graph should have an extra tensor from the added dequantize
// op.
ASSERT_EQ(quantized_graph->tensors()->size(),
float_graph->tensors()->size() + 1);
// Check that a dequantize op exists.
int32_t dequant_input_idx = -1;
int32_t dequant_output_idx = -1;
for (size_t i = 0; i < quantized_graph->operators()->size(); ++i) {
const auto op = quantized_graph->operators()->Get(i);
const uint32_t op_code_idx = op->opcode_index();
if (GetBuiltinCode(output_model->operator_codes()->Get(op_code_idx)) ==
BuiltinOperator_DEQUANTIZE) {
dequant_input_idx = op->inputs()->Get(0);
dequant_output_idx = op->outputs()->Get(0);
}
}
ASSERT_GT(dequant_input_idx, -1);
ASSERT_GT(dequant_output_idx, -1);
std::vector<int> used_tensors;
for (size_t i = 0; i < quantized_graph->tensors()->size(); ++i) {
const auto quant_tensor = quantized_graph->tensors()->Get(i);
const auto float_tensor = FindMatchingExpectedTensor(
/*quantized_model=*/output_model, /*expected_model=*/model_,
/*quantized_tensor=*/quant_tensor,
/*possible_tensors=*/float_graph->tensors(),
/*used_tensors=*/used_tensors);
// If the tensor is a weight, it should have type INT8.
// If the tensor is a bias, it should have type FLOAT32.
// If the tensor is an input or output it should have type FLOAT32.
// The input to dequantize should be INT8, and all other tensors should be
// FLOAT32.
if (i == dequant_input_idx) {
EXPECT_EQ(quant_tensor->type(), TensorType_INT8);
// The dequantize should still be quantized per-channel
EXPECT_EQ(quant_tensor->quantization()->scale()->size(), 5);
EXPECT_EQ(quant_tensor->quantization()->quantized_dimension(), 0);
} else if (i == dequant_output_idx) {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
} else if (IsModelInputOrOutput(output_model, i)) {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
} else if (float_tensor != nullptr &&
float_tensor->name()->str() == "conv_bias") {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
} else if ((!CreateMutableModelFromFile(output_model)
->buffers[quant_tensor->buffer()]
.get()
->data.empty())) {
// If it's a non-bias constant tensor, it must be the weight.
EXPECT_EQ(quant_tensor->type(), TensorType_INT8);
} else {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
}
}
}
}
} // namespace
} // namespace optimize
} // namespace tflite
int main(int argc, char** argv) {
std::string model_file;
const std::vector<tensorflow::Flag> flag_list = {
tensorflow::Flag("test_model_file", &model_file,
"Path to test tflite model file."),
};
const bool parse_result = tensorflow::Flags::Parse(&argc, argv, flag_list);
if (!parse_result) {
std::cerr << "Required test_model_file\n";
std::abort();
}
g_test_model_dir = new std::string(tensorflow::io::Dirname(model_file));
::tensorflow::port::InitMain(argv[0], &argc, &argv);
return RUN_ALL_TESTS();
}
tensorflow/compiler/mlir/lite/quantization/lite/test_util.cc
+96
View File
@@ -0,0 +1,96 @@
/* Copyright 2018 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/lite/test_util.h"
#include <cstdarg>
#include <cstdio>
#include <gtest/gtest.h>
namespace mlir {
namespace lite {
namespace internal {
const char* kConvModelWithMinus128Plus127Weights =
"single_conv_weights_min_minus_127_max_plus_127.bin";
const char* kConvModelWith0Plus10Weights =
"single_conv_weights_min_0_max_plus_10.bin";
const char* kConvModelWithNoBias = "single_conv_no_bias.bin";
const char* kSingleSoftmaxModelMinMinus5MaxPlus5 =
"single_softmax_min_minus_5_max_plus_5.bin";
const char* kSingleAvgPoolModelMinMinus5MaxPlus5 =
"single_avg_pool_min_minus_5_max_plus_5.bin";
const char* kModelWithSharedWeights = "weight_shared_between_convs.bin";
const char* kMultiInputAddWithReshape = "multi_input_add_reshape.bin";
const char* kQuantizedWithGather = "quantized_with_gather.bin";
const char* kConstInputAddModel = "add_with_const_input.bin";
const char* kFloatConcatMax5Max10Max10 = "concat.bin";
const char* kModelWithBroadcastToOp = "broadcast_to.bin";
const char* kModelWithCustomOp = "custom_op.bin";
const char* kModelWithArgMaxOp = "argmax.bin";
const char* kModelWithFCOp = "fc.bin";
const char* kModelWithGatherNDOp = "gather_nd.bin";
const char* kModelWithWhereOp = "where.bin";
const char* kModelMixed = "mixed.bin";
const char* kModelMixed16x8 = "mixed16x8.bin";
const char* kModelSplit = "split.bin";
const char* kModelPack = "pack.bin";
const char* kLstmCalibrated = "lstm_calibrated.bin";
const char* kLstmQuantized = "lstm_quantized.bin";
const char* kUnidirectionalSequenceLstmCalibrated =
"unidirectional_sequence_lstm_calibrated.bin";
const char* kUnidirectionalSequenceLstmQuantized =
"unidirectional_sequence_lstm_quantized.bin";
const char* kModelWithMinimumOp = "minimum.bin";
const char* kModelWithMaximumOp = "maximum.bin";
const char* kLstmCalibrated2 = "lstm_calibrated2.bin";
const char* kLstmQuantized2 = "lstm_quantized2.bin";
const char* kModelWithTranspose = "transpose.bin";
const char* kSvdfCalibrated = "svdf_calibrated.bin";
const char* kSvdfQuantized = "svdf_quantized.bin";
const char* kModelWithUnpack = "unpack.bin";
const char* kQatModelWithFc = "fc_qat.bin";
const char* kModelWithResourceVarsCalibrated = "resource_vars_calibrated.bin";
int FailOnErrorReporter::Report(const char* format, va_list args) {
char buf[1024];
vsnprintf(buf, sizeof(buf), format, args);
EXPECT_TRUE(false) << "Error happened: " << buf;
return 0;
}
} // namespace internal
} // namespace lite
} // namespace mlir
tensorflow/compiler/mlir/lite/quantization/lite/test_util.h
+145
View File
@@ -0,0 +1,145 @@
/* Copyright 2018 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_LITE_TEST_UTIL_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_LITE_TEST_UTIL_H_
#include <cstdarg>
#include "tensorflow/compiler/mlir/lite/core/api/error_reporter.h"
namespace mlir {
namespace lite {
namespace internal {
// Test model with a single convolution.
// Floating point weights of the model are all integers and lie in
// range[-127, 127]. The weights have been put in such a way that each
// channel has at least one weight as -127 and one weight as 127.
// The activations are all in range: [-128, 127]
// This means all bias computations should result in 1.0 scale.
extern const char* kConvModelWithMinus128Plus127Weights;
// Test model with single convolution where all weights are integers between
// [0, 10] weights are randomly distributed. It is not guaranteed that min max
// for weights are going to appear in each channel.
// Activations have min = 0, max = 10.
extern const char* kConvModelWith0Plus10Weights;
// Test model where no bias is in the conv.
extern const char* kConvModelWithNoBias;
// A floating point model with a single softmax. The input tensor has min
// and max in range [-5, 5], not necessarily -5 or +5.
extern const char* kSingleSoftmaxModelMinMinus5MaxPlus5;
// A floating point model with a single average pool. The input tensor has min
// and max in range [-5, 5], not necessarily -5 or +5.
extern const char* kSingleAvgPoolModelMinMinus5MaxPlus5;
// Test model with a weights variable that is shared between a convolution layer
// and an add operation.
extern const char* kModelWithSharedWeights;
// Test model with Add followed by a reshape. Model has 2 inputs for add.
extern const char* kMultiInputAddWithReshape;
// Test gather operation with quantized input.
extern const char* kQuantizedWithGather;
// Test model with a tf.constant input to tf.add. Model has 2 inputs one
// constant and other placeholder.
extern const char* kConstInputAddModel;
// A float test model with concat that has [0, 5] and [0, 10] for inputs and [0,
// 10] as output.
extern const char* kFloatConcatMax5Max10Max10;
// Test model with broadcast_to op.
extern const char* kModelWithBroadcastToOp;
// Test model with a custom op.
extern const char* kModelWithCustomOp;
// Test model with a argmax op.
extern const char* kModelWithArgMaxOp;
// Test model with a fully connected op.
extern const char* kModelWithFCOp;
// Test model with a gather_nd op.
extern const char* kModelWithGatherNDOp;
// Test model with a Where op.
extern const char* kModelWithWhereOp;
// Test model with mixed quantizable and un-quantizable ops.
// reshape->custom->custom->squeeze.
extern const char* kModelMixed;
// Test model with mixed quantizable and
// and un-quantizable ops for
// activations in 16-bit.
extern const char* kModelMixed16x8;
// Test model with split op.
extern const char* kModelSplit;
// Test model with pack op.
extern const char* kModelPack;
// Test model with LSTM op that has layer norm, has projection, without
// peephole, without cifg.
extern const char* kLstmCalibrated;
extern const char* kLstmQuantized;
// Test model with LSTM op that has peephole, without layer norm, without
// projection, without cifg.
extern const char* kLstmCalibrated2;
extern const char* kLstmQuantized2;
extern const char* kUnidirectionalSequenceLstmCalibrated;
extern const char* kUnidirectionalSequenceLstmQuantized;
// Test model with a minimum op.
extern const char* kModelWithMinimumOp;
// Test model with a maximum op.
extern const char* kModelWithMaximumOp;
// Test model with a transpose op.
extern const char* kModelWithTranspose;
// Test model with SVDF op.
extern const char* kSvdfCalibrated;
extern const char* kSvdfQuantized;
// Test model with an unpack op.
extern const char* kModelWithUnpack;
// Test QAT model with fc op.
extern const char* kQatModelWithFc;
// Test calibrated model with resource variables.
extern const char* kModelWithResourceVarsCalibrated;
// An error reporter that fails on testing.
class FailOnErrorReporter : public tflite::ErrorReporter {
public:
int Report(const char* format, va_list args) override;
};
} // namespace internal
} // namespace lite
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_LITE_TEST_UTIL_H_
tensorflow/compiler/mlir/lite/quantization/lite/testdata/README.md
Vendored
+30
View File
@@ -0,0 +1,30 @@
# Test models for testing quantization
This directory contains test models for testing quantization.
## Models
* `single_conv_weights_min_0_max_plus_10.bin` \
A floating point model with single convolution where all weights are
integers between [0, 10] weights are randomly distributed. It is not
guaranteed that min max for weights are going to appear in each channel.
All activations have min maxes and activations are in range [0,10].
* `single_conv_weights_min_minus_127_max_plus_127.bin` \
A floating point model with a single convolution where weights of the model
are all integers that lie in range[-127, 127]. The weights have been put in
such a way that each channel has at least one weight as -127 and one weight
as 127. The activations are all in range: [-128, 127].
This means all bias computations should result in 1.0 scale.
* `single_softmax_min_minus_5_max_5.bin` \
A floating point model with a single softmax. The input tensor has min
and max in range [-5, 5], not necessarily -5 or +5.
* `single_avg_pool_input_min_minus_5_max_5.bin` \
A floating point model with a single average pool. The input tensor has min
and max in range [-5, 5], not necessarily -5 or +5.
* `weight_shared_between_convs.bin` \
A floating point model with two convs that have a use the same weight tensor.
* `multi_input_add_reshape.bin` \
A floating point model with two inputs with an add followed by a reshape.
* `quantized_with_gather.bin` \
A floating point model with an input with a gather, modeling a situation
of mapping categorical input to embeddings.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/add_with_const_input.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/argmax.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/broadcast_to.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/concat.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/custom_op.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/fc.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/fc_qat.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/gather_nd.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/lstm_calibrated.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/lstm_calibrated2.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/lstm_quantized.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/lstm_quantized2.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/maximum.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/minimum.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/mixed.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/mixed16x8.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/multi_input_add_reshape.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/pack.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/quantized_with_gather.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/resource_vars_calibrated.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/single_avg_pool_min_minus_5_max_plus_5.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/single_conv_no_bias.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/single_conv_weights_min_0_max_plus_10.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/single_conv_weights_min_minus_127_max_plus_127.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/single_softmax_min_minus_5_max_plus_5.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/split.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/svdf_calibrated.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/svdf_quantized.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/transpose.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/unidirectional_sequence_lstm_calibrated.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/unidirectional_sequence_lstm_quantized.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/unpack.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/weight_shared_between_convs.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/testdata/where.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/compiler/mlir/lite/quantization/lite/tfl_quantizer.cc
+71
View File
@@ -0,0 +1,71 @@
/* 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 <iostream>
#include <string>
#include <system_error>
#include "absl/status/status.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/InitLLVM.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/PrettyStackTrace.h"
#include "llvm/Support/raw_ostream.h"
#include "tensorflow/compiler/mlir/lite/quantization/lite/quantize_model.h"
#include "tensorflow/compiler/mlir/lite/schema/schema_generated.h"
using llvm::cl::opt;
// NOLINTNEXTLINE
static opt<std::string> inputFileName(llvm::cl::Positional,
llvm::cl::desc("<input file>"),
llvm::cl::init("-"));
namespace mlir {
namespace {
absl::Status QuantizeAnnotatedModel(llvm::StringRef buffer,
std::string& output_buffer) {
return mlir::lite::QuantizeModel(
buffer, tflite::TensorType_INT8, tflite::TensorType_INT8,
tflite::TensorType_INT8, {}, /*disable_per_channel=*/false,
/*fully_quantize=*/true, output_buffer);
}
} // namespace
} // namespace mlir
int main(int argc, char** argv) {
llvm::InitLLVM y(argc, argv);
llvm::cl::ParseCommandLineOptions(argc, argv);
auto file_or_err = llvm::MemoryBuffer::getFileOrSTDIN(inputFileName.c_str());
if (std::error_code error = file_or_err.getError()) {
llvm::errs() << argv[0] << ": could not open input file '" << inputFileName
<< "': " << error.message() << "\n";
return 1;
}
auto buffer = file_or_err->get();
std::string output_buffer;
if (auto status = mlir::QuantizeAnnotatedModel(buffer->getBuffer().str(),
output_buffer);
!status.ok()) {
llvm::errs() << status.message() << "\n";
return 1;
}
std::cout << output_buffer << "\n";
return 0;
}
tensorflow/compiler/mlir/lite/quantization/lite/tfl_to_std.cc
+86
View File
@@ -0,0 +1,86 @@
/* Copyright 2020 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/lite/tfl_to_std.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/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
#include "tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_utils.h"
#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
#include "tensorflow/compiler/mlir/lite/utils/utils.h"
namespace mlir {
namespace TFL {
void ConvertTFLQuantOpsToMlirQuantOps(func::FuncOp func) {
OpBuilder b(func);
func.walk([&](Operation* op) {
b.setInsertionPoint(op);
if (auto dq = llvm::dyn_cast<DequantizeOp>(op)) {
auto dcast = quantfork::DequantizeCastOp::create(
b, dq.getLoc(), dq.getOutput().getType(), dq.getInput());
dq.getOutput().replaceAllUsesWith(dcast);
dq.erase();
} else if (auto q = llvm::dyn_cast<QuantizeOp>(op)) {
auto qcast = quantfork::QuantizeCastOp::create(
b, q.getLoc(), q.getOutput().getType(), q.getInput());
q.getOutput().replaceAllUsesWith(qcast);
q.erase();
} else if (auto q = llvm::dyn_cast<ConstOp>(op)) {
auto value = q.getValue();
auto type = q.getResult().getType();
if (arith::ConstantOp::isBuildableWith(value, type)) {
auto c = arith::ConstantOp::create(b, q.getLoc(), q.getValue());
q.getOutput().replaceAllUsesWith(c);
q.erase();
} else if (TFL::NoValueOp::isBuildableWith(value, type)) {
auto c = TFL::NoValueOp::create(b, q.getLoc(), type, mlir::UnitAttr());
q.getOutput().replaceAllUsesWith(c);
q.erase();
}
}
});
}
void ConvertMlirQuantOpsToTFLQuantOps(func::FuncOp func) {
OpBuilder b(func);
func.walk([&](Operation* op) {
b.setInsertionPoint(op);
if (auto dq = llvm::dyn_cast<quantfork::DequantizeCastOp>(op)) {
auto dcast = DequantizeOp::create(b, dq.getLoc(),
dq.getResult().getType(), dq.getArg());
dq.getResult().replaceAllUsesWith(dcast);
if (auto extra_attr = op->getAttr(kVolatileOpAttrName)) {
dcast->setAttr(kVolatileOpAttrName, extra_attr);
}
dq.erase();
} else if (auto q = llvm::dyn_cast<quantfork::QuantizeCastOp>(op)) {
auto out_type = q.getResult().getType();
auto qcast = QuantizeOp::create(b, q.getLoc(), out_type, q.getArg(),
TypeAttr::get(out_type));
q.getResult().replaceAllUsesWith(qcast);
if (auto extra_attr = op->getAttr(kVolatileOpAttrName)) {
qcast->setAttr(kVolatileOpAttrName, extra_attr);
}
q.erase();
}
});
}
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/quantization/lite/tfl_to_std.h
+64
View File
@@ -0,0 +1,64 @@
/* Copyright 2020 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_LITE_TFL_TO_STD_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_LITE_TFL_TO_STD_H_
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
namespace mlir {
namespace TFL {
// Converts all the tfl.quantize/tfl.dequantize ops to the ops in the mlir.quant
// dialect ones in the function.
void ConvertTFLQuantOpsToMlirQuantOps(func::FuncOp func);
// Converts all the mlir.quant dialect ops to the tfl.quantize/tfl.dequantize
// ops in the function.
void ConvertMlirQuantOpsToTFLQuantOps(func::FuncOp func);
// A helper class to convert target function to another representation using
// `ConvertForward` function during construction and convert target function
// back to the original representation using `ConvertBackward` function during
// deconstruction.
template <void (*ConvertForward)(func::FuncOp),
void (*ConvertBackward)(func::FuncOp)>
class ScopedOpsConverter {
public:
explicit ScopedOpsConverter(func::FuncOp func) : func_(func) {
ConvertForward(func_);
}
ScopedOpsConverter(const ScopedOpsConverter&) = delete;
ScopedOpsConverter operator=(const ScopedOpsConverter&) = delete;
ScopedOpsConverter(const ScopedOpsConverter&&) = delete;
ScopedOpsConverter operator=(const ScopedOpsConverter&&) = delete;
~ScopedOpsConverter() { ConvertBackward(func_); }
private:
func::FuncOp func_;
};
using ScopedTFLQuantOpsToMlirQuantOpsConverter =
ScopedOpsConverter<ConvertTFLQuantOpsToMlirQuantOps,
ConvertMlirQuantOpsToTFLQuantOps>;
using ScopedMlirQuantOpsToTFLQuantOpsConverter =
ScopedOpsConverter<ConvertMlirQuantOpsToTFLQuantOps,
ConvertTFLQuantOpsToMlirQuantOps>;
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_LITE_TFL_TO_STD_H_
tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/BUILD
+143
View File
@@ -0,0 +1,143 @@
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"],
default_visibility = [
"//visibility:private",
],
licenses = ["notice"],
)
cc_library(
name = "portable_tensor_utils",
srcs = ["portable_tensor_utils.cc"],
hdrs = ["portable_tensor_utils.h"],
visibility = ["//tensorflow/compiler/mlir/lite/quantization/common/quantization_lib:__pkg__"],
)
cc_library(
name = "model_utils",
srcs = ["model_utils.cc"],
hdrs = ["model_utils.h"],
deps = [
"//tensorflow/compiler/mlir/lite/schema:schema_conversion_utils",
"//tensorflow/compiler/mlir/lite/schema:schema_fbs",
"//tensorflow/compiler/mlir/lite/schema:schema_utils",
"@com_google_absl//absl/strings",
],
)
tf_cc_test(
name = "model_utils_test",
srcs = ["model_utils_test.cc"],
deps = [
":model_utils",
"//tensorflow/compiler/mlir/lite/schema:schema_fbs",
"@com_google_googletest//:gtest_main",
],
)
cc_library(
name = "quantization_utils",
srcs = ["quantization_utils.cc"],
hdrs = ["quantization_utils.h"],
deps = [
":model_utils",
":portable_tensor_utils",
"//tensorflow/compiler/mlir/lite/kernels/internal:runtime_shape",
"//tensorflow/compiler/mlir/lite/schema:schema_fbs",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@eigen_archive//:eigen3",
],
)
tf_cc_test(
name = "quantization_utils_test",
srcs = ["quantization_utils_test.cc"],
args = [
"--test_model_file=$(location //tensorflow/compiler/mlir/lite/quantization/lite:testdata/single_conv_weights_min_0_max_plus_10.bin)",
],
data = [
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/single_conv_weights_min_0_max_plus_10.bin",
],
deps = [
":quantization_utils",
"//tensorflow/compiler/mlir/lite/core:absl_error_model_builder",
"//tensorflow/compiler/mlir/lite/quantization/lite:test_util",
"//tensorflow/compiler/mlir/lite/schema:schema_fbs",
"//tensorflow/compiler/mlir/lite/schema:schema_utils",
"@com_google_absl//absl/status",
"@com_google_googletest//:gtest",
"@eigen_archive//:eigen3",
"@tsl//tsl/platform:path",
"@tsl//tsl/platform:platform_port",
"@xla//xla/tsl/lib/core:status_test_util",
"@xla//xla/tsl/util:command_line_flags",
],
)
cc_library(
name = "quantize_weights",
srcs = select({
"//tensorflow:ios": ["quantize_weights_portable.cc"],
"//tensorflow:android": ["quantize_weights_portable.cc"],
"//conditions:default": ["quantize_weights.cc"],
}),
hdrs = ["quantize_weights.h"],
compatible_with = get_compatible_with_portable(),
visibility = ["//visibility:public"],
deps = [
":model_utils",
":portable_tensor_utils",
":quantization_utils",
"//tensorflow/compiler/mlir/lite/schema:schema_fbs",
"//tensorflow/compiler/mlir/lite/schema:schema_utils",
"//tensorflow/core/platform:logging",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/log",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
"@flatbuffers//:runtime_cc",
] + select({
"//tensorflow:android": [],
"//tensorflow:ios": [],
"//conditions:default": [
"//tensorflow/compiler/mlir/lite/quantization/lite:quantize_weights",
],
}),
)
tf_cc_test(
name = "quantize_weights_test",
srcs = ["quantize_weights_test.cc"],
args = [
"--test_model_file=$(location //tensorflow/compiler/mlir/lite/quantization/lite:testdata/single_conv_weights_min_0_max_plus_10.bin)",
],
data = [
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/custom_op.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/quantized_with_gather.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/single_conv_weights_min_0_max_plus_10.bin",
"//tensorflow/compiler/mlir/lite/quantization/lite:testdata/weight_shared_between_convs.bin",
],
tags = [
"tflite_not_portable_android",
"tflite_not_portable_ios",
],
deps = [
":quantize_weights",
"//tensorflow/compiler/mlir/lite/core:absl_error_model_builder",
"//tensorflow/compiler/mlir/lite/quantization/lite:test_util",
"//tensorflow/compiler/mlir/lite/schema:schema_fbs",
"//tensorflow/compiler/mlir/lite/schema:schema_utils",
"//tensorflow/core:framework_internal",
"@com_google_googletest//:gtest",
"@flatbuffers",
"@tsl//tsl/platform:path",
"@tsl//tsl/platform:platform_port",
"@xla//xla/tsl/util:command_line_flags",
],
)
tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/README
+7
View File
@@ -0,0 +1,7 @@
This folder organizes legacy toco quantize_weights and its dependencies for easy
deletion, per go/tfl_split_quantize_weights.
When the MLIR quantizer replaces the legacy TOCO quantizer entirely, this entire
folder should be deleted. The partial copies of utils in this folder should not
be used anywhere but to support the legacy TOCO quantizer in this folder for
a clean and easy deletion.
tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/model_utils.cc
+108
View File
@@ -0,0 +1,108 @@
/* 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 is the MLIR copy of part of
// third_party/tensorflow/lite/tools/optimize/model_utils.cc as part of the
// effort to decouple TFLite from MLIR.
#include "tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/model_utils.h"
#include <cstddef>
#include <cstdint>
#include <memory>
#include <string>
#include <vector>
#include "tensorflow/compiler/mlir/lite/schema/schema_conversion_utils.h"
#include "tensorflow/compiler/mlir/lite/schema/schema_generated.h"
#include "tensorflow/compiler/mlir/lite/schema/schema_utils.h"
namespace mlir {
namespace lite {
namespace toco_legacy {
using std::string;
using tflite::BuiltinOperator;
using tflite::BuiltinOperator_DEQUANTIZE;
using tflite::ModelT;
using tflite::OperatorCodeT;
using tflite::OperatorT;
using tflite::TensorT;
using tflite::TensorType;
// LINT.IfChange(GetOrInsertOpCodeIndex)
// Returns the index of the OpCode.
// If a OpCode doesn't exist, adds it and returns its index.
int32_t GetOrInsertOpCodeIndex(ModelT* model, const BuiltinOperator& op_code,
int32_t version) {
for (size_t i = 0; i < model->operator_codes.size(); ++i) {
if (tflite::GetBuiltinCode(model->operator_codes[i].get()) == op_code) {
return i;
}
}
model->operator_codes.push_back(std::make_unique<OperatorCodeT>());
int op_code_idx = model->operator_codes.size() - 1;
model->operator_codes[op_code_idx]->builtin_code = op_code;
model->operator_codes[op_code_idx]->deprecated_builtin_code =
tflite::ConvertBuiltinCodeToDeprecatedBuiltinCode(op_code);
// Version 2 and onwards supports INT8 inputs.
model->operator_codes[op_code_idx]->version = version;
// Return the index of the newly placed OperatorCodeT.
return op_code_idx;
}
// LINT.ThenChange(//tensorflow/lite/tools/optimize/model_utils.cc:GetOrInsertOpCodeIndex)
// LINT.IfChange(MakeDequantizeOperator)
// Creates a Dequantize OperatorT object.
void MakeDequantizeOperator(ModelT* model, std::unique_ptr<OperatorT>* op,
int32_t input, int32_t output) {
OperatorT* op_raw = new OperatorT;
// Version 2 and onwards supports INT8 inputs.
op_raw->opcode_index =
GetOrInsertOpCodeIndex(model, BuiltinOperator_DEQUANTIZE, 2);
op_raw->inputs = {input};
op_raw->outputs = {output};
op->reset(op_raw);
}
// LINT.ThenChange(//tensorflow/lite/tools/optimize/model_utils.cc:MakeDequantizeOperator)
// LINT.IfChange(MakeTensor)
// Create a new TensorT object without quantization parameters.
void MakeTensor(const string& name, const std::vector<int32_t>& shape,
const std::vector<int32_t>& shape_signature,
const TensorType& type, std::unique_ptr<TensorT>* tensor) {
TensorT* tensor_raw = new TensorT;
tensor_raw->name = name;
tensor_raw->shape = shape;
if (!shape_signature.empty()) {
tensor_raw->shape_signature = shape_signature;
}
tensor_raw->type = type;
tensor->reset(tensor_raw);
}
// LINT.ThenChange(//tensorflow/lite/tools/optimize/model_utils.cc:MakeTensor)
// LINT.IfChange(HasMinMax)
bool HasMinMax(const TensorT* tensor) {
return tensor->quantization && !tensor->quantization->min.empty() &&
!tensor->quantization->max.empty();
}
// LINT.ThenChange(//tensorflow/lite/tools/optimize/model_utils.cc:HasMinMax)
} // namespace toco_legacy
} // namespace lite
} // namespace mlir
tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/model_utils.h
+60
View File
@@ -0,0 +1,60 @@
/* 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 is the MLIR copy of part of
// third_party/tensorflow/lite/tools/optimize/model_utils.h as part of the
// effort to decouple TFLite from MLIR.
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_LITE_TOCO_LEGACY_MODEL_UTILS_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_LITE_TOCO_LEGACY_MODEL_UTILS_H_
#include <cstdint>
#include <memory>
#include <string>
#include <vector>
#include "tensorflow/compiler/mlir/lite/schema/schema_generated.h"
namespace mlir {
namespace lite {
namespace toco_legacy {
using std::string;
using tflite::ModelT;
using tflite::OperatorT;
using tflite::TensorT;
using tflite::TensorType;
// LINT.IfChange(MakeDequantizeOperator)
// Creates a Dequantize OperatorT object.
void MakeDequantizeOperator(ModelT* model, std::unique_ptr<OperatorT>* op,
int32_t input, int32_t output);
// LINT.ThenChange(//tensorflow/lite/tools/optimize/model_utils.h:MakeDequantizeOperator)
// LINT.IfChange(MakeTensor)
// Create a new TensorT object without quantization parameters.
void MakeTensor(const string& name, const std::vector<int32_t>& shape,
const std::vector<int32_t>& shape_signature,
const TensorType& type, std::unique_ptr<TensorT>* tensor);
// LINT.ThenChange(//tensorflow/lite/tools/optimize/model_utils.h:MakeTensor)
// LINT.IfChange(HasMinMax)
bool HasMinMax(const TensorT* tensor);
// LINT.ThenChange(//tensorflow/lite/tools/optimize/model_utils.h:HasMinMax)
} // namespace toco_legacy
} // namespace lite
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_LITE_TOCO_LEGACY_MODEL_UTILS_H_
tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/model_utils_test.cc
+50
View File
@@ -0,0 +1,50 @@
/* Copyright 2018 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 is the MLIR copy of part of
// third_party/tensorflow/lite/tools/optimize/model_utils_test.cc as part of the
// effort to decouple TFLite from MLIR.
#include "tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/model_utils.h"
#include <memory>
#include <string>
#include <gtest/gtest.h>
#include "tensorflow/compiler/mlir/lite/schema/schema_generated.h"
namespace mlir {
namespace lite {
namespace toco_legacy {
namespace {
using std::string;
// LINT.IfChange(HasMinMaxTest)
TEST(ModelUtilsTest, HasMinMax) {
tflite::TensorT tensor;
tensor.quantization = std::make_unique<tflite::QuantizationParametersT>();
tensor.quantization->min.push_back(0.5);
EXPECT_FALSE(mlir::lite::toco_legacy::HasMinMax(&tensor));
tensor.quantization->max.push_back(1.5);
EXPECT_TRUE(mlir::lite::toco_legacy::HasMinMax(&tensor));
}
// LINT.ThenChange(//tensorflow/lite/tools/optimize/model_utils_test.cc:HasMinMaxTest)
} // namespace
} // namespace toco_legacy
} // namespace lite
} // namespace mlir
tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/portable_tensor_utils.cc
+66
View File
@@ -0,0 +1,66 @@
/* 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 is the MLIR copy of part of
// third_party/tensorflow/lite/kernels/internal/reference/portable_tensor_utils.cc
// as part of the effort to decouple TFLite from MLIR.
#include "tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/portable_tensor_utils.h"
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <cstring>
namespace mlir {
namespace lite {
namespace toco_legacy {
// LINT.IfChange(portable_symmetric_quantize_floats)
void PortableSymmetricQuantizeFloats(const float* values, const int size,
int8_t* quantized_values, float* min_value,
float* max_value, float* scaling_factor) {
auto minmax = std::minmax_element(values, values + size);
*min_value = *minmax.first;
*max_value = *minmax.second;
PortableSymmetricQuantizeFloats(values, size, quantized_values, *min_value,
*max_value, scaling_factor);
}
void PortableSymmetricQuantizeFloats(const float* values, const int size,
int8_t* quantized_values, float min_value,
float max_value, float* scaling_factor) {
const int32_t kScale = 127;
const float range = std::max(std::abs(min_value), std::abs(max_value));
if (range == 0) {
memset(quantized_values, 0, size * sizeof(int8_t));
*scaling_factor = 1;
return;
}
*scaling_factor = range / kScale;
const float scaling_factor_inv = kScale / range;
for (int i = 0; i < size; ++i) {
const int32_t quantized_value =
static_cast<int32_t>(round(values[i] * scaling_factor_inv));
// Clamp: just in case some odd numeric offset.
quantized_values[i] = static_cast<int8_t>(
std::min(kScale, std::max(-kScale, quantized_value)));
}
}
// LINT.ThenChange(//tensorflow/lite/kernels/internal/reference/portable_tensor_utils.cc:portable_symmetric_quantize_floats)
} // namespace toco_legacy
} // namespace lite
} // namespace mlir
tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/portable_tensor_utils.h
+42
View File
@@ -0,0 +1,42 @@
/* Copyright 2017 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 is the MLIR copy of part of
// third_party/tensorflow/lite/kernels/internal/reference/portable_tensor_utils.h
// as part of the effort to decouple TFLite from MLIR.
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_LITE_TOCO_LEGACY_PORTABLE_TENSOR_UTILS_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_LITE_TOCO_LEGACY_PORTABLE_TENSOR_UTILS_H_
#include <cstdint>
namespace mlir {
namespace lite {
namespace toco_legacy {
// LINT.IfChange(portable_symmetric_quantize_floats)
void PortableSymmetricQuantizeFloats(const float* values, int size,
int8_t* quantized_values, float* min_value,
float* max_value, float* scaling_factor);
void PortableSymmetricQuantizeFloats(const float* values, int size,
int8_t* quantized_values, float min_value,
float max_value, float* scaling_factor);
// LINT.ThenChange(//tensorflow/lite/kernels/internal/reference/portable_tensor_utils.h:portable_symmetric_quantize_floats)
} // namespace toco_legacy
} // namespace lite
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_LITE_TOCO_LEGACY_PORTABLE_TENSOR_UTILS_H_
tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/quantization_utils.cc
+371
View File
@@ -0,0 +1,371 @@
/* Copyright 2018 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 is the MLIR copy of part of
// third_party/tensorflow/lite/tools/optimize/quantization_utils.cc as part of
// the effort to decouple TFLite from MLIR.
#include "tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/quantization_utils.h"
#include <algorithm>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <cstdlib>
#include <memory>
#include <vector>
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "Eigen/Core" // from @eigen_archive
#include "tensorflow/compiler/mlir/lite/kernels/internal/runtime_shape.h"
#include "tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/model_utils.h"
#include "tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/portable_tensor_utils.h"
#include "tensorflow/compiler/mlir/lite/schema/schema_generated.h"
namespace mlir {
namespace lite {
namespace toco_legacy {
namespace {
// LINT.IfChange(QuantizationUtilsConstants)
const int8_t kMinQuantizedValue8bit = -127;
const int8_t kMaxQuantizedValue8bit = 127;
const int8_t kMinQuantizedValue4bit = -7;
const int8_t kMaxQuantizedValue4bit = 7;
// The maximum number of dimensions supported in per-channel quantization.
constexpr int kPerChannelMaxDim = 4;
// LINT.ThenChange(//tensorflow/lite/tools/optimize/quantization_utils.cc:QuantizationUtilsConstants)
} // namespace
using absl::InternalError;
using mlir::RuntimeShape;
using tflite::BufferT;
using tflite::QuantizationParametersT;
using tflite::TensorT;
using tflite::TensorType;
using tflite::TensorType_INT8;
// LINT.IfChange(NumElements)
absl::Status NumElements(const TensorT& tensor, uint64_t* num_elements) {
*num_elements = 1;
for (const int64_t dim : tensor.shape) {
if (dim <= 0 || *num_elements > UINT64_MAX / static_cast<uint64_t>(dim)) {
return InternalError("Invalid tensor shape.");
}
*num_elements *= dim;
}
return absl::OkStatus();
}
// LINT.ThenChange(//tensorflow/lite/tools/optimize/quantization_utils.cc:NumElements)
// LINT.IfChange(FillPerChannelMinMax)
absl::Status FillPerChannelMinMax(
const float* const input, const std::vector<int32_t>& dimension,
int32_t channel_dim_index, QuantizationParametersT* quantization_params) {
if (!quantization_params->min.empty() || !quantization_params->max.empty()) {
return absl::InvalidArgumentError(
"Min or max already present in tensor quantization params.");
}
if (dimension.size() > kPerChannelMaxDim) {
return absl::InvalidArgumentError(absl::StrCat(
"Expected tensor with less than %d dimensions, but got %d.",
kPerChannelMaxDim + 1, dimension.size()));
}
if (channel_dim_index >= dimension.size()) {
return absl::InvalidArgumentError(absl::StrCat(
"Expected channel_dim_index to be less than %d, but got %d.",
dimension.size(), channel_dim_index));
}
const int32_t channel_dim_size = dimension[channel_dim_index];
quantization_params->quantized_dimension = channel_dim_index;
quantization_params->min = std::vector<float>(channel_dim_size);
quantization_params->max = std::vector<float>(channel_dim_size);
std::vector<bool> has_min_max_value(channel_dim_size, false);
int indices[kPerChannelMaxDim];
RuntimeShape unextended_tensor_dims(dimension.size(), dimension.data());
RuntimeShape tensor_dims =
RuntimeShape::ExtendedShape(kPerChannelMaxDim, unextended_tensor_dims);
channel_dim_index +=
kPerChannelMaxDim - unextended_tensor_dims.DimensionsCount();
// Compute min max ranges per channel
for (indices[0] = 0; indices[0] < tensor_dims.Dims(0); indices[0]++) {
for (indices[1] = 0; indices[1] < tensor_dims.Dims(1); indices[1]++) {
for (indices[2] = 0; indices[2] < tensor_dims.Dims(2); indices[2]++) {
for (indices[3] = 0; indices[3] < tensor_dims.Dims(3); indices[3]++) {
int channel_idx = indices[channel_dim_index];
const float val = input[Offset(tensor_dims, indices)];
if (has_min_max_value[channel_idx]) {
if (quantization_params->min[channel_idx] > val) {
quantization_params->min[channel_idx] = val;
} else if (quantization_params->max[channel_idx] < val) {
quantization_params->max[channel_idx] = val;
}
} else {
quantization_params->min[channel_idx] = val;
quantization_params->max[channel_idx] = val;
has_min_max_value[channel_idx] = true;
}
}
}
}
}
return absl::OkStatus();
}
// LINT.ThenChange(//tensorflow/lite/tools/optimize/quantization_utils.cc:FillPerChannelMinMax)
// LINT.IfChange(SymmetricPerChannelQuantization)
// Per-channel quantize a tensor at the given index and fills both scales and
// quantized values.
absl::Status SymmetricPerChannelQuantization(
TensorT* tensor, const float* const input, int32_t channel_dim_index,
std::vector<float>* output_scales, std::vector<int8_t>* output_value) {
if (tensor == nullptr) {
return absl::InvalidArgumentError("Cannot quantize. Tensor is null.");
}
const int32_t channel_dim_size = tensor->shape[channel_dim_index];
// Fill per channel max and min values if needed
if (tensor->quantization == nullptr) {
tensor->quantization = std::make_unique<QuantizationParametersT>();
}
if (!HasMinMax(tensor)) {
absl::Status status = FillPerChannelMinMax(
input, tensor->shape, channel_dim_index, tensor->quantization.get());
if (!status.ok()) {
return status;
}
}
// Calculate scales per channel using max and min values from tensor.
std::vector<float> scale_invs(channel_dim_size);
const float half_scale = kMaxQuantizedValue8bit;
for (int channel_idx = 0; channel_idx < channel_dim_size; channel_idx++) {
const float half_range =
std::max(std::abs(tensor->quantization->min[channel_idx]),
std::abs(tensor->quantization->max[channel_idx]));
output_scales->at(channel_idx) = half_range / half_scale;
if (half_range == 0) {
scale_invs[channel_idx] = 0;
} else {
scale_invs[channel_idx] = half_scale / half_range;
}
}
// Quantize the input values.
SymmetricPerChannelQuantizeValues(input, scale_invs, tensor->shape,
channel_dim_index, output_value);
return absl::OkStatus();
}
// LINT.ThenChange(//tensorflow/lite/tools/optimize/quantization_utils.cc:SymmetricPerChannelQuantization)
// LINT.IfChange(SymmetricPerChannelQuantizeValues)
void SymmetricPerChannelQuantizeValues(const float* const input,
const std::vector<float>& scales_inv,
const std::vector<int32_t>& dimension,
int32_t channel_dim_index,
std::vector<int8_t>* output_value) {
// Quantize the values.
int indices[kPerChannelMaxDim];
RuntimeShape unextended_tensor_dims(dimension.size(), dimension.data());
RuntimeShape tensor_dims =
RuntimeShape::ExtendedShape(kPerChannelMaxDim, unextended_tensor_dims);
channel_dim_index +=
kPerChannelMaxDim - unextended_tensor_dims.DimensionsCount();
for (indices[0] = 0; indices[0] < tensor_dims.Dims(0); indices[0]++) {
for (indices[1] = 0; indices[1] < tensor_dims.Dims(1); indices[1]++) {
for (indices[2] = 0; indices[2] < tensor_dims.Dims(2); indices[2]++) {
for (indices[3] = 0; indices[3] < tensor_dims.Dims(3); indices[3]++) {
int channel_idx = indices[channel_dim_index];
int index = Offset(tensor_dims, indices);
const float val = input[index];
const int32_t quantized_value =
static_cast<int32_t>(round(val * scales_inv[channel_idx]));
output_value->at(index) = std::min<int8_t>(
kMaxQuantizedValue8bit,
std::max<int8_t>(kMinQuantizedValue8bit, quantized_value));
}
}
}
}
}
// LINT.ThenChange(//tensorflow/lite/tools/optimize/quantization_utils.cc:SymmetricPerChannelQuantizeValues)
// LINT.IfChange(SymmetricQuantizeTensor)
absl::Status SymmetricQuantizeTensor(ModelT* model, TensorT* tensor) {
if (model == nullptr || tensor == nullptr) {
return absl::InvalidArgumentError("No tensor to quantize.");
}
BufferT* buffer = model->buffers[tensor->buffer].get();
if (buffer == nullptr) {
return absl::InvalidArgumentError("Missing buffer.");
}
const float* float_data = reinterpret_cast<const float*>(buffer->data.data());
uint64_t num_elements;
absl::Status status = NumElements(*tensor, &num_elements);
if (!status.ok()) {
return status;
}
std::vector<int8_t> quantized_buffer;
quantized_buffer.resize(num_elements);
float min_value, max_value, scaling_factor;
mlir::lite::toco_legacy::PortableSymmetricQuantizeFloats(
float_data, num_elements, quantized_buffer.data(), &min_value, &max_value,
&scaling_factor);
if (tensor->quantization == nullptr) {
tensor->quantization = std::make_unique<QuantizationParametersT>();
}
tensor->quantization->scale = std::vector<float>(1, scaling_factor);
tensor->quantization->zero_point = std::vector<int64_t>(1, 0);
uint8_t* uint8_buffer = reinterpret_cast<uint8_t*>(quantized_buffer.data());
model->buffers[tensor->buffer]->data.assign(uint8_buffer,
uint8_buffer + num_elements);
// Update the tensor type.
tensor->type = TensorType_INT8;
return absl::OkStatus();
}
// LINT.ThenChange(//tensorflow/lite/tools/optimize/quantization_utils.cc:SymmetricQuantizeTensor)
// LINT.IfChange(QuantizeTensorFloat16)
absl::Status QuantizeTensorFloat16(ModelT* model, TensorT* tensor) {
if (model == nullptr || tensor == nullptr) {
return absl::InvalidArgumentError("No tensor to quantize.");
}
BufferT* buffer = model->buffers[tensor->buffer].get();
if (buffer == nullptr) {
return absl::InvalidArgumentError("Missing buffer.");
}
uint64_t num_elements;
absl::Status status = NumElements(*tensor, &num_elements);
if (!status.ok()) {
return status;
}
// Copy single byte buffer data to float vector to guard against misalignment.
std::vector<float> float_vector(num_elements);
uint8_t* first = buffer->data.data();
std::copy(first, first + buffer->data.size(),
reinterpret_cast<uint8_t*>(float_vector.data()));
// Transform float data to float16.
std::vector<Eigen::half> quantized_buffer;
quantized_buffer.resize(num_elements);
constexpr float kMaxFloat16Value = 65504.f;
constexpr float kMinFloat16Value = -65504.f;
std::transform(float_vector.begin(), float_vector.end(),
quantized_buffer.begin(), [=](float a) {
float clamped = std::min(std::max(a, kMinFloat16Value),
kMaxFloat16Value);
return static_cast<Eigen::half>(clamped);
});
char* half_buffer = reinterpret_cast<char*>(quantized_buffer.data());
model->buffers[tensor->buffer]->data.assign(
half_buffer, half_buffer + sizeof(Eigen::half) * num_elements);
// Update the tensor type.
tensor->type = tflite::TensorType_FLOAT16;
return absl::OkStatus();
}
// LINT.ThenChange(//tensorflow/lite/tools/optimize/quantization_utils.cc:QuantizeTensorFloat16)
// LINT.IfChange(AddQuantizationParams)
absl::Status AddQuantizationParams(const std::vector<float>& scales,
const std::vector<int64_t>& zero_point,
int quantized_dimension,
const uint8_t* buffer_data,
size_t buffer_size, TensorType output_type,
ModelT* model, TensorT* tensor) {
if (tensor->quantization == nullptr) {
tensor->quantization = std::make_unique<QuantizationParametersT>();
}
tensor->quantization->scale.assign(scales.begin(), scales.end());
if (zero_point.size() != scales.size()) {
return absl::InvalidArgumentError(
absl::StrCat("Received zero_point of size %d and scales of size %d. "
"These sizes should match.",
zero_point.size(), scales.size()));
}
tensor->quantization->zero_point.assign(zero_point.begin(), zero_point.end());
tensor->quantization->quantized_dimension = quantized_dimension;
model->buffers[tensor->buffer]->data.assign(buffer_data,
buffer_data + buffer_size);
// Update the tensor type.
tensor->type = output_type;
return absl::OkStatus();
}
// LINT.ThenChange(//tensorflow/lite/tools/optimize/quantization_utils.cc:AddQuantizationParams)
// LINT.IfChange(SymmetricQuantizeTensorPerChannel)
absl::Status SymmetricQuantizeTensorPerChannel(ModelT* model, TensorT* tensor,
int32_t channel_dim_index) {
if (tensor->shape.size() > kPerChannelMaxDim) {
return absl::InvalidArgumentError(absl::StrCat(
"SymmetricQuantizeTensorPerChannel requires tensor with less than %d "
"dimensions, but got %d dimension(s).",
kPerChannelMaxDim + 1, tensor->shape.size()));
}
// Get dimensions.
uint64_t num_elements;
absl::Status status = NumElements(*tensor, &num_elements);
if (!status.ok()) {
return status;
}
const int32_t channel_dim_size = tensor->shape[channel_dim_index];
// Get input float data.
const BufferT* buffer = model->buffers[tensor->buffer].get();
const float* float_input_data =
reinterpret_cast<const float*>(buffer->data.data());
// Create container for output scale and output data.
std::vector<float> scales(channel_dim_size);
std::vector<int8_t> final_buffer(num_elements);
// Quantize the input data with respect to channel_dim_index.
status = SymmetricPerChannelQuantization(
tensor, float_input_data, channel_dim_index, &scales, &final_buffer);
if (!status.ok()) {
return status;
}
// Set the buffers and output type.
uint8_t* uint8_buffer = reinterpret_cast<uint8_t*>(final_buffer.data());
const size_t buffer_size = num_elements * sizeof(int8_t);
std::vector<int64_t> zero_point(scales.size(), 0);
return AddQuantizationParams(scales, zero_point, channel_dim_index,
uint8_buffer, buffer_size, TensorType_INT8,
model, tensor);
}
// LINT.ThenChange(//tensorflow/lite/tools/optimize/quantization_utils.cc:SymmetricQuantizeTensorPerChannel)
} // namespace toco_legacy
} // namespace lite
} // namespace mlir
tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/quantization_utils.h
+110
View File
@@ -0,0 +1,110 @@
/* Copyright 2018 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 is the MLIR copy of part of
// third_party/tensorflow/lite/tools/optimize/quantization_utils.h as part of
// the effort to decouple TFLite from MLIR.
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_LITE_TOCO_LEGACY_QUANTIZATION_UTILS_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_LITE_TOCO_LEGACY_QUANTIZATION_UTILS_H_
#include <cstddef>
#include <cstdint>
#include <vector>
#include "absl/status/status.h"
#include "tensorflow/compiler/mlir/lite/schema/schema_generated.h"
namespace mlir {
namespace lite {
namespace toco_legacy {
using tflite::ModelT;
using tflite::QuantizationParametersT;
using tflite::TensorT;
using tflite::TensorType;
// LINT.IfChange(num_elements)
// Returns the number of elements in the given tensor.
absl::Status NumElements(const TensorT& tensor, uint64_t* num_elements);
// LINT.ThenChange(//tensorflow/lite/tools/optimize/quantization_utils.h:num_elements)
// LINT.IfChange(fill_per_channel_min_max)
// Populates the max and min values for per channel quantization.
absl::Status FillPerChannelMinMax(const float* input,
const std::vector<int32_t>& dimension,
int32_t channel_dim_index,
QuantizationParametersT* quantization_params);
// LINT.ThenChange(//tensorflow/lite/tools/optimize/quantization_utils.h:fill_per_channel_min_max)
// LINT.IfChange(symmetric_per_channel_quantization)
// Per-channel quantize a tensor at the given index and returns both scales and
// quantized values.
// Parameters:
// - tensor is the tensor to be quantized, needed to access associated
// quantization parameters
// - input is the float input data to be quantized.
// - channel_dim_index is the channel index within "dimension".
// dimension[channel_dim_index] gives the number of channels.
// - output_scale is the output scale, the size of which equals the number of
// channels.
// - output_value is the output data, the size of which equals the number of
// inputs.
absl::Status SymmetricPerChannelQuantization(TensorT* tensor,
const float* input,
int32_t channel_dim_index,
std::vector<float>* output_scales,
std::vector<int8_t>* output_value);
// LINT.ThenChange(//tensorflow/lite/tools/optimize/quantization_utils.h:symmetric_per_channel_quantization)
// LINT.IfChange(symmetric_per_channel_quantize_values)
// Quantize the values given an array of scales.
void SymmetricPerChannelQuantizeValues(const float* input,
const std::vector<float>& scales_inv,
const std::vector<int32_t>& dimension,
int32_t channel_dim_index,
std::vector<int8_t>* output_value);
// LINT.ThenChange(//tensorflow/lite/tools/optimize/quantization_utils.h:symmetric_per_channel_quantize_values)
// LINT.IfChange(symmetric_quantize_tensor)
// Quantizes tensor using symmetric quantization with the min and max elements
// of the tensor.
absl::Status SymmetricQuantizeTensor(ModelT* model, TensorT* tensor);
// LINT.ThenChange(//tensorflow/lite/tools/optimize/quantization_utils.h:symmetric_quantize_tensor)
// LINT.IfChange(symmetric_quantize_tensor_per_channel)
// Quantizes tensor with per channel.
absl::Status SymmetricQuantizeTensorPerChannel(ModelT* model, TensorT* tensor,
int32_t channel_dim_index);
// LINT.ThenChange(//tensorflow/lite/tools/optimize/quantization_utils.h:symmetric_quantize_tensor_per_channel)
// LINT.IfChange(quantize_tensor_float16)
// Quantizes tensor to float16.
absl::Status QuantizeTensorFloat16(ModelT* model, TensorT* tensor);
// LINT.ThenChange(//tensorflow/lite/tools/optimize/quantization_utils.h:quantize_tensor_float16)
// LINT.IfChange(add_quantization_params)
absl::Status AddQuantizationParams(const std::vector<float>& scales,
const std::vector<int64_t>& zero_point,
int quantized_dimension,
const uint8_t* buffer_data,
size_t buffer_size, TensorType output_type,
ModelT* model, TensorT* tensor);
// LINT.ThenChange(//tensorflow/lite/tools/optimize/quantization_utils.h:add_quantization_params)
} // namespace toco_legacy
} // namespace lite
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_LITE_TOCO_LEGACY_QUANTIZATION_UTILS_H_
tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/quantization_utils_test.cc
+496
View File
@@ -0,0 +1,496 @@
/* Copyright 2018 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 is the MLIR copy of part of
// third_party/tensorflow/lite/tools/optimize/quantization_utils_test.cc as part
// of the effort to decouple TFLite from MLIR.
#include "tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/quantization_utils.h"
#include <algorithm>
#include <cstdint>
#include <cstdlib>
#include <iostream>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "Eigen/Core" // from @eigen_archive
#include "tensorflow/compiler/mlir/lite/core/absl_error_model_builder.h"
#include "tensorflow/compiler/mlir/lite/quantization/lite/test_util.h"
#include "tensorflow/compiler/mlir/lite/schema/schema_generated.h"
#include "tensorflow/compiler/mlir/lite/schema/schema_utils.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "xla/tsl/util/command_line_flags.h"
#include "tsl/platform/init_main.h"
#include "tsl/platform/path.h"
namespace {
std::string* g_test_model_dir = nullptr;
} // namespace
namespace mlir {
namespace lite {
namespace toco_legacy {
namespace {
using mlir::TFL::FlatBufferModelAbslError;
using tflite::BuiltinOperator_CONV_2D;
using tflite::QuantizationParametersT;
using tflite::SubGraphT;
using tflite::TensorT;
using tflite::TensorType_FLOAT16;
using tflite::TensorType_FLOAT32;
using tflite::TensorType_INT8;
std::unique_ptr<FlatBufferModelAbslError> ReadModel(const char* model) {
auto model_path = tsl::io::JoinPath(*g_test_model_dir, model);
return FlatBufferModelAbslError::BuildFromFile(model_path.c_str());
}
std::unique_ptr<FlatBufferModelAbslError> ReadConvModel() {
return ReadModel(mlir::lite::internal::kConvModelWith0Plus10Weights);
}
using ::testing::ElementsAreArray;
class QuantizationUtilsTest : public testing::Test {};
TEST_F(QuantizationUtilsTest, NumElements) {
TensorT tensor;
tensor.shape = {1, 2, 3, 4};
uint64_t num_elements;
TF_EXPECT_OK(NumElements(tensor, &num_elements));
EXPECT_EQ(num_elements, 1 * 2 * 3 * 4);
tensor.shape = {5};
TF_EXPECT_OK(NumElements(tensor, &num_elements));
EXPECT_EQ(num_elements, 5);
tensor.shape = {};
TF_EXPECT_OK(NumElements(tensor, &num_elements));
// Scalars with empty shape have 1 element.
EXPECT_EQ(num_elements, 1);
tensor.shape = {1, 2, 3, -1};
EXPECT_EQ(NumElements(tensor, &num_elements).code(),
absl::StatusCode::kInternal);
}
TEST_F(QuantizationUtilsTest, SymmetricPerChannelQuantizationWithNullQParams) {
// Set up an input with [3, 2, 2, 2] size and 0 is the channel index.
const std::vector<float> input = {
3.0, 2.0, 5.0, -2.0, 3.0, 2.0, 5.0, -2.0, // Channel 1.
1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, // Channel 2.
1.0, 0.0, -1.0, -2.0, -3.0, -4.0, -5.0, -6.0, // Channel 3.
};
const int channel_index = 0;
// Create holder for output scale and data.
std::vector<float> output_scales(3);
std::vector<int8_t> output_data(3 * 2 * 2 * 2);
// Call SymmetricPerChannelQuantization with quant_params as a null pointer
// and verify the result.
TensorT tensor = TensorT();
tensor.quantization = nullptr;
tensor.shape = {3, 2, 2, 2};
TF_EXPECT_OK(mlir::lite::toco_legacy::SymmetricPerChannelQuantization(
&tensor, input.data(), channel_index, &output_scales, &output_data));
const std::vector<float> expected_output_scales = {0.0393700786, 0.0629921257,
0.0472440943};
const std::vector<int8_t> expected_output_data = {
76, 51, 127, -51, 76, 51, 127, -51, // Channel 1.
16, 32, 48, 64, 79, 95, 111, 127, // Channel 2.
21, 0, -21, -42, -64, -85, -106, -127, // Channel 3.
};
EXPECT_THAT(output_scales, ElementsAreArray(expected_output_scales));
EXPECT_THAT(output_data, ElementsAreArray(expected_output_data));
}
TEST_F(QuantizationUtilsTest, SymmetricPerChannelQuantization) {
// Set up an input with [3, 2, 2, 2] size and 0 is the channel index.
const std::vector<float> input = {
3.0, 2.0, 5.0, -2.0, 3.0, 2.0, 5.0, -2.0, // Channel 1.
1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, // Channel 2.
1.0, 0.0, -1.0, -2.0, -3.0, -4.0, -5.0, -6.0, // Channel 3.
};
const int32_t channel_index = 0;
// Create holder for output scale and data.
std::vector<float> output_scales(3);
std::vector<int8_t> output_data(3 * 2 * 2 * 2);
// Initialize pointer to quantization parameters
TensorT tensor = TensorT();
tensor.quantization = std::make_unique<QuantizationParametersT>();
tensor.shape = {3, 2, 2, 2};
TF_EXPECT_OK(mlir::lite::toco_legacy::FillPerChannelMinMax(
input.data(), tensor.shape, channel_index, tensor.quantization.get()));
// Test that FillPerChanneMinMax worked
const std::vector<float> expected_mins = {-2.0, 1.0, -6.0};
const std::vector<float> expected_maxs = {5.0, 8.0, 1.0};
EXPECT_THAT(tensor.quantization->min, ElementsAreArray(expected_mins));
EXPECT_THAT(tensor.quantization->max, ElementsAreArray(expected_maxs));
// Call SymmetricPerChannelQuantization with quant_params as a null pointer
// and verify the result.
TF_EXPECT_OK(mlir::lite::toco_legacy::SymmetricPerChannelQuantization(
&tensor, input.data(), channel_index, &output_scales, &output_data));
const std::vector<float> expected_output_scales = {0.0393700786, 0.0629921257,
0.0472440943};
const std::vector<int8_t> expected_output_data = {
76, 51, 127, -51, 76, 51, 127, -51, // Channel 1.
16, 32, 48, 64, 79, 95, 111, 127, // Channel 2.
21, 0, -21, -42, -64, -85, -106, -127, // Channel 3.
};
EXPECT_THAT(output_scales, ElementsAreArray(expected_output_scales));
EXPECT_THAT(output_data, ElementsAreArray(expected_output_data));
}
TEST_F(QuantizationUtilsTest, SymmetricPerChannelQuantization2DTensor) {
// Set up an input with [3, 8] size and 0 is the channel index.
const std::vector<float> input = {
3.0, 2.0, 5.0, -2.0, 3.0, 2.0, 5.0, -2.0, // Batch 1.
1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, // Batch 2.
1.0, 0.0, -1.0, -2.0, -3.0, -4.0, -5.0, -6.0, // Batch 3.
};
const int32_t channel_index = 1;
// Create holder for output scale and data.
std::vector<float> output_scales(8);
std::vector<int8_t> output_data(3 * 8);
// Initialize pointer to quantization parameters
TensorT tensor = TensorT();
tensor.quantization = std::make_unique<QuantizationParametersT>();
tensor.shape = {3, 8};
TF_EXPECT_OK(mlir::lite::toco_legacy::FillPerChannelMinMax(
input.data(), tensor.shape, channel_index, tensor.quantization.get()));
// Test that FillPerChanneMinMax worked
const std::vector<float> expected_mins = {1.0, 0.0, -1.0, -2.0,
-3.0, -4.0, -5.0, -6.0};
const std::vector<float> expected_maxs = {3.0, 2.0, 5.0, 4.0,
5.0, 6.0, 7.0, 8.0};
EXPECT_THAT(tensor.quantization->min, ElementsAreArray(expected_mins));
EXPECT_THAT(tensor.quantization->max, ElementsAreArray(expected_maxs));
// Call SymmetricPerChannelQuantization with quant_params as a null pointer
// and verify the result.
TF_EXPECT_OK(mlir::lite::toco_legacy::SymmetricPerChannelQuantization(
&tensor, input.data(), channel_index, &output_scales, &output_data));
const std::vector<float> expected_output_scales = {
0.02362204724, 0.01574803149, 0.03937007874, 0.03149606299,
0.03937007874, 0.04724409448, 0.05511811023, 0.06299212598};
const std::vector<int8_t> expected_output_data = {
127, 127, 127, -64, 76, 42, 91, -32, // Batch 1.
42, 127, 76, 127, 127, 127, 127, 127, // Batch 2.
42, 0, -25, -64, -76, -85, -91, -95, // Batch 3.
};
EXPECT_THAT(output_scales, ElementsAreArray(expected_output_scales));
EXPECT_THAT(output_data, ElementsAreArray(expected_output_data));
}
TEST_F(QuantizationUtilsTest, SymmetricPerChannelQuantizeValues) {
// Set up an input with [3, 1, 1, 2] size and 0 is the channel index.
const std::vector<float> input = {
13.0, 21.0, // Channel 1.
21.0, 22.0, // Channel 2.
31.0, 40.0, // Channel 3.
};
const std::vector<float> scales_inv = {2, 0.5, 3};
const std::vector<int32_t> dimension = {3, 1, 1, 2};
const int channel_index = 0;
// Create holder for output data.
std::vector<int8_t> output_data(3 * 1 * 1 * 2);
// Call SymmetricPerChannelQuantizeValues and verify the result.
SymmetricPerChannelQuantizeValues(input.data(), scales_inv, dimension,
channel_index, &output_data);
const std::vector<int8_t> expected_output_data = {
26, 42, // Channel 1.
11, 11, // Channel 2.
93, 120, // Channel 3.
};
EXPECT_THAT(output_data, ElementsAreArray(expected_output_data));
}
TEST_F(QuantizationUtilsTest, FillPerChannelMinMax) {
// Set up an input with [3, 1, 1, 2] size.
const std::vector<float> input = {
13.0, 21.0, // Channel 1.
21.0, 22.0, // Channel 2.
31.0, 40.0, // Channel 3.
};
// Initialize pointer to quantization parameters.
QuantizationParametersT quantization_params = QuantizationParametersT();
std::vector<int> dimension = {3, 1, 1, 2};
int32_t channel_dim_idx = 0;
const std::vector<float> expected_mins = {13.0, 21.0, 31.0};
const std::vector<float> expected_maxs = {21.0, 22.0, 40.0};
TF_EXPECT_OK(mlir::lite::toco_legacy::FillPerChannelMinMax(
input.data(), dimension, channel_dim_idx, &quantization_params));
EXPECT_EQ(quantization_params.min, expected_mins);
EXPECT_EQ(quantization_params.max, expected_maxs);
EXPECT_EQ(quantization_params.quantized_dimension, channel_dim_idx);
}
TEST_F(QuantizationUtilsTest, FillPerChannelMinMaxFillDim3) {
// Set up an input with [3, 1, 1, 2] size.
const std::vector<float> input = {
// Channel 1, Channel 2
13.0, 21.0, 21.0, 22.0, 31.0, 40.0,
};
// Initialize pointer to quantization parameters.
QuantizationParametersT quantization_params = QuantizationParametersT();
std::vector<int> dimension = {3, 1, 1, 2};
int32_t channel_dim_idx = 3;
const std::vector<float> expected_mins = {13.0, 21.0};
const std::vector<float> expected_maxs = {31.0, 40.0};
TF_EXPECT_OK(mlir::lite::toco_legacy::FillPerChannelMinMax(
input.data(), dimension, channel_dim_idx, &quantization_params));
EXPECT_EQ(quantization_params.min, expected_mins);
EXPECT_EQ(quantization_params.max, expected_maxs);
EXPECT_EQ(quantization_params.quantized_dimension, channel_dim_idx);
}
TEST_F(QuantizationUtilsTest, FillPerChannelMinMax2DTensor) {
// Set up an input with [3, 2] size.
const std::vector<float> input = {
// Channel 1, Channel 2
13.0, 21.0, 21.0, 22.0, 31.0, 40.0,
};
// Initialize pointer to quantization parameters.
QuantizationParametersT quantization_params = QuantizationParametersT();
std::vector<int> dimension = {3, 2};
int32_t channel_dim_idx = 1;
const std::vector<float> expected_mins = {13.0, 21.0};
const std::vector<float> expected_maxs = {31.0, 40.0};
TF_EXPECT_OK(mlir::lite::toco_legacy::FillPerChannelMinMax(
input.data(), dimension, channel_dim_idx, &quantization_params));
EXPECT_EQ(quantization_params.min, expected_mins);
EXPECT_EQ(quantization_params.max, expected_maxs);
EXPECT_EQ(quantization_params.quantized_dimension, channel_dim_idx);
}
TEST_F(QuantizationUtilsTest, SymmetricQuantizeTensorNullInputs) {
EXPECT_EQ(SymmetricQuantizeTensor(nullptr, nullptr).code(),
absl::StatusCode::kInvalidArgument);
}
TEST_F(QuantizationUtilsTest, SymmetricQuantizeTensorNullQuantParams) {
// Conv model has weights between 0 and 10.
// Quantize the weights tensor.
ASSERT_TRUE(g_test_model_dir);
ASSERT_FALSE(g_test_model_dir->empty());
auto test_model = ReadConvModel();
ASSERT_TRUE(test_model);
auto readonly_model = test_model->GetModel();
ASSERT_TRUE(readonly_model);
ASSERT_TRUE(readonly_model->subgraphs());
ASSERT_GE(readonly_model->subgraphs()->size(), 1);
tflite::ModelT model;
readonly_model->UnPackTo(&model);
auto subgraph = model.subgraphs[0].get();
auto conv_op = subgraph->operators.at(0).get();
ASSERT_EQ(
GetBuiltinCode(model.operator_codes.at(conv_op->opcode_index).get()),
BuiltinOperator_CONV_2D);
int32_t weights_tensor_idx = conv_op->inputs[1];
TensorT* weights_tensor = subgraph->tensors.at(weights_tensor_idx).get();
// Empty quantization parameters.
weights_tensor->quantization = std::make_unique<QuantizationParametersT>();
EXPECT_EQ(weights_tensor->type, TensorType_FLOAT32);
size_t float_buffer_size =
model.buffers.at(weights_tensor->buffer)->data.size();
TF_EXPECT_OK(SymmetricQuantizeTensor(&model, weights_tensor));
size_t quant_buffer_size =
model.buffers.at(weights_tensor->buffer)->data.size();
EXPECT_EQ(weights_tensor->type, TensorType_INT8);
EXPECT_EQ(quant_buffer_size * 4, float_buffer_size);
}
TEST_F(QuantizationUtilsTest, SymmetricQuantizeTensor) {
// Conv model has weights between 0 and 10.
// Quantize the weights tensor.
ASSERT_TRUE(g_test_model_dir);
ASSERT_FALSE(g_test_model_dir->empty());
auto test_model = ReadConvModel();
ASSERT_TRUE(test_model);
auto readonly_model = test_model->GetModel();
ASSERT_TRUE(readonly_model);
ASSERT_TRUE(readonly_model->subgraphs());
ASSERT_GE(readonly_model->subgraphs()->size(), 1);
tflite::ModelT model;
readonly_model->UnPackTo(&model);
auto subgraph = model.subgraphs[0].get();
auto conv_op = subgraph->operators.at(0).get();
ASSERT_EQ(
GetBuiltinCode(model.operator_codes.at(conv_op->opcode_index).get()),
BuiltinOperator_CONV_2D);
int32_t weights_tensor_idx = conv_op->inputs[1];
TensorT* weights_tensor = subgraph->tensors.at(weights_tensor_idx).get();
EXPECT_EQ(weights_tensor->type, TensorType_FLOAT32);
size_t float_buffer_size =
model.buffers.at(weights_tensor->buffer)->data.size();
TF_EXPECT_OK(SymmetricQuantizeTensor(&model, weights_tensor));
size_t quant_buffer_size =
model.buffers.at(weights_tensor->buffer)->data.size();
EXPECT_EQ(weights_tensor->type, TensorType_INT8);
EXPECT_EQ(quant_buffer_size * 4, float_buffer_size);
}
TEST_F(QuantizationUtilsTest, QuantizeFloat16Clamp) {
// Create data.
auto model = std::make_unique<ModelT>();
auto subgraph = std::make_unique<tflite::SubGraphT>();
auto tensor = std::make_unique<TensorT>();
auto buffer = std::make_unique<tflite::BufferT>();
constexpr int kNumElements = 6;
const std::vector<float> weights = {2.0, 1.0, 65504., 65505, -65504., -99999};
auto weights_reinterpreted_data =
reinterpret_cast<const unsigned char*>(weights.data());
buffer->data.assign(weights_reinterpreted_data,
weights_reinterpreted_data + weights.size() * 4);
tensor->buffer = 0;
tensor->shape = {1, kNumElements};
// Wire the model.
model->subgraphs.push_back(std::move(subgraph));
model->subgraphs[0]->tensors.push_back(std::move(tensor));
model->buffers.push_back(std::move(buffer));
// Call and verify.
TF_EXPECT_OK(QuantizeTensorFloat16(model.get(),
model->subgraphs[0]->tensors[0].get()));
auto weightsf16 = reinterpret_cast<Eigen::half*>(
model->buffers[model->subgraphs[0]->tensors[0]->buffer]->data.data());
std::vector<float> wf32(kNumElements);
std::transform(weightsf16, weightsf16 + 6, wf32.begin(),
[](Eigen::half a) { return static_cast<float>(a); });
EXPECT_THAT(wf32,
ElementsAreArray({2.0, 1.0, 65504., 65504., -65504., -65504.}));
EXPECT_EQ(model->subgraphs[0]->tensors[0]->type, TensorType_FLOAT16);
}
TEST_F(QuantizationUtilsTest, QuantizeFloat16) {
// Conv model has weights between 0 and 10.
// Quantize the weights tensor.
ASSERT_TRUE(g_test_model_dir != nullptr);
ASSERT_FALSE(g_test_model_dir->empty());
auto test_model = ReadConvModel();
ASSERT_TRUE(test_model);
auto readonly_model = test_model->GetModel();
ASSERT_TRUE(readonly_model);
ASSERT_TRUE(readonly_model->subgraphs());
ASSERT_GE(readonly_model->subgraphs()->size(), 1);
tflite::ModelT model;
readonly_model->UnPackTo(&model);
auto subgraph = model.subgraphs[0].get();
auto conv_op = subgraph->operators.at(0).get();
ASSERT_EQ(
GetBuiltinCode(model.operator_codes.at(conv_op->opcode_index).get()),
BuiltinOperator_CONV_2D);
int32_t weights_tensor_idx = conv_op->inputs[1];
TensorT* weights_tensor = subgraph->tensors.at(weights_tensor_idx).get();
EXPECT_EQ(weights_tensor->type, TensorType_FLOAT32);
size_t float_buffer_size =
model.buffers.at(weights_tensor->buffer)->data.size();
TF_EXPECT_OK(QuantizeTensorFloat16(&model, weights_tensor));
size_t quant_buffer_size =
model.buffers.at(weights_tensor->buffer)->data.size();
EXPECT_EQ(weights_tensor->type, TensorType_FLOAT16);
EXPECT_EQ(quant_buffer_size * 2, float_buffer_size);
}
TEST_F(QuantizationUtilsTest, AddQuantizationParams) {
// Create data.
auto model = std::make_unique<ModelT>();
auto subgraph = std::make_unique<tflite::SubGraphT>();
auto tensor = std::make_unique<TensorT>();
auto buffer = std::make_unique<tflite::BufferT>();
const std::vector<float> scales = {0.5, 1.0, 1.5};
const std::vector<int64_t> zero_points = {5, 10, 15};
const int32_t quantizated_dimension = 3;
const std::vector<uint8_t> buffer_data = {1, 2, 3, 4};
const int32_t buffer_size = 4;
tensor->buffer = 0;
// Wire the model.
model->subgraphs.push_back(std::move(subgraph));
model->subgraphs[0]->tensors.push_back(std::move(tensor));
model->buffers.push_back(std::move(buffer));
// Call and verify.
TF_EXPECT_OK(AddQuantizationParams(scales, zero_points, quantizated_dimension,
buffer_data.data(), buffer_size,
TensorType_INT8, model.get(),
model->subgraphs[0]->tensors[0].get()));
EXPECT_THAT(model->subgraphs[0]->tensors[0]->quantization->scale,
ElementsAreArray(scales));
EXPECT_THAT(model->subgraphs[0]->tensors[0]->quantization->zero_point,
ElementsAreArray(zero_points));
EXPECT_THAT(model->buffers[model->subgraphs[0]->tensors[0]->buffer]->data,
ElementsAreArray(buffer_data));
EXPECT_EQ(model->subgraphs[0]->tensors[0]->type, TensorType_INT8);
}
} // namespace
} // namespace toco_legacy
} // namespace lite
} // namespace mlir
int main(int argc, char** argv) {
std::string model_file;
const std::vector<tsl::Flag> flag_list = {
tsl::Flag("test_model_file", &model_file,
"Path to test tflite model file."),
};
const bool parse_result = tsl::Flags::Parse(&argc, argv, flag_list);
if (!parse_result) {
std::cerr << "Required test_model_file\n";
std::abort();
}
g_test_model_dir = new std::string(tsl::io::Dirname(model_file));
::tsl::port::InitMain(argv[0], &argc, &argv);
return RUN_ALL_TESTS();
}
tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/quantize_weights.cc
+752
View File
@@ -0,0 +1,752 @@
/* Copyright 2018 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/lite/toco_legacy/quantize_weights.h"
#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <memory>
#include <string>
#include <tuple>
#include <utility>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "flatbuffers/buffer.h" // from @flatbuffers
#include "flatbuffers/flatbuffer_builder.h" // from @flatbuffers
#include "tensorflow/compiler/mlir/lite/quantization/lite/quantize_weights.h"
#include "tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/model_utils.h"
#include "tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/quantization_utils.h"
#include "tensorflow/compiler/mlir/lite/schema/schema_generated.h"
#include "tensorflow/compiler/mlir/lite/schema/schema_utils.h"
#include "tensorflow/core/platform/logging.h"
namespace mlir {
namespace lite {
namespace toco_legacy {
namespace {
using absl::flat_hash_set;
using mlir::lite::toco_legacy::
CustomOpMap; // Use this instead of mlir::lite::CustomOpMap because that
// uses mlir::lite::CustomOpInfo in
// tensorflow/compiler/mlir/lite/quantization/lite/quantize_weights.h,
// and we need mlir::lite::toco_legacy::CustomOpInfo, in
// tensorflow/compiler/mlir/lite/quantization/lite/optimize/quantize_weights.h
using tflite::BufferT;
using tflite::BuiltinOperator;
using tflite::BuiltinOperator_BATCH_MATMUL;
using tflite::BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM;
using tflite::BuiltinOperator_BIDIRECTIONAL_SEQUENCE_RNN;
using tflite::BuiltinOperator_CONV_2D;
using tflite::BuiltinOperator_CUSTOM;
using tflite::BuiltinOperator_DEPTHWISE_CONV_2D;
using tflite::BuiltinOperator_EMBEDDING_LOOKUP;
using tflite::BuiltinOperator_FULLY_CONNECTED;
using tflite::BuiltinOperator_GATHER;
using tflite::BuiltinOperator_LSTM;
using tflite::BuiltinOperator_RNN;
using tflite::BuiltinOperator_SVDF;
using tflite::BuiltinOperator_TRANSPOSE_CONV;
using tflite::BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM;
using tflite::BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN;
using tflite::FinishModelBuffer;
using tflite::GetBuiltinCode;
using tflite::Model;
using tflite::ModelT;
using tflite::OperatorCodeT;
using tflite::OperatorT;
using tflite::SubGraphT;
using tflite::TensorT;
using tflite::TensorType_FLOAT32;
using tflite::TensorType_INT8;
struct ConsumerOpInfo {
OperatorT* op;
// The index of the op in the operators vector.
int32_t op_idx;
// The index of the tensor to quantize in subgraph->tensors.
int32_t op_input_idx;
};
struct TensorPerChannel {
TensorT* t;
bool is_per_channel;
int channel_dim;
};
// The default minimum number of elements a weights array must have to be
// quantized by this transformation.
const int kWeightsMinNumElementsDefault = 1024;
// Redefined from tensorflow/lite/core/c/common.h as local const int instead of
// discouraged #define macro.
const int kTfLiteOptionalTensor = -1;
// Convert the MLIR CustomOpMap from the TFlite CustomOpMap as their member
// variables differ.
void ConstructMLIRCustomOpMap(mlir::lite::CustomOpMap& mlir_map,
const CustomOpMap& tflite_map) {
for (const auto& entry : tflite_map) {
mlir_map[entry.first].quantizable_input_indices =
entry.second.quantizable_input_indices;
mlir_map[entry.first].is_weight_only = !entry.second.is_hybrid;
mlir_map[entry.first].no_side_effect = true;
}
}
// Gets the operators that consume tensor_idx.
std::vector<ConsumerOpInfo> GetTensorConsumers(const ModelT* model,
const SubGraphT* subgraph,
int32_t tensor_idx) {
// TODO(suharshs): If this proves to be too slow, avoid calling it per tensor,
// instead doing one sweep for the entire model.
std::vector<ConsumerOpInfo> consumer_ops;
for (size_t op_idx = 0; op_idx < subgraph->operators.size(); ++op_idx) {
OperatorT* op = subgraph->operators[op_idx].get();
if (op == nullptr) {
continue;
}
for (size_t i = 0; i < op->inputs.size(); ++i) {
if (op->inputs[i] == tensor_idx) {
consumer_ops.push_back(
{op, static_cast<int32_t>(op_idx), static_cast<int32_t>(i)});
}
}
}
return consumer_ops;
}
// Gets the list of op->inputs indices of the weights inputs to be quantized for
// the provided op.
std::vector<int32_t> GetWeightInputIndices(const OperatorCodeT* op_code,
const CustomOpMap& custom_op_map) {
const BuiltinOperator builtin_op_code = GetBuiltinCode(op_code);
if (builtin_op_code == BuiltinOperator_CUSTOM) {
const std::string custom_code = op_code->custom_code;
const auto& custom_op_info = custom_op_map.find(custom_code);
if (custom_op_info != custom_op_map.end()) {
return custom_op_info->second.quantizable_input_indices;
}
} else if (builtin_op_code == BuiltinOperator_CONV_2D ||
builtin_op_code == BuiltinOperator_DEPTHWISE_CONV_2D ||
builtin_op_code == BuiltinOperator_FULLY_CONNECTED ||
builtin_op_code == BuiltinOperator_BATCH_MATMUL ||
builtin_op_code == BuiltinOperator_EMBEDDING_LOOKUP ||
builtin_op_code == BuiltinOperator_TRANSPOSE_CONV) {
return {1};
} else if (builtin_op_code == BuiltinOperator_SVDF) {
// tensorflow/lite/kernels/svdf.cc
return {1, 2};
} else if (builtin_op_code == BuiltinOperator_LSTM ||
builtin_op_code == BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM) {
// tensorflow/lite/kernels/lstm.cc
// tensorflow/lite/kernels/unidirectional_sequence_lstm.cc
return {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 16};
} else if (builtin_op_code == BuiltinOperator_RNN ||
builtin_op_code == BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN) {
// tensorflow/lite/kernels/basic_rnn.cc
// tensorflow/lite/kernels/unidirectional_sequence_rnn.cc
return {1, 2};
} else if (builtin_op_code == BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM) {
// tensorflow/lite/kernels/bidirectional_sequence_lstm.cc
return {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 16, 18, 19, 20, 21,
22, 23, 24, 25, 26, 27, 28, 33, 40, 41, 42, 43, 44, 45, 46, 47};
} else if (builtin_op_code == BuiltinOperator_BIDIRECTIONAL_SEQUENCE_RNN) {
// tensorflow/lite/kernels/bidirectional_sequence_rnn.cc
return {1, 2, 4, 5, 6, 8, 9, 10, 11};
} else if (builtin_op_code == BuiltinOperator_GATHER) {
// tensorflow/lite/kernels/gather.cc
return {0};
}
return {};
}
// Checks that a specific input can be quantized.
bool IsQuantizedInput(const OperatorCodeT* op_code,
const CustomOpMap& custom_op_map, int op_input_idx) {
const auto quantized_input_indices =
GetWeightInputIndices(op_code, custom_op_map);
return std::find(std::begin(quantized_input_indices),
std::end(quantized_input_indices),
op_input_idx) != std::end(quantized_input_indices);
}
// Returns true if the operator supports hybrid evaluation.
bool IsHybridEvaluationOp(const OperatorT* op, const OperatorCodeT* op_code,
const CustomOpMap& custom_op_map,
bool use_updated_hybrid_scheme) {
const BuiltinOperator builtin_op_code = GetBuiltinCode(op_code);
// Operations that support hybrid evaluation.
bool eval_hybrid = false;
if (builtin_op_code == BuiltinOperator_CUSTOM) {
const std::string custom_code = op_code->custom_code;
const auto custom_op_info = custom_op_map.find(custom_code);
if (custom_op_info == custom_op_map.end()) {
return {};
} else {
return custom_op_info->second.is_hybrid;
}
} else if (builtin_op_code == BuiltinOperator_FULLY_CONNECTED ||
builtin_op_code == BuiltinOperator_BATCH_MATMUL ||
builtin_op_code == BuiltinOperator_CONV_2D ||
builtin_op_code == BuiltinOperator_SVDF ||
builtin_op_code == BuiltinOperator_RNN ||
builtin_op_code == BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM ||
builtin_op_code == BuiltinOperator_BIDIRECTIONAL_SEQUENCE_RNN ||
builtin_op_code == BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM ||
builtin_op_code == BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN) {
eval_hybrid = true;
} else if (builtin_op_code == BuiltinOperator_LSTM) {
const tflite::LSTMOptionsT* options = op->builtin_options.AsLSTMOptions();
// Only lstm kernel_type full supports hybrid evaluation.
if (options->kernel_type == tflite::LSTMKernelType_FULL) {
eval_hybrid = true;
}
} else if (builtin_op_code == BuiltinOperator_DEPTHWISE_CONV_2D) {
eval_hybrid = use_updated_hybrid_scheme;
}
return eval_hybrid;
}
// Returns true if all of the op's inputs are quantized.
bool CheckAllOpInputsQuantized(const SubGraphT* subgraph, const OperatorT* op,
const OperatorCodeT* op_code,
const CustomOpMap& custom_op_map) {
std::vector<int32_t> op_input_indices =
GetWeightInputIndices(op_code, custom_op_map);
for (const int32_t op_input_idx : op_input_indices) {
int32_t tensor_idx = op->inputs[op_input_idx];
if (tensor_idx == -1) {
// Optional tensor.
continue;
}
TensorT* tensor = subgraph->tensors[tensor_idx].get();
if (tensor->type != TensorType_INT8) {
return false;
}
}
return true;
}
// Inserts Tensors for each input tensor of op that should be
// quantized into tensor_map.
absl::Status InsertQuantizableInputTensorsFromOperator(
const ModelT* model, OperatorT* op, uint64_t weights_min_num_elements,
const CustomOpMap& custom_op_map,
absl::flat_hash_map<int32_t, TensorPerChannel>* tensor_map,
int subgraph_index, bool use_updated_hybrid_scheme) {
SubGraphT* subgraph = model->subgraphs.at(subgraph_index).get();
const OperatorCodeT* op_code = model->operator_codes[op->opcode_index].get();
auto builtin_code = GetBuiltinCode(op_code);
std::vector<int32_t> op_input_indices =
GetWeightInputIndices(op_code, custom_op_map);
for (const int32_t op_input_idx : op_input_indices) {
int32_t tensor_idx = op->inputs[op_input_idx];
if (tensor_idx == -1) {
LOG(INFO) << "Skipping optional tensor input " << op_input_idx
<< " of operation " << EnumNameBuiltinOperator(builtin_code);
continue;
}
TensorT* tensor = subgraph->tensors[tensor_idx].get();
if (tensor->type != TensorType_FLOAT32) {
LOG(INFO) << "Skipping quantization of tensor " << tensor->name
<< " that is not type float.";
continue;
}
uint64_t num_elements;
if (!mlir::lite::toco_legacy::NumElements(*tensor, &num_elements).ok()) {
return absl::InternalError("Error in quantization_utils NumElements");
}
if (num_elements < weights_min_num_elements) {
LOG(INFO) << "Skipping quantization of tensor " << tensor->name
<< " because it has fewer than " << weights_min_num_elements
<< " elements (" << num_elements << ").";
continue;
}
// Some tensors may have a null buffer vector, indicating an intermediate
// array.
if (model->buffers[tensor->buffer]->data.data() == nullptr) {
LOG(INFO) << "Skipping quantization of tensor " << tensor->name
<< " because it has no allocated buffer.";
continue;
}
if (builtin_code == BuiltinOperator_DEPTHWISE_CONV_2D) {
tensor_map->insert({tensor_idx,
{tensor, /*is_per_channel=*/use_updated_hybrid_scheme,
/*dim=*/3}});
} else if (builtin_code == BuiltinOperator_CONV_2D) {
tensor_map->insert({tensor_idx,
{tensor, /*is_per_channel=*/use_updated_hybrid_scheme,
/*dim=*/0}});
} else {
switch (builtin_code) {
case BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM:
op->builtin_options.AsBidirectionalSequenceLSTMOptions()
->asymmetric_quantize_inputs = use_updated_hybrid_scheme;
break;
case BuiltinOperator_BIDIRECTIONAL_SEQUENCE_RNN:
op->builtin_options.AsBidirectionalSequenceRNNOptions()
->asymmetric_quantize_inputs = use_updated_hybrid_scheme;
break;
case BuiltinOperator_FULLY_CONNECTED:
op->builtin_options.AsFullyConnectedOptions()
->asymmetric_quantize_inputs = use_updated_hybrid_scheme;
break;
case BuiltinOperator_BATCH_MATMUL:
op->builtin_options.AsBatchMatMulOptions()
->asymmetric_quantize_inputs = use_updated_hybrid_scheme;
break;
case BuiltinOperator_LSTM:
op->builtin_options.AsLSTMOptions()->asymmetric_quantize_inputs =
use_updated_hybrid_scheme;
break;
case BuiltinOperator_RNN:
op->builtin_options.AsRNNOptions()->asymmetric_quantize_inputs =
use_updated_hybrid_scheme;
break;
case BuiltinOperator_SVDF:
op->builtin_options.AsSVDFOptions()->asymmetric_quantize_inputs =
use_updated_hybrid_scheme;
break;
case BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM:
op->builtin_options.AsUnidirectionalSequenceLSTMOptions()
->asymmetric_quantize_inputs = use_updated_hybrid_scheme;
break;
case BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN:
op->builtin_options.AsSequenceRNNOptions()
->asymmetric_quantize_inputs = use_updated_hybrid_scheme;
break;
default:
break;
}
tensor_map->insert({tensor_idx, {tensor, /*is_per_channel=*/false}});
}
}
return absl::OkStatus();
}
// Updates operator code versions for the operators with INT8 inputs.
void UpdateInt8OperatorVersions(ModelT* model, bool use_updated_hybrid_scheme) {
for (int i = 0, end = model->operator_codes.size(); i < end; ++i) {
const BuiltinOperator& op_code =
GetBuiltinCode(model->operator_codes[i].get());
if (op_code == BuiltinOperator_RNN ||
op_code == BuiltinOperator_BIDIRECTIONAL_SEQUENCE_RNN ||
op_code == BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM ||
op_code == BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN) {
model->operator_codes[i]->version = use_updated_hybrid_scheme ? 3 : 2;
} else if (op_code == BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM ||
op_code == BuiltinOperator_EMBEDDING_LOOKUP) {
model->operator_codes[i]->version = 3;
} else if (op_code == BuiltinOperator_LSTM) {
model->operator_codes[i]->version = use_updated_hybrid_scheme ? 4 : 3;
} else if (op_code == BuiltinOperator_CONV_2D) {
model->operator_codes[i]->version = use_updated_hybrid_scheme ? 5 : 2;
} else if (op_code == BuiltinOperator_FULLY_CONNECTED) {
model->operator_codes[i]->version = use_updated_hybrid_scheme ? 9 : 3;
} else if (op_code == BuiltinOperator_BATCH_MATMUL) {
model->operator_codes[i]->version = use_updated_hybrid_scheme ? 4 : 1;
} else if (op_code == BuiltinOperator_SVDF) {
model->operator_codes[i]->version = use_updated_hybrid_scheme ? 4 : 2;
} else if (op_code == BuiltinOperator_DEPTHWISE_CONV_2D) {
model->operator_codes[i]->version = 6;
}
}
}
// Returns true if the op in consumer_op_infos can pass through quantization.
bool IsQuantizationPassThroughOps(
const ModelT* model, const std::vector<ConsumerOpInfo>& consumer_op_infos) {
if (consumer_op_infos.size() != 1) {
return false;
}
const OperatorT* consumer_op = consumer_op_infos.front().op;
const BuiltinOperator op_code =
GetBuiltinCode(model->operator_codes[consumer_op->opcode_index].get());
return op_code == BuiltinOperator_GATHER ||
op_code == BuiltinOperator_EMBEDDING_LOOKUP;
}
// Copies quantization parameters from input to output and returns consumers of
// the output tensor as a tuple with values:
// - index of the output tensor
// - pointer to the output tensor
// - vector of consumers ops.
std::tuple<int32_t, TensorT*, std::vector<ConsumerOpInfo>>
PassQuantizationAndGetConsumers(
const ModelT* model, const SubGraphT* subgraph,
const std::vector<ConsumerOpInfo>& consumer_op_infos,
const CustomOpMap& custom_op_map) {
const OperatorT* op = consumer_op_infos.front().op;
const OperatorCodeT* op_code = model->operator_codes[op->opcode_index].get();
if (op->outputs.size() != 1) {
LOG(ERROR)
<< "An op that passes quantization has more than one quantized output";
return std::make_tuple(-1, nullptr, std::vector<ConsumerOpInfo>());
}
const int32_t output_tensor_idx = op->outputs.front();
const auto input_idx = GetWeightInputIndices(op_code, custom_op_map);
if (input_idx.size() != 1) {
LOG(ERROR)
<< "An op that passes quantization has more than one quantized input";
return std::make_tuple(-1, nullptr, std::vector<ConsumerOpInfo>());
}
const int32_t input_tensor_idx = op->inputs[input_idx.front()];
// Propagate quantization params.
const TensorT* input_tensor = subgraph->tensors[input_tensor_idx].get();
TensorT* output_tensor = subgraph->tensors[output_tensor_idx].get();
if (!output_tensor->quantization) {
output_tensor->quantization =
std::make_unique<tflite::QuantizationParametersT>();
}
*output_tensor->quantization = *input_tensor->quantization;
output_tensor->type = TensorType_INT8;
return std::make_tuple(
output_tensor_idx, output_tensor,
GetTensorConsumers(model, subgraph, output_tensor_idx));
}
inline bool IsOpDenylisted(const flat_hash_set<BuiltinOperator>& op_denylist,
const BuiltinOperator op_code) {
return op_denylist.find(op_code) != op_denylist.end();
}
absl::Status QuantizeWeightsInt8(
flatbuffers::FlatBufferBuilder* builder, const Model* input_model,
bool use_hybrid_evaluation, uint64_t weights_min_num_elements,
const CustomOpMap& custom_op_map, bool use_updated_hybrid_scheme,
const absl::flat_hash_set<BuiltinOperator>& op_denylist = {}) {
std::unique_ptr<ModelT> model;
model.reset(input_model->UnPack());
for (int subgraph_index = 0, end = model->subgraphs.size();
subgraph_index < end; ++subgraph_index) {
SubGraphT* subgraph = model->subgraphs.at(subgraph_index).get();
absl::flat_hash_map<int32_t, TensorPerChannel> tensor_map;
for (int i = 0; i < subgraph->operators.size(); ++i) {
OperatorT* op = subgraph->operators[i].get();
absl::Status status = InsertQuantizableInputTensorsFromOperator(
model.get(), op, weights_min_num_elements, custom_op_map, &tensor_map,
subgraph_index, use_updated_hybrid_scheme);
if (!status.ok()) return status;
}
for (std::pair<int32_t, TensorPerChannel> tensor_pair : tensor_map) {
// Quantize the tensor.
if (tensor_pair.second.is_per_channel) {
if (!mlir::lite::toco_legacy::SymmetricQuantizeTensorPerChannel(
model.get(), tensor_pair.second.t,
tensor_pair.second.channel_dim)
.ok()) {
return absl::InternalError(
"SymmetricQuantizeTensorPerChannel failed");
}
} else {
if (!mlir::lite::toco_legacy::SymmetricQuantizeTensor(
model.get(), tensor_pair.second.t)
.ok()) {
return absl::InternalError("SymmetricQuantizeTensor failed");
}
}
}
// Examine the tensor consumers to determine which require dequantize ops.
for (const auto& tensor_pair : tensor_map) {
int32_t tensor_idx = tensor_pair.first;
TensorT* tensor = tensor_pair.second.t;
std::vector<ConsumerOpInfo> consumer_op_infos =
GetTensorConsumers(model.get(), subgraph, tensor_idx);
if (IsQuantizationPassThroughOps(model.get(), consumer_op_infos)) {
std::tie(tensor_idx, tensor, consumer_op_infos) =
PassQuantizationAndGetConsumers(model.get(), subgraph,
consumer_op_infos, custom_op_map);
if (tensor_idx < 0) {
// Error message is already logged by PassQuantizationAndGetConsumers.
return absl::InternalError("PassQuantizationAndGetConsumers failed");
}
}
std::vector<ConsumerOpInfo> dequant_op_infos; // Ops that need dequants.
for (ConsumerOpInfo& consumer_op_info : consumer_op_infos) {
OperatorT* consumer_op = consumer_op_info.op;
const OperatorCodeT* consumer_op_code =
model->operator_codes[consumer_op->opcode_index].get();
// If the op is a hybrid op and all the required tensors are quantized,
// we have no further work to do, but for all ops that require
// dequantization we need to add a Dequantize op.
bool eval_hybrid =
use_hybrid_evaluation &&
!IsOpDenylisted(op_denylist, GetBuiltinCode(consumer_op_code)) &&
IsHybridEvaluationOp(consumer_op, consumer_op_code, custom_op_map,
use_updated_hybrid_scheme) &&
CheckAllOpInputsQuantized(subgraph, consumer_op, consumer_op_code,
custom_op_map) &&
IsQuantizedInput(consumer_op_code, custom_op_map,
consumer_op_info.op_input_idx);
if (!eval_hybrid) {
dequant_op_infos.push_back(consumer_op_info);
}
}
// Check if this tensor is an output tensor.
int32_t output_index = -1;
for (int32_t i = 0; i < subgraph->outputs.size(); ++i) {
if (subgraph->outputs[i] == tensor_idx) {
output_index = i;
break;
}
}
// If no ops require dequant and it is not output, we are done for this
// tensor.
if (dequant_op_infos.empty() && output_index < 0) {
continue;
}
// Create a new tensor to be the output of the dequantize op.
std::unique_ptr<TensorT> dequantize_output;
const std::string dequant_name = tensor->name + "_dequantize";
mlir::lite::toco_legacy::MakeTensor(
dequant_name, tensor->shape, tensor->shape_signature,
TensorType_FLOAT32, &dequantize_output);
const int32_t dequantize_output_idx = subgraph->tensors.size();
subgraph->tensors.push_back(std::move(dequantize_output));
// Create the Dequantize operation.
std::unique_ptr<OperatorT> dequantize_op;
mlir::lite::toco_legacy::MakeDequantizeOperator(
model.get(), &dequantize_op, tensor_idx, dequantize_output_idx);
// Update the op_input of all the ops that need the created dequantize
// operation.
int32_t min_op_idx = subgraph->operators.size();
for (ConsumerOpInfo& dequant_op_info : dequant_op_infos) {
dequant_op_info.op->inputs[dequant_op_info.op_input_idx] =
dequantize_output_idx;
min_op_idx = std::min(dequant_op_info.op_idx, min_op_idx);
}
// Update output name.
if (output_index >= 0) {
subgraph->outputs[output_index] = dequantize_output_idx;
}
// Insert the newly created Dequantize operation before the earliest
// consumer, since TFLite requires operators to be topo-sorted.
subgraph->operators.insert(subgraph->operators.begin() + min_op_idx,
std::move(dequantize_op));
}
}
// Update the modified operator code versions.
UpdateInt8OperatorVersions(model.get(), use_updated_hybrid_scheme);
flatbuffers::Offset<Model> output_model_location =
Model::Pack(*builder, model.get());
FinishModelBuffer(*builder, output_model_location);
return absl::OkStatus();
}
absl::Status QuantizeWeightsFloat16(flatbuffers::FlatBufferBuilder* builder,
const Model* input_model) {
std::unique_ptr<ModelT> model;
model.reset(input_model->UnPack());
for (int subgraph_index = 0, end = model->subgraphs.size();
subgraph_index < end; ++subgraph_index) {
SubGraphT* subgraph = model->subgraphs.at(subgraph_index).get();
absl::flat_hash_map<int32_t, TensorT*> tensor_map;
for (int i = 0, sub_end = subgraph->operators.size(); i < sub_end; ++i) {
OperatorT* op = subgraph->operators[i].get();
for (auto tensor_idx : op->inputs) {
// Skip optional tensors.
if (tensor_idx == kTfLiteOptionalTensor) {
continue;
}
TensorT* tensor = subgraph->tensors[tensor_idx].get();
BufferT* buffer = model->buffers[tensor->buffer].get();
if (buffer == nullptr) {
return absl::InternalError("Buffer is null");
}
// Quantize tensors that have data to quantize.
bool is_constant = !model->buffers[tensor->buffer].get()->data.empty();
if (tensor->type == TensorType_FLOAT32 && is_constant) {
tensor_map.insert({tensor_idx, tensor});
}
}
}
// The hash map ensures that we quantize each tensor exactly once.
for (std::pair<int32_t, TensorT*> tensor_pair : tensor_map) {
// Quantize the tensor.
if (!mlir::lite::toco_legacy::QuantizeTensorFloat16(model.get(),
tensor_pair.second)
.ok()) {
return absl::InternalError("QuantizeTensorFloat16 failed");
}
int32_t tensor_idx = tensor_pair.first;
TensorT* tensor = tensor_pair.second;
std::vector<ConsumerOpInfo> dequant_op_infos =
GetTensorConsumers(model.get(), subgraph, tensor_idx);
// Create a new tensor to be the output of the dequantize op.
std::unique_ptr<TensorT> dequantize_output;
const std::string dequant_name = tensor->name + "_dequantize";
mlir::lite::toco_legacy::MakeTensor(
dequant_name, tensor->shape, tensor->shape_signature,
TensorType_FLOAT32, &dequantize_output);
const int32_t dequantize_output_idx = subgraph->tensors.size();
subgraph->tensors.push_back(std::move(dequantize_output));
// Create the Dequantize operation.
std::unique_ptr<OperatorT> dequantize_op;
mlir::lite::toco_legacy::MakeDequantizeOperator(
model.get(), &dequantize_op, tensor_idx, dequantize_output_idx);
// Update the op_input of all the ops that need the created dequantize
// operation.
int32_t min_op_idx = subgraph->operators.size();
for (ConsumerOpInfo& dequant_op_info : dequant_op_infos) {
dequant_op_info.op->inputs[dequant_op_info.op_input_idx] =
dequantize_output_idx;
min_op_idx = std::min(dequant_op_info.op_idx, min_op_idx);
}
// Insert the newly created Dequantize operation before the earliest
// consumer, since TFLite requires operators to be topo-sorted.
subgraph->operators.insert(subgraph->operators.begin() + min_op_idx,
std::move(dequantize_op));
}
}
flatbuffers::Offset<Model> output_model_location =
Model::Pack(*builder, model.get());
FinishModelBuffer(*builder, output_model_location);
return absl::OkStatus();
}
} // namespace
namespace internal {
absl::Status QuantizeWeights(flatbuffers::FlatBufferBuilder* builder,
const Model* input_model,
uint64_t weights_min_num_elements,
bool use_hybrid_evaluation,
QuantizerType quantizer_type) {
// By default we require that only weights with more than
// kWeightsMinSizeDefault elements are quantized.
if (quantizer_type == QuantizerType::MLIR_QUANTIZER) {
return mlir::lite::QuantizeWeights(
builder, input_model, weights_min_num_elements, use_hybrid_evaluation);
}
CustomOpMap custom_op_map;
return QuantizeWeightsInt8(builder, input_model, use_hybrid_evaluation,
weights_min_num_elements, custom_op_map,
kUseUpdatedHybridSchemeDefault);
}
} // namespace internal
absl::Status QuantizeWeights(flatbuffers::FlatBufferBuilder* builder,
const Model* input_model,
uint64_t weights_min_num_elements,
QuantizerType quantizer_type) {
if (quantizer_type == QuantizerType::MLIR_QUANTIZER) {
return mlir::lite::QuantizeWeights(builder, input_model,
weights_min_num_elements);
}
CustomOpMap custom_op_map;
return QuantizeWeightsInt8(builder, input_model, true,
weights_min_num_elements, custom_op_map,
kUseUpdatedHybridSchemeDefault);
}
absl::Status QuantizeWeights(flatbuffers::FlatBufferBuilder* builder,
const Model* input_model, BufferType quant_type,
bool use_updated_hybrid_scheme,
QuantizerType quantizer_type) {
// By default we require that only weights with more than
// kWeightsMinSizeDefault elements are quantized.
if (quantizer_type == QuantizerType::MLIR_QUANTIZER) {
return mlir::lite::QuantizeWeights(builder, input_model,
(mlir::lite::BufferType)quant_type,
use_updated_hybrid_scheme);
}
switch (quant_type) {
case BufferType::QUANTIZED_INT8: {
mlir::lite::toco_legacy::CustomOpMap custom_op_map;
return QuantizeWeightsInt8(builder, input_model, true,
kWeightsMinNumElementsDefault, custom_op_map,
use_updated_hybrid_scheme);
}
case BufferType::QUANTIZED_FLOAT16:
return QuantizeWeightsFloat16(builder, input_model);
}
}
absl::Status QuantizeWeights(flatbuffers::FlatBufferBuilder* builder,
const Model* input_model,
uint64_t weights_min_num_elements,
const CustomOpMap& custom_op_map,
QuantizerType quantizer_type) {
if (quantizer_type == QuantizerType::MLIR_QUANTIZER) {
mlir::lite::CustomOpMap mlir_custom_op_map;
ConstructMLIRCustomOpMap(mlir_custom_op_map, custom_op_map);
return mlir::lite::QuantizeWeights(
builder, input_model, weights_min_num_elements, mlir_custom_op_map);
}
return QuantizeWeightsInt8(builder, input_model, true,
weights_min_num_elements, custom_op_map,
kUseUpdatedHybridSchemeDefault);
}
absl::Status QuantizeWeights(flatbuffers::FlatBufferBuilder* builder,
const Model* input_model,
uint64_t weights_min_num_elements,
const CustomOpMap& custom_op_map,
bool use_updated_hybrid_scheme,
const flat_hash_set<BuiltinOperator>& op_denylist,
QuantizerType quantizer_type) {
if (quantizer_type == QuantizerType::MLIR_QUANTIZER) {
mlir::lite::CustomOpMap mlir_custom_op_map;
ConstructMLIRCustomOpMap(mlir_custom_op_map, custom_op_map);
return mlir::lite::QuantizeWeights(
builder, input_model, weights_min_num_elements, mlir_custom_op_map,
use_updated_hybrid_scheme, op_denylist);
}
return QuantizeWeightsInt8(builder, input_model,
/*use_hybrid_evaluation=*/true,
weights_min_num_elements, custom_op_map,
use_updated_hybrid_scheme, op_denylist);
}
} // namespace toco_legacy
} // namespace lite
} // namespace mlir
tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/quantize_weights.h
+109
View File
@@ -0,0 +1,109 @@
/* Copyright 2018 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_LITE_TOCO_LEGACY_QUANTIZE_WEIGHTS_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_LITE_TOCO_LEGACY_QUANTIZE_WEIGHTS_H_
#include <cstdint>
#include <string>
#include <unordered_map>
#include <vector>
#include "absl/container/flat_hash_set.h"
#include "absl/status/status.h"
#include "flatbuffers/flatbuffer_builder.h" // from @flatbuffers
#include "tensorflow/compiler/mlir/lite/schema/schema_generated.h"
namespace mlir {
namespace lite {
namespace toco_legacy {
using ::tflite::BuiltinOperator;
using ::tflite::Model;
// Supported resulting types from quantization process.
enum class BufferType { QUANTIZED_INT8, QUANTIZED_FLOAT16 };
enum class QuantizerType { OLD_QUANTIZER, MLIR_QUANTIZER };
// Stores information about how to quantize a user-specified custom operation.
struct CustomOpInfo {
std::vector<std::int32_t> quantizable_input_indices;
bool is_hybrid;
};
// Map from custom op code to custom op quantization information.
using CustomOpMap = std::unordered_map<std::string, CustomOpInfo>;
// This macro is for internal use for conversions requiring previous behavior.
#ifdef TFLITE_USE_PREVIOUS_HYBRID_SCHEME
// Use asymmetric quantized activations and per-channel quantized weights.
constexpr bool kUseUpdatedHybridSchemeDefault = false;
#else
// Use symmetric quantized activations and per-channel quantized weights.
constexpr bool kUseUpdatedHybridSchemeDefault = true;
#endif
// Quantizes input_model and populates the provided builder with the new model.
// By default only weights tensors weight more than 1024 elements will be
// quantized.
//
// A tflite::Model can be obtained from the builder with:
// const uint8_t* buffer = builder->GetBufferPointer();
// tflite::Model* model = GetModel(buffer);
absl::Status QuantizeWeights(
flatbuffers::FlatBufferBuilder* builder, const Model* input_model,
BufferType quant_type = BufferType::QUANTIZED_INT8,
bool use_updated_hybrid_scheme = kUseUpdatedHybridSchemeDefault,
QuantizerType quantizer_type = QuantizerType::OLD_QUANTIZER);
// Same as above, but only weights with greater than or equal
// weights_min_num_elements elements will be quantized.
absl::Status QuantizeWeights(
flatbuffers::FlatBufferBuilder* builder, const Model* input_model,
uint64_t weights_min_num_elements,
QuantizerType quantizer_type = QuantizerType::OLD_QUANTIZER);
// Same as above, but with entry point of quantizing custom ops.
absl::Status QuantizeWeights(
flatbuffers::FlatBufferBuilder* builder, const Model* input_model,
uint64_t weights_min_num_elements, const CustomOpMap& custom_op_map,
QuantizerType quantizer_type = QuantizerType::OLD_QUANTIZER);
// Same as above, but if use updated_hybrid_scheme is false,
// use previous quantization scheme. Optional op_denylist argument
// disables hybrid evaluation for provided BuiltinOperators.
absl::Status QuantizeWeights(
flatbuffers::FlatBufferBuilder* builder, const Model* input_model,
uint64_t weights_min_num_elements, const CustomOpMap& custom_op_map,
bool use_updated_hybrid_scheme,
const absl::flat_hash_set<BuiltinOperator>& op_denylist = {},
QuantizerType quantizer_type = QuantizerType::OLD_QUANTIZER);
namespace internal {
// If use_hybrid_evaluation is false, will disable using hybrid eval for
// operations that support it.
//
// We use this internal QuantizeWeights call to test models with hybrid
// evaluation disabled.
absl::Status QuantizeWeights(
flatbuffers::FlatBufferBuilder* builder, const Model* input_model,
uint64_t weights_min_num_elements, bool use_hybrid_evaluation,
QuantizerType quantizer_type = QuantizerType::OLD_QUANTIZER);
} // namespace internal
} // namespace toco_legacy
} // namespace lite
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_LITE_TOCO_LEGACY_QUANTIZE_WEIGHTS_H_
tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/quantize_weights_portable.cc
+739
View File
@@ -0,0 +1,739 @@
/* Copyright 2022 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.
==============================================================================*/
// clang-format off
#include "tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/quantize_weights.h"
// clang-format on
#include <algorithm>
#include <memory>
#include <string>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "absl/memory/memory.h"
#include "absl/status/status.h"
#include "flatbuffers/buffer.h" // from @flatbuffers
#include "flatbuffers/flatbuffer_builder.h" // from @flatbuffers
#include "tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/model_utils.h"
#include "tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/portable_tensor_utils.h"
#include "tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/quantization_utils.h"
#include "tensorflow/compiler/mlir/lite/schema/schema_generated.h"
#include "tensorflow/compiler/mlir/lite/schema/schema_utils.h"
#include "tensorflow/core/platform/logging.h"
namespace mlir {
namespace lite {
namespace toco_legacy {
namespace {
using absl::flat_hash_set;
using mlir::lite::toco_legacy::
CustomOpMap; // Use this instead of mlir::lite::CustomOpMap because that
// uses mlir::lite::CustomOpInfo in
// tensorflow/compiler/mlir/lite/quantization/lite/quantize_weights.h,
// and we need mlir::lite::toco_legacy::CustomOpInfo, in
// tensorflow/compiler/mlir/lite/quantization/lite/optimize/quantize_weights.h
using tflite::BufferT;
using tflite::BuiltinOperator;
using tflite::BuiltinOperator_BATCH_MATMUL;
using tflite::BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM;
using tflite::BuiltinOperator_BIDIRECTIONAL_SEQUENCE_RNN;
using tflite::BuiltinOperator_CONV_2D;
using tflite::BuiltinOperator_CUSTOM;
using tflite::BuiltinOperator_DEPTHWISE_CONV_2D;
using tflite::BuiltinOperator_EMBEDDING_LOOKUP;
using tflite::BuiltinOperator_FULLY_CONNECTED;
using tflite::BuiltinOperator_GATHER;
using tflite::BuiltinOperator_LSTM;
using tflite::BuiltinOperator_RNN;
using tflite::BuiltinOperator_SVDF;
using tflite::BuiltinOperator_TRANSPOSE_CONV;
using tflite::BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM;
using tflite::BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN;
using tflite::FinishModelBuffer;
using tflite::GetBuiltinCode;
using tflite::Model;
using tflite::ModelT;
using tflite::OperatorCodeT;
using tflite::OperatorT;
using tflite::SubGraphT;
using tflite::TensorT;
using tflite::TensorType_FLOAT32;
using tflite::TensorType_INT8;
struct ConsumerOpInfo {
OperatorT* op;
// The index of the op in the operators vector.
int32_t op_idx;
// The index of the tensor to quantize in subgraph->tensors.
int32_t op_input_idx;
};
struct TensorPerChannel {
TensorT* t;
bool is_per_channel;
int channel_dim;
};
// The default minimum number of elements a weights array must have to be
// quantized by this transformation.
const int kWeightsMinNumElementsDefault = 1024;
// Redefined from tensorflow/lite/core/c/common.h as local const int instead of
// discouraged #define macro.
const int kTfLiteOptionalTensor = -1;
// Gets the operators that consume tensor_idx.
std::vector<ConsumerOpInfo> GetTensorConsumers(const ModelT* model,
const SubGraphT* subgraph,
int32_t tensor_idx) {
// TODO(suharshs): If this proves to be too slow, avoid calling it per tensor,
// instead doing one sweep for the entire model.
std::vector<ConsumerOpInfo> consumer_ops;
for (size_t op_idx = 0; op_idx < subgraph->operators.size(); ++op_idx) {
OperatorT* op = subgraph->operators[op_idx].get();
if (op == nullptr) {
continue;
}
for (size_t i = 0; i < op->inputs.size(); ++i) {
if (op->inputs[i] == tensor_idx) {
consumer_ops.push_back(
{op, static_cast<int32_t>(op_idx), static_cast<int32_t>(i)});
}
}
}
return consumer_ops;
}
// Gets the list of op->inputs indices of the weights inputs to be quantized for
// the provided op.
std::vector<int32_t> GetWeightInputIndices(const OperatorCodeT* op_code,
const CustomOpMap& custom_op_map) {
const BuiltinOperator builtin_op_code = GetBuiltinCode(op_code);
if (builtin_op_code == BuiltinOperator_CUSTOM) {
const std::string custom_code = op_code->custom_code;
const auto& custom_op_info = custom_op_map.find(custom_code);
if (custom_op_info != custom_op_map.end()) {
return custom_op_info->second.quantizable_input_indices;
}
} else if (builtin_op_code == BuiltinOperator_CONV_2D ||
builtin_op_code == BuiltinOperator_DEPTHWISE_CONV_2D ||
builtin_op_code == BuiltinOperator_FULLY_CONNECTED ||
builtin_op_code == BuiltinOperator_BATCH_MATMUL ||
builtin_op_code == BuiltinOperator_EMBEDDING_LOOKUP ||
builtin_op_code == BuiltinOperator_TRANSPOSE_CONV) {
return {1};
} else if (builtin_op_code == BuiltinOperator_SVDF) {
// tensorflow/lite/kernels/svdf.cc
return {1, 2};
} else if (builtin_op_code == BuiltinOperator_LSTM ||
builtin_op_code == BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM) {
// tensorflow/lite/kernels/lstm.cc
// tensorflow/lite/kernels/unidirectional_sequence_lstm.cc
return {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 16};
} else if (builtin_op_code == BuiltinOperator_RNN ||
builtin_op_code == BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN) {
// tensorflow/lite/kernels/basic_rnn.cc
// tensorflow/lite/kernels/unidirectional_sequence_rnn.cc
return {1, 2};
} else if (builtin_op_code == BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM) {
// tensorflow/lite/kernels/bidirectional_sequence_lstm.cc
return {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 16, 18, 19, 20, 21,
22, 23, 24, 25, 26, 27, 28, 33, 40, 41, 42, 43, 44, 45, 46, 47};
} else if (builtin_op_code == BuiltinOperator_BIDIRECTIONAL_SEQUENCE_RNN) {
// tensorflow/lite/kernels/bidirectional_sequence_rnn.cc
return {1, 2, 4, 5, 6, 8, 9, 10, 11};
} else if (builtin_op_code == BuiltinOperator_GATHER) {
// tensorflow/lite/kernels/gather.cc
return {0};
}
return {};
}
// Checks that a specific input can be quantized.
bool IsQuantizedInput(const OperatorCodeT* op_code,
const CustomOpMap& custom_op_map, int op_input_idx) {
const auto quantized_input_indices =
GetWeightInputIndices(op_code, custom_op_map);
return std::find(std::begin(quantized_input_indices),
std::end(quantized_input_indices),
op_input_idx) != std::end(quantized_input_indices);
}
// Returns true if the operator supports hybrid evaluation.
bool IsHybridEvaluationOp(const OperatorT* op, const OperatorCodeT* op_code,
const CustomOpMap& custom_op_map,
bool use_updated_hybrid_scheme) {
const BuiltinOperator builtin_op_code = GetBuiltinCode(op_code);
// Operations that support hybrid evaluation.
bool eval_hybrid = false;
if (builtin_op_code == BuiltinOperator_CUSTOM) {
const std::string custom_code = op_code->custom_code;
const auto custom_op_info = custom_op_map.find(custom_code);
if (custom_op_info == custom_op_map.end()) {
return {};
} else {
return custom_op_info->second.is_hybrid;
}
} else if (builtin_op_code == BuiltinOperator_FULLY_CONNECTED ||
builtin_op_code == BuiltinOperator_BATCH_MATMUL ||
builtin_op_code == BuiltinOperator_CONV_2D ||
builtin_op_code == BuiltinOperator_SVDF ||
builtin_op_code == BuiltinOperator_RNN ||
builtin_op_code == BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM ||
builtin_op_code == BuiltinOperator_BIDIRECTIONAL_SEQUENCE_RNN ||
builtin_op_code == BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM ||
builtin_op_code == BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN) {
eval_hybrid = true;
} else if (builtin_op_code == BuiltinOperator_LSTM) {
const tflite::LSTMOptionsT* options = op->builtin_options.AsLSTMOptions();
// Only lstm kernel_type full supports hybrid evaluation.
if (options->kernel_type == tflite::LSTMKernelType_FULL) {
eval_hybrid = true;
}
} else if (builtin_op_code == BuiltinOperator_DEPTHWISE_CONV_2D) {
eval_hybrid = use_updated_hybrid_scheme;
}
return eval_hybrid;
}
// Returns true if all of the op's inputs are quantized.
bool CheckAllOpInputsQuantized(const SubGraphT* subgraph, const OperatorT* op,
const OperatorCodeT* op_code,
const CustomOpMap& custom_op_map) {
std::vector<int32_t> op_input_indices =
GetWeightInputIndices(op_code, custom_op_map);
for (const int32_t op_input_idx : op_input_indices) {
int32_t tensor_idx = op->inputs[op_input_idx];
if (tensor_idx == -1) {
// Optional tensor.
continue;
}
TensorT* tensor = subgraph->tensors[tensor_idx].get();
if (tensor->type != TensorType_INT8) {
return false;
}
}
return true;
}
// Inserts Tensors for each input tensor of op that should be
// quantized into tensor_map.
absl::Status InsertQuantizableInputTensorsFromOperator(
const ModelT* model, OperatorT* op, uint64_t weights_min_num_elements,
const CustomOpMap& custom_op_map,
absl::flat_hash_map<int32_t, TensorPerChannel>* tensor_map,
int subgraph_index, bool use_updated_hybrid_scheme) {
SubGraphT* subgraph = model->subgraphs.at(subgraph_index).get();
const OperatorCodeT* op_code = model->operator_codes[op->opcode_index].get();
auto builtin_code = GetBuiltinCode(op_code);
std::vector<int32_t> op_input_indices =
GetWeightInputIndices(op_code, custom_op_map);
for (const int32_t op_input_idx : op_input_indices) {
int32_t tensor_idx = op->inputs[op_input_idx];
if (tensor_idx == -1) {
LOG(INFO) << "Skipping optional tensor input " << op_input_idx
<< " of operation " << EnumNameBuiltinOperator(builtin_code);
continue;
}
TensorT* tensor = subgraph->tensors[tensor_idx].get();
if (tensor->type != TensorType_FLOAT32) {
LOG(INFO) << "Skipping quantization of tensor " << tensor->name
<< " that is not type float.";
continue;
}
uint64_t num_elements;
if (!mlir::lite::toco_legacy::NumElements(*tensor, &num_elements).ok()) {
return absl::InternalError("Error in quantization_utils NumElements");
}
if (num_elements < weights_min_num_elements) {
LOG(INFO) << "Skipping quantization of tensor " << tensor->name
<< " because it has fewer than " << weights_min_num_elements
<< " elements (" << num_elements << ").";
continue;
}
// Some tensors may have a null buffer vector, indicating an intermediate
// array.
if (model->buffers[tensor->buffer]->data.data() == nullptr) {
LOG(INFO) << "Skipping quantization of tensor " << tensor->name
<< " because it has no allocated buffer.";
continue;
}
if (builtin_code == BuiltinOperator_DEPTHWISE_CONV_2D) {
tensor_map->insert({tensor_idx,
{tensor, /*is_per_channel=*/use_updated_hybrid_scheme,
/*dim=*/3}});
} else if (builtin_code == BuiltinOperator_CONV_2D) {
tensor_map->insert({tensor_idx,
{tensor, /*is_per_channel=*/use_updated_hybrid_scheme,
/*dim=*/0}});
} else {
switch (builtin_code) {
case BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM:
op->builtin_options.AsBidirectionalSequenceLSTMOptions()
->asymmetric_quantize_inputs = use_updated_hybrid_scheme;
break;
case BuiltinOperator_BIDIRECTIONAL_SEQUENCE_RNN:
op->builtin_options.AsBidirectionalSequenceRNNOptions()
->asymmetric_quantize_inputs = use_updated_hybrid_scheme;
break;
case BuiltinOperator_FULLY_CONNECTED:
op->builtin_options.AsFullyConnectedOptions()
->asymmetric_quantize_inputs = use_updated_hybrid_scheme;
break;
case BuiltinOperator_BATCH_MATMUL:
op->builtin_options.AsBatchMatMulOptions()
->asymmetric_quantize_inputs = use_updated_hybrid_scheme;
break;
case BuiltinOperator_LSTM:
op->builtin_options.AsLSTMOptions()->asymmetric_quantize_inputs =
use_updated_hybrid_scheme;
break;
case BuiltinOperator_RNN:
op->builtin_options.AsRNNOptions()->asymmetric_quantize_inputs =
use_updated_hybrid_scheme;
break;
case BuiltinOperator_SVDF:
op->builtin_options.AsSVDFOptions()->asymmetric_quantize_inputs =
use_updated_hybrid_scheme;
break;
case BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM:
op->builtin_options.AsUnidirectionalSequenceLSTMOptions()
->asymmetric_quantize_inputs = use_updated_hybrid_scheme;
break;
case BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN:
op->builtin_options.AsSequenceRNNOptions()
->asymmetric_quantize_inputs = use_updated_hybrid_scheme;
break;
default:
break;
}
tensor_map->insert({tensor_idx, {tensor, /*is_per_channel=*/false}});
}
}
return absl::OkStatus();
}
// Updates operator code versions for the operators with INT8 inputs.
void UpdateInt8OperatorVersions(ModelT* model, bool use_updated_hybrid_scheme) {
for (int i = 0, end = model->operator_codes.size(); i < end; ++i) {
const BuiltinOperator& op_code =
GetBuiltinCode(model->operator_codes[i].get());
if (op_code == BuiltinOperator_RNN ||
op_code == BuiltinOperator_BIDIRECTIONAL_SEQUENCE_RNN ||
op_code == BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM ||
op_code == BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN) {
model->operator_codes[i]->version = use_updated_hybrid_scheme ? 3 : 2;
} else if (op_code == BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM ||
op_code == BuiltinOperator_EMBEDDING_LOOKUP) {
model->operator_codes[i]->version = 3;
} else if (op_code == BuiltinOperator_LSTM) {
model->operator_codes[i]->version = use_updated_hybrid_scheme ? 4 : 3;
} else if (op_code == BuiltinOperator_CONV_2D) {
model->operator_codes[i]->version = use_updated_hybrid_scheme ? 5 : 2;
} else if (op_code == BuiltinOperator_FULLY_CONNECTED) {
model->operator_codes[i]->version = use_updated_hybrid_scheme ? 9 : 3;
} else if (op_code == BuiltinOperator_BATCH_MATMUL) {
model->operator_codes[i]->version = use_updated_hybrid_scheme ? 4 : 1;
} else if (op_code == BuiltinOperator_SVDF) {
model->operator_codes[i]->version = use_updated_hybrid_scheme ? 4 : 2;
} else if (op_code == BuiltinOperator_DEPTHWISE_CONV_2D) {
model->operator_codes[i]->version = 6;
}
}
}
// Returns true if the op in consumer_op_infos can pass through quantization.
bool IsQuantizationPassThroughOps(
const ModelT* model, const std::vector<ConsumerOpInfo>& consumer_op_infos) {
if (consumer_op_infos.size() != 1) {
return false;
}
const OperatorT* consumer_op = consumer_op_infos.front().op;
const BuiltinOperator op_code =
GetBuiltinCode(model->operator_codes[consumer_op->opcode_index].get());
return op_code == BuiltinOperator_GATHER ||
op_code == BuiltinOperator_EMBEDDING_LOOKUP;
}
// Copies quantization parameters from input to output and returns consumers of
// the output tensor as a tuple with values:
// - index of the output tensor
// - pointer to the output tensor
// - vector of consumers ops.
std::tuple<int32_t, TensorT*, std::vector<ConsumerOpInfo>>
PassQuantizationAndGetConsumers(
const ModelT* model, const SubGraphT* subgraph,
const std::vector<ConsumerOpInfo>& consumer_op_infos,
const CustomOpMap& custom_op_map) {
const OperatorT* op = consumer_op_infos.front().op;
const OperatorCodeT* op_code = model->operator_codes[op->opcode_index].get();
if (op->outputs.size() != 1) {
LOG(ERROR)
<< "An op that passes quantization has more than one quantized output";
return std::make_tuple(-1, nullptr, std::vector<ConsumerOpInfo>());
}
const int32_t output_tensor_idx = op->outputs.front();
const auto input_idx = GetWeightInputIndices(op_code, custom_op_map);
if (input_idx.size() != 1) {
LOG(ERROR)
<< "An op that passes quantization has more than one quantized input";
return std::make_tuple(-1, nullptr, std::vector<ConsumerOpInfo>());
}
const int32_t input_tensor_idx = op->inputs[input_idx.front()];
// Propagate quantization params.
const TensorT* input_tensor = subgraph->tensors[input_tensor_idx].get();
TensorT* output_tensor = subgraph->tensors[output_tensor_idx].get();
if (!output_tensor->quantization) {
output_tensor->quantization =
std::make_unique<tflite::QuantizationParametersT>();
}
*output_tensor->quantization = *input_tensor->quantization;
output_tensor->type = TensorType_INT8;
return std::make_tuple(
output_tensor_idx, output_tensor,
GetTensorConsumers(model, subgraph, output_tensor_idx));
}
inline bool IsOpDenylisted(const flat_hash_set<BuiltinOperator>& op_denylist,
const BuiltinOperator op_code) {
return op_denylist.find(op_code) != op_denylist.end();
}
absl::Status QuantizeWeightsInt8(
flatbuffers::FlatBufferBuilder* builder, const Model* input_model,
bool use_hybrid_evaluation, uint64_t weights_min_num_elements,
const CustomOpMap& custom_op_map, bool use_updated_hybrid_scheme,
const flat_hash_set<BuiltinOperator>& op_denylist = {}) {
std::unique_ptr<ModelT> model;
model.reset(input_model->UnPack());
for (int subgraph_index = 0, end = model->subgraphs.size();
subgraph_index < end; ++subgraph_index) {
SubGraphT* subgraph = model->subgraphs.at(subgraph_index).get();
absl::flat_hash_map<int32_t, TensorPerChannel> tensor_map;
for (int i = 0; i < subgraph->operators.size(); ++i) {
OperatorT* op = subgraph->operators[i].get();
if (!InsertQuantizableInputTensorsFromOperator(
model.get(), op, weights_min_num_elements, custom_op_map,
&tensor_map, subgraph_index, use_updated_hybrid_scheme)
.ok()) {
return absl::InternalError(
"Failed to insert quantizable input tensors from operator");
}
}
for (std::pair<int32_t, TensorPerChannel> tensor_pair : tensor_map) {
// Quantize the tensor.
if (tensor_pair.second.is_per_channel) {
if (!mlir::lite::toco_legacy::SymmetricQuantizeTensorPerChannel(
model.get(), tensor_pair.second.t,
tensor_pair.second.channel_dim)
.ok()) {
return absl::InternalError("Failed to quantize tensor per channel");
}
} else {
if (!mlir::lite::toco_legacy::SymmetricQuantizeTensor(
model.get(), tensor_pair.second.t)
.ok()) {
return absl::InternalError("Failed to quantize tensor");
}
}
}
// Examine the tensor consumers to determine which require dequantize ops.
for (const auto& tensor_pair : tensor_map) {
int32_t tensor_idx = tensor_pair.first;
TensorT* tensor = tensor_pair.second.t;
std::vector<ConsumerOpInfo> consumer_op_infos =
GetTensorConsumers(model.get(), subgraph, tensor_idx);
if (IsQuantizationPassThroughOps(model.get(), consumer_op_infos)) {
std::tie(tensor_idx, tensor, consumer_op_infos) =
PassQuantizationAndGetConsumers(model.get(), subgraph,
consumer_op_infos, custom_op_map);
if (tensor_idx < 0) {
// Error message is already logged by PassQuantizationAndGetConsumers.
return absl::InternalError(
"Failed to pass quantization and get consumers");
}
}
std::vector<ConsumerOpInfo> dequant_op_infos; // Ops that need dequants.
for (ConsumerOpInfo& consumer_op_info : consumer_op_infos) {
OperatorT* consumer_op = consumer_op_info.op;
const OperatorCodeT* consumer_op_code =
model->operator_codes[consumer_op->opcode_index].get();
// If the op is a hybrid op and all the required tensors are quantized,
// we have no further work to do, but for all ops that require
// dequantization we need to add a Dequantize op.
bool eval_hybrid =
use_hybrid_evaluation &&
!IsOpDenylisted(op_denylist, GetBuiltinCode(consumer_op_code)) &&
IsHybridEvaluationOp(consumer_op, consumer_op_code, custom_op_map,
use_updated_hybrid_scheme) &&
CheckAllOpInputsQuantized(subgraph, consumer_op, consumer_op_code,
custom_op_map) &&
IsQuantizedInput(consumer_op_code, custom_op_map,
consumer_op_info.op_input_idx);
if (!eval_hybrid) {
dequant_op_infos.push_back(consumer_op_info);
}
}
// Check if this tensor is an output tensor.
int32_t output_index = -1;
for (int32_t i = 0; i < subgraph->outputs.size(); ++i) {
if (subgraph->outputs[i] == tensor_idx) {
output_index = i;
break;
}
}
// If no ops require dequant and it is not output, we are done for this
// tensor.
if (dequant_op_infos.empty() && output_index < 0) {
continue;
}
// Create a new tensor to be the output of the dequantize op.
std::unique_ptr<TensorT> dequantize_output;
const std::string dequant_name = tensor->name + "_dequantize";
mlir::lite::toco_legacy::MakeTensor(
dequant_name, tensor->shape, tensor->shape_signature,
TensorType_FLOAT32, &dequantize_output);
const int32_t dequantize_output_idx = subgraph->tensors.size();
subgraph->tensors.push_back(std::move(dequantize_output));
// Create the Dequantize operation.
std::unique_ptr<OperatorT> dequantize_op;
mlir::lite::toco_legacy::MakeDequantizeOperator(
model.get(), &dequantize_op, tensor_idx, dequantize_output_idx);
// Update the op_input of all the ops that need the created dequantize
// operation.
int32_t min_op_idx = subgraph->operators.size();
for (ConsumerOpInfo& dequant_op_info : dequant_op_infos) {
dequant_op_info.op->inputs[dequant_op_info.op_input_idx] =
dequantize_output_idx;
min_op_idx = std::min(dequant_op_info.op_idx, min_op_idx);
}
// Update output name.
if (output_index >= 0) {
subgraph->outputs[output_index] = dequantize_output_idx;
}
// Insert the newly created Dequantize operation before the earliest
// consumer, since TFLite requires operators to be topo-sorted.
subgraph->operators.insert(subgraph->operators.begin() + min_op_idx,
std::move(dequantize_op));
}
}
// Update the modified operator code versions.
UpdateInt8OperatorVersions(model.get(), use_updated_hybrid_scheme);
flatbuffers::Offset<Model> output_model_location =
Model::Pack(*builder, model.get());
FinishModelBuffer(*builder, output_model_location);
return absl::OkStatus();
}
absl::Status QuantizeWeightsFloat16(flatbuffers::FlatBufferBuilder* builder,
const Model* input_model) {
std::unique_ptr<ModelT> model;
model.reset(input_model->UnPack());
for (int subgraph_index = 0, end = model->subgraphs.size();
subgraph_index < end; ++subgraph_index) {
SubGraphT* subgraph = model->subgraphs.at(subgraph_index).get();
absl::flat_hash_map<int32_t, TensorT*> tensor_map;
for (int i = 0, sub_end = subgraph->operators.size(); i < sub_end; ++i) {
OperatorT* op = subgraph->operators[i].get();
for (auto tensor_idx : op->inputs) {
// Skip optional tensors.
if (tensor_idx == kTfLiteOptionalTensor) {
continue;
}
TensorT* tensor = subgraph->tensors[tensor_idx].get();
BufferT* buffer = model->buffers[tensor->buffer].get();
if (buffer == nullptr) {
return absl::InternalError("Buffer is null");
}
// Quantize tensors that have data to quantize.
bool is_constant = !model->buffers[tensor->buffer].get()->data.empty();
if (tensor->type == TensorType_FLOAT32 && is_constant) {
tensor_map.insert({tensor_idx, tensor});
}
}
}
// The hash map ensures that we quantize each tensor exactly once.
for (std::pair<int32_t, TensorT*> tensor_pair : tensor_map) {
// Quantize the tensor.
if (!mlir::lite::toco_legacy::QuantizeTensorFloat16(model.get(),
tensor_pair.second)
.ok()) {
return absl::InternalError("QuantizeTensorFloat16 failed");
}
int32_t tensor_idx = tensor_pair.first;
TensorT* tensor = tensor_pair.second;
std::vector<ConsumerOpInfo> dequant_op_infos =
GetTensorConsumers(model.get(), subgraph, tensor_idx);
// Create a new tensor to be the output of the dequantize op.
std::unique_ptr<TensorT> dequantize_output;
const std::string dequant_name = tensor->name + "_dequantize";
mlir::lite::toco_legacy::MakeTensor(
dequant_name, tensor->shape, tensor->shape_signature,
TensorType_FLOAT32, &dequantize_output);
const int32_t dequantize_output_idx = subgraph->tensors.size();
subgraph->tensors.push_back(std::move(dequantize_output));
// Create the Dequantize operation.
std::unique_ptr<OperatorT> dequantize_op;
mlir::lite::toco_legacy::MakeDequantizeOperator(
model.get(), &dequantize_op, tensor_idx, dequantize_output_idx);
// Update the op_input of all the ops that need the created dequantize
// operation.
int32_t min_op_idx = subgraph->operators.size();
for (ConsumerOpInfo& dequant_op_info : dequant_op_infos) {
dequant_op_info.op->inputs[dequant_op_info.op_input_idx] =
dequantize_output_idx;
min_op_idx = std::min(dequant_op_info.op_idx, min_op_idx);
}
// Insert the newly created Dequantize operation before the earliest
// consumer, since TFLite requires operators to be topo-sorted.
subgraph->operators.insert(subgraph->operators.begin() + min_op_idx,
std::move(dequantize_op));
}
}
flatbuffers::Offset<Model> output_model_location =
Model::Pack(*builder, model.get());
FinishModelBuffer(*builder, output_model_location);
return absl::OkStatus();
}
} // namespace
namespace internal {
absl::Status QuantizeWeights(flatbuffers::FlatBufferBuilder* builder,
const Model* input_model,
uint64_t weights_min_num_elements,
bool use_hybrid_evaluation,
QuantizerType quantizer_type) {
if (quantizer_type == QuantizerType::MLIR_QUANTIZER) {
LOG(ERROR) << "Portable targets cannot use the MLIR quantizer.";
return absl::InternalError(
"Portable targets cannot use the MLIR quantizer.");
}
// By default we require that only weights with more than
// kWeightsMinSizeDefault elements are quantized.
CustomOpMap custom_op_map;
return QuantizeWeightsInt8(builder, input_model, use_hybrid_evaluation,
weights_min_num_elements, custom_op_map,
kUseUpdatedHybridSchemeDefault);
}
} // namespace internal
absl::Status QuantizeWeights(flatbuffers::FlatBufferBuilder* builder,
const Model* input_model,
uint64_t weights_min_num_elements,
QuantizerType quantizer_type) {
if (quantizer_type == QuantizerType::MLIR_QUANTIZER) {
LOG(ERROR) << "Portable targets cannot use the MLIR quantizer.";
return absl::InternalError(
"Portable targets cannot use the MLIR quantizer.");
}
CustomOpMap custom_op_map;
return QuantizeWeightsInt8(builder, input_model, true,
weights_min_num_elements, custom_op_map,
kUseUpdatedHybridSchemeDefault);
}
absl::Status QuantizeWeights(flatbuffers::FlatBufferBuilder* builder,
const Model* input_model, BufferType quant_type,
bool use_updated_hybrid_scheme,
QuantizerType quantizer_type) {
if (quantizer_type == QuantizerType::MLIR_QUANTIZER) {
LOG(ERROR) << "Portable targets cannot use the MLIR quantizer.";
return absl::InternalError(
"Portable targets cannot use the MLIR quantizer.");
}
switch (quant_type) {
case BufferType::QUANTIZED_INT8: {
// By default we require that only weights with more than
// kWeightsMinSizeDefault elements are quantized.
CustomOpMap custom_op_map;
return QuantizeWeightsInt8(builder, input_model, true,
kWeightsMinNumElementsDefault, custom_op_map,
use_updated_hybrid_scheme);
}
case BufferType::QUANTIZED_FLOAT16:
return QuantizeWeightsFloat16(builder, input_model);
}
}
absl::Status QuantizeWeights(flatbuffers::FlatBufferBuilder* builder,
const Model* input_model,
uint64_t weights_min_num_elements,
const CustomOpMap& custom_op_map,
QuantizerType quantizer_type) {
if (quantizer_type == QuantizerType::MLIR_QUANTIZER) {
LOG(ERROR) << "Portable targets cannot use the MLIR quantizer.";
return absl::InternalError(
"Portable targets cannot use the MLIR quantizer.");
}
return QuantizeWeightsInt8(builder, input_model, true,
weights_min_num_elements, custom_op_map,
kUseUpdatedHybridSchemeDefault);
}
absl::Status QuantizeWeights(flatbuffers::FlatBufferBuilder* builder,
const Model* input_model,
uint64_t weights_min_num_elements,
const CustomOpMap& custom_op_map,
bool use_updated_hybrid_scheme,
const flat_hash_set<BuiltinOperator>& op_denylist,
QuantizerType quantizer_type) {
if (quantizer_type == QuantizerType::MLIR_QUANTIZER) {
LOG(ERROR) << "Portable targets cannot use the MLIR quantizer.";
return absl::InternalError(
"Portable targets cannot use the MLIR quantizer.");
}
return QuantizeWeightsInt8(builder, input_model,
/*use_hybrid_evaluation=*/true,
weights_min_num_elements, custom_op_map,
use_updated_hybrid_scheme, op_denylist);
}
} // namespace toco_legacy
} // namespace lite
} // namespace mlir
tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/quantize_weights_test.cc
+701
View File
@@ -0,0 +1,701 @@
/* Copyright 2018 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/lite/toco_legacy/quantize_weights.h"
#include <cstddef>
#include <cstdint>
#include <cstdlib>
#include <iostream>
#include <memory>
#include <string>
#include <vector>
#include <gtest/gtest.h>
#include "flatbuffers/flatbuffer_builder.h" // from @flatbuffers
#include "flatbuffers/vector.h" // from @flatbuffers
#include "tensorflow/compiler/mlir/lite/core/absl_error_model_builder.h"
#include "tensorflow/compiler/mlir/lite/quantization/lite/test_util.h"
#include "tensorflow/compiler/mlir/lite/schema/schema_generated.h"
#include "tensorflow/compiler/mlir/lite/schema/schema_utils.h"
#include "xla/tsl/util/command_line_flags.h"
#include "tsl/platform/init_main.h"
#include "tsl/platform/path.h"
namespace {
std::string* g_test_model_dir = nullptr;
} // namespace
namespace mlir {
namespace lite {
namespace toco_legacy {
namespace {
using mlir::TFL::FlatBufferModelAbslError;
using tflite::BuiltinOperator_CONV_2D;
using tflite::BuiltinOperator_CUSTOM;
using tflite::BuiltinOperator_DEQUANTIZE;
using tflite::GetModel;
using tflite::Model;
using tflite::TensorType_FLOAT16;
using tflite::TensorType_FLOAT32;
using tflite::TensorType_INT8;
std::unique_ptr<FlatBufferModelAbslError> ReadTestModel() {
auto model_path = tsl::io::JoinPath(
*g_test_model_dir, ::mlir::lite::internal::kConvModelWith0Plus10Weights);
return FlatBufferModelAbslError::BuildFromFile(model_path.c_str());
}
std::unique_ptr<FlatBufferModelAbslError> ReadSharedWeightsTestModel() {
auto model_path = tsl::io::JoinPath(
*g_test_model_dir, ::mlir::lite::internal::kModelWithSharedWeights);
return FlatBufferModelAbslError::BuildFromFile(model_path.c_str());
}
std::unique_ptr<FlatBufferModelAbslError> ReadGatherTestModel() {
auto model_path = tsl::io::JoinPath(
*g_test_model_dir, ::mlir::lite::internal::kQuantizedWithGather);
return FlatBufferModelAbslError::BuildFromFile(model_path.c_str());
}
std::unique_ptr<FlatBufferModelAbslError> ReadCustomOpTestModel() {
auto model_path = tsl::io::JoinPath(
*g_test_model_dir, ::mlir::lite::internal::kModelWithCustomOp);
return FlatBufferModelAbslError::BuildFromFile(model_path.c_str());
}
template <typename T>
std::vector<T> GetAsVector(const flatbuffers::Vector<T>* vec) {
return std::vector<T>(vec->begin(), vec->end());
}
class QuantizeWeightsTest : public testing::Test {
protected:
QuantizeWeightsTest() = default;
void LoadBasicModel() {
input_model_ = ReadTestModel();
model_ = input_model_->GetModel();
}
void LoadSharedWeightsModel() {
input_model_ = ReadSharedWeightsTestModel();
model_ = input_model_->GetModel();
}
void LoadGatherTestModel() {
input_model_ = ReadGatherTestModel();
model_ = input_model_->GetModel();
}
void LoadCustomOpTestModel() {
input_model_ = ReadCustomOpTestModel();
model_ = input_model_->GetModel();
}
std::unique_ptr<FlatBufferModelAbslError> input_model_;
const Model* model_;
bool IsModelInputOrOutput(const Model* model, uint32_t tensor_idx) {
for (size_t subgraph_idx = 0; subgraph_idx < model_->subgraphs()->size();
++subgraph_idx) {
const auto subgraph = model->subgraphs()->Get(subgraph_idx);
for (size_t i = 0; i < subgraph->inputs()->size(); ++i) {
if (subgraph->inputs()->Get(i) == tensor_idx) {
return true;
}
}
for (size_t i = 0; i < subgraph->outputs()->size(); ++i) {
if (subgraph->outputs()->Get(i) == tensor_idx) {
return true;
}
}
}
return false;
}
// Returns the producer op code of the specified tensor_idx.
bool GetProducerOpCode(const Model* model, uint32_t subgraph_idx,
uint32_t tensor_idx,
tflite::BuiltinOperator* op_code) {
const auto subgraph = model->subgraphs()->Get(subgraph_idx);
for (size_t op_idx = 0; op_idx < subgraph->operators()->size(); ++op_idx) {
const auto op = subgraph->operators()->Get(op_idx);
for (size_t i = 0; i < op->outputs()->size(); ++i) {
if (op->outputs()->Get(i) == tensor_idx) {
const uint32_t op_code_idx = op->opcode_index();
*op_code = GetBuiltinCode(model->operator_codes()->Get(op_code_idx));
return true;
}
}
}
return false;
}
};
TEST_F(QuantizeWeightsTest, QuantizationSucceeds) {
LoadBasicModel();
flatbuffers::FlatBufferBuilder builder;
ASSERT_TRUE(
QuantizeWeights(&builder, model_, 0, QuantizerType::OLD_QUANTIZER).ok());
const uint8_t* buffer = builder.GetBufferPointer();
const Model* output_model = GetModel(buffer);
ASSERT_TRUE(output_model);
}
TEST_F(QuantizeWeightsTest, WeightsMinNumElements) {
LoadBasicModel();
// Make weights_min_size sufficiently large such that no quantization should
// happen, i.e. the original model is the same size as the old one.
flatbuffers::FlatBufferBuilder builder;
const uint64_t kWeightsMinNumElements = 1000000;
ASSERT_TRUE(QuantizeWeights(&builder, model_, kWeightsMinNumElements,
QuantizerType::OLD_QUANTIZER)
.ok());
const uint8_t* buffer = builder.GetBufferPointer();
const Model* output_model = GetModel(buffer);
ASSERT_TRUE(output_model);
for (size_t subgraph_idx = 0; subgraph_idx < model_->subgraphs()->size();
subgraph_idx++) {
const auto quantized_graph = output_model->subgraphs()->Get(subgraph_idx);
const auto float_graph = model_->subgraphs()->Get(subgraph_idx);
ASSERT_EQ(quantized_graph->tensors()->size(),
float_graph->tensors()->size());
for (size_t i = 0; i < quantized_graph->tensors()->size(); i++) {
const auto quant_tensor = quantized_graph->tensors()->Get(i);
const auto float_tensor = float_graph->tensors()->Get(i);
// Everything should remain equal between the two graphs.
EXPECT_EQ(quant_tensor->buffer(), float_tensor->buffer());
EXPECT_EQ(quant_tensor->is_variable(), float_tensor->is_variable());
EXPECT_EQ(GetAsVector(quant_tensor->shape()),
GetAsVector(float_tensor->shape()));
EXPECT_EQ(quant_tensor->name()->str(), float_tensor->name()->str());
EXPECT_EQ(quant_tensor->type(), float_tensor->type());
}
}
}
TEST_F(QuantizeWeightsTest, HybridConv) {
LoadBasicModel();
flatbuffers::FlatBufferBuilder builder;
ASSERT_TRUE(
QuantizeWeights(&builder, model_, 0, QuantizerType::OLD_QUANTIZER).ok());
const uint8_t* buffer = builder.GetBufferPointer();
const Model* output_model = GetModel(buffer);
ASSERT_TRUE(output_model);
// Nothing should change.
ASSERT_EQ(output_model->subgraphs()->size(), model_->subgraphs()->size());
for (size_t subgraph_idx = 0; subgraph_idx < model_->subgraphs()->size();
subgraph_idx++) {
const auto quantized_graph = output_model->subgraphs()->Get(subgraph_idx);
const auto float_graph = model_->subgraphs()->Get(subgraph_idx);
ASSERT_EQ(quantized_graph->tensors()->size(),
float_graph->tensors()->size());
// Make sure the graph only has one Conv operation.
ASSERT_EQ(quantized_graph->operators()->size(), 1);
const auto op = quantized_graph->operators()->Get(0);
const uint32_t op_code_idx = op->opcode_index();
ASSERT_EQ(GetBuiltinCode(output_model->operator_codes()->Get(op_code_idx)),
BuiltinOperator_CONV_2D);
for (size_t i = 0; i < quantized_graph->tensors()->size(); i++) {
const auto quant_tensor = quantized_graph->tensors()->Get(i);
const auto float_tensor = float_graph->tensors()->Get(i);
EXPECT_EQ(quant_tensor->buffer(), float_tensor->buffer());
EXPECT_EQ(quant_tensor->is_variable(), float_tensor->is_variable());
EXPECT_EQ(GetAsVector(quant_tensor->shape()),
GetAsVector(float_tensor->shape()));
EXPECT_EQ(quant_tensor->name()->str(), float_tensor->name()->str());
// If the tensor is a weight, it should have type INT8, otherwise it
// should stay with type FLOAT32.
// If the tensor is a bias, it should have type FLOAT32.
if (quant_tensor->name()->str() == "conv_bias") {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
} else if (IsModelInputOrOutput(output_model, i)) {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
} else if (quant_tensor->buffer() != 0) {
EXPECT_EQ(quant_tensor->type(), TensorType_INT8)
<< quant_tensor->name()->str();
auto shape = GetAsVector(quant_tensor->shape());
if (kUseUpdatedHybridSchemeDefault) {
EXPECT_EQ(quant_tensor->quantization()->scale()->size(), shape[0]);
} else {
EXPECT_EQ(quant_tensor->quantization()->scale()->size(), 1);
}
} else {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
}
}
}
}
TEST_F(QuantizeWeightsTest, DequantizeConv) {
LoadBasicModel();
flatbuffers::FlatBufferBuilder builder;
ASSERT_TRUE(internal::QuantizeWeights(&builder, model_, 0,
/*use_hybrid_evaluation=*/false,
QuantizerType::OLD_QUANTIZER)
.ok());
const uint8_t* buffer = builder.GetBufferPointer();
const Model* output_model = GetModel(buffer);
ASSERT_TRUE(output_model);
ASSERT_EQ(output_model->subgraphs()->size(), model_->subgraphs()->size());
for (size_t subgraph_idx = 0; subgraph_idx < model_->subgraphs()->size();
++subgraph_idx) {
const auto quantized_graph = output_model->subgraphs()->Get(subgraph_idx);
const auto float_graph = model_->subgraphs()->Get(subgraph_idx);
// The output graph should have an extra tensor from the added dequantize
// op.
ASSERT_EQ(quantized_graph->tensors()->size(),
float_graph->tensors()->size() + 1);
// Check that a dequantize op exists.
int32_t dequant_input_idx = -1;
int32_t dequant_output_idx = -1;
for (size_t i = 0; i < quantized_graph->operators()->size(); ++i) {
const auto op = quantized_graph->operators()->Get(i);
const uint32_t op_code_idx = op->opcode_index();
if (GetBuiltinCode(output_model->operator_codes()->Get(op_code_idx)) ==
BuiltinOperator_DEQUANTIZE) {
dequant_input_idx = op->inputs()->Get(0);
dequant_output_idx = op->outputs()->Get(0);
}
}
ASSERT_GT(dequant_input_idx, -1);
ASSERT_GT(dequant_output_idx, -1);
for (size_t i = 0; i < quantized_graph->tensors()->size(); ++i) {
const auto quant_tensor = quantized_graph->tensors()->Get(i);
// If the tensor is a weight, it should have type INT8.
// If the tensor is a bias, it should have type FLOAT32.
// If the tensor is an input or output it should have type FLOAT32.
// The input to dequantize should be INT8, and all other tensors should be
// FLOAT32.
if (i == dequant_input_idx) {
EXPECT_EQ(quant_tensor->type(), TensorType_INT8);
} else if (i == dequant_output_idx) {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
} else if (IsModelInputOrOutput(output_model, i)) {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
} else if (quant_tensor->name()->str() == "conv_bias") {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
} else if (quant_tensor->buffer() != 0) {
// If it's a non-bias constant tensor, it must be the weight.
EXPECT_EQ(quant_tensor->type(), TensorType_INT8);
} else {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
}
}
}
}
TEST_F(QuantizeWeightsTest, DequantizeConvFloat16) {
LoadBasicModel();
flatbuffers::FlatBufferBuilder builder;
ASSERT_TRUE(QuantizeWeights(&builder, model_, BufferType::QUANTIZED_FLOAT16,
kUseUpdatedHybridSchemeDefault,
QuantizerType::OLD_QUANTIZER)
.ok());
const uint8_t* buffer = builder.GetBufferPointer();
const Model* output_model = GetModel(buffer);
ASSERT_TRUE(output_model);
ASSERT_EQ(output_model->subgraphs()->size(), model_->subgraphs()->size());
for (size_t subgraph_idx = 0; subgraph_idx < model_->subgraphs()->size();
++subgraph_idx) {
const auto quantized_graph = output_model->subgraphs()->Get(subgraph_idx);
const auto float_graph = model_->subgraphs()->Get(subgraph_idx);
// The output graph should have two extra tensors from the added dequantize
// op.
ASSERT_EQ(quantized_graph->tensors()->size(),
float_graph->tensors()->size() + 2);
// Check that a dequantize op exists.
int32_t dequant_input_idx = -1;
int32_t dequant_output_idx = -1;
for (size_t i = 0; i < quantized_graph->operators()->size(); ++i) {
const auto op = quantized_graph->operators()->Get(i);
const uint32_t op_code_idx = op->opcode_index();
if (GetBuiltinCode(output_model->operator_codes()->Get(op_code_idx)) ==
BuiltinOperator_DEQUANTIZE) {
dequant_input_idx = op->inputs()->Get(0);
dequant_output_idx = op->outputs()->Get(0);
}
}
ASSERT_GT(dequant_input_idx, -1);
ASSERT_GT(dequant_output_idx, -1);
for (size_t i = 0; i < quantized_graph->tensors()->size(); ++i) {
const auto quant_tensor = quantized_graph->tensors()->Get(i);
// If the tensor is a weight, it should have type FLOAT16.
// If the tensor is a bias, it should have type FLOAT16.
// If the tensor is an input or output it should have type FLOAT32.
// The input to dequantize should be FLOAT16, and all other tensors should
// be FLOAT32.
if (i == dequant_input_idx) {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT16);
} else if (i == dequant_output_idx) {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
} else if (IsModelInputOrOutput(output_model, i)) {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
} else if (quant_tensor->name()->str() == "conv_bias") {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT16);
} else if (quant_tensor->buffer() != 0) {
// If it's a non-bias constant tensor, it must be the weight.
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT16);
} else {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
}
}
}
}
TEST_F(QuantizeWeightsTest, SharedWeights_Hybrid) {
LoadSharedWeightsModel();
flatbuffers::FlatBufferBuilder builder;
ASSERT_TRUE(
QuantizeWeights(&builder, model_, 0, QuantizerType::OLD_QUANTIZER).ok());
const uint8_t* buffer = builder.GetBufferPointer();
const Model* output_model = GetModel(buffer);
ASSERT_TRUE(output_model);
ASSERT_EQ(output_model->subgraphs()->size(), model_->subgraphs()->size());
uint32_t num_conv_ops = 0;
for (size_t subgraph_idx = 0; subgraph_idx < model_->subgraphs()->size();
++subgraph_idx) {
const auto quantized_graph = output_model->subgraphs()->Get(subgraph_idx);
for (size_t i = 0; i < quantized_graph->operators()->size(); ++i) {
const auto op = quantized_graph->operators()->Get(i);
const uint32_t op_code_idx = op->opcode_index();
const auto op_code =
GetBuiltinCode(output_model->operator_codes()->Get(op_code_idx));
if (op_code == BuiltinOperator_CONV_2D) {
num_conv_ops++;
// Ensure that each convolution's weights tensor is now INT8.
const auto weights_tensor =
quantized_graph->tensors()->Get(op->inputs()->Get(1));
EXPECT_EQ(weights_tensor->type(), TensorType_INT8);
}
}
}
// Ensure that there were exactly two convolutions in the model.
EXPECT_EQ(num_conv_ops, 2);
}
TEST_F(QuantizeWeightsTest, SharedWeights_Dequantize) {
LoadSharedWeightsModel();
flatbuffers::FlatBufferBuilder builder;
ASSERT_TRUE(internal::QuantizeWeights(&builder, model_, 0,
/*use_hybrid_evaluation*/ false,
QuantizerType::OLD_QUANTIZER)
.ok());
const uint8_t* buffer = builder.GetBufferPointer();
const Model* output_model = GetModel(buffer);
ASSERT_TRUE(output_model);
ASSERT_EQ(output_model->subgraphs()->size(), model_->subgraphs()->size());
uint32_t num_conv_ops = 0;
for (size_t subgraph_idx = 0; subgraph_idx < model_->subgraphs()->size();
++subgraph_idx) {
const auto quantized_graph = output_model->subgraphs()->Get(subgraph_idx);
for (size_t i = 0; i < quantized_graph->operators()->size(); ++i) {
const auto op = quantized_graph->operators()->Get(i);
const uint32_t op_code_idx = op->opcode_index();
const auto op_code =
GetBuiltinCode(output_model->operator_codes()->Get(op_code_idx));
if (op_code == BuiltinOperator_CONV_2D) {
num_conv_ops++;
// Ensure that each convolution's weights tensor is still FLOAT
// (the output of the dequantize).
uint32_t weights_tensor_index = op->inputs()->Get(1);
const auto weights_tensor =
quantized_graph->tensors()->Get(weights_tensor_index);
EXPECT_EQ(weights_tensor->type(), TensorType_FLOAT32);
// Check that it comes from a dequantize operation.
BuiltinOperator producer_op_code;
ASSERT_TRUE(GetProducerOpCode(output_model, subgraph_idx,
weights_tensor_index, &producer_op_code));
EXPECT_EQ(producer_op_code, BuiltinOperator_DEQUANTIZE);
}
}
}
// Ensure that there were exactly two convolutions in the model.
EXPECT_EQ(num_conv_ops, 2);
}
TEST_F(QuantizeWeightsTest, VerifyGatherQuantization) {
LoadGatherTestModel();
flatbuffers::FlatBufferBuilder builder;
ASSERT_TRUE(
QuantizeWeights(&builder, model_, 0, QuantizerType::OLD_QUANTIZER).ok());
const uint8_t* buffer = builder.GetBufferPointer();
const Model* output_model = GetModel(buffer);
ASSERT_TRUE(output_model);
ASSERT_EQ(output_model->subgraphs()->size(), model_->subgraphs()->size());
for (size_t subgraph_idx = 0; subgraph_idx < model_->subgraphs()->size();
++subgraph_idx) {
const auto quantized_graph = output_model->subgraphs()->Get(subgraph_idx);
for (size_t i = 0; i < quantized_graph->operators()->size(); ++i) {
const auto op = quantized_graph->operators()->Get(i);
const uint32_t op_code_idx = op->opcode_index();
const auto op_code =
GetBuiltinCode(output_model->operator_codes()->Get(op_code_idx));
if (op_code == tflite::BuiltinOperator_GATHER) {
uint32_t input_tensor_index = op->inputs()->Get(0);
const auto weights_tensor =
quantized_graph->tensors()->Get(input_tensor_index);
EXPECT_EQ(weights_tensor->type(), TensorType_INT8);
}
}
}
}
TEST_F(QuantizeWeightsTest, VerifyCustomOpQuantizationDequantize) {
LoadCustomOpTestModel();
// The custom op is not hybrid, and the second input is a constant that can
// be quantized.
CustomOpMap custom_op_map;
custom_op_map["CustomTestOp"] = {
.quantizable_input_indices = {1},
.is_hybrid = false,
};
flatbuffers::FlatBufferBuilder builder;
ASSERT_TRUE(QuantizeWeights(&builder, model_, 0, custom_op_map,
QuantizerType::OLD_QUANTIZER)
.ok());
const uint8_t* buffer = builder.GetBufferPointer();
const Model* output_model = GetModel(buffer);
ASSERT_TRUE(output_model);
ASSERT_EQ(output_model->subgraphs()->size(), model_->subgraphs()->size());
const auto quantized_graph = output_model->subgraphs()->Get(0);
// A dequantize op should be added.
ASSERT_EQ(quantized_graph->operators()->size(),
model_->subgraphs()->Get(0)->operators()->size() + 1);
int num_custom_ops_found = 0;
for (size_t i = 0; i < quantized_graph->operators()->size(); ++i) {
const auto op = quantized_graph->operators()->Get(i);
const uint32_t op_code_idx = op->opcode_index();
const auto op_code =
GetBuiltinCode(output_model->operator_codes()->Get(op_code_idx));
if (op_code == BuiltinOperator_CUSTOM) {
uint32_t weights_tensor_index = op->inputs()->Get(1);
const auto weights_tensor =
quantized_graph->tensors()->Get(weights_tensor_index);
EXPECT_EQ(weights_tensor->type(), TensorType_FLOAT32);
// Check that it comes from a dequantize operation.
BuiltinOperator producer_op_code;
ASSERT_TRUE(GetProducerOpCode(output_model, 0, weights_tensor_index,
&producer_op_code));
EXPECT_EQ(producer_op_code, BuiltinOperator_DEQUANTIZE);
num_custom_ops_found++;
}
}
EXPECT_EQ(num_custom_ops_found, 1);
}
TEST_F(QuantizeWeightsTest, VerifyCustomOpQuantizationHybrid) {
LoadCustomOpTestModel();
// The custom op is hybrid, and the second input is a constant that can
// be quantized.
CustomOpMap custom_op_map;
custom_op_map["CustomTestOp"] = {
.quantizable_input_indices = {1},
.is_hybrid = true,
};
flatbuffers::FlatBufferBuilder builder;
ASSERT_TRUE(QuantizeWeights(&builder, model_, 0, custom_op_map,
QuantizerType::OLD_QUANTIZER)
.ok());
const uint8_t* buffer = builder.GetBufferPointer();
const Model* output_model = GetModel(buffer);
ASSERT_TRUE(output_model);
ASSERT_EQ(output_model->subgraphs()->size(), model_->subgraphs()->size());
const auto quantized_graph = output_model->subgraphs()->Get(0);
ASSERT_EQ(quantized_graph->operators()->size(),
model_->subgraphs()->Get(0)->operators()->size());
int num_custom_ops_found = 0;
for (size_t i = 0; i < quantized_graph->operators()->size(); ++i) {
const auto op = quantized_graph->operators()->Get(i);
const uint32_t op_code_idx = op->opcode_index();
const auto op_code =
GetBuiltinCode(output_model->operator_codes()->Get(op_code_idx));
if (op_code == BuiltinOperator_CUSTOM) {
uint32_t weights_tensor_index = op->inputs()->Get(1);
const auto weights_tensor =
quantized_graph->tensors()->Get(weights_tensor_index);
EXPECT_EQ(weights_tensor->type(), TensorType_INT8);
num_custom_ops_found++;
}
}
EXPECT_EQ(num_custom_ops_found, 1);
}
TEST_F(QuantizeWeightsTest, VerifyUpdatedHybridSchemeFalseQuantizationHybrid) {
LoadBasicModel();
flatbuffers::FlatBufferBuilder builder;
const CustomOpMap custom_op_map;
ASSERT_TRUE(QuantizeWeights(&builder, model_, 0, custom_op_map,
/*use_updated_hybrid_scheme=*/false,
/*op_denylist=*/{}, QuantizerType::OLD_QUANTIZER)
.ok());
const uint8_t* buffer = builder.GetBufferPointer();
const Model* output_model = GetModel(buffer);
ASSERT_TRUE(output_model);
// Nothing should change.
ASSERT_EQ(output_model->subgraphs()->size(), model_->subgraphs()->size());
for (size_t subgraph_idx = 0; subgraph_idx < model_->subgraphs()->size();
subgraph_idx++) {
const auto quantized_graph = output_model->subgraphs()->Get(subgraph_idx);
const auto float_graph = model_->subgraphs()->Get(subgraph_idx);
ASSERT_EQ(quantized_graph->tensors()->size(),
float_graph->tensors()->size());
// Make sure the graph only has one Conv operation.
ASSERT_EQ(quantized_graph->operators()->size(), 1);
const auto op = quantized_graph->operators()->Get(0);
const uint32_t op_code_idx = op->opcode_index();
ASSERT_EQ(GetBuiltinCode(output_model->operator_codes()->Get(op_code_idx)),
BuiltinOperator_CONV_2D);
for (size_t i = 0; i < quantized_graph->tensors()->size(); i++) {
const auto quant_tensor = quantized_graph->tensors()->Get(i);
const auto float_tensor = float_graph->tensors()->Get(i);
EXPECT_EQ(quant_tensor->buffer(), float_tensor->buffer());
EXPECT_EQ(quant_tensor->is_variable(), float_tensor->is_variable());
EXPECT_EQ(GetAsVector(quant_tensor->shape()),
GetAsVector(float_tensor->shape()));
EXPECT_EQ(quant_tensor->name()->str(), float_tensor->name()->str());
// If the tensor is a weight, it should have type INT8, otherwise it
// should stay with type FLOAT32.
// If the tensor is a bias, it should have type FLOAT32.
if (quant_tensor->name()->str() == "conv_bias") {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
} else if (IsModelInputOrOutput(output_model, i)) {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
} else if (quant_tensor->buffer() != 0) {
EXPECT_EQ(quant_tensor->type(), TensorType_INT8)
<< quant_tensor->name()->str();
auto shape = GetAsVector(quant_tensor->shape());
EXPECT_EQ(quant_tensor->quantization()->scale()->size(), 1);
} else {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
}
}
}
}
TEST_F(QuantizeWeightsTest, DequantizeConvBlocklisted) {
LoadBasicModel();
flatbuffers::FlatBufferBuilder builder;
const CustomOpMap custom_op_map;
ASSERT_TRUE(QuantizeWeights(&builder, model_, 0, custom_op_map,
/*use_updated_hybrid_scheme=*/true,
/*op_denylist*/ {BuiltinOperator_CONV_2D},
QuantizerType::OLD_QUANTIZER)
.ok());
const uint8_t* buffer = builder.GetBufferPointer();
const Model* output_model = GetModel(buffer);
ASSERT_TRUE(output_model);
ASSERT_EQ(output_model->subgraphs()->size(), model_->subgraphs()->size());
for (size_t subgraph_idx = 0; subgraph_idx < model_->subgraphs()->size();
++subgraph_idx) {
const auto quantized_graph = output_model->subgraphs()->Get(subgraph_idx);
const auto float_graph = model_->subgraphs()->Get(subgraph_idx);
// The output graph should have an extra tensor from the added dequantize
// op.
ASSERT_EQ(quantized_graph->tensors()->size(),
float_graph->tensors()->size() + 1);
// Check that a dequantize op exists.
int32_t dequant_input_idx = -1;
int32_t dequant_output_idx = -1;
for (size_t i = 0; i < quantized_graph->operators()->size(); ++i) {
const auto op = quantized_graph->operators()->Get(i);
const uint32_t op_code_idx = op->opcode_index();
if (GetBuiltinCode(output_model->operator_codes()->Get(op_code_idx)) ==
BuiltinOperator_DEQUANTIZE) {
dequant_input_idx = op->inputs()->Get(0);
dequant_output_idx = op->outputs()->Get(0);
}
}
ASSERT_GT(dequant_input_idx, -1);
ASSERT_GT(dequant_output_idx, -1);
for (size_t i = 0; i < quantized_graph->tensors()->size(); ++i) {
const auto quant_tensor = quantized_graph->tensors()->Get(i);
// If the tensor is a weight, it should have type INT8.
// If the tensor is a bias, it should have type FLOAT32.
// If the tensor is an input or output it should have type FLOAT32.
// The input to dequantize should be INT8, and all other tensors should be
// FLOAT32.
if (i == dequant_input_idx) {
EXPECT_EQ(quant_tensor->type(), TensorType_INT8);
// The dequantize should still be quantized per-channel
EXPECT_EQ(quant_tensor->quantization()->scale()->size(), 5);
EXPECT_EQ(quant_tensor->quantization()->quantized_dimension(), 0);
} else if (i == dequant_output_idx) {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
} else if (IsModelInputOrOutput(output_model, i)) {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
} else if (quant_tensor->name()->str() == "conv_bias") {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
} else if (quant_tensor->buffer() != 0) {
// If it's a non-bias constant tensor, it must be the weight.
EXPECT_EQ(quant_tensor->type(), TensorType_INT8);
} else {
EXPECT_EQ(quant_tensor->type(), TensorType_FLOAT32);
}
}
}
}
} // namespace
} // namespace toco_legacy
} // namespace lite
} // namespace mlir
int main(int argc, char** argv) {
std::string model_file;
const std::vector<tsl::Flag> flag_list = {
tsl::Flag("test_model_file", &model_file,
"Path to test tflite model file."),
};
const bool parse_result = tsl::Flags::Parse(&argc, argv, flag_list);
if (!parse_result) {
std::cerr << "Required test_model_file\n";
std::abort();
}
g_test_model_dir = new std::string(tsl::io::Dirname(model_file));
::tsl::port::InitMain(argv[0], &argc, &argv);
return RUN_ALL_TESTS();
}
tensorflow/compiler/mlir/lite/quantization/numerical_utils.cc
+97
View File
@@ -0,0 +1,97 @@
/* Copyright 2020 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/numerical_utils.h"
#include <assert.h>
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <limits>
#include <optional>
namespace mlir {
namespace quant {
// Converts a double-precision floating-point multiplier to a quantized
// multiplier.
//
// Args:
// double_multiplier: The double-precision floating-point multiplier.
//
// Returns:
// A quantized multiplier, represented as a pair of integers: the quantized
// multiplier and the shift amount. The shift amount is the number of bits
// that the quantized multiplier should be shifted to the right before being
// used.
QuantizedMultiplier QuantizeMultiplier(double double_multiplier) {
if (double_multiplier < 1e-6) {
return {0, 0};
}
int32_t shift;
const double q = frexp(double_multiplier, &shift);
int64_t quantized_multiplier = round(q * (1LL << 31));
assert(quantized_multiplier <= (1LL << 31));
if (quantized_multiplier == (1LL << 31)) {
quantized_multiplier /= 2;
++shift;
}
assert(quantized_multiplier <= std::numeric_limits<int32_t>::max());
// Check that the shift amount is not greater than 31 or less than -31.
if (shift > 31 || shift < -31) {
return {0, 0};
}
return {static_cast<int32_t>(quantized_multiplier), shift};
}
// Calculates the quantized range for a given scale, zero point, minimum and
// maximum values, and quantization range.
//
// Args:
// scale: The scale factor for the quantized values.
// zero_point: The zero point for the quantized values.
// rmin: The minimum value of the quantized values.
// rmax: The maximum value of the quantized values.
// qmin: The minimum value of the quantization range.
// qmax: The maximum value of the quantization range.
//
// Returns:
// A quantized range, represented as a pair of integers: the minimum and
// maximum quantized values.
QuantizedRange CalculateQuantizedRange(double scale, int32_t zero_point,
std::optional<double> rmin,
std::optional<double> rmax, int32_t qmin,
int32_t qmax) {
auto quantize = [scale, zero_point](float f) {
return zero_point + static_cast<int32_t>(std::round(f / scale));
};
if (rmin.has_value() && rmax.has_value()) {
return {std::max(qmin, quantize(rmin.value())),
std::min(qmax, quantize(rmax.value()))};
} else if (rmin.has_value()) {
return {std::max(qmin, quantize(rmin.value())), qmax};
} else if (rmax.has_value()) {
return {qmin, std::min(qmax, quantize(rmax.value()))};
} else {
return {qmin, qmax};
}
}
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/lite/quantization/numerical_utils.h
+46
View File
@@ -0,0 +1,46 @@
/* Copyright 2020 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_NUMERICAL_UTILS_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_NUMERICAL_UTILS_H_
#include <cstdint>
#include <optional>
#include <utility>
#include "absl/types/optional.h"
namespace mlir {
namespace quant {
using QuantizedMultiplier = std::pair<int32_t, int32_t>;
using QuantizedRange = std::pair<int32_t, int32_t>;
// Decompose double precision multiplier to integer multiplier and exponent.
// double_multiplier = int_multiplier * 2 ^ (-31 + exponent)
// int_multiplier will be range of (2^31, 2^30].
QuantizedMultiplier QuantizeMultiplier(double double_multiplier);
// Calculate the effective quantized value range for the scale, zero point. The
// range is the minimum range defined by [rmin, rmax] and [qmin, qmax].
QuantizedRange CalculateQuantizedRange(double scale, int32_t zero_point,
std::optional<double> rmin,
std::optional<double> rmax, int32_t qmin,
int32_t qmax);
} // namespace quant
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_NUMERICAL_UTILS_H_
tensorflow/compiler/mlir/lite/quantization/numerical_utils_test.cc
+113
View File
@@ -0,0 +1,113 @@
/* Copyright 2020 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/numerical_utils.h"
#include <cmath>
#include <optional>
#include <gtest/gtest.h>
namespace mlir {
namespace quant {
namespace {
double ComposeScale(const QuantizedMultiplier& input) {
return input.first * exp2(-31 + input.second);
}
TEST(NumericalUtils, QuantizeMultiplier) {
// Decompose multiplier larger than 1.
ASSERT_FLOAT_EQ(ComposeScale(QuantizeMultiplier(1.0e6)), 1.0e6);
ASSERT_FLOAT_EQ(ComposeScale(QuantizeMultiplier(1.0e3)), 1.0e3);
ASSERT_FLOAT_EQ(ComposeScale(QuantizeMultiplier(10.)), 10.);
ASSERT_FLOAT_EQ(ComposeScale(QuantizeMultiplier(5.)), 5.);
ASSERT_FLOAT_EQ(ComposeScale(QuantizeMultiplier(2.)), 2.);
// Decompose multiplier between 1.0 and 1e-6.
ASSERT_FLOAT_EQ(ComposeScale(QuantizeMultiplier(0.0)), 0.0);
ASSERT_FLOAT_EQ(ComposeScale(QuantizeMultiplier(1.0)), 1.0);
ASSERT_FLOAT_EQ(ComposeScale(QuantizeMultiplier(1.0e-1)), 1.0e-1);
ASSERT_FLOAT_EQ(ComposeScale(QuantizeMultiplier(1.0e-2)), 1.0e-2);
ASSERT_FLOAT_EQ(ComposeScale(QuantizeMultiplier(1.0e-3)), 1.0e-3);
ASSERT_FLOAT_EQ(ComposeScale(QuantizeMultiplier(1.0e-4)), 1.0e-4);
ASSERT_FLOAT_EQ(ComposeScale(QuantizeMultiplier(1.0e-5)), 1.0e-5);
ASSERT_FLOAT_EQ(ComposeScale(QuantizeMultiplier(1.0e-6)), 1.0e-6);
// When scale is smaller than 1.0e-6, it is decomposed to {0, 0}.
ASSERT_FLOAT_EQ(ComposeScale(QuantizeMultiplier(1.0e-7)), 0.0);
ASSERT_FLOAT_EQ(ComposeScale(QuantizeMultiplier(1.0e-8)), 0.0);
}
TEST(NumericalUtils, ActivationRange) {
// zero point = 0
auto a =
CalculateQuantizedRange(1e-6, 0, std::nullopt, std::nullopt, -128, 127);
ASSERT_EQ(a.first, -128);
ASSERT_EQ(a.second, 127);
auto b = CalculateQuantizedRange(1e-6, 0, 0.0, std::nullopt, -128, 127);
ASSERT_EQ(b.first, 0);
ASSERT_EQ(b.second, 127);
auto c = CalculateQuantizedRange(1e-6, 0, -1.0, 1.0, -128, 127);
ASSERT_EQ(c.first, -128);
ASSERT_EQ(c.second, 127);
auto d = CalculateQuantizedRange(1e-6, 0, 0.0, 6.0, -128, 127);
ASSERT_EQ(d.first, 0);
ASSERT_EQ(d.second, 127);
// zero point = 100
auto e =
CalculateQuantizedRange(1e-6, 100, std::nullopt, std::nullopt, -128, 127);
ASSERT_EQ(e.first, -128);
ASSERT_EQ(e.second, 127);
auto f = CalculateQuantizedRange(1e-6, 100, 0.0, std::nullopt, -128, 127);
ASSERT_EQ(f.first, 100);
ASSERT_EQ(f.second, 127);
auto g = CalculateQuantizedRange(1e-6, 100, -1.0, 1.0, -128, 127);
ASSERT_EQ(g.first, -128);
ASSERT_EQ(g.second, 127);
auto h = CalculateQuantizedRange(1e-6, 100, 0.0, 6.0, -128, 127);
ASSERT_EQ(h.first, 100);
ASSERT_EQ(h.second, 127);
// zero point = -100
auto i = CalculateQuantizedRange(1e-6, -100, std::nullopt, std::nullopt, -128,
127);
ASSERT_EQ(i.first, -128);
ASSERT_EQ(i.second, 127);
auto j = CalculateQuantizedRange(1e-6, -100, 0.0, std::nullopt, -128, 127);
ASSERT_EQ(j.first, -100);
ASSERT_EQ(j.second, 127);
auto k = CalculateQuantizedRange(1e-6, -100, -1.0, 1.0, -128, 127);
ASSERT_EQ(k.first, -128);
ASSERT_EQ(k.second, 127);
auto l = CalculateQuantizedRange(1e-6, -100, 0.0, 6.0, -128, 127);
ASSERT_EQ(l.first, -100);
ASSERT_EQ(l.second, 127);
}
} // namespace
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/lite/quantization/quantization_context.cc
+367
View File
@@ -0,0 +1,367 @@
/* Copyright 2020 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/quantization_context.h"
#include <vector>
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/ErrorHandling.h"
#include "llvm/Support/raw_ostream.h"
#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/BuiltinTypeInterfaces.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/MLIRContext.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 "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_utils.h"
#include "tensorflow/compiler/mlir/lite/quantization/device_target.h"
#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
#define DEBUG_TYPE "quantization-context"
namespace mlir {
namespace quant {
QuantizeContext::QuantizeContext(func::FuncOp func, const DeviceTarget &spec)
: func_(func), target_spec_(spec) {
llvm::DenseMap<Value, int> value_to_state;
func.walk([&](quantfork::QuantizeRegionOp op) {
for (int i = 0, e = op.getNumOperands(); i != e; ++i) {
states_manager_.InitializeOperandState(op, i, &value_to_state);
}
for (int res = 0, e = op.getNumResults(); res != e; ++res) {
states_manager_.InitializeResultState(op, res, &value_to_state);
}
});
}
std::vector<quantfork::QuantizeRegionOp> QuantizeContext::GetAllOps() {
std::vector<quantfork::QuantizeRegionOp> all_ops;
all_ops.reserve(128);
func_.walk([&](quantfork::QuantizeRegionOp op) { all_ops.push_back(op); });
return all_ops;
}
KernelSpecs::Signature QuantizeContext::GetSignature(
quantfork::QuantizeRegionOp op) {
KernelSpecs::Signature signature;
signature.reserve(op.getInputSpecs().size() + op.getOutputSpecs().size());
for (int i = 0; i < op.getNumOperands(); ++i) {
DeviceTarget::AppendToSignature(GetOperandParams(op, i), &signature);
}
for (int i = 0; i < op.getNumResults(); ++i) {
DeviceTarget::AppendToSignature(GetResultParams(op, i), &signature);
}
return signature;
}
LogicalResult QuantizeContext::Handle(
quantfork::QuantizeRegionOp op,
llvm::SmallVectorImpl<Operation *> *new_items, bool *changed) {
auto signature = GetSignature(op);
auto spec = target_spec_.GetKernelSpec(op.getLogicalKernel(), signature);
if (!spec.has_value()) {
op.emitWarning(
"Couldn't find kernel from the registration for quantization.");
return success();
}
switch (spec->type) {
case ScaleConstraintType::OutputInputFreeScale: {
// no propagation.
*changed |= false;
break;
}
case ScaleConstraintType::CustomScale: {
if (failed(spec->scale_fn(this, op, new_items, changed))) {
return failure();
}
break;
}
case ScaleConstraintType::OutputInputSameScale: {
auto params = GetQuantParamsForSameScaleConstraint(op);
if (EmptyParams(params)) {
*changed |= false;
break;
}
// propagate this params to all the quantizable ports.
if (failed(PropagateQuantParams(op, params, new_items, changed))) {
return failure();
}
break;
}
default: {
// TODO(fengliuai): implement the other types.
llvm_unreachable("no implementation.");
return failure();
}
}
return success();
}
LogicalResult QuantizeContext::Finalize() {
MLIRContext *context = func_.getContext();
func_.walk([&](quantfork::QuantizeRegionOp op) {
llvm::SmallVector<Attribute, 4> input_specs;
auto original_input_specs = op.getInputSpecs().getValue();
for (int i = 0, e = op.getNumOperands(); i != e; ++i) {
auto &state = states_manager_.GetOperandQuantState(op, i);
auto &requantize = states_manager_.GetOperandRequantizeState(op, i);
if (state.IsEmpty() && requantize.pos == RequantizeState::NO_REQUANTIZE) {
input_specs.push_back(original_input_specs[i]);
} else if (requantize.pos == RequantizeState::ON_OUTPUT) {
input_specs.push_back(TypeAttr::get(requantize.params));
} else {
input_specs.push_back(TypeAttr::get(state.params));
}
}
op->setAttr("input_specs", ArrayAttr::get(context, input_specs));
llvm::SmallVector<Attribute, 4> output_specs;
auto original_output_specs = op.getOutputSpecs().getValue();
for (int res = 0, e = op.getNumResults(); res != e; ++res) {
auto &state = states_manager_.GetResultQuantState(op, res);
auto &requantize = states_manager_.GetResultRequantizeState(op, res);
if (state.IsEmpty() && requantize.pos == RequantizeState::NO_REQUANTIZE) {
output_specs.push_back(original_output_specs[res]);
} else if (requantize.pos == RequantizeState::ON_INPUT) {
output_specs.push_back(TypeAttr::get(requantize.params));
} else {
output_specs.push_back(TypeAttr::get(state.params));
}
}
op->setAttr("output_specs", ArrayAttr::get(context, output_specs));
});
return success();
}
void QuantizeContext::DumpStates(quantfork::QuantizeRegionOp current_op) {
if (current_op) {
llvm::errs() << "\n\n\n" << current_op.getLogicalKernel() << "\n";
}
func_.walk([&](quantfork::QuantizeRegionOp op) {
if (current_op == op) llvm::errs() << "===>>>";
llvm::errs() << op.getLogicalKernel() << " : (";
for (auto i = 0; i < op.getNumOperands(); ++i) {
if (auto params = GetOperandParams(op, i))
params.print(llvm::errs());
else
llvm::errs() << "_";
llvm::errs() << ",";
}
llvm::errs() << ") -> (";
for (auto i = 0; i < op.getNumResults(); ++i) {
if (auto params = GetResultParams(op, i))
params.print(llvm::errs());
else
llvm::errs() << "_";
llvm::errs() << ",";
}
llvm::errs() << ")\n";
});
}
// 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.
TFL::QuantParams QuantizeContext::GetQuantParamsForSameScaleConstraint(
Operation *op) {
// Two vector to collect Non-empty operands and results states.
std::vector<quant::QuantState *> mutable_states, immutable_states;
for (int i = 0, e = op->getNumOperands(); i != e; ++i) {
auto &state = states_manager_.GetOperandQuantState(op, i);
if (state.immutable) {
immutable_states.push_back(&state);
} else if (!state.IsEmpty()) {
mutable_states.push_back(&state);
}
}
int immutable_operands_num = immutable_states.size();
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, e = op->getNumResults(); i != e; ++i) {
auto &state = states_manager_.GetResultQuantState(op, i);
if (state.immutable) {
immutable_states.push_back(&state);
} else if (!state.IsEmpty()) {
mutable_states.push_back(&state);
}
}
int immutable_results_num = immutable_states.size() - immutable_operands_num;
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;
}
LLVM_DEBUG(llvm::dbgs()
<< "Quantization parameters are not collected in an ideal place. "
"Has to fallback values which might introduce errors.\n");
// 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 {};
}
LogicalResult QuantizeContext::PropagateQuantParams(
Operation *op, const TFL::QuantParams params,
quant::AdjacentOperations *new_items, bool *changed) {
// Use the final state to set all the operands' parameters.
for (int i = 0, e = op->getNumOperands(); i != e; ++i) {
auto ele =
llvm::cast<ShapedType>(op->getOperand(i).getType()).getElementType();
if (isa<FloatType>(ele) && SetOperandParams(op, i, params)) {
*changed |= true;
new_items->push_back(op->getOperand(i).getDefiningOp());
}
}
// Use the final state to set all the results' parameters.
for (int res = 0, e = op->getNumResults(); res != e; ++res) {
auto ele =
llvm::cast<ShapedType>(op->getResult(res).getType()).getElementType();
if (isa<FloatType>(ele) && SetResultParams(op, res, params)) {
auto users = op->getResult(res).getUsers();
*changed |= !users.empty();
new_items->append(users.begin(), users.end());
}
}
return success();
}
int QuantizeContext::StatesManager::InitializeState(
quantfork::QuantizeRegionOp op, int index, bool as_result) {
Attribute params_attr;
if (as_result) {
params_attr = op.getOutputSpecs()[index];
} else {
params_attr = op.getInputSpecs()[index];
}
TFL::QuantParams params =
dyn_cast<QuantizedType>(cast<TypeAttr>(params_attr).getValue());
bool immutable = !EmptyParams(params);
int next_state_index = states_.size();
states_.push_back({params, immutable});
if (as_result) {
result_states_.insert({{op, index}, next_state_index});
} else {
operand_states_.insert({{op, index}, next_state_index});
}
return next_state_index;
}
void QuantizeContext::StatesManager::InitializeOperandState(
quantfork::QuantizeRegionOp op, int index,
llvm::DenseMap<Value, int> *cache) {
Value in = op.getOperand(index);
auto cached = cache->insert({in, 0});
if (!cached.second) {
operand_states_.insert({{op, index}, cached.first->second});
return;
}
cached.first->second = InitializeState(op, index, /*as_result=*/false);
}
void QuantizeContext::StatesManager::InitializeResultState(
quantfork::QuantizeRegionOp op, int index,
llvm::DenseMap<Value, int> *cache) {
auto res = op.getResult(index);
auto cached = cache->insert({res, 0});
if (!cached.second) {
result_states_.insert({{op, index}, cached.first->second});
return;
}
cached.first->second = InitializeState(op, index, /*as_result=*/true);
}
bool QuantizeContext::StatesManager::SetConstantResultParams(Operation *op) {
llvm_unreachable("no implementation.");
return false;
}
bool QuantizeContext::StatesManager::SetResultParams(Operation *op,
int res_index,
TFL::QuantParams params) {
auto &state = GetResultQuantState(op, res_index);
if (state.params == params) {
return false;
}
if (!state.IsEmpty()) {
auto &rescale = GetResultRequantizeState(op, res_index);
rescale.params = params;
rescale.pos = RequantizeState::ON_INPUT;
return false;
}
state.params = params;
return true;
}
bool QuantizeContext::StatesManager::SetOperandParams(Operation *op, int index,
TFL::QuantParams params) {
auto &state = GetOperandQuantState(op, index);
if (state.params == params) {
return false;
}
if (!state.IsEmpty()) {
auto &rescale = GetOperandRequantizeState(op, index);
rescale.params = params;
rescale.pos = RequantizeState::ON_OUTPUT;
return false;
}
state.params = params;
return true;
}
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/lite/quantization/quantization_context.h
+247
View File
@@ -0,0 +1,247 @@
/* Copyright 2020 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_QUANTIZATION_CONTEXT_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_QUANTIZATION_CONTEXT_H_
#include <unordered_map>
#include <utility>
#include <vector>
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/SmallVector.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/Quant/IR/QuantTypes.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_utils.h"
#include "tensorflow/compiler/mlir/lite/quantization/device_target.h"
#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
namespace mlir {
namespace quant {
static bool EmptyParams(TFL::QuantParams p) {
return p == quant::QuantizedType();
}
// The state for each op result during the quantization parameters propagation.
struct QuantState {
// Quantization parameters propagated to an op result.
TFL::QuantParams 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() { return EmptyParams(params); }
};
// 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.
TFL::QuantParams params;
};
// This class manages all the intermediate quantization states.
class QuantizeContext {
public:
QuantizeContext(func::FuncOp func, const DeviceTarget &spec);
// Returns all the quant region ops.
std::vector<quantfork::QuantizeRegionOp> GetAllOps();
// For each quant region op, propagates its quantization parameters according
// to the kernel specification and also returns the adjacent quant region ops
// which get the new quantization parameters propagated.
LogicalResult Handle(quantfork::QuantizeRegionOp op,
llvm::SmallVectorImpl<Operation *> *new_items,
bool *changed);
// Updates the port quantization specifications of all the quant region ops
// with the propagation results.
LogicalResult Finalize();
// Dumps the states stores in the state manager.
void DumpStates(quantfork::QuantizeRegionOp current_op = {});
// Update the quantization parameter for certain result of the op. By this
// method, the quantization parameter is propagated to all the users of the
// result as well.
bool SetResultParams(Operation *op, int index, TFL::QuantParams params) {
return states_manager_.SetResultParams(op, index, params);
}
// Update the quantization parameter for certain operand of the op. By this
// method, the quantization parameter is propagated to the defining op of
// operand as well.
bool SetOperandParams(Operation *op, int index, TFL::QuantParams params) {
return states_manager_.SetOperandParams(op, index, params);
}
// Return the quantization parameter of certain result of the op.
TFL::QuantParams GetResultParams(Operation *op, int index) {
return states_manager_.GetResultParams(op, index);
}
// Return the quantization parameter of certain operand of the op.
TFL::QuantParams GetOperandParams(Operation *op, int index) {
return states_manager_.GetOperandParams(op, index);
}
// Return the signature of the op.
KernelSpecs::Signature GetSignature(quantfork::QuantizeRegionOp op);
// 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.
TFL::QuantParams GetQuantParamsForSameScaleConstraint(Operation *op);
// Propagate `params` to all the quantizable port of the `op`. The adjacent
// ops, which have the parameters propagated to, are collected by `new_items`,
// so they can be added to the working queue. `changed` is set to true if
// there are any new elements being added to `new_items`.
LogicalResult PropagateQuantParams(Operation *op, TFL::QuantParams params,
AdjacentOperations *new_items,
bool *changed);
private:
class StatesManager {
public:
// Sets the quantization parameters of the constant result according to its
// content.
//
// Always returns true.
bool SetConstantResultParams(Operation *op);
// Sets the quantization parameters of the result to a fixed value. If any
// quantization parameters have been propagated, a `requantize` will happen
// on the input of propagated quantization.
//
// Returns true, if the users of the result needs to be added to the
// worklist.
bool SetResultParams(Operation *op, int index, TFL::QuantParams params);
// Sets the quantization parameters of the operand to a fixed value. If any
// quantization parameters have been propagated, a `requantize` will happen
// on the output of propagated quantization.
//
// Returns true, if the defining op of the operand needs to be added to the
// worklist.
bool SetOperandParams(Operation *op, int index, TFL::QuantParams params);
// Returns the quantization parameters of the index-th result of the op.
TFL::QuantParams GetResultParams(Operation *op, int index) {
return states_[result_states_[{op, index}]].params;
}
// Returns the quantization parameters of the index-th operand of the op.
TFL::QuantParams GetOperandParams(Operation *op, int index) {
return states_[operand_states_[{op, index}]].params;
}
private:
friend class QuantizeContext;
// Uses the type of `val` 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.
int InitializeState(quantfork::QuantizeRegionOp op, int index,
bool as_result);
// 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(quantfork::QuantizeRegionOp op, int index,
llvm::DenseMap<Value, int> *cache);
// 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(quantfork::QuantizeRegionOp op, int index,
llvm::DenseMap<Value, int> *cache);
// Returns the state of the index-th operand of the op.
QuantState &GetOperandQuantState(Operation *op, int index) {
return states_[operand_states_[{op, index}]];
}
// Returns the state of the index-th result of the op.
QuantState &GetResultQuantState(Operation *op, int index) {
return states_[result_states_[{op, index}]];
}
// Returns the state of the index-th operand of the op.
RequantizeState &GetOperandRequantizeState(Operation *op, int index) {
return rescale_states_[operand_states_[{op, index}]];
}
// Returns the state of the index-th result of the op.
RequantizeState &GetResultRequantizeState(Operation *op, int index) {
return rescale_states_[result_states_[{op, index}]];
}
private:
// This is used to identify an operand or result of an op. The second
// element of this pair is the index of the operand or result.
using OpValue = std::pair<mlir::Operation *, int>;
// 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_`.
std::unordered_map<int, RequantizeState> rescale_states_;
// Maps of indexes to the propagation state vector from the ops operands,
// results and arguments.
llvm::DenseMap<OpValue, int> operand_states_;
llvm::DenseMap<OpValue, int> result_states_;
};
func::FuncOp func_;
DeviceTarget target_spec_;
StatesManager states_manager_;
};
} // namespace quant
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_QUANTIZATION_CONTEXT_H_
tensorflow/compiler/mlir/lite/quantization/quantization_info.proto
+89
View File
@@ -0,0 +1,89 @@
// Copyright 2026 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.
// ==============================================================================
syntax = "proto3";
package mlir.quant;
// Represents the quantization parameters for a list of named tensors.
message QuantizationInfo {
// min/max of the per axis value range. To quantize the value, the metadata
// of the target properties should be specified or read from the ops
// quantization specification.
message MinMax {
float min = 1;
float max = 2;
}
// Affine parameters to quantize the per axis value. The metadata of the
// target properties should be specified as well.
message AffineParams {
float scale = 1;
int32 zero_point = 2;
}
// Params to quantize the axis. Only one of the field can be used.
message PerAxisParams {
oneof params_oneof {
// min/max of the ranges.
MinMax min_max = 1;
// affine parameters to quantize the per axis value.
AffineParams affine_params = 2;
}
}
// The metadata defines the target properties.
message Metadata {
// Bit number of fixed-point data the target kernel supports.
int32 num_bits = 1;
// The quantized axis index if it is per-axis quantization.
int32 quantize_axis = 2;
// The minimum allowed value of the fixed-point data range.
// This can also be used to derive the sign of storage type.
int32 range_min = 3;
// The minimum allowed value of the fixed-point data range.
int32 range_max = 4;
}
// The quantization parameters for a named tensor.
message QuantParams {
// The tensor name has the following convention:
// tensor_name := op_name | op_name ’:’ port_number.
// If the op has only one port, op_name can be used.
// If the op has internal states, such as fused LSTM, the port_number should
// follow a predefined convention.
oneof name_oneof {
string name = 1;
// An regex can be used to match multiple tensors.
string name_regex = 2;
}
// The quantization parameters for the tensor. If it is for per-axis, the
// parameters should be defined for each axis, otherwise, if it is for
// per-tensor, this repeated field should only contain a single element.
repeated PerAxisParams params = 3;
// Metadata about the quantization parameters.
Metadata meta = 5;
}
// List of quantization parameters for tensors.
repeated QuantParams entries = 1;
}
tensorflow/compiler/mlir/lite/quantization/quantization_passes.h
+46
View File
@@ -0,0 +1,46 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_QUANTIZATION_PASSES_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_QUANTIZATION_PASSES_H_
#include <functional>
#include <string>
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
namespace mlir {
namespace quant {
using OperationToName = std::function<llvm::StringRef(Operation* op)>;
// Creates an instance pass to import quantization stats to the operations in
// the function. A custom method to get the name from the op is used because
// different dialect ops might have different ways to assign the name.
std::unique_ptr<OperationPass<func::FuncOp>> CreateImportQuantStatsPass(
OperationToName op_to_name, const std::string& stats_str);
// Creates an instance pass to import quantization stats to the operations in
// the function. A custom method to get the name from the op is used because
// different dialect ops might have different ways to assign the name.
std::unique_ptr<OperationPass<func::FuncOp>>
CreateImportQuantStatsPassForTFControlDialect(const std::string& stats_str);
} // namespace quant
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_QUANTIZATION_PASSES_H_
tensorflow/compiler/mlir/lite/quantization/tensorflow/BUILD
+77
View File
@@ -0,0 +1,77 @@
load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library")
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"],
default_visibility = [
":friends",
],
licenses = ["notice"],
)
package_group(
name = "friends",
packages = [
"//learning/brain/mobile/model_optimization/...",
"//tensorflow/compiler/mlir/...",
"//tensorflow/compiler/mlir/lite/...",
],
)
td_library(
name = "ptq_td_files",
srcs = [
"fallback_to_flex_patterns.td",
],
compatible_with = get_compatible_with_portable(),
visibility = ["//visibility:private"],
deps = [
"//tensorflow/compiler/mlir/tensorflow:tensorflow_ops_td_files",
"@llvm-project//mlir:ArithOpsTdFiles",
"@llvm-project//mlir:FuncTdFiles",
],
)
gentbl_cc_library(
name = "ptq_fallback_to_flex_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {"fallback_to_flex_patterns.inc": ["-gen-rewriters"]},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "fallback_to_flex_patterns.td",
deps = [":ptq_td_files"],
)
cc_library(
name = "tf_quantization_passes",
srcs = [
"fallback_to_flex_ops.cc",
"fallback_to_flex_patterns.inc",
"tf_to_quant.cc",
],
hdrs = [
"passes.h",
],
deps = [
":ptq_fallback_to_flex_inc_gen",
"//tensorflow/compiler/mlir/lite:tensorflow_lite",
"//tensorflow/compiler/mlir/lite/quantization/common/quantization_lib",
"//tensorflow/compiler/mlir/lite/quantization/ir:QuantOps",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:translate_utils",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/platform:statusor",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:string_view",
"@flatbuffers",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:QuantOps",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
],
alwayslink = 1,
)

Some files were not shown because too many files have changed in this diff Show More