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/quantization/tensorflow/BUILD
+539
View File
@@ -0,0 +1,539 @@
load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library")
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("@rules_python//python:proto.bzl", "py_proto_library")
load("@xla//third_party/rules_python/python:py_binary.bzl", "py_binary")
load("//tensorflow:tensorflow.bzl", "tf_cc_binary")
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
load("//tensorflow/compiler/mlir/quantization/tensorflow:internal_visibility_allowlist.bzl", "internal_visibility_allowlist")
load("//tensorflow/core/platform:build_config.bzl", "tf_proto_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
":internal_visibility_allowlist_package",
"//tensorflow:__pkg__",
],
licenses = ["notice"],
)
# TODO(b/249181641): Restructure source codes for TF quantization.
package_group(
name = "internal_visibility_allowlist_package",
packages = [
"//tensorflow/compiler/mlir/lite/...",
"//tensorflow/compiler/mlir/quantization/...",
"//tensorflow/compiler/mlir/tensorflow_to_stablehlo/...",
] + internal_visibility_allowlist(),
)
py_binary(
name = "gen_quantized_function_library",
srcs = ["gen_quantized_function_library.py"],
strict_deps = True,
deps = [
"@absl_py//absl:app",
"@absl_py//absl/flags",
],
)
genrule(
name = "quantized_function_library",
srcs = [
"passes/quantized_function_library_uniform_quantized.mlir",
"passes/quantized_function_library.mlir",
"passes/quantized_function_library_uniform_quantized_drq.mlir",
"passes/quantized_function_library_tf_drq.mlir",
"passes/quantized_function_library_xla_weight_only.mlir",
],
outs = [
"passes/quantized_function_library.h",
],
cmd = "$(location gen_quantized_function_library) --output_file $(RULEDIR)/passes/quantized_function_library.h --src '$(SRCS)'",
compatible_with = get_compatible_with_portable(),
tools = ["gen_quantized_function_library"],
)
cc_library(
name = "manipulate_model_attr",
srcs = [
"passes/manipulate_model_attr.cc",
],
hdrs = [
"passes/manipulate_model_attr.h",
],
compatible_with = get_compatible_with_portable(),
deps = [
"//tensorflow/compiler/mlir/tensorflow",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
],
)
cc_library(
name = "remove_identity_op_pattern",
srcs = [
"passes/remove_identity_op_pattern.cc",
],
hdrs = [
"passes/remove_identity_op_pattern.h",
],
compatible_with = get_compatible_with_portable(),
deps = [
"//tensorflow/compiler/mlir/tensorflow",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
],
)
td_library(
name = "quant_td_files",
srcs = [
"passes/cast_bf16_ops_to_f32.td",
"passes/lift_quantizable_spots_as_functions.td",
"passes/lift_quantizable_spots_as_functions_drq.td",
"passes/optimize.td",
"passes/post_quantize.td",
"passes/prepare_lifting.td",
"passes/prepare_quantize.td",
"passes/preprocess_op.td",
"passes/quantize_composite_functions.td",
"passes/replace_cast_hacks_with_tf_xla_ops.td",
"passes/tf_quant_ops.td",
],
compatible_with = get_compatible_with_portable(),
deps = [
"//tensorflow/compiler/mlir/quantization/common:quant_td_files",
"//tensorflow/compiler/mlir/quantization/common/ir:QuantizationOpsTdFiles",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_ops_td_files",
"@llvm-project//mlir:ArithOpsTdFiles",
"@llvm-project//mlir:FuncTdFiles",
],
)
gentbl_cc_library(
name = "convert_tf_xla_op_to_tf_op_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {"passes/convert_tf_xla_op_to_tf_op.inc": ["-gen-rewriters"]},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "passes/convert_tf_xla_op_to_tf_op.td",
deps = [":quant_td_files"],
)
gentbl_cc_library(
name = "cast_bf16_ops_to_f32_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {"passes/cast_bf16_ops_to_f32.inc": ["-gen-rewriters"]},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "passes/cast_bf16_ops_to_f32.td",
deps = [":quant_td_files"],
)
gentbl_cc_library(
name = "prepare_lifting_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {"passes/prepare_lifting.inc": ["-gen-rewriters"]},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "passes/prepare_lifting.td",
deps = [":quant_td_files"],
)
gentbl_cc_library(
name = "lift_quantizable_spots_as_functions_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {"passes/lift_quantizable_spots_as_functions.inc": ["-gen-rewriters"]},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "passes/lift_quantizable_spots_as_functions.td",
deps = [":quant_td_files"],
)
gentbl_cc_library(
name = "lift_quantizable_spots_as_functions_drq_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {"passes/lift_quantizable_spots_as_functions_drq.inc": ["-gen-rewriters"]},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "passes/lift_quantizable_spots_as_functions_drq.td",
deps = [":quant_td_files"],
)
gentbl_cc_library(
name = "prepare_quantize_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {"passes/prepare_quantize.inc": ["-gen-rewriters"]},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "passes/prepare_quantize.td",
deps = [":quant_td_files"],
)
gentbl_cc_library(
name = "quantize_composite_functions_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {"passes/quantize_composite_functions.inc": ["-gen-rewriters"]},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "passes/quantize_composite_functions.td",
deps = [":quant_td_files"],
)
gentbl_cc_library(
name = "tf_quant_ops_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {
"passes/tf_quant_ops.h.inc": ["-gen-op-decls"],
"passes/tf_quant_ops.cc.inc": ["-gen-op-defs"],
},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "passes/tf_quant_ops.td",
deps = [
":quant_td_files",
],
)
gentbl_cc_library(
name = "optimize_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {"passes/optimize.inc": ["-gen-rewriters"]},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "passes/optimize.td",
deps = [":quant_td_files"],
)
gentbl_cc_library(
name = "convert_tpu_model_to_cpu_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {"passes/convert_tpu_model_to_cpu.inc": ["-gen-rewriters"]},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "passes/convert_tpu_model_to_cpu.td",
deps = [":quant_td_files"],
)
gentbl_cc_library(
name = "post_quantize_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {"passes/post_quantize.inc": ["-gen-rewriters"]},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "passes/post_quantize.td",
deps = [":quant_td_files"],
)
gentbl_cc_library(
name = "preprocess_op_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {"passes/preprocess_op.inc": ["-gen-rewriters"]},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "passes/preprocess_op.td",
deps = [":quant_td_files"],
)
cc_library(
name = "tf_quant_ops",
srcs = [
"passes/tf_quant_ops.cc",
"passes/tf_quant_ops.cc.inc",
"passes/tf_quant_ops.h.inc",
],
hdrs = ["passes/tf_quant_ops.h"],
compatible_with = get_compatible_with_portable(),
visibility = ["//visibility:public"],
deps = [
"//tensorflow/compiler/mlir/tensorflow:tensorflow_attributes",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_op_interfaces",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_side_effects",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_structs",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_traits",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:CallOpInterfaces",
"@llvm-project//mlir:ControlFlowInterfaces",
"@llvm-project//mlir:DerivedAttributeOpInterface",
"@llvm-project//mlir:Dialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:InferTypeOpInterface",
"@llvm-project//mlir:LoopLikeInterface",
"@llvm-project//mlir:Parser",
"@llvm-project//mlir:SideEffectInterfaces",
"@llvm-project//mlir:Support",
],
)
gentbl_cc_library(
name = "replace_cast_hacks_with_tf_xla_ops_inc_gen",
compatible_with = get_compatible_with_portable(),
tbl_outs = {"passes/replace_cast_hacks_with_tf_xla_ops.inc": ["-gen-rewriters"]},
tblgen = "@llvm-project//mlir:mlir-tblgen",
td_file = "passes/replace_cast_hacks_with_tf_xla_ops.td",
deps = [":quant_td_files"],
)
cc_library(
name = "passes",
srcs = [
"passes/add_dump_tensor_op.cc",
"passes/add_quantization_unit_loc.cc",
"passes/cast_bf16_ops_to_f32.cc",
"passes/cast_bf16_ops_to_f32.inc",
"passes/convert_custom_aggregation_op_to_quant_stats.cc",
"passes/convert_fake_quant_to_qdq.cc",
"passes/convert_tf_xla_op_to_tf_op.cc",
"passes/convert_tf_xla_op_to_tf_op.inc",
"passes/convert_tpu_model_to_cpu.cc",
"passes/convert_tpu_model_to_cpu.inc",
"passes/duplicate_shape_determining_constants.cc",
"passes/insert_custom_aggregation_ops.cc",
"passes/insert_main_function.cc",
"passes/insert_quantized_functions.cc",
"passes/insert_restore_op.cc",
"passes/insert_save_op.cc",
"passes/lift_hashtable_ops_as_args.cc",
"passes/lift_quantizable_spots_as_functions.cc",
"passes/lift_quantizable_spots_as_functions.inc",
"passes/lift_quantizable_spots_as_functions_drq.cc",
"passes/lift_quantizable_spots_as_functions_drq.inc",
"passes/mark_functions_noinline.cc",
"passes/merge_duplicate_resource_ops.cc",
"passes/merge_initializer_function_ops_to_main.cc",
"passes/merge_save_function_ops_to_main.cc",
"passes/optimize.cc",
"passes/optimize.inc",
"passes/post_quantize.cc",
"passes/post_quantize.inc",
"passes/prepare_lifting.cc",
"passes/prepare_lifting.inc",
"passes/prepare_quantize.cc",
"passes/prepare_quantize.inc",
"passes/prepare_quantize_drq.cc",
"passes/preprocess_op.cc",
"passes/propagate_quantize_type.cc",
"passes/quantize.cc",
"passes/quantize_composite_functions.cc",
"passes/quantize_composite_functions.inc",
"passes/quantize_weights.cc",
"passes/quantized_function_library.h",
"passes/remove_var_init_by_const.cc",
"passes/replace_cast_hacks_with_tf_xla_ops.cc",
"passes/replace_cast_hacks_with_tf_xla_ops.inc",
"passes/unfreeze_constants.cc",
],
hdrs = [
"passes/constants.h",
"passes/passes.h",
],
compatible_with = get_compatible_with_portable(),
deps = [
":cast_bf16_ops_to_f32_inc_gen",
":convert_tf_xla_op_to_tf_op_inc_gen",
":convert_tpu_model_to_cpu_inc_gen",
":lift_quantizable_spots_as_functions_drq_inc_gen",
":lift_quantizable_spots_as_functions_inc_gen",
":manipulate_model_attr",
":optimize_inc_gen",
":post_quantize_inc_gen",
":prepare_lifting_inc_gen",
":prepare_quantize_inc_gen",
":preprocess_op_gen",
":quantization_options_proto_cc",
":quantize_composite_functions_inc_gen",
":remove_identity_op_pattern",
":replace_cast_hacks_with_tf_xla_ops_inc_gen",
":tf_quant_ops",
"//tensorflow/compiler/mlir/quantization/common:attrs_and_constraints",
"//tensorflow/compiler/mlir/quantization/common:func",
"//tensorflow/compiler/mlir/quantization/common:lift_as_function_call",
"//tensorflow/compiler/mlir/quantization/common/ir:QuantOps",
"//tensorflow/compiler/mlir/quantization/common/quantization_lib",
"//tensorflow/compiler/mlir/quantization/common/quantization_lib:quantization_config",
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
"//tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration:calibration_parameters",
"//tensorflow/compiler/mlir/quantization/tensorflow/cc:const_op_size",
"//tensorflow/compiler/mlir/quantization/tensorflow/cc:constant_fold",
"//tensorflow/compiler/mlir/quantization/tensorflow/cc:quantization_unit_loc",
"//tensorflow/compiler/mlir/quantization/tensorflow/cc:run_passes",
"//tensorflow/compiler/mlir/quantization/tensorflow/ops:tf_op_quant_spec",
"//tensorflow/compiler/mlir/quantization/tensorflow/ops:tf_quantize_op",
"//tensorflow/compiler/mlir/quantization/tensorflow/utils:fake_quant_utils",
"//tensorflow/compiler/mlir/quantization/tensorflow/utils:tf_to_uniform_attribute_utils",
"//tensorflow/compiler/mlir/quantization/tensorflow/utils:tf_to_xla_attribute_utils",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:error_util",
"//tensorflow/compiler/mlir/tensorflow:import_model",
"//tensorflow/compiler/mlir/tensorflow:mangling_util",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_attributes",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/compiler/mlir/tensorflow:xla_call_module_attrs",
"//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
"//tensorflow/compiler/mlir/tensorflow/transforms:tf_dialect_passes",
"//tensorflow/compiler/mlir/tensorflow/transforms:tf_saved_model_freeze_variables",
"//tensorflow/compiler/mlir/tensorflow/transforms:tf_saved_model_passes",
"//tensorflow/compiler/mlir/utils:name_utils",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/ir/importexport:convert_tensor",
"//tensorflow/core/platform:macros",
"//tensorflow/core/platform:path",
"//tensorflow/core/tpu:tpu_defs",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/cleanup",
"@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/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
"@com_googlesource_code_re2//:re2",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:Dialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:FunctionInterfaces",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:QuantOps",
"@llvm-project//mlir:Rewrite",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
"@llvm-project//mlir:Transforms",
"@llvm-project//mlir:UBDialect",
"@xla//xla:xla_data_proto_cc",
],
# Alwayslink is required for registering the MLIR passes.
# TODO(b/255530126): Split the pass registration from the definitions to avoid binary size bloat.
alwayslink = True,
)
cc_library(
name = "quantize_preprocess",
srcs = [
"quantize_preprocess.cc",
],
hdrs = [
"quantize_preprocess.h",
],
compatible_with = get_compatible_with_portable(),
visibility = ["//visibility:public"],
deps = [
":passes",
"//tensorflow/compiler/mlir/quantization/stablehlo:bridge_passes",
"//tensorflow/compiler/mlir/quantization/stablehlo/cc:pass_pipeline",
"//tensorflow/compiler/mlir/quantization/tensorflow/cc:run_passes",
"//tensorflow/compiler/mlir/stablehlo:fold_broadcast_pass",
"//tensorflow/compiler/mlir/stablehlo:fuse_convolution_pass",
"//tensorflow/compiler/mlir/stablehlo:legalize_tf_xla_call_module_to_stablehlo_pass",
"//tensorflow/compiler/mlir/stablehlo:rename_entrypoint_to_main",
"//tensorflow/compiler/mlir/stablehlo:tf_stablehlo",
"//tensorflow/compiler/mlir/stablehlo:unfuse_batch_norm_pass",
"//tensorflow/compiler/mlir/tensorflow:error_util",
"//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
"//tensorflow/compiler/mlir/tensorflow/transforms:tf_dialect_passes",
"//tensorflow/compiler/mlir/tensorflow/transforms:tf_saved_model_freeze_variables",
"//tensorflow/compiler/mlir/tensorflow/transforms:tf_saved_model_passes",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core/platform:path",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Transforms",
"@xla//xla/mlir_hlo:all_passes",
],
)
cc_library(
name = "quantize_passes",
srcs = ["quantize_passes.cc"],
hdrs = ["quantize_passes.h"],
compatible_with = get_compatible_with_portable(),
visibility = ["//visibility:public"],
deps = [
":passes",
":quantization_options_proto_cc",
"//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
"//tensorflow/compiler/mlir/tensorflow/transforms:tf_dialect_passes",
"//tensorflow/core/platform:path",
"@com_google_absl//absl/strings",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Transforms",
"@xla//xla/mlir_hlo:mhlo_passes",
],
)
# OSS only: This target is header-only. Link `quantization_options_proto_impl` only to
# `libtensorflow_framework.so` via `lib_internal_impl`. Do NOT link
# `quantization_options_proto_impl` directly unless the target does not link
# `libtensorflow_framework.so`.
tf_proto_library(
name = "quantization_options_proto",
srcs = [
"quantization_options.proto",
],
make_default_target_header_only = True,
protodeps = [
"//tensorflow/core:protos_all",
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto",
],
visibility = ["//visibility:public"],
)
# copybara:uncomment_begin(google-only)
# py_proto_library(
# name = "quantization_options_py_pb2",
# visibility = ["//visibility:public"],
# deps = [":quantization_options_proto"],
# )
# copybara:uncomment_end
# OSS only: This target is header-only. Link `exported_model_proto_impl` only to
# `libtensorflow_framework.so` via `lib_internal_impl`. Do NOT link
# `exported_model_proto_impl` directly unless the target does not link
# `libtensorflow_framework.so`.
tf_proto_library(
name = "exported_model_proto",
srcs = ["exported_model.proto"],
make_default_target_header_only = True,
protodeps = [
"//tensorflow/core:protos_all",
],
visibility = [
":internal_visibility_allowlist_package",
# To be visible from `lib_internal_impl`.
"//tensorflow/core:__pkg__",
"//tensorflow/python:__pkg__",
],
)
# copybara:uncomment_begin(google-only)
# py_proto_library(
# name = "exported_model_py_pb2",
# deps = [":exported_model_proto"],
# )
# copybara:uncomment_end
tf_cc_binary(
name = "tf-quant-opt",
srcs = ["passes/tf_quant_opt.cc"],
deps = [
":passes",
"//tensorflow/compiler/mlir:init_mlir",
"//tensorflow/compiler/mlir/quantization/common/ir:QuantOps",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
"//tensorflow/compiler/mlir/tensorflow/transforms:tf_dialect_passes",
"//tensorflow/core/ir/types:Dialect",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:AllPassesAndDialects",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:FuncExtensions",
"@llvm-project//mlir:MlirOptLib",
"@llvm-project//mlir:QuantOps",
"@llvm-project//mlir:RegisterAllPasses",
"@llvm-project//mlir:SCFDialect",
"@llvm-project//mlir:ShapeDialect",
"@stablehlo//:stablehlo_ops",
],
)
tensorflow/compiler/mlir/quantization/tensorflow/calibrator/BUILD
+229
View File
@@ -0,0 +1,229 @@
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("@rules_python//python:proto.bzl", "py_proto_library")
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:pytype.default.bzl", "pytype_strict_contrib_test", "pytype_strict_library")
load(
"//tensorflow:tensorflow.bzl",
"tf_cc_test",
"tf_gen_op_wrapper_py",
)
load(
"//tensorflow:tensorflow.default.bzl",
"get_compatible_with_portable",
"tf_kernel_library",
"tf_py_strict_test",
)
load("//tensorflow/core/platform:build_config.bzl", "tf_proto_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//tensorflow/compiler/mlir/quantization/tensorflow:internal_visibility_allowlist_package",
"//tensorflow/core:__pkg__",
"//tensorflow/tools/pip_package:__subpackages__",
],
licenses = ["notice"],
)
cc_library(
name = "calibration_statistics_collector_base",
hdrs = ["calibration_statistics_collector_base.h"],
compatible_with = get_compatible_with_portable(),
deps = [
":calibration_statistics_proto_cc",
"@com_google_absl//absl/types:span",
],
)
cc_library(
name = "calibration_statistics_collector_min_max",
srcs = ["calibration_statistics_collector_min_max.cc"],
hdrs = ["calibration_statistics_collector_min_max.h"],
compatible_with = get_compatible_with_portable(),
deps = [
":calibration_statistics_collector_base",
":calibration_statistics_proto_cc",
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
"//tensorflow/compiler/mlir/quantization/tensorflow:quantization_options_proto_cc",
"@com_google_absl//absl/types:span",
],
)
cc_library(
name = "calibration_statistics_collector_average_min_max",
srcs = ["calibration_statistics_collector_average_min_max.cc"],
hdrs = ["calibration_statistics_collector_average_min_max.h"],
compatible_with = get_compatible_with_portable(),
deps = [
":calibration_statistics_collector_base",
":calibration_statistics_proto_cc",
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
"//tensorflow/compiler/mlir/quantization/tensorflow:quantization_options_proto_cc",
"@com_google_absl//absl/types:span",
],
)
cc_library(
name = "calibration_statistics_collector_histogram",
srcs = ["calibration_statistics_collector_histogram.cc"],
hdrs = ["calibration_statistics_collector_histogram.h"],
compatible_with = get_compatible_with_portable(),
deps = [
":calibration_statistics_collector_base",
":calibration_statistics_proto_cc",
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
"//tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration:calibration_parameters",
"//tensorflow/compiler/mlir/quantization/tensorflow:quantization_options_proto_cc",
"@com_google_absl//absl/types:span",
],
)
pytype_strict_library(
name = "calibration_algorithm",
srcs = ["calibration_algorithm.py"],
deps = [
":calibration_statistics_proto_py",
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_py",
"//third_party/py/numpy",
],
)
pytype_strict_contrib_test(
name = "calibration_algorithm_test",
srcs = ["calibration_algorithm_test.py"],
deps = [
":calibration_algorithm",
":calibration_statistics_proto_py",
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_py",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
tf_proto_library(
name = "calibration_statistics_proto",
srcs = ["calibration_statistics.proto"],
make_default_target_header_only = True,
)
# copybara:uncomment_begin(google-only)
# py_proto_library(
# name = "calibration_statistics_py_pb2",
# deps = [
# ":calibration_statistics_proto",
# ],
# )
# copybara:uncomment_end
tf_cc_test(
name = "calibration_statistics_collector_test",
size = "small",
srcs = ["calibration_statistics_collector_test.cc"],
deps = [
":calibration_statistics_collector_average_min_max",
":calibration_statistics_collector_histogram",
":calibration_statistics_collector_min_max",
":calibration_statistics_proto_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"@com_google_googletest//:gtest_main",
],
)
tf_kernel_library(
name = "custom_aggregator_op",
srcs = ["custom_aggregator_op.cc"],
compatible_with = get_compatible_with_portable(),
visibility = [
"//tensorflow:__pkg__",
"//tensorflow/compiler/mlir/quantization/tensorflow/python:__pkg__",
],
deps = [
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
"//tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration:calibration_parameters",
"//tensorflow/compiler/mlir/quantization/tensorflow:quantization_options_proto_cc",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/status",
"@xla//xla/tsl/platform:errors",
],
)
tf_gen_op_wrapper_py(
name = "gen_custom_aggregator_op_wrapper",
out = "custom_aggregator_op_wrapper.py",
extra_py_deps = [
"//tensorflow/python:pywrap_tfe",
"//tensorflow/python/util:dispatch",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
],
# Prevent unintentionally generating Python wrappers for all TF ops.
op_allowlist = ["CustomAggregator"],
py_lib_rule = py_library,
visibility = ["//visibility:private"],
deps = [":custom_aggregator_op"],
)
tf_py_strict_test(
name = "custom_aggregator_op_test",
size = "small",
srcs = ["integration_test/custom_aggregator_op_test.py"],
deps = [
":calibration_statistics_proto_py",
":gen_custom_aggregator_op_wrapper",
"//tensorflow:tensorflow_py",
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_py",
"//tensorflow/python:pywrap_tensorflow",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/platform:client_testlib",
],
)
tf_kernel_library(
name = "calibration_statistics_saver_op",
srcs = ["calibration_statistics_saver_op.cc"],
compatible_with = get_compatible_with_portable(),
visibility = ["//tensorflow/compiler/mlir/quantization/tensorflow/python:__pkg__"],
deps = [
":calibration_statistics_collector_average_min_max",
":calibration_statistics_collector_base",
":calibration_statistics_collector_histogram",
":calibration_statistics_collector_min_max",
":calibration_statistics_proto_cc",
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/platform:logging",
"@com_google_absl//absl/base:nullability",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:span",
"@xla//xla/tsl/platform:env",
],
)
tf_cc_test(
name = "calibration_statistics_saver_op_test",
srcs = ["calibration_statistics_saver_op_test.cc"],
deps = [
":calibration_statistics_proto_cc",
":calibration_statistics_saver_op",
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core:testlib",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status:status_matchers",
"@com_google_googletest//:gtest",
"@xla//xla/tsl/platform:errors",
],
)
tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_algorithm.py
+395
View File
@@ -0,0 +1,395 @@
# 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.
# ==============================================================================
"""Defines CalibrationAlgorithm for calculating min and max values calculated by calibration method."""
import abc
import itertools
import logging
import numpy as np
from tensorflow.compiler.mlir.quantization.stablehlo import quantization_config_pb2 as stablehlo_quant_config_pb2
from tensorflow.compiler.mlir.quantization.tensorflow.calibrator import calibration_statistics_pb2 as calib_stats_pb2
_CalibrationMethod = (
stablehlo_quant_config_pb2.CalibrationOptions.CalibrationMethod
)
_REGISTRY = {}
def _implements(calib_method: _CalibrationMethod):
def decorator(cls):
assert calib_method not in _REGISTRY
_REGISTRY[calib_method] = cls
return cls
return decorator
class _CalibrationAlgorithmBase(abc.ABC):
"""Abstract base class for calibration algorithm."""
def __init__(
self,
statistics: calib_stats_pb2.CalibrationStatistics,
calib_opts: stablehlo_quant_config_pb2.CalibrationOptions,
):
self._statistics = statistics
self._calib_opts = calib_opts
@abc.abstractmethod
def get_min_max_value(self) -> tuple[float, float]:
pass
class _HistogramCalibrationAlgorithmBase(_CalibrationAlgorithmBase):
"""Base class for histogram calibrators."""
def __init__(
self,
statistics: calib_stats_pb2.CalibrationStatistics,
calib_opts: stablehlo_quant_config_pb2.CalibrationOptions,
):
"""Builds histogram using statistics.histogram_statistics.
lower_bound hist_mid
v v
|=========|=========|=========|=========|=========|
bin width
Args:
statistics: Collected calibration statistics.
calib_opts: Calibration options used for calculating min and max.
"""
super().__init__(statistics, calib_opts)
hist_stats = statistics.histogram_statistics
self._bin_width = hist_stats.bin_width
self._lower_bound = hist_stats.lower_bound
self._hist_freq = np.array(hist_stats.hist_freq)
self._num_bins = len(self._hist_freq)
self._num_bits = 8
# i-th bin has a range [bins[i], bins[i + 1]).
# bins[i] = lower_bound + i * bin_width
# bins[i + 1] = lower_bound + (i + 1) * bin_width
# So hist_mids[i] = (lower_bound + bin_width / 2) + bin_width * i
first_mid = self._lower_bound + self._bin_width / 2
last_mid = first_mid + (self._num_bins - 1) * self._bin_width
self._hist_mids = np.linspace(first_mid, last_mid, self._num_bins)
def _get_dequantized_hist_mids_after_quantize(
self, quant_min: float, quant_max: float
) -> np.ndarray:
"""Quantizes and dequantizes hist_mids using quant_min and quant_max.
Quantization converts the range of numbers from [quant_min, quant_max] to
[0, 2^num_bits - 1]. Values less than quant_min are converted to 0, and
values greater than quant_max are converted to 2^num_bits - 1.
The histogram represents the distribution of the data, and our goal is to
find the quant_min and quant_max that best describe this distribution. To do
this, we quantize hist_mids using quant_min and quant_max and dequantize
them again. Then the difference between hist_mids and dequantized hist_mids
equates to quantization error when using quant_min and quant_max.
Args:
quant_min: The minimum real value that can be represented by a quantized
value.
quant_max: The maximum real value that can be represented by a quantized
value.
Returns:
dequantized hist_mids after quantizing by quant_min and quant_max
"""
maxbound = 2**self._num_bits - 1
minbound = 0
scale = (quant_max - quant_min) / maxbound
zero_point = -quant_min / scale
# Limit the range of zero_point and scale in case (quant_max - quant_min)
# is unusually small.
if abs(zero_point) > 9e9:
zero_point = 9e9
if abs(scale) < 1e-9:
scale = 1e-9
zero_point = round(zero_point)
quantized_hist_mids = np.clip(
np.round(self._hist_mids / scale) + zero_point, minbound, maxbound
)
dequantized_hist_mids = scale * (quantized_hist_mids - zero_point)
return dequantized_hist_mids
def _get_weighted_mean_squared_error(
self, quant_min, quant_max
) -> tuple[float, float, float]:
"""Gets mean squared error between hist_mids and dequantized hist_mids.
Quantization converts the range of numbers from [quant_min, quant_max] to
[0, 2^num_bits - 1]. Values less than quant_min are converted to 0, and
values greater than quant_max are converted to 2^num_bits - 1.
Args:
quant_min: The minimum real value that can be represented by a quantized
value.
quant_max: The maximum real value that can be represented by a quantized
value.
Returns:
(error, quant_min, quant_max): Tuple of weighted mean squared error.
error = (hist_mids - dequantized_hist_mids)**2 * hist_freq
"""
dequantized_hist_mids = self._get_dequantized_hist_mids_after_quantize(
quant_min, quant_max
)
squared_error = (self._hist_mids - dequantized_hist_mids) ** 2
weighted_error = np.sum(squared_error * self._hist_freq)
return (weighted_error, quant_min, quant_max)
def _get_min_max_value_by_expanding_range(
self, start_idx: int
) -> tuple[float, float]:
"""Starting from start_idx, expand left and right alternately to find the min value of mse loss.
Args:
start_idx: Index to start quantization.
Returns:
(min_value, max_value): Min and max calculated.
"""
# Tuple of (mse_error, quant_min, quant_max).
mse_min = (float('inf'), float('inf'), float('inf'))
left, right = start_idx, start_idx
# If this value is true, it moves left, otherwise it moves right.
move_left = True
while not (left == 0 and right == self._num_bins - 1):
# Decrease left if right can't be moved or move_left is true.
if (move_left and left > 0) or (right == self._num_bins - 1):
left = max(left - 1, 0)
# Else increase right.
else:
right = min(right + 1, self._num_bins - 1)
# Toogle the move_left.
move_left = not move_left
quant_min, quant_max = self._hist_mids[left], self._hist_mids[right]
mse_tuple = self._get_weighted_mean_squared_error(quant_min, quant_max)
mse_min = min(mse_tuple, mse_min)
# Extract (quant_min, quant_max) from (mse_error, quant_min, quant_max).
min_value, max_value = mse_min[1], mse_min[2]
return min_value, max_value
@_implements(_CalibrationMethod.CALIBRATION_METHOD_MIN_MAX)
class _MinMax(_CalibrationAlgorithmBase):
"""MinMaxCalibrationAlgorithm for calculating min and max values of calibration result.
MinMax calibration calculates the global min and global max values.
global min = min of given sample inputs
global max = max of given sample inputs
"""
def get_min_max_value(self) -> tuple[float, float]:
"""Calculates the global min and max values.
Returns:
(min_value, max_value): Min and max calculated using MinMax
"""
return (
self._statistics.min_max_statistics.global_min,
self._statistics.min_max_statistics.global_max,
)
@_implements(_CalibrationMethod.CALIBRATION_METHOD_AVERAGE_MIN_MAX)
class _AverageMinMax(_CalibrationAlgorithmBase):
"""AverageMinMaxCalibrationAlgorithm for calculating min and max values of calibration result.
AverageMinMax calibration calculates the average of min and max values.
average of min = sum of min values / number of samples
average of max = sum of max values / number of samples
"""
def get_min_max_value(self) -> tuple[float, float]:
"""Calculates the average of min and max values.
Returns:
(min_value, max_value): Min and max calculated using AverageMinMax
Raises:
ValueError: num_samples is 0.
"""
average_min_max_statistics = self._statistics.average_min_max_statistics
# num_samples is guaranteed to be larger than 0 because
# get_statistics_from_calibrator throws an exception if num_samples == 0.
num_samples = average_min_max_statistics.num_samples
if num_samples == 0:
raise ValueError(
'num_samples must not be 0 when calibration method is'
f' AverageMinMax: {self._calib_opts}'
)
min_value, max_value = (
average_min_max_statistics.min_sum / num_samples,
average_min_max_statistics.max_sum / num_samples,
)
return min_value, max_value
@_implements(_CalibrationMethod.CALIBRATION_METHOD_HISTOGRAM_PERCENTILE)
class _HistogramPercentile(_HistogramCalibrationAlgorithmBase):
"""HistogramPercentile for calculating min and max values of calibration result."""
def get_min_max_value(self) -> tuple[float, float]:
"""Calculates min and max from statistics using calibration options.
A "percentile" is a statistical concept that represents the value below
which a given percentage of data falls in a dataset. It involves sorting the
data from smallest to largest and then finding the value at a specified
percentage position. For example, the 0.01 percentile represents the value
in a given data set that corresponds to the lowest 0.01% of the data.
HistogramPercentile calibration uses min_percentile and max_percentile to
find min and max.
min_percentile and max_percentile must be in range [0, 100].
min_percentile is 0.001 by default.
max_percentile is 99.999 by default.
Returns:
(min_value, max_value): Min and max calculated using HistogramPercentile
"""
total_freq = sum(self._hist_freq)
# hist_freq_cumsum is dividing cumulative sum of hist_freq by total_freq
# hist_freq_cumsum's value is in range [0, 1] by its definition
hist_freq_cumsum = np.cumsum(self._hist_freq) / total_freq
# min_percentile and max_percentile are converted from [0, 100] to [0, 1].
min_quantile, max_quantile = (
self._calib_opts.calibration_parameters.min_percentile / 100.0,
self._calib_opts.calibration_parameters.max_percentile / 100.0,
)
# Get index of min/max quantile.
min_quantile_idx, max_quantile_idx = (
np.searchsorted(hist_freq_cumsum, min_quantile, side='right'),
np.searchsorted(hist_freq_cumsum, max_quantile, side='left'),
)
# Get value of min/max quantile index.
min_value, max_value = (
self._hist_mids[min_quantile_idx],
self._hist_mids[max_quantile_idx],
)
return min_value, max_value
@_implements(_CalibrationMethod.CALIBRATION_METHOD_HISTOGRAM_MSE_BRUTEFORCE)
class _HistogramMseBruteforce(_HistogramCalibrationAlgorithmBase):
"""HistogramMseBruteforce for calculating min and max values of calibration result."""
def get_min_max_value(self) -> tuple[float, float]:
"""Finds the optimal quant_min and quant_max by testing all possible cases.
It guarantees optimal quant_min and quant_max for the representative
dataset, but not for the test dataset.
Returns:
(min_value, max_value): Min and max calculated using
HistogramMseBruteforce.
"""
if self._num_bins > 512:
logging.warning(
'num_bins=%d is too large. The HISTOGRAM_MSE_BRUTEFORCE method tests'
' all histogram mid value pairs, so it may take a long time.',
self._num_bins,
)
# Tuple of (mse_error, quant_min, quant_max).
mse_min = (float('inf'), float('inf'), float('inf'))
# Calculate the error for all hist_mid pairs.
for left, right in itertools.combinations(range(self._num_bins), 2):
quant_min, quant_max = self._hist_mids[left], self._hist_mids[right]
mse_tuple = self._get_weighted_mean_squared_error(quant_min, quant_max)
mse_min = min(mse_tuple, mse_min)
min_value, max_value = mse_min[1], mse_min[2]
return min_value, max_value
@_implements(_CalibrationMethod.CALIBRATION_METHOD_HISTOGRAM_MSE_MAX_FREQUENCY)
class _HistogramMseMaxFrequency(_HistogramCalibrationAlgorithmBase):
"""HistogramMseMaxFrequency for calculating min and max values of calibration result."""
def get_min_max_value(self) -> tuple[float, float]:
"""Finds min and max starting from the index of the max frequency.
The HistogramMseMaxFrequency method starts from the bin with the highest
frequency and expands the range to both sides. This performs well when data
is well spread on both sides of the max frequency.
Returns:
(min_value, max_value): Min and max calculated using method to expand the
range based on max frequency.
"""
# Find the index of max frequency.
freq_max_idx = np.argmax(self._hist_freq)
return self._get_min_max_value_by_expanding_range(freq_max_idx)
@_implements(_CalibrationMethod.CALIBRATION_METHOD_HISTOGRAM_MSE_SYMMETRIC)
class _HistogramMseSymmetric(_HistogramCalibrationAlgorithmBase):
"""HistogramMseSymmetric for calculating min and max values of calibration result."""
def get_min_max_value(self) -> tuple[float, float]:
"""Finds min and max starting from the center index.
The HistogramMseSymmetric method starts from the center bin and expands the
range to both sides. This works better when the data is well-centered.
Returns:
(min_value, max_value): Min and max calculated using the method starting
from center and expanding.
"""
# This function is currently only called in this method, but will be used in
# other methods in the future.
return self._get_min_max_value_by_expanding_range(self._num_bins // 2)
def get_min_max_value(
statistics: calib_stats_pb2.CalibrationStatistics,
calib_opts: stablehlo_quant_config_pb2.CalibrationOptions,
) -> tuple[float, float]:
"""Calculates min and max from statistics using calibration options.
Args:
statistics: Collected calibration statistics.
calib_opts: Calibration options used for calculating min and max.
Returns:
(min_value, max_value): Min and max calculated using calib_opts.
Raises:
ValueError: Unsupported calibration method is given.
"""
calib_method = calib_opts.calibration_method
if calib_method not in _REGISTRY:
raise ValueError(f'Unsupported calibration method: {calib_method}')
calibration_algorithm = _REGISTRY[calib_method](statistics, calib_opts)
return calibration_algorithm.get_min_max_value()
tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_algorithm_test.py
+122
View File
@@ -0,0 +1,122 @@
# 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.
# ==============================================================================
"""Tests for CalibrationAlgorithm."""
from absl.testing import parameterized
import numpy as np
from tensorflow.compiler.mlir.quantization.stablehlo import quantization_config_pb2 as stablehlo_quant_config_pb2
from tensorflow.compiler.mlir.quantization.tensorflow.calibrator import calibration_algorithm
from tensorflow.compiler.mlir.quantization.tensorflow.calibrator import calibration_statistics_pb2 as calib_stats_pb2
from tensorflow.python.platform import test
_CalibrationMethod = (
stablehlo_quant_config_pb2.CalibrationOptions.CalibrationMethod
)
class CalibrationAlgorithmTest(test.TestCase, parameterized.TestCase):
def test_min_max_max(self):
calib_opts = stablehlo_quant_config_pb2.CalibrationOptions(
calibration_method=_CalibrationMethod.CALIBRATION_METHOD_MIN_MAX
)
statistics = calib_stats_pb2.CalibrationStatistics()
statistics.min_max_statistics.global_min = 1.0
statistics.min_max_statistics.global_max = 5.0
min_value, max_value = calibration_algorithm.get_min_max_value(
statistics, calib_opts
)
self.assertAllEqual((min_value, max_value), (1.0, 5.0))
def test_average_min_max(self):
calib_opts = stablehlo_quant_config_pb2.CalibrationOptions(
calibration_method=_CalibrationMethod.CALIBRATION_METHOD_AVERAGE_MIN_MAX
)
statistics = calib_stats_pb2.CalibrationStatistics()
statistics.average_min_max_statistics.min_sum = 5.0
statistics.average_min_max_statistics.max_sum = 50.0
statistics.average_min_max_statistics.num_samples = 5
min_value, max_value = calibration_algorithm.get_min_max_value(
statistics, calib_opts
)
self.assertAllEqual((min_value, max_value), (1.0, 10.0))
@parameterized.named_parameters(
{
"testcase_name": "with_histogram_percentile",
"calibration_options": stablehlo_quant_config_pb2.CalibrationOptions(
calibration_method=_CalibrationMethod.CALIBRATION_METHOD_HISTOGRAM_PERCENTILE,
calibration_parameters=stablehlo_quant_config_pb2.CalibrationOptions.CalibrationParameters(
min_percentile=0.001, max_percentile=99.999
),
),
},
{
"testcase_name": "with_histogram_mse_bruteforce",
"calibration_options": stablehlo_quant_config_pb2.CalibrationOptions(
calibration_method=_CalibrationMethod.CALIBRATION_METHOD_HISTOGRAM_MSE_BRUTEFORCE,
calibration_parameters=stablehlo_quant_config_pb2.CalibrationOptions.CalibrationParameters(),
),
},
{
"testcase_name": "with_histogram_mse_max_frequency",
"calibration_options": stablehlo_quant_config_pb2.CalibrationOptions(
calibration_method=_CalibrationMethod.CALIBRATION_METHOD_HISTOGRAM_MSE_MAX_FREQUENCY,
calibration_parameters=stablehlo_quant_config_pb2.CalibrationOptions.CalibrationParameters(),
),
},
{
"testcase_name": "with_histogram_mse_symmetric",
"calibration_options": stablehlo_quant_config_pb2.CalibrationOptions(
calibration_method=_CalibrationMethod.CALIBRATION_METHOD_HISTOGRAM_MSE_SYMMETRIC,
calibration_parameters=stablehlo_quant_config_pb2.CalibrationOptions.CalibrationParameters(),
),
},
)
def test_histogram_calibration_methods(self, calibration_options):
statistics = calib_stats_pb2.CalibrationStatistics()
statistics.histogram_statistics.lower_bound = 0.0
statistics.histogram_statistics.bin_width = 1.0
hist_freq = np.zeros(501, dtype=np.int32)
# Advanced calibration methods that use histograms detect outliers, so they
# don't use the outliers as min/max values.
hist_freq[0] = 1
hist_freq[-1] = 1
# The majority of the data exists around the center.
hist_freq[250] = 1000
for i in range(1, 201):
hist_freq[250 - i] = 1000 - i
hist_freq[250 + i] = 1000 - i
statistics.histogram_statistics.hist_freq.extend(hist_freq.tolist())
# Histogram calibration methods should remove outliers.
min_value, max_value = calibration_algorithm.get_min_max_value(
statistics, calibration_options
)
# Since the min/max values may differ slightly for each calibration method,
# also check the nearby values.
self.assertAllInRange(min_value, 49, 51)
self.assertAllInRange(max_value, 449, 451)
if __name__ == "__main__":
test.main()
tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics.proto
+66
View File
@@ -0,0 +1,66 @@
// 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 tensorflow.calibrator;
option cc_enable_arenas = true;
// Calibration algorithm's collecting statistics.
// NEXT_ID: 4
message CalibrationStatistics {
message MinMaxStatistics {
// global minimum of all sample datasets.
float global_min = 1;
// global maximum of all sample datasets.
float global_max = 2;
}
message AverageMinMaxStatistics {
// sum of batch's minimum in each sample dataset.
float min_sum = 1;
// sum of batch's maximum in each sample dataset.
float max_sum = 2;
// number of sample datasets
int32 num_samples = 3;
}
message HistogramStatistics {
// width of bin
float bin_width = 1;
// lower_bound is the first bin's min value.
// lower_bound and bin_width can be used to restore the histogram.
float lower_bound = 2;
// hist_freq[i] saves frequency of range [bins[i], bins[i + 1]).
// bins[i] = lower_bound + bin_width * i
// bins[i + 1] = lower_bound + bin_width * (i + 1)
repeated float hist_freq = 3;
}
MinMaxStatistics min_max_statistics = 1;
AverageMinMaxStatistics average_min_max_statistics = 2;
HistogramStatistics histogram_statistics = 3;
}
message CalibrationStatisticsMap {
// A map from the id of CustomAggregator op to its collected statistics.
map<string, CalibrationStatistics> statistics = 1;
}
tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_collector_average_min_max.cc
+55
View File
@@ -0,0 +1,55 @@
/* 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.
==============================================================================*/
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_collector_average_min_max.h"
#include <cstdint>
#include <optional>
#include "absl/types/span.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics.pb.h"
namespace tensorflow {
namespace calibrator {
void CalibrationStatisticsCollectorAverageMinMax::ClearData() {
average_min_max_statistics_.set_min_sum(0.0);
average_min_max_statistics_.set_max_sum(0.0);
average_min_max_statistics_.set_num_samples(0);
}
void CalibrationStatisticsCollectorAverageMinMax::Collect(
const float min, const float max, absl::Span<const int64_t> histogram) {
const float current_min_sum = average_min_max_statistics_.min_sum();
const float current_max_sum = average_min_max_statistics_.max_sum();
const int current_num_samples = average_min_max_statistics_.num_samples();
average_min_max_statistics_.set_min_sum(current_min_sum + min);
average_min_max_statistics_.set_max_sum(current_max_sum + max);
average_min_max_statistics_.set_num_samples(current_num_samples + 1);
}
std::optional<CalibrationStatistics>
CalibrationStatisticsCollectorAverageMinMax::GetStatistics() const {
if (average_min_max_statistics_.num_samples() == 0) return std::nullopt;
CalibrationStatistics statistics;
statistics.mutable_average_min_max_statistics()->CopyFrom(
average_min_max_statistics_);
return statistics;
}
} // namespace calibrator
} // namespace tensorflow
tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_collector_average_min_max.h
+52
View File
@@ -0,0 +1,52 @@
/* 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_QUANTIZATION_TENSORFLOW_CALIBRATOR_CALIBRATION_STATISTICS_COLLECTOR_AVERAGE_MIN_MAX_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_CALIBRATOR_CALIBRATION_STATISTICS_COLLECTOR_AVERAGE_MIN_MAX_H_
#include <optional>
#include "absl/types/span.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_collector_base.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
namespace tensorflow {
namespace calibrator {
using ::stablehlo::quantization::CalibrationOptions;
// AverageMinMax calibration calculates the average of min and max values.
// average of min = sum of min values / number of samples
// average of max = sum of max values / number of samples
class CalibrationStatisticsCollectorAverageMinMax
: public CalibrationStatisticsCollectorBase {
public:
explicit CalibrationStatisticsCollectorAverageMinMax() { ClearData(); }
void ClearData() override;
void Collect(float min, float max,
absl::Span<const int64_t> histogram) override;
std::optional<CalibrationStatistics> GetStatistics() const override;
private:
CalibrationStatistics::AverageMinMaxStatistics average_min_max_statistics_;
};
} // namespace calibrator
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_CALIBRATOR_CALIBRATION_STATISTICS_COLLECTOR_AVERAGE_MIN_MAX_H_
tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_collector_base.h
+45
View File
@@ -0,0 +1,45 @@
/* 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_QUANTIZATION_TENSORFLOW_CALIBRATOR_CALIBRATION_STATISTICS_COLLECTOR_BASE_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_CALIBRATOR_CALIBRATION_STATISTICS_COLLECTOR_BASE_H_
#include <cstdint>
#include <optional>
#include "absl/types/span.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics.pb.h"
namespace tensorflow {
namespace calibrator {
// Abstract base class for CalibrationStatisticsCollcetor such as
// CalibrationStatisticsCollectorMinMax. Each class collects different
// statistics based on the calibration methods.
class CalibrationStatisticsCollectorBase {
public:
// Collect data for calibration.
virtual void Collect(float min, float max,
absl::Span<const int64_t> histogram) = 0;
virtual void ClearData() = 0;
// Return the statistics needed for a given calibration method.
virtual std::optional<CalibrationStatistics> GetStatistics() const = 0;
virtual ~CalibrationStatisticsCollectorBase() = default;
};
} // namespace calibrator
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_CALIBRATOR_CALIBRATION_STATISTICS_COLLECTOR_BASE_H_
tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_collector_histogram.cc
+134
View File
@@ -0,0 +1,134 @@
/* 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.
==============================================================================*/
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_collector_histogram.h"
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <deque>
#include <optional>
#include <utility>
#include "absl/types/span.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/calibration_parameters.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
namespace tensorflow {
namespace calibrator {
namespace {
using ::stablehlo::quantization::CalculateBinIndex;
using ::stablehlo::quantization::CalculateBinWidth;
using ::stablehlo::quantization::CalculateLowerBound;
// Gets the histogram frequencies for the given range.
float GetRangeFrequencies(absl::Span<const int64_t> histogram,
const float bin_width, const float lower_bound,
const float range_start, const float range_end) {
float freq_sum = 0.f;
for (float range = std::max(range_start, lower_bound); range < range_end;
range += bin_width) {
const int32_t idx = CalculateBinIndex(range, lower_bound, bin_width);
if (idx >= histogram.size()) break;
// If the range is smaller than bin width, add the proportional value of
// that bin.
const float proportion = std::min(range_end - range, bin_width) / bin_width;
freq_sum += histogram[idx] * proportion;
}
return freq_sum;
}
} // namespace
void CalibrationStatisticsCollectorHistogram::ClearData() {
hist_freq_.clear();
}
void CalibrationStatisticsCollectorHistogram::Collect(
const float min, const float max, absl::Span<const int64_t> histogram) {
if (histogram.empty()) return;
// Reconstruct the bin width, lower and upper bound from the collected data.
const float collected_bin_width =
CalculateBinWidth(min, max, histogram.size());
const float collected_lower_bound =
CalculateLowerBound(min, collected_bin_width);
const float collected_upper_bound =
std::ceil(max / collected_bin_width) * collected_bin_width;
// When histogram is not initialized.
if (hist_freq_.empty()) {
bin_width_ = collected_bin_width;
lower_bound_ = collected_lower_bound;
}
const auto [lower_idx, upper_idx] =
ExpandHistogramIfNeeded(collected_lower_bound, collected_upper_bound);
for (int32_t idx = lower_idx; idx <= upper_idx; ++idx) {
// Calculate the range covered by this index then add with the collected
// frequency associated to that range.
const float range_start = lower_bound_ + idx * bin_width_;
hist_freq_[idx] += GetRangeFrequencies(histogram, collected_bin_width,
collected_lower_bound, range_start,
range_start + bin_width_);
}
}
std::optional<CalibrationStatistics>
CalibrationStatisticsCollectorHistogram::GetStatistics() const {
if (hist_freq_.empty()) return std::nullopt;
CalibrationStatistics::HistogramStatistics hist_stats;
// Skip trailing zeros in the histogram.
int32_t real_size = hist_freq_.size();
for (; real_size > 0; --real_size) {
if (hist_freq_[real_size - 1] != 0) break;
}
hist_stats.set_lower_bound(lower_bound_);
hist_stats.set_bin_width(bin_width_);
hist_stats.mutable_hist_freq()->Assign(hist_freq_.begin(),
hist_freq_.begin() + real_size);
CalibrationStatistics statistics;
statistics.mutable_histogram_statistics()->CopyFrom(hist_stats);
return statistics;
}
std::pair<int32_t, int32_t>
CalibrationStatisticsCollectorHistogram::ExpandHistogramIfNeeded(
const float lower_bound, const float upper_bound) {
int32_t lower_idx = CalculateBinIndex(lower_bound, lower_bound_, bin_width_);
// If lower_idx < 0, then expand the histogram to the left.
if (lower_idx < 0) {
hist_freq_.insert(hist_freq_.begin(), -lower_idx, 0);
lower_bound_ -= bin_width_ * (-lower_idx);
lower_idx = 0;
}
int32_t upper_idx = CalculateBinIndex(upper_bound, lower_bound_, bin_width_);
// If upper_idx >= hist_freq_.size(), then expand the histogram to the right.
if (upper_idx >= hist_freq_.size()) {
hist_freq_.resize(upper_idx + 1, 0);
}
return std::make_pair(lower_idx, upper_idx);
}
} // namespace calibrator
} // namespace tensorflow
tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_collector_histogram.h
+66
View File
@@ -0,0 +1,66 @@
/* 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_QUANTIZATION_TENSORFLOW_CALIBRATOR_CALIBRATION_STATISTICS_COLLECTOR_HISTOGRAM_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_CALIBRATOR_CALIBRATION_STATISTICS_COLLECTOR_HISTOGRAM_H_
#include <cstdint>
#include <deque>
#include <optional>
#include <utility>
#include "absl/types/span.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_collector_base.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
namespace tensorflow {
namespace calibrator {
class CalibrationStatisticsCollectorHistogram
: public CalibrationStatisticsCollectorBase {
public:
explicit CalibrationStatisticsCollectorHistogram() { ClearData(); }
void ClearData() override;
void Collect(float min, float max,
absl::Span<const int64_t> histogram) override;
std::optional<CalibrationStatistics> GetStatistics() const override;
private:
// Expands the histogram so the lower_bound and upper_bound can fit in the
// histogram. Returns the indexes associated to those values.
std::pair<int32_t, int32_t> ExpandHistogramIfNeeded(float lower_bound,
float upper_bound);
// hist_freq_[i] saves frequency of range [bins[i], bins[i + 1]).
// bins[i] = lower_bound_ + bin_width_ * i
// bins[i + 1] = lower_bound_ + bin_width_ * (i + 1)
std::deque<float> hist_freq_;
// Width of bin
float bin_width_;
// The first bin's left value. [left, right)
float lower_bound_;
};
} // namespace calibrator
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_CALIBRATOR_CALIBRATION_STATISTICS_COLLECTOR_HISTOGRAM_H_
tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_collector_min_max.cc
+58
View File
@@ -0,0 +1,58 @@
/* 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.
==============================================================================*/
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_collector_min_max.h"
#include <algorithm>
#include <cstdint>
#include <limits>
#include <optional>
#include "absl/types/span.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics.pb.h"
namespace tensorflow {
namespace calibrator {
void CalibrationStatisticsCollectorMinMax::ClearData() {
// global_min will be updated by std::min(global_min, input_value) so
// it is initialized with the value numeric_limits<float>::max().
min_max_statistics_.set_global_min(std::numeric_limits<float>::max());
// global_max will be updated by std::max(global_max, input_value) so it
// is initialized with the value numeric_limits<float>::lowest().
min_max_statistics_.set_global_max(std::numeric_limits<float>::lowest());
}
void CalibrationStatisticsCollectorMinMax::Collect(
const float min, const float max, absl::Span<const int64_t> histogram) {
min_max_statistics_.set_global_min(
std::min(min_max_statistics_.global_min(), min));
min_max_statistics_.set_global_max(
std::max(min_max_statistics_.global_max(), max));
}
std::optional<CalibrationStatistics>
CalibrationStatisticsCollectorMinMax::GetStatistics() const {
if (min_max_statistics_.global_min() == std::numeric_limits<float>::max())
return std::nullopt;
CalibrationStatistics statistics;
statistics.mutable_min_max_statistics()->CopyFrom(min_max_statistics_);
return statistics;
}
} // namespace calibrator
} // namespace tensorflow
tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_collector_min_max.h
+53
View File
@@ -0,0 +1,53 @@
/* 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_QUANTIZATION_TENSORFLOW_CALIBRATOR_CALIBRATION_STATISTICS_COLLECTOR_MIN_MAX_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_CALIBRATOR_CALIBRATION_STATISTICS_COLLECTOR_MIN_MAX_H_
#include <cstdint>
#include <optional>
#include "absl/types/span.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_collector_base.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
namespace tensorflow {
namespace calibrator {
using ::stablehlo::quantization::CalibrationOptions;
// MinMax calibration calculates the global min and global max values.
// global min = min of given sample inputs
// global max = max of given sample inputs
class CalibrationStatisticsCollectorMinMax
: public CalibrationStatisticsCollectorBase {
public:
explicit CalibrationStatisticsCollectorMinMax() { ClearData(); }
void ClearData() override;
void Collect(float min, float max,
absl::Span<const int64_t> histogram) override;
std::optional<CalibrationStatistics> GetStatistics() const override;
private:
CalibrationStatistics::MinMaxStatistics min_max_statistics_;
};
} // namespace calibrator
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_CALIBRATOR_CALIBRATION_STATISTICS_COLLECTOR_MIN_MAX_H_
tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_collector_test.cc
+304
View File
@@ -0,0 +1,304 @@
/* 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 <optional>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_collector_average_min_max.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_collector_histogram.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_collector_min_max.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace calibrator {
namespace {
using ::testing::ElementsAre;
TEST(CalibrationStatisticsCollectorTest, SimpleMinMax) {
auto collector = CalibrationStatisticsCollectorMinMax();
collector.Collect(
/*min=*/1.0f, /*max=*/10.f, /*histogram=*/{});
collector.Collect(
/*min=*/-5.0f, /*max=*/5.f, /*histogram=*/{});
std::optional<CalibrationStatistics> statistics = collector.GetStatistics();
EXPECT_TRUE(statistics.has_value());
EXPECT_EQ(statistics.value().min_max_statistics().global_min(), -5.0f);
EXPECT_EQ(statistics.value().min_max_statistics().global_max(), 10.0f);
}
TEST(CalibrationStatisticsCollectorTest, SimpleAverageMinMax) {
auto collector = CalibrationStatisticsCollectorAverageMinMax();
collector.Collect(
/*min=*/1.0f, /*max=*/10.f, /*histogram=*/{});
collector.Collect(
/*min=*/-5.0f, /*max=*/5.f, /*histogram=*/{});
std::optional<CalibrationStatistics> statistics = collector.GetStatistics();
EXPECT_TRUE(statistics.has_value());
EXPECT_EQ(statistics.value().average_min_max_statistics().min_sum(), -4.0f);
EXPECT_EQ(statistics.value().average_min_max_statistics().max_sum(), 15.0f);
EXPECT_EQ(statistics.value().average_min_max_statistics().num_samples(), 2);
}
TEST(CalibrationStatisticsCollectorTest, ClearDataAndGetResultsMinMax) {
auto collector = CalibrationStatisticsCollectorMinMax();
collector.Collect(
/*min=*/1.0f, /*max=*/10.f, /*histogram=*/{});
collector.Collect(
/*min=*/-5.0f, /*max=*/5.f, /*histogram=*/{});
std::optional<CalibrationStatistics> statistics = collector.GetStatistics();
EXPECT_TRUE(statistics.has_value());
EXPECT_EQ(statistics.value().min_max_statistics().global_min(), -5.0f);
EXPECT_EQ(statistics.value().min_max_statistics().global_max(), 10.0f);
collector.ClearData();
statistics = collector.GetStatistics();
EXPECT_FALSE(statistics.has_value());
collector.Collect(
/*min=*/1.0f, /*max=*/10.f, /*histogram=*/{});
collector.Collect(
/*min=*/2.0f, /*max=*/5.f, /*histogram=*/{});
statistics = collector.GetStatistics();
EXPECT_TRUE(statistics.has_value());
EXPECT_EQ(statistics.value().min_max_statistics().global_min(), 1.0f);
EXPECT_EQ(statistics.value().min_max_statistics().global_max(), 10.0f);
}
TEST(CalibrationStatisticsCollectorTest, ClearDataAndGetResultsAverageMinMax) {
auto collector = CalibrationStatisticsCollectorAverageMinMax();
collector.Collect(
/*min=*/1.0f, /*max=*/10.f, /*histogram=*/{});
collector.Collect(
/*min=*/-5.0f, /*max=*/5.f, /*histogram=*/{});
std::optional<CalibrationStatistics> statistics = collector.GetStatistics();
EXPECT_TRUE(statistics.has_value());
EXPECT_EQ(statistics.value().average_min_max_statistics().min_sum(), -4.0f);
EXPECT_EQ(statistics.value().average_min_max_statistics().max_sum(), 15.0f);
EXPECT_EQ(statistics.value().average_min_max_statistics().num_samples(), 2);
collector.ClearData();
statistics = collector.GetStatistics();
EXPECT_FALSE(statistics.has_value());
collector.Collect(
/*min=*/1.0f, /*max=*/10.f, /*histogram=*/{});
statistics = collector.GetStatistics();
EXPECT_TRUE(statistics.has_value());
EXPECT_EQ(statistics.value().average_min_max_statistics().min_sum(), 1.0f);
EXPECT_EQ(statistics.value().average_min_max_statistics().max_sum(), 10.0f);
EXPECT_EQ(statistics.value().average_min_max_statistics().num_samples(), 1);
}
TEST(HistogramStatisticsCollectorTest, SingleBatchSimple) {
CalibrationOptions calib_opts;
calib_opts.set_calibration_method(
CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_MSE_MAX_FREQUENCY);
auto collector = CalibrationStatisticsCollectorHistogram();
collector.Collect(
/*min=*/1.f, /*max=*/16.f, /*histogram=*/{1, 0, 3, 5, 7, 6, 5, 0});
std::optional<CalibrationStatistics> statistics = collector.GetStatistics();
EXPECT_TRUE(statistics.has_value());
EXPECT_EQ(statistics.value().histogram_statistics().lower_bound(), 0.f);
EXPECT_EQ(statistics.value().histogram_statistics().bin_width(), 2.f);
// Trailing zeros should be removed.
EXPECT_THAT(statistics.value().histogram_statistics().hist_freq(),
ElementsAre(1, 0, 3, 5, 7, 6, 5));
}
TEST(HistogramStatisticsCollectorTest, AggregateSameBatchSize) {
CalibrationOptions calib_opts;
calib_opts.set_calibration_method(
CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_MSE_MAX_FREQUENCY);
auto collector = CalibrationStatisticsCollectorHistogram();
collector.Collect(
/*min=*/1.f, /*max=*/16.f, /*histogram=*/{1, 0, 3, 5, 7, 6, 5, 1});
std::optional<CalibrationStatistics> statistics = collector.GetStatistics();
EXPECT_TRUE(statistics.has_value());
EXPECT_EQ(statistics.value().histogram_statistics().lower_bound(), 0.f);
EXPECT_EQ(statistics.value().histogram_statistics().bin_width(), 2.f);
EXPECT_THAT(statistics.value().histogram_statistics().hist_freq(),
ElementsAre(1, 0, 3, 5, 7, 6, 5, 1));
collector.Collect(
/*min=*/-1.f, /*max=*/12.f, /*histogram=*/{1, 0, 1, 2, 2, 1, 1, 0});
statistics = collector.GetStatistics();
EXPECT_TRUE(statistics.has_value());
EXPECT_EQ(statistics.value().histogram_statistics().lower_bound(), -2.f);
EXPECT_EQ(statistics.value().histogram_statistics().bin_width(), 2.f);
EXPECT_THAT(statistics.value().histogram_statistics().hist_freq(),
ElementsAre(1, 1, 1, 5, 7, 8, 7, 5, 1));
}
TEST(HistogramStatisticsCollectorTest, AggregateSmallerBatchSizeExpandLeft) {
CalibrationOptions calib_opts;
calib_opts.set_calibration_method(
CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_MSE_MAX_FREQUENCY);
auto collector = CalibrationStatisticsCollectorHistogram();
collector.Collect(
/*min=*/1.f, /*max=*/16.f, /*histogram=*/{1, 0, 3, 5, 7, 6, 5, 1});
std::optional<CalibrationStatistics> statistics = collector.GetStatistics();
EXPECT_TRUE(statistics.has_value());
EXPECT_EQ(statistics.value().histogram_statistics().lower_bound(), 0.f);
EXPECT_EQ(statistics.value().histogram_statistics().bin_width(), 2.f);
EXPECT_THAT(statistics.value().histogram_statistics().hist_freq(),
ElementsAre(1, 0, 3, 5, 7, 6, 5, 1));
collector.Collect(
/*min=*/-1.f, /*max=*/5.f, /*histogram=*/{1, 0, 1, 2, 2, 1, 1, 0});
statistics = collector.GetStatistics();
EXPECT_TRUE(statistics.has_value());
EXPECT_EQ(statistics.value().histogram_statistics().lower_bound(), -2.f);
EXPECT_EQ(statistics.value().histogram_statistics().bin_width(), 2.f);
EXPECT_THAT(statistics.value().histogram_statistics().hist_freq(),
ElementsAre(1, 2, 4, 5, 5, 7, 6, 5, 1));
}
TEST(HistogramStatisticsCollectorTest, AggregateSmallerBatchSizeExpandRight) {
CalibrationOptions calib_opts;
calib_opts.set_calibration_method(
CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_MSE_MAX_FREQUENCY);
auto collector = CalibrationStatisticsCollectorHistogram();
collector.Collect(
/*min=*/1.f, /*max=*/16.f, /*histogram=*/{1, 0, 3, 5, 7, 6, 5, 1});
std::optional<CalibrationStatistics> statistics = collector.GetStatistics();
EXPECT_TRUE(statistics.has_value());
EXPECT_EQ(statistics.value().histogram_statistics().lower_bound(), 0.f);
EXPECT_EQ(statistics.value().histogram_statistics().bin_width(), 2.f);
EXPECT_THAT(statistics.value().histogram_statistics().hist_freq(),
ElementsAre(1, 0, 3, 5, 7, 6, 5, 1));
collector.Collect(
/*min=*/13.f, /*max=*/19.f, /*histogram=*/{1, 0, 1, 2, 2, 1, 1, 0});
statistics = collector.GetStatistics();
EXPECT_TRUE(statistics.has_value());
EXPECT_EQ(statistics.value().histogram_statistics().lower_bound(), 0.f);
EXPECT_EQ(statistics.value().histogram_statistics().bin_width(), 2.f);
EXPECT_THAT(statistics.value().histogram_statistics().hist_freq(),
ElementsAre(1, 0, 3, 5, 7, 6, 6, 2, 4, 2));
}
TEST(HistogramStatisticsCollectorTest, AggregateTinyBinWidth) {
CalibrationOptions calib_opts;
calib_opts.set_calibration_method(
CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_MSE_MAX_FREQUENCY);
auto collector = CalibrationStatisticsCollectorHistogram();
collector.Collect(
/*min=*/1.f, /*max=*/16.f, /*histogram=*/{1, 0, 3, 5, 7, 6, 5, 1});
std::optional<CalibrationStatistics> statistics = collector.GetStatistics();
EXPECT_TRUE(statistics.has_value());
EXPECT_EQ(statistics.value().histogram_statistics().lower_bound(), 0.f);
EXPECT_EQ(statistics.value().histogram_statistics().bin_width(), 2.f);
EXPECT_THAT(statistics.value().histogram_statistics().hist_freq(),
ElementsAre(1, 0, 3, 5, 7, 6, 5, 1));
collector.Collect(
/*min=*/-1.f, /*max=*/-0.99998f, /*histogram=*/{1, 0, 1, 2, 2, 1, 1, 0});
statistics = collector.GetStatistics();
EXPECT_TRUE(statistics.has_value());
EXPECT_EQ(statistics.value().histogram_statistics().lower_bound(), -2.f);
EXPECT_EQ(statistics.value().histogram_statistics().bin_width(), 2.f);
EXPECT_THAT(statistics.value().histogram_statistics().hist_freq(),
ElementsAre(8, 1, 0, 3, 5, 7, 6, 5, 1));
}
TEST(HistogramStatisticsCollectorTest, AggregateLargerBatchSizeExpandLeft) {
CalibrationOptions calib_opts;
calib_opts.set_calibration_method(
CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_MSE_MAX_FREQUENCY);
auto collector = CalibrationStatisticsCollectorHistogram();
collector.Collect(
/*min=*/1.f, /*max=*/16.f, /*histogram=*/{1, 0, 3, 5, 7, 6, 5, 1});
std::optional<CalibrationStatistics> statistics = collector.GetStatistics();
EXPECT_TRUE(statistics.has_value());
EXPECT_EQ(statistics.value().histogram_statistics().lower_bound(), 0.f);
EXPECT_EQ(statistics.value().histogram_statistics().bin_width(), 2.f);
EXPECT_THAT(statistics.value().histogram_statistics().hist_freq(),
ElementsAre(1, 0, 3, 5, 7, 6, 5, 1));
collector.Collect(
/*min=*/-5.f, /*max=*/5.f, /*histogram=*/{1, 2, 2, 1});
statistics = collector.GetStatistics();
EXPECT_TRUE(statistics.has_value());
EXPECT_EQ(statistics.value().histogram_statistics().lower_bound(), -8.f);
EXPECT_EQ(statistics.value().histogram_statistics().bin_width(), 2.f);
EXPECT_THAT(statistics.value().histogram_statistics().hist_freq(),
ElementsAre(0.5, 0.5, 1, 1, 2, 1, 3.5, 5.5, 7, 6, 5, 1));
}
TEST(HistogramStatisticsCollectorTest, AggregateLargerBatchSizeExpandRight) {
CalibrationOptions calib_opts;
calib_opts.set_calibration_method(
CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_MSE_MAX_FREQUENCY);
auto collector = CalibrationStatisticsCollectorHistogram();
collector.Collect(
/*min=*/1.f, /*max=*/16.f, /*histogram=*/{1, 0, 3, 5, 7, 6, 5, 1});
std::optional<CalibrationStatistics> statistics = collector.GetStatistics();
EXPECT_TRUE(statistics.has_value());
EXPECT_EQ(statistics.value().histogram_statistics().lower_bound(), 0.f);
EXPECT_EQ(statistics.value().histogram_statistics().bin_width(), 2.f);
EXPECT_THAT(statistics.value().histogram_statistics().hist_freq(),
ElementsAre(1, 0, 3, 5, 7, 6, 5, 1));
collector.Collect(
/*min=*/10.f, /*max=*/21.f, /*histogram=*/{1, 2, 2, 1});
statistics = collector.GetStatistics();
EXPECT_TRUE(statistics.has_value());
EXPECT_EQ(statistics.value().histogram_statistics().lower_bound(), 0.f);
EXPECT_EQ(statistics.value().histogram_statistics().bin_width(), 2.f);
EXPECT_THAT(statistics.value().histogram_statistics().hist_freq(),
ElementsAre(1, 0, 3, 5, 7.5, 6.5, 6, 2, 1, 1, 0.5, 0.5));
}
} // namespace
} // namespace calibrator
} // namespace tensorflow
tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_saver_op.cc
+190
View File
@@ -0,0 +1,190 @@
/* 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 <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>
#include "absl/base/nullability.h"
#include "absl/container/flat_hash_map.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_collector_average_min_max.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_collector_base.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_collector_histogram.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_collector_min_max.h"
#include "xla/tsl/platform/file_system.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/op_requires.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/platform/logging.h"
namespace tensorflow {
namespace {
using ::stablehlo::quantization::CalibrationOptions;
using CalibrationMethod =
::stablehlo::quantization::CalibrationOptions_CalibrationMethod;
using ::tensorflow::calibrator::CalibrationStatistics;
using ::tensorflow::calibrator::CalibrationStatisticsCollectorAverageMinMax;
using ::tensorflow::calibrator::CalibrationStatisticsCollectorBase;
using ::tensorflow::calibrator::CalibrationStatisticsCollectorHistogram;
using ::tensorflow::calibrator::CalibrationStatisticsCollectorMinMax;
using ::tensorflow::calibrator::CalibrationStatisticsMap;
} // namespace
REGISTER_OP("CalibrationStatisticsSaver")
.Input("args: Tin")
.Attr("Tin: list(type) >= 0")
.Attr("ids: list(string) >= 1")
.Attr("calibration_methods: list(int) >= 1")
.Attr("output_file_path: string")
.SetIsStateful()
.Doc(R"doc(
Aggregates and saves the calibration statistics data.
This op collects outputs of multiples CustomAggregator ops, which includes
`min`, `max` and `histogram`. Then it aggregates them according to the
calibration method and save the result to the given file path as a binary
proto file.)doc");
class CalibrationStatisticsSaverOp : public OpKernel {
public:
explicit CalibrationStatisticsSaverOp(
OpKernelConstruction* absl_nonnull context)
: OpKernel(context) {
std::string output_file_path;
OP_REQUIRES_OK(context,
context->GetAttr("output_file_path", &output_file_path));
OP_REQUIRES_OK(context, context->env()->NewWritableFile(output_file_path,
&output_file_));
OP_REQUIRES_OK(context, context->GetAttr("ids", &ids_));
OP_REQUIRES_OK(context, context->GetAttr("calibration_methods",
&calibration_methods_));
OP_REQUIRES(
context, ids_.size() == calibration_methods_.size(),
absl::AbortedError(
"The `ids` and `calibration_methods` must have the same size."));
// Check the number and type of inputs.
OP_REQUIRES(context, context->num_inputs() == ids_.size() * 3,
absl::AbortedError("The number of inputs must be three times "
"the size of the `ids` list."));
for (int i = 0; i < ids_.size(); ++i) {
OP_REQUIRES(context, context->input_type(i * 3) == DT_FLOAT,
absl::AbortedError("The input `min` must have float type."));
OP_REQUIRES(context, context->input_type(i * 3 + 1) == DT_FLOAT,
absl::AbortedError("The input `max` must have float type."));
OP_REQUIRES(
context, context->input_type(i * 3 + 2) == DT_INT64,
absl::AbortedError("The input `histogram` must have int64 type."));
}
}
~CalibrationStatisticsSaverOp() override {
// Save to file during destruction so we only save it once.
// TODO - b/335044516 : Find a way to flush outside of the destructor.
CalibrationStatisticsMap statistics_map;
for (const auto& [id, collector] : id_to_collector_) {
std::optional<CalibrationStatistics> statistics =
collector->GetStatistics();
if (!statistics.has_value()) continue;
statistics_map.mutable_statistics()->emplace(id, std::move(*statistics));
}
if (auto status = output_file_->Append(statistics_map.SerializeAsString());
!status.ok()) {
LOG(ERROR) << "Failed to write calibration statistics: "
<< status.message();
}
if (auto status = output_file_->Close(); !status.ok()) {
LOG(ERROR) << "Failed to close calibration statistics file: "
<< status.message();
}
}
void Compute(OpKernelContext* absl_nonnull context) override {
for (int idx = 0; idx < ids_.size(); ++idx) {
AssignIfNotExists(
ids_[idx], static_cast<CalibrationMethod>(calibration_methods_[idx]));
const Tensor& min_tensor = context->input(3 * idx);
const Tensor& max_tensor = context->input(3 * idx + 1);
const Tensor& histogram_tensor = context->input(3 * idx + 2);
const float min_value = min_tensor.scalar<float>()();
const float max_value = max_tensor.scalar<float>()();
auto histogram_flat = histogram_tensor.flat<int64_t>();
absl::Span<const int64_t> histogram_data =
absl::MakeSpan(histogram_flat.data(), histogram_flat.size());
id_to_collector_[ids_[idx]]->Collect(min_value, max_value,
histogram_data);
}
}
private:
// The path to save calibration statistics data.
std::unique_ptr<tsl::WritableFile> output_file_;
// The id and calibration method of preceding CustomAggregator ops.
std::vector<std::string> ids_;
std::vector<int32_t> calibration_methods_;
// Map from id to its collector instance.
absl::flat_hash_map<std::string,
std::unique_ptr<CalibrationStatisticsCollectorBase>>
id_to_collector_;
void AssignIfNotExists(absl::string_view id,
const CalibrationMethod calibration_method) {
std::unique_ptr<CalibrationStatisticsCollectorBase>& collector =
id_to_collector_[id];
if (collector != nullptr) return;
switch (calibration_method) {
case CalibrationOptions::CALIBRATION_METHOD_AVERAGE_MIN_MAX:
collector =
std::make_unique<CalibrationStatisticsCollectorAverageMinMax>();
break;
case CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_PERCENTILE:
case CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_MSE_BRUTEFORCE:
case CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_MSE_SYMMETRIC:
case CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_MSE_MAX_FREQUENCY:
collector = std::make_unique<CalibrationStatisticsCollectorHistogram>();
break;
case CalibrationOptions::CALIBRATION_METHOD_MIN_MAX:
default:
collector = std::make_unique<CalibrationStatisticsCollectorMinMax>();
}
}
};
REGISTER_KERNEL_BUILDER(Name("CalibrationStatisticsSaver").Device(DEVICE_CPU),
CalibrationStatisticsSaverOp);
} // namespace tensorflow
tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_saver_op_test.cc
+289
View File
@@ -0,0 +1,289 @@
/* 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 <cstdint>
#include <string>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/log/check.h"
#include "absl/status/status_matchers.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics.pb.h"
#include "xla/tsl/platform/errors.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/kernels/ops_testutil.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace {
using ::stablehlo::quantization::CalibrationOptions;
using ::tensorflow::calibrator::CalibrationStatistics;
using ::tensorflow::calibrator::CalibrationStatisticsMap;
using ::testing::Contains;
using ::testing::ElementsAre;
using ::testing::HasSubstr;
using ::testing::Key;
using ::testing::SizeIs;
class CalibrationStatisticsSaverTest : public OpsTestBase {};
TEST_F(CalibrationStatisticsSaverTest, MissingOutputPath) {
std::vector<std::string> ids{"1"};
std::vector<int32_t> calibration_methods{
CalibrationOptions::CALIBRATION_METHOD_AVERAGE_MIN_MAX};
std::vector<NodeDefBuilder::NodeOut> inputs;
inputs.emplace_back("min", 0, DT_FLOAT);
inputs.emplace_back("max", 0, DT_FLOAT);
CHECK_OK(NodeDefBuilder("op", "CalibrationStatisticsSaver")
.Input(inputs)
.Attr("ids", ids)
.Attr("calibration_methods", calibration_methods)
.Finalize(node_def()));
ASSERT_THAT(InitOp(),
absl_testing::StatusIs(
tsl::error::INVALID_ARGUMENT,
HasSubstr("NodeDef missing attr 'output_file_path'")));
}
TEST_F(CalibrationStatisticsSaverTest, WrongNumInputs) {
std::vector<std::string> ids{"1"};
std::vector<int32_t> calibration_methods{
CalibrationOptions::CALIBRATION_METHOD_AVERAGE_MIN_MAX};
std::vector<NodeDefBuilder::NodeOut> inputs;
inputs.emplace_back("min", 0, DT_FLOAT);
inputs.emplace_back("max", 0, DT_FLOAT);
CHECK_OK(NodeDefBuilder("op", "CalibrationStatisticsSaver")
.Input(inputs)
.Attr("ids", ids)
.Attr("calibration_methods", calibration_methods)
.Attr("output_file_path", "/tmp/statistics.pbtxt")
.Finalize(node_def()));
ASSERT_THAT(InitOp(),
absl_testing::StatusIs(
tsl::error::ABORTED,
HasSubstr("The number of inputs must be three times "
"the size of the `ids` list.")));
}
TEST_F(CalibrationStatisticsSaverTest, WrongInputTypes) {
std::vector<std::string> ids{"1"};
std::vector<int32_t> calibration_methods{
CalibrationOptions::CALIBRATION_METHOD_AVERAGE_MIN_MAX};
std::vector<NodeDefBuilder::NodeOut> inputs;
inputs.emplace_back("min", 0, DT_FLOAT);
inputs.emplace_back("max", 0, DT_FLOAT);
inputs.emplace_back("histogram", 0, DT_FLOAT);
CHECK_OK(NodeDefBuilder("op", "CalibrationStatisticsSaver")
.Input(inputs)
.Attr("ids", ids)
.Attr("calibration_methods", calibration_methods)
.Attr("output_file_path", "/tmp/statistics.pbtxt")
.Finalize(node_def()));
ASSERT_THAT(InitOp(),
absl_testing::StatusIs(
tsl::error::ABORTED,
HasSubstr("The input `histogram` must have int64 type")));
}
TEST_F(CalibrationStatisticsSaverTest, SimpleMinMax) {
std::vector<std::string> ids{"1"};
std::vector<int32_t> calibration_methods{
CalibrationOptions::CALIBRATION_METHOD_MIN_MAX};
std::vector<NodeDefBuilder::NodeOut> inputs;
inputs.emplace_back("min", 0, DT_FLOAT);
inputs.emplace_back("max", 0, DT_FLOAT);
inputs.emplace_back("histogram", 0, DT_INT64);
const std::string dir = testing::TmpDir();
const std::string output_file_path = io::JoinPath(dir, "statistics.pbtxt");
CHECK_OK(NodeDefBuilder("op", "CalibrationStatisticsSaver")
.Input(inputs)
.Attr("ids", ids)
.Attr("calibration_methods", calibration_methods)
.Attr("output_file_path", output_file_path)
.Finalize(node_def()));
CHECK_OK(InitOp());
AddInputFromArray<float>(TensorShape({}), {1.f});
AddInputFromArray<float>(TensorShape({}), {5.f});
AddInputFromArray<int64_t>(TensorShape({0}), {});
CHECK_OK(RunOpKernel());
kernel_.reset();
CalibrationStatisticsMap statistics_map;
CHECK_OK(ReadBinaryProto(Env::Default(), output_file_path, &statistics_map));
ASSERT_THAT(statistics_map.statistics(), SizeIs(1));
ASSERT_THAT(statistics_map.statistics(), ElementsAre(Key("1")));
const CalibrationStatistics& stats = statistics_map.statistics().at("1");
ASSERT_TRUE(stats.has_min_max_statistics());
EXPECT_FLOAT_EQ(stats.min_max_statistics().global_min(), 1.f);
EXPECT_FLOAT_EQ(stats.min_max_statistics().global_max(), 5.f);
}
TEST_F(CalibrationStatisticsSaverTest, SimpleAverageMinMax) {
std::vector<std::string> ids{"1"};
std::vector<int32_t> calibration_methods{
CalibrationOptions::CALIBRATION_METHOD_AVERAGE_MIN_MAX};
std::vector<NodeDefBuilder::NodeOut> inputs;
inputs.emplace_back("min", 0, DT_FLOAT);
inputs.emplace_back("max", 0, DT_FLOAT);
inputs.emplace_back("histogram", 0, DT_INT64);
const std::string dir = testing::TmpDir();
const std::string output_file_path = io::JoinPath(dir, "statistics.pbtxt");
CHECK_OK(NodeDefBuilder("op", "CalibrationStatisticsSaver")
.Input(inputs)
.Attr("ids", ids)
.Attr("calibration_methods", calibration_methods)
.Attr("output_file_path", output_file_path)
.Finalize(node_def()));
CHECK_OK(InitOp());
AddInputFromArray<float>(TensorShape({}), {1.f});
AddInputFromArray<float>(TensorShape({}), {5.f});
AddInputFromArray<int64_t>(TensorShape({0}), {});
CHECK_OK(RunOpKernel());
kernel_.reset();
CalibrationStatisticsMap statistics_map;
CHECK_OK(ReadBinaryProto(Env::Default(), output_file_path, &statistics_map));
ASSERT_THAT(statistics_map.statistics(), SizeIs(1));
ASSERT_THAT(statistics_map.statistics(), ElementsAre(Key("1")));
const CalibrationStatistics& stats = statistics_map.statistics().at("1");
ASSERT_TRUE(stats.has_average_min_max_statistics());
EXPECT_FLOAT_EQ(stats.average_min_max_statistics().min_sum(), 1.f);
EXPECT_FLOAT_EQ(stats.average_min_max_statistics().max_sum(), 5.f);
EXPECT_EQ(stats.average_min_max_statistics().num_samples(), 1);
}
TEST_F(CalibrationStatisticsSaverTest, SimpleHistogram) {
std::vector<std::string> ids{"1"};
std::vector<int32_t> calibration_methods{
CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_MSE_BRUTEFORCE};
std::vector<NodeDefBuilder::NodeOut> inputs;
inputs.emplace_back("min", 0, DT_FLOAT);
inputs.emplace_back("max", 0, DT_FLOAT);
inputs.emplace_back("histogram", 0, DT_INT64);
const std::string dir = testing::TmpDir();
const std::string output_file_path = io::JoinPath(dir, "statistics.pbtxt");
CHECK_OK(NodeDefBuilder("op", "CalibrationStatisticsSaver")
.Input(inputs)
.Attr("ids", ids)
.Attr("calibration_methods", calibration_methods)
.Attr("output_file_path", output_file_path)
.Finalize(node_def()));
CHECK_OK(InitOp());
AddInputFromArray<float>(TensorShape({}), {1.f});
AddInputFromArray<float>(TensorShape({}), {5.f});
AddInputFromArray<int64_t>(TensorShape({8}), {1, 4, 6, 7, 3, 2, 1, 0});
CHECK_OK(RunOpKernel());
kernel_.reset();
CalibrationStatisticsMap statistics_map;
CHECK_OK(ReadBinaryProto(Env::Default(), output_file_path, &statistics_map));
ASSERT_THAT(statistics_map.statistics(), SizeIs(1));
ASSERT_THAT(statistics_map.statistics(), ElementsAre(Key("1")));
const CalibrationStatistics& stats = statistics_map.statistics().at("1");
ASSERT_TRUE(stats.has_histogram_statistics());
EXPECT_FLOAT_EQ(stats.histogram_statistics().bin_width(), 0.5f);
EXPECT_FLOAT_EQ(stats.histogram_statistics().lower_bound(), 1.f);
EXPECT_THAT(stats.histogram_statistics().hist_freq(),
ElementsAre(1, 4, 6, 7, 3, 2, 1));
}
TEST_F(CalibrationStatisticsSaverTest, MultipleStats) {
std::vector<std::string> ids{"1", "2"};
std::vector<int32_t> calibration_methods{
CalibrationOptions::CALIBRATION_METHOD_AVERAGE_MIN_MAX,
CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_MSE_BRUTEFORCE};
std::vector<NodeDefBuilder::NodeOut> inputs;
inputs.emplace_back("min", 0, DT_FLOAT);
inputs.emplace_back("max", 0, DT_FLOAT);
inputs.emplace_back("histogram", 0, DT_INT64);
inputs.emplace_back("min", 0, DT_FLOAT);
inputs.emplace_back("max", 0, DT_FLOAT);
inputs.emplace_back("histogram", 0, DT_INT64);
const std::string dir = testing::TmpDir();
const std::string output_file_path = io::JoinPath(dir, "statistics.pbtxt");
CHECK_OK(NodeDefBuilder("op", "CalibrationStatisticsSaver")
.Input(inputs)
.Attr("ids", ids)
.Attr("calibration_methods", calibration_methods)
.Attr("output_file_path", output_file_path)
.Finalize(node_def()));
CHECK_OK(InitOp());
AddInputFromArray<float>(TensorShape({}), {1.f});
AddInputFromArray<float>(TensorShape({}), {5.f});
AddInputFromArray<int64_t>(TensorShape({0}), {});
AddInputFromArray<float>(TensorShape({}), {1.f});
AddInputFromArray<float>(TensorShape({}), {5.f});
AddInputFromArray<int64_t>(TensorShape({8}), {1, 4, 6, 7, 3, 2, 1, 0});
CHECK_OK(RunOpKernel());
kernel_.reset();
CalibrationStatisticsMap statistics_map;
CHECK_OK(ReadBinaryProto(Env::Default(), output_file_path, &statistics_map));
ASSERT_THAT(statistics_map.statistics(), SizeIs(2));
ASSERT_THAT(statistics_map.statistics(), Contains(Key("1")));
ASSERT_THAT(statistics_map.statistics(), Contains(Key("2")));
const CalibrationStatistics& stats_1 = statistics_map.statistics().at("1");
ASSERT_TRUE(stats_1.has_average_min_max_statistics());
EXPECT_FLOAT_EQ(stats_1.average_min_max_statistics().min_sum(), 1.f);
EXPECT_FLOAT_EQ(stats_1.average_min_max_statistics().max_sum(), 5.f);
EXPECT_EQ(stats_1.average_min_max_statistics().num_samples(), 1);
const CalibrationStatistics& stats_2 = statistics_map.statistics().at("2");
ASSERT_TRUE(stats_2.has_histogram_statistics());
EXPECT_FLOAT_EQ(stats_2.histogram_statistics().bin_width(), 0.5f);
EXPECT_FLOAT_EQ(stats_2.histogram_statistics().lower_bound(), 1.f);
EXPECT_THAT(stats_2.histogram_statistics().hist_freq(),
ElementsAre(1, 4, 6, 7, 3, 2, 1));
}
} // namespace
} // namespace tensorflow
tensorflow/compiler/mlir/quantization/tensorflow/calibrator/custom_aggregator_op.cc
+157
View File
@@ -0,0 +1,157 @@
/* 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.
==============================================================================*/
#define EIGEN_USE_THREADS
#include <cstdint>
#include <string>
#include "absl/status/status.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/calibration_parameters.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
#include "xla/tsl/platform/errors.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/op_requires.h"
#include "tensorflow/core/framework/shape_inference.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/types.h"
namespace tensorflow {
namespace {
using ::stablehlo::quantization::CalculateBinIndexSafe;
using ::stablehlo::quantization::CalculateBinWidth;
using ::stablehlo::quantization::CalculateLowerBound;
using ::stablehlo::quantization::CalibrationOptions;
using CPUDevice = ::Eigen::ThreadPoolDevice;
using CalibrationMethod =
::stablehlo::quantization::CalibrationOptions_CalibrationMethod;
} // namespace
REGISTER_OP("CustomAggregator")
.Input("input: float")
.Output("output: float")
.Output("min: float")
.Output("max: float")
.Output("histogram: int64")
.Attr("id: string")
.Attr("calibration_method: int = 0")
.Attr("num_bins: int = 0")
.Attr("min_percentile: float = 0.0")
.Attr("max_percentile: float = 0.0")
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
c->set_output(0, c->input(0));
c->set_output(1, c->Scalar());
c->set_output(2, c->Scalar());
const tensorflow::AttrValue* num_bins_attr;
TF_RETURN_IF_ERROR(c->GetAttr("num_bins", &num_bins_attr));
c->set_output(3, c->MakeShape({num_bins_attr->i()}));
return absl::OkStatus();
});
class CustomAggregatorOp : public OpKernel {
public:
explicit CustomAggregatorOp(OpKernelConstruction* context)
: OpKernel(context) {
OP_REQUIRES_OK(context, context->GetAttr("id", &id_));
int calibration_method_value;
int num_bins;
float min_percentile;
float max_percentile;
OP_REQUIRES_OK(context, context->GetAttr("calibration_method",
&calibration_method_value));
OP_REQUIRES_OK(context, context->GetAttr("num_bins", &num_bins));
OP_REQUIRES_OK(context,
context->GetAttr("min_percentile", &min_percentile));
OP_REQUIRES_OK(context,
context->GetAttr("max_percentile", &max_percentile));
auto calibration_method =
static_cast<CalibrationMethod>(calibration_method_value);
OP_REQUIRES(
context,
calibration_method !=
CalibrationOptions::CALIBRATION_METHOD_UNSPECIFIED,
absl::AbortedError("The calibration method must be specified."));
calib_opts_.set_calibration_method(calibration_method);
calib_opts_.mutable_calibration_parameters()->set_num_bins(num_bins);
calib_opts_.mutable_calibration_parameters()->set_min_percentile(
min_percentile);
calib_opts_.mutable_calibration_parameters()->set_max_percentile(
max_percentile);
}
void Compute(OpKernelContext* context) override {
const Tensor& input_tensor = context->input(0);
// Use the same input for the first output.
context->set_output(0, input_tensor);
// Calculate min/max statistics.
const auto input_flat = input_tensor.flat<float>();
Tensor *min_output = nullptr, *max_output = nullptr;
OP_REQUIRES_OK(context, context->allocate_output("min", {}, &min_output));
OP_REQUIRES_OK(context, context->allocate_output("max", {}, &max_output));
min_output->scalar<float>().device(
context->template eigen_device<CPUDevice>()) = input_flat.minimum();
max_output->scalar<float>().device(
context->template eigen_device<CPUDevice>()) = input_flat.maximum();
// Calculate histogram statistics.
const int32_t num_bins = calib_opts_.calibration_parameters().num_bins();
Tensor* histogram_output = nullptr;
OP_REQUIRES_OK(context, context->allocate_output("histogram", {num_bins},
&histogram_output));
if (num_bins > 0) {
const float min_value = min_output->scalar<float>()();
const float max_value = max_output->scalar<float>()();
CalculateHistogramStatistics(context, input_tensor, min_value, max_value,
num_bins, histogram_output);
}
}
private:
std::string id_;
CalibrationOptions calib_opts_;
void CalculateHistogramStatistics(OpKernelContext* context,
const Tensor& input_tensor,
const float min_value,
const float max_value,
const int32_t num_bins,
Tensor* histogram_tensor) {
const auto input_flat = input_tensor.flat<float>();
auto histogram_flat = histogram_tensor->flat<int64_t>();
histogram_flat.setZero();
const float bin_width = CalculateBinWidth(min_value, max_value, num_bins);
const float lower_bound = CalculateLowerBound(min_value, bin_width);
for (int i = 0; i < input_flat.size(); ++i) {
int32_t bin_index = CalculateBinIndexSafe(
input_flat.data()[i], lower_bound, bin_width, num_bins);
histogram_flat.data()[bin_index] += 1;
}
}
};
REGISTER_KERNEL_BUILDER(Name("CustomAggregator").Device(DEVICE_CPU),
CustomAggregatorOp);
} // namespace tensorflow
tensorflow/compiler/mlir/quantization/tensorflow/calibrator/integration_test/custom_aggregator_op_test.py
+148
View File
@@ -0,0 +1,148 @@
# 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.
# ==============================================================================
"""Tests for Custom Aggregator op."""
import tensorflow # pylint: disable=unused-import
from tensorflow.compiler.mlir.quantization.stablehlo import quantization_config_pb2 as stablehlo_quant_config_pb2
from tensorflow.compiler.mlir.quantization.tensorflow.calibrator import custom_aggregator_op_wrapper
from tensorflow.python import pywrap_tensorflow # pylint: disable=unused-import
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.platform import test
_CalibrationMethod = (
stablehlo_quant_config_pb2.CalibrationOptions.CalibrationMethod
)
class CustomAggregatorTest(test.TestCase):
def setUp(self):
super(CustomAggregatorTest, self).setUp()
ops.disable_eager_execution()
def testBypassAndMinMax(self):
with self.session():
input_tensor = array_ops.constant(
[1.0, 2.0, 3.0, 4.0, 5.0], dtypes.float32
)
aggregator = custom_aggregator_op_wrapper.custom_aggregator(
input_tensor,
id='1',
calibration_method=_CalibrationMethod.CALIBRATION_METHOD_MIN_MAX,
)
aggregator_output = self.evaluate(aggregator)
self.assertAllEqual(aggregator_output.output, [1.0, 2.0, 3.0, 4.0, 5.0])
self.assertEqual(aggregator_output.min, 1.0)
self.assertEqual(aggregator_output.max, 5.0)
self.assertEmpty(aggregator_output.histogram)
def testTwoIdentities(self):
with self.session():
input_tensor1 = array_ops.constant(
[1.0, 2.0, 3.0, 4.0, 5.0], dtypes.float32
)
aggregator1 = custom_aggregator_op_wrapper.custom_aggregator(
input_tensor1,
'2',
calibration_method=_CalibrationMethod.CALIBRATION_METHOD_MIN_MAX,
)
aggregator1_output = self.evaluate(aggregator1)
self.assertAllEqual(aggregator1_output.output, [1.0, 2.0, 3.0, 4.0, 5.0])
self.assertEqual(aggregator1_output.min, 1.0)
self.assertEqual(aggregator1_output.max, 5.0)
self.assertEmpty(aggregator1_output.histogram)
input_tensor2 = array_ops.constant(
[-1.0, -2.0, -3.0, -4.0, -5.0], dtypes.float32
)
aggregator2 = custom_aggregator_op_wrapper.custom_aggregator(
input_tensor2,
'3',
calibration_method=_CalibrationMethod.CALIBRATION_METHOD_MIN_MAX,
)
aggregator2_output = self.evaluate(aggregator2)
self.assertAllEqual(
aggregator2_output.output, [-1.0, -2.0, -3.0, -4.0, -5.0]
)
self.assertEqual(aggregator2_output.min, -5.0)
self.assertEqual(aggregator2_output.max, -1.0)
self.assertEmpty(aggregator2_output.histogram)
def testBypassAndAverageMinMax(self):
with self.session():
input_tensor1 = array_ops.constant(
[-50.0, -25.0, 0.0, 25.0, 50.0], dtypes.float32
)
aggregator1 = custom_aggregator_op_wrapper.custom_aggregator(
input_tensor1,
'6',
calibration_method=_CalibrationMethod.CALIBRATION_METHOD_AVERAGE_MIN_MAX,
)
aggregator1_output = self.evaluate(aggregator1)
self.assertAllEqual(
aggregator1_output.output,
[-50.0, -25.0, 0.0, 25.0, 50.0],
)
self.assertEqual(aggregator1_output.min, -50.0)
self.assertEqual(aggregator1_output.max, 50.0)
self.assertEmpty(aggregator1_output.histogram)
input_tensor2 = array_ops.constant(
[-100.0, -50.0, 0.0, 50.0, 100.0], dtypes.float32
)
aggregator2 = custom_aggregator_op_wrapper.custom_aggregator(
input_tensor2,
'6',
calibration_method=_CalibrationMethod.CALIBRATION_METHOD_AVERAGE_MIN_MAX,
)
aggregator2_output = self.evaluate(aggregator2)
self.assertAllEqual(
aggregator2_output.output, [-100.0, -50.0, 0.0, 50.0, 100.0]
)
self.assertEqual(aggregator2_output.min, -100.0)
self.assertEqual(aggregator2_output.max, 100.0)
self.assertEmpty(aggregator2_output.histogram)
def testHistogramCalibration(self):
with self.session():
input_tensor = array_ops.constant(
[1.0, 1.0, 3.0, 4.0, 6.0], dtypes.float32
)
aggregator = custom_aggregator_op_wrapper.custom_aggregator(
input_tensor,
id='7',
calibration_method=_CalibrationMethod.CALIBRATION_METHOD_HISTOGRAM_MSE_BRUTEFORCE,
num_bins=512,
)
aggregator_output = self.evaluate(aggregator)
self.assertAllEqual(aggregator_output.output, [1.0, 1.0, 3.0, 4.0, 6.0])
self.assertEqual(aggregator_output.min, 1.0)
self.assertEqual(aggregator_output.max, 6.0)
self.assertLen(aggregator_output.histogram, 512)
self.assertEqual(sum(aggregator_output.histogram), 5)
self.assertEqual(aggregator_output.histogram[0], 2)
self.assertEqual(aggregator_output.histogram[128], 1)
self.assertEqual(aggregator_output.histogram[192], 1)
self.assertEqual(aggregator_output.histogram[320], 1)
if __name__ == '__main__':
test.main()
tensorflow/compiler/mlir/quantization/tensorflow/calibrator/pywrap_calibration.pyi
+32
View File
@@ -0,0 +1,32 @@
# 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.
# ==============================================================================
from tensorflow.compiler.mlir.quantization.tensorflow.calibrator import calibration_statistics_pb2
# LINT.IfChange(clear_calibrator)
def clear_calibrator() -> None: ...
# LINT.ThenChange()
# LINT.IfChange(clear_data_from_calibrator)
def clear_data_from_calibrator(id: bytes) -> None: ...
# LINT.ThenChange()
# LINT.IfChange(get_statistics_from_calibrator)
def get_statistics_from_calibrator(
id: bytes,
) -> calibration_statistics_pb2.CalibrationStatistics: ...
# LINT.ThenChange()
tensorflow/compiler/mlir/quantization/tensorflow/cc/BUILD
+210
View File
@@ -0,0 +1,210 @@
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load(
"//tensorflow:tensorflow.bzl",
"tf_cc_test",
)
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/quantization:__subpackages__",
],
licenses = ["notice"],
)
cc_library(
name = "save_variables",
srcs = ["save_variables.cc"],
hdrs = ["save_variables.h"],
compatible_with = get_compatible_with_portable(),
deps = [
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/core:framework",
"//tensorflow/core/ir/importexport:convert_tensor",
"//tensorflow/core/util/tensor_bundle",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:logging",
"@xla//xla/tsl/platform:status",
],
)
tf_cc_test(
name = "save_variables_test",
srcs = ["save_variables_test.cc"],
deps = [
":save_variables",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/framework:tensor_testutil",
"//tensorflow/core/util/tensor_bundle",
"@com_google_absl//absl/cleanup",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:status_matchers",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
],
)
cc_library(
name = "const_op_size",
srcs = ["const_op_size.cc"],
hdrs = ["const_op_size.h"],
compatible_with = get_compatible_with_portable(),
deps = [
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_remaining_ops",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"@com_google_absl//absl/algorithm:container",
"@llvm-project//mlir:IR",
],
)
tf_cc_test(
name = "const_op_size_test",
srcs = ["const_op_size_test.cc"],
deps = [
":const_op_size",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest_main",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
"@llvm-project//mlir:Support",
],
)
cc_library(
name = "convert_asset_args",
srcs = ["convert_asset_args.cc"],
hdrs = ["convert_asset_args.h"],
compatible_with = get_compatible_with_portable(),
deps = [
"//tensorflow/compiler/mlir/quantization/common:func",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:import_model",
"//tensorflow/core/protobuf:for_core_protos_cc",
"@com_google_absl//absl/algorithm:container",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
)
tf_cc_test(
name = "convert_asset_args_test",
srcs = ["convert_asset_args_test.cc"],
deps = [
":convert_asset_args",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/protobuf:for_core_protos_cc",
"@com_google_absl//absl/strings:string_view",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
"@llvm-project//mlir:Support",
],
)
cc_library(
name = "run_passes",
srcs = ["run_passes.cc"],
hdrs = ["run_passes.h"],
compatible_with = get_compatible_with_portable(),
deps = [
"//tensorflow/compiler/mlir/quantization/tensorflow/debugging:mlir_dump",
"//tensorflow/compiler/mlir/tensorflow:error_util",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/strings:string_view",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:statusor",
],
)
cc_library(
name = "constant_fold",
srcs = [
"constant_fold.cc",
],
hdrs = [
"constant_fold.h",
],
compatible_with = get_compatible_with_portable(),
deps = [
"//tensorflow/compiler/mlir/quantization/common:lift_as_function_call",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow/transforms:constant_fold_utils",
"@com_google_absl//absl/container:flat_hash_set",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
],
)
tf_cc_test(
name = "constant_fold_test",
srcs = ["constant_fold_test.cc"],
deps = [
":constant_fold",
"//tensorflow/compiler/mlir/quantization/common:attrs_and_constraints",
"//tensorflow/compiler/mlir/quantization/common:test_base",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/core:tensorflow",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"@com_google_absl//absl/strings:string_view",
"@com_google_googletest//:gtest",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TransformUtils",
],
)
cc_library(
name = "quantization_unit_loc",
srcs = ["quantization_unit_loc.cc"],
hdrs = ["quantization_unit_loc.h"],
compatible_with = get_compatible_with_portable(),
deps = [
"//tensorflow/compiler/mlir/quantization/tensorflow:quantization_options_proto_cc",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:string_view",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
)
tensorflow/compiler/mlir/quantization/tensorflow/cc/const_op_size.cc
+81
View File
@@ -0,0 +1,81 @@
/* 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/quantization/tensorflow/cc/const_op_size.h"
#include <climits>
#include <cstdint>
#include "absl/algorithm/container.h"
#include "mlir/IR/BuiltinAttributeInterfaces.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
namespace mlir {
namespace quant {
namespace {
// For types that have varying sizes or difficult to determine the size of, each
// element is arbitrarily considered to be 4 bytes.
constexpr int64_t kAssumedNumBytesPerElem = 4;
int64_t GetSizeOfIntOrFloatConst(TF::ConstOp const_op) {
const Type dtype = const_op.getDtype();
const ElementsAttr const_value = const_op.getValue();
const auto bytes_per_elem =
static_cast<int64_t>(dtype.getIntOrFloatBitWidth() / CHAR_BIT);
return bytes_per_elem * const_value.getNumElements();
}
int64_t GetSizeOfStringConst(TF::ConstOp const_op) {
const ElementsAttr const_value = const_op.getValue();
// This cast is guaranteed to succeed. See `ConvertToTensorProto` from
// tensorflow/core/ir/importexport/convert_tensor.cc.
const auto str_attr = cast<DenseStringElementsAttr>(const_value);
// Sum the sizes of each string.
return absl::c_accumulate(
str_attr.getRawStringData(), 0,
[](int64_t acc, const StringRef str_value) -> int64_t {
return acc + str_value.size();
});
}
// Arbitrarily calculate the size of const of type whose size is unkown or
// varying. Each element of such a type is considered to have
// `kAssumedNumBytesPerElem` bytes.
int64_t GetSizeOfUnsupportedTypeConst(TF::ConstOp const_op) {
return kAssumedNumBytesPerElem * const_op.getValue().getNumElements();
}
} // namespace
int64_t GetSizeInBytes(TF::ConstOp const_op) {
const Type dtype = const_op.getDtype();
if (dtype.isIntOrFloat()) {
return GetSizeOfIntOrFloatConst(const_op);
} else if (isa<TF::StringType>(dtype)) {
return GetSizeOfStringConst(const_op);
} else {
return GetSizeOfUnsupportedTypeConst(const_op);
}
}
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/cc/const_op_size.h
+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.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_CC_CONST_OP_SIZE_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_CC_CONST_OP_SIZE_H_
#include <cstdint>
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
namespace mlir {
namespace quant {
// Returns the size in bytes of the underlying data of `const_op`. If the
// underlying type's size cannot be determined, it assumes 4 bytes per element.
int64_t GetSizeInBytes(TF::ConstOp const_op);
} // namespace quant
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_CC_CONST_OP_SIZE_H_
tensorflow/compiler/mlir/quantization/tensorflow/cc/const_op_size_test.cc
+154
View File
@@ -0,0 +1,154 @@
/* 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/quantization/tensorflow/cc/const_op_size.h"
#include <cstdint>
#include <gmock/gmock.h>
#include "absl/strings/string_view.h"
#include "llvm/Support/Casting.h"
#include "mlir/IR/AsmState.h" // from @llvm-project
#include "mlir/IR/Block.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
#include "mlir/Parser/Parser.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/core/platform/test.h"
namespace mlir {
namespace quant {
namespace {
using ::testing::Eq;
class GetSizeInBytesTest : public ::testing::Test {
protected:
GetSizeInBytesTest() : ctx_() { ctx_.loadDialect<TF::TensorFlowDialect>(); }
MLIRContext ctx_;
};
TF::ConstOp ParseConstOp(const absl::string_view const_op_str, Block& block,
MLIRContext& ctx) {
const LogicalResult parse_result =
parseSourceString(const_op_str, &block, ParserConfig(&ctx));
EXPECT_TRUE(succeeded(parse_result));
auto const_op = dyn_cast_or_null<TF::ConstOp>(block.front());
EXPECT_TRUE(const_op);
return const_op;
}
TEST_F(GetSizeInBytesTest, Int32ScalarConstOpSizeInBytes) {
constexpr absl::string_view kConstOpExpr =
R"mlir(%cst = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>)mlir";
Block block{};
TF::ConstOp int_tensor_const_op = ParseConstOp(kConstOpExpr, block, ctx_);
const int64_t num_bytes = GetSizeInBytes(int_tensor_const_op);
EXPECT_THAT(num_bytes, Eq(4));
}
TEST_F(GetSizeInBytesTest, Int32ConstOpSizeInBytes) {
constexpr absl::string_view kConstOpExpr =
R"mlir(%cst = "tf.Const"() {value = dense<1> : tensor<2xi32>} : () -> tensor<2xi32>)mlir";
Block block{};
TF::ConstOp int_tensor_const_op = ParseConstOp(kConstOpExpr, block, ctx_);
const int64_t num_bytes = GetSizeInBytes(int_tensor_const_op);
EXPECT_THAT(num_bytes, Eq(8));
}
TEST_F(GetSizeInBytesTest, Int8ConstOpSizeInBytes) {
constexpr absl::string_view kConstOpExpr =
R"mlir(%cst = "tf.Const"() {value = dense<2> : tensor<3xi8>} : () -> tensor<3xi8>)mlir";
Block block{};
TF::ConstOp int_tensor_const_op = ParseConstOp(kConstOpExpr, block, ctx_);
const int64_t num_bytes = GetSizeInBytes(int_tensor_const_op);
EXPECT_THAT(num_bytes, Eq(3));
}
TEST_F(GetSizeInBytesTest, Float32ConstOpSizeInBytes) {
constexpr absl::string_view kConstOpExpr =
R"mlir(%cst = "tf.Const"() {value = dense<3.0> : tensor<4xf32>} : () -> tensor<4xf32>)mlir";
Block block{};
TF::ConstOp int_tensor_const_op = ParseConstOp(kConstOpExpr, block, ctx_);
const int64_t num_bytes = GetSizeInBytes(int_tensor_const_op);
EXPECT_THAT(num_bytes, Eq(16));
}
TEST_F(GetSizeInBytesTest, Float64ConstOpSizeInBytes) {
constexpr absl::string_view kConstOpExpr =
R"mlir(%cst = "tf.Const"() {value = dense<3.0> : tensor<2xf64>} : () -> tensor<2xf64>)mlir";
Block block{};
TF::ConstOp int_tensor_const_op = ParseConstOp(kConstOpExpr, block, ctx_);
const int64_t num_bytes = GetSizeInBytes(int_tensor_const_op);
EXPECT_THAT(num_bytes, Eq(16));
}
TEST_F(GetSizeInBytesTest, Bfloat16ConstOpSizeInBytes) {
constexpr absl::string_view kConstOpExpr = R"mlir(
%cst = "tf.Const"() {value = dense<1.0> : tensor<7xbf16>} : () -> tensor<7xbf16>
)mlir";
Block block{};
TF::ConstOp int_tensor_const_op = ParseConstOp(kConstOpExpr, block, ctx_);
const int64_t num_bytes = GetSizeInBytes(int_tensor_const_op);
EXPECT_THAT(num_bytes, Eq(14));
}
TEST_F(GetSizeInBytesTest, TfStringConstOpSizeInBytes) {
constexpr absl::string_view kConstOpExpr = R"mlir(
%cst = "tf.Const"() {value = dense<["Hello World", "Quantization"]> : tensor<2x!tf_type.string>} : () -> tensor<2x!tf_type.string>
)mlir";
Block block{};
TF::ConstOp int_tensor_const_op = ParseConstOp(kConstOpExpr, block, ctx_);
// Sum of the number of characters in "Hello World" and "Quantization".
const int64_t num_bytes = GetSizeInBytes(int_tensor_const_op);
EXPECT_THAT(num_bytes, Eq(23));
}
TEST_F(GetSizeInBytesTest, ConstOpWithUnknownSizeAssumes4BytesPerElement) {
constexpr absl::string_view kConstOpExpr = R"mlir(
%cst = "tf.Const"() {value = #tf_type<tensor_proto : "0xDEADBAAD"> : tensor<!tf_type.variant>} : () -> tensor<!tf_type.variant>
)mlir";
Block block{};
TF::ConstOp int_tensor_const_op = ParseConstOp(kConstOpExpr, block, ctx_);
// For non-fixed size like tf_type.variant, the size of each element is
// assumed to be 4 bytes.
const int64_t num_bytes = GetSizeInBytes(int_tensor_const_op);
EXPECT_THAT(num_bytes, Eq(4));
}
} // namespace
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/cc/constant_fold.cc
+150
View File
@@ -0,0 +1,150 @@
/* 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.
==============================================================================*/
#include "tensorflow/compiler/mlir/quantization/tensorflow/cc/constant_fold.h"
#include "absl/container/flat_hash_set.h"
#include "llvm/ADT/SmallVector.h"
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/PatternMatch.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/quantization/common/lift_as_function_call.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/constant_fold_utils.h"
namespace mlir {
namespace quant {
namespace {
// Folds the operation recursively and return the results.
LogicalResult FoldOperation(OpBuilder& builder, Operation* op,
SmallVector<Value>& results) {
SmallVector<ElementsAttr> inputs;
for (auto operand : op->getOperands()) {
auto preceding_const_op = operand.getDefiningOp<TF::ConstOp>();
if (preceding_const_op) {
inputs.push_back(preceding_const_op.getValue());
continue;
}
Operation* preceding_op = operand.getDefiningOp();
int preceding_result_id = -1;
for (auto preceding_result : preceding_op->getResults()) {
if (operand == preceding_result) {
preceding_result_id = preceding_result.getResultNumber();
break;
}
}
SmallVector<Value> preceding_results;
if (failed(FoldOperation(builder, preceding_op, preceding_results))) {
return failure();
}
auto preceding_result = preceding_results[preceding_result_id];
preceding_const_op = preceding_result.getDefiningOp<TF::ConstOp>();
inputs.push_back(preceding_const_op.getValue());
}
SmallVector<Attribute> result_values;
if (failed(TF::EvaluateOperation(op, inputs, result_values))) {
return failure();
}
results.clear();
builder.setInsertionPointAfter(op);
for (const auto& result_value : result_values) {
results.push_back(TF::ConstOp::create(builder, op->getLoc(), result_value));
}
return success();
}
bool IsOperationFoldable(Operation* op) {
if (isa<TF::ConstOp>(op)) return true;
if (op->getDialect()->getNamespace() != "tf" || !TF::CanBeFolded(op)) {
return false;
}
// Check if the operands are foldable as well.
for (auto operand : op->getOperands()) {
auto preceding_op = operand.getDefiningOp();
if (!preceding_op || !IsOperationFoldable(preceding_op)) {
return false;
}
}
return true;
}
// TODO: b/289744814 - Refactor to have a single source of truth of TF Quant
// specs.
absl::flat_hash_set<int> GetQuantizableOperands(Operation* op) {
absl::flat_hash_set<int> quantizable_operands;
if (isa<TF::DepthwiseConv2dNativeOp, TF::Conv2DOp, TF::Conv3DOp, TF::MatMulOp,
TF::BatchMatMulOp>(op)) {
quantizable_operands.insert(1);
} else if (isa<TF::GatherOp>(op)) {
quantizable_operands.insert(0);
} else if (auto einsum_op = dyn_cast<TF::EinsumOp>(op)) {
if (IsEinsumSupportedByXlaDotV2(einsum_op.getEquationAttr())) {
quantizable_operands.insert(1);
}
}
return quantizable_operands;
}
} // namespace
SmallVector<Value> ConstantFoldOpIfPossible(Operation* op) {
if (!IsOperationFoldable(op)) return op->getResults();
OpBuilder builder(op);
SmallVector<Value> results;
if (failed(FoldOperation(builder, op, results))) {
return op->getResults();
}
return results;
}
LogicalResult ConstantFoldQuantizableOperands::matchAndRewrite(
Operation* op, PatternRewriter& rewriter) const {
absl::flat_hash_set<int> quantizable_operands = GetQuantizableOperands(op);
if (quantizable_operands.empty()) return failure();
bool has_change = false;
for (auto operand_idx : quantizable_operands) {
Value operand = op->getOperand(operand_idx);
Operation* preceding_op = operand.getDefiningOp();
if (!preceding_op || isa<TF::ConstOp>(preceding_op)) continue;
int preceding_result_idx = -1;
for (auto preceding_result : preceding_op->getResults()) {
if (operand == preceding_result) {
preceding_result_idx = preceding_result.getResultNumber();
break;
}
}
has_change = has_change || IsOperationFoldable(preceding_op);
SmallVector<Value> folded_results = ConstantFoldOpIfPossible(preceding_op);
op->setOperand(operand_idx, folded_results[preceding_result_idx]);
}
return success(/*isSuccess=*/has_change);
}
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/cc/constant_fold.h
+43
View File
@@ -0,0 +1,43 @@
/* 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_QUANTIZATION_TENSORFLOW_CC_CONSTANT_FOLD_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_CC_CONSTANT_FOLD_H_
#include "llvm/ADT/SmallVector.h"
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
namespace mlir {
namespace quant {
// Applies constant folding recursively if the operation and all of its operands
// are foldable. Returns the constants generated by constant-folding or the
// original operation's outputs if not folded.
SmallVector<Value> ConstantFoldOpIfPossible(Operation* op);
// This pattern tries to constant-fold the quantizable operands of supported
// TF operations.
struct ConstantFoldQuantizableOperands : public RewritePattern {
public:
explicit ConstantFoldQuantizableOperands(MLIRContext* context)
: RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}
LogicalResult matchAndRewrite(Operation* op,
PatternRewriter& rewriter) const override;
};
} // namespace quant
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_CC_CONSTANT_FOLD_H_
tensorflow/compiler/mlir/quantization/tensorflow/cc/constant_fold_test.cc
+201
View File
@@ -0,0 +1,201 @@
/* 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.
==============================================================================*/
#include "tensorflow/compiler/mlir/quantization/tensorflow/cc/constant_fold.h"
#include <utility>
#include <gmock/gmock.h>
#include "absl/strings/string_view.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "mlir/IR/PatternMatch.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 "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.h"
#include "tensorflow/compiler/mlir/quantization/common/test_base.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/core/platform/test.h"
namespace mlir {
namespace quant {
namespace {
using ::testing::NotNull;
using ::testing::SizeIs;
using ConstantFoldingTest = QuantizationTestBase;
TEST_F(ConstantFoldingTest, FoldLargeConstant) {
constexpr absl::string_view kModuleCode = R"mlir(
module {
func.func @test_fold_constant() -> (tensor<1024x24x24x3xf32>) {
%zp = "tf.Const"() {value = dense<2> : tensor<i32>} : () -> tensor<i32>
%scale = "tf.Const"() {value = dense<2.0> : tensor<f32>} : () -> tensor<f32>
%weight = "tf.Const"() {value = dense<1> : tensor<1024x24x24x3xi8>} : () -> tensor<1024x24x24x3xi8>
%input_i32 = "tf.Cast"(%weight) : (tensor<1024x24x24x3xi8>) -> tensor<1024x24x24x3xi32>
%output = "tf.Sub"(%input_i32, %zp) : (tensor<1024x24x24x3xi32>, tensor<i32>) -> tensor<1024x24x24x3xi32>
%cast = "tf.Cast"(%output) : (tensor<1024x24x24x3xi32>) -> tensor<1024x24x24x3xf32>
%mul = "tf.Mul"(%cast, %scale) : (tensor<1024x24x24x3xf32>, tensor<f32>) -> tensor<1024x24x24x3xf32>
func.return %mul : tensor<1024x24x24x3xf32>
}
}
)mlir";
OwningOpRef<ModuleOp> module_op_ref = ParseModuleOpString(kModuleCode);
const auto test_func =
module_op_ref->lookupSymbol<func::FuncOp>("test_fold_constant");
ASSERT_THAT(test_func, NotNull());
Operation* mul_op = FindOperationOfType<TF::MulOp>(test_func);
SmallVector<Value> results = ConstantFoldOpIfPossible(mul_op);
EXPECT_THAT(results, SizeIs(1));
EXPECT_TRUE(isa<TF::ConstOp>(results[0].getDefiningOp()));
}
TEST_F(ConstantFoldingTest, NotFoldingIdentity) {
constexpr absl::string_view kModuleCode = R"mlir(
module {
func.func @test_fold_constant() -> (tensor<1024x24x24x3xf32>) {
%zp = "tf.Const"() {value = dense<2> : tensor<i32>} : () -> tensor<i32>
%scale = "tf.Const"() {value = dense<2.0> : tensor<f32>} : () -> tensor<f32>
%weight = "tf.Const"() {value = dense<1> : tensor<1024x24x24x3xi8>} : () -> tensor<1024x24x24x3xi8>
%input_i32 = "tf.Cast"(%weight) : (tensor<1024x24x24x3xi8>) -> tensor<1024x24x24x3xi32>
%output = "tf.Sub"(%input_i32, %zp) : (tensor<1024x24x24x3xi32>, tensor<i32>) -> tensor<1024x24x24x3xi32>
%cast = "tf.Cast"(%output) : (tensor<1024x24x24x3xi32>) -> tensor<1024x24x24x3xf32>
%identity = "tf.Identity"(%scale) : (tensor<f32>) -> tensor<f32>
%mul = "tf.Mul"(%cast, %identity) : (tensor<1024x24x24x3xf32>, tensor<f32>) -> tensor<1024x24x24x3xf32>
func.return %mul : tensor<1024x24x24x3xf32>
}
}
)mlir";
OwningOpRef<ModuleOp> module_op_ref = ParseModuleOpString(kModuleCode);
const auto test_func =
module_op_ref->lookupSymbol<func::FuncOp>("test_fold_constant");
ASSERT_THAT(test_func, NotNull());
Operation* op_to_fold = FindOperationOfType<TF::MulOp>(test_func);
SmallVector<Value> results = ConstantFoldOpIfPossible(op_to_fold);
EXPECT_THAT(results, SizeIs(1));
// No constant-folding since the IdentityOp has `TF_NoConstantFold` trait.
auto mul_op = dyn_cast_or_null<TF::MulOp>(results[0].getDefiningOp());
EXPECT_THAT(mul_op, NotNull());
// Even though the preceding CastOp is foldable, it shouldn't be folded since
// we are calling from the MulOp.
EXPECT_TRUE(isa<TF::CastOp>(mul_op.getX().getDefiningOp()));
}
TEST_F(ConstantFoldingTest, NotFoldingArgument) {
constexpr absl::string_view kModuleCode = R"mlir(
module {
func.func @test_fold_constant(%arg0: tensor<f32>) -> (tensor<1024x24x24x3xf32>) {
%zp = "tf.Const"() {value = dense<2> : tensor<i32>} : () -> tensor<i32>
%weight = "tf.Const"() {value = dense<1> : tensor<1024x24x24x3xi8>} : () -> tensor<1024x24x24x3xi8>
%input_i32 = "tf.Cast"(%weight) : (tensor<1024x24x24x3xi8>) -> tensor<1024x24x24x3xi32>
%output = "tf.Sub"(%input_i32, %zp) : (tensor<1024x24x24x3xi32>, tensor<i32>) -> tensor<1024x24x24x3xi32>
%cast = "tf.Cast"(%output) : (tensor<1024x24x24x3xi32>) -> tensor<1024x24x24x3xf32>
%mul = "tf.Mul"(%cast, %arg0) : (tensor<1024x24x24x3xf32>, tensor<f32>) -> tensor<1024x24x24x3xf32>
func.return %mul : tensor<1024x24x24x3xf32>
}
}
)mlir";
OwningOpRef<ModuleOp> module_op_ref = ParseModuleOpString(kModuleCode);
const auto test_func =
module_op_ref->lookupSymbol<func::FuncOp>("test_fold_constant");
ASSERT_THAT(test_func, NotNull());
Operation* op_to_fold = FindOperationOfType<TF::MulOp>(test_func);
SmallVector<Value> results = ConstantFoldOpIfPossible(op_to_fold);
EXPECT_THAT(results, SizeIs(1));
// No constant-folding since the second operand is an argument.
TF::MulOp mul_op = dyn_cast_or_null<TF::MulOp>(results[0].getDefiningOp());
EXPECT_THAT(mul_op, NotNull());
// Even though the preceding CastOp is foldable, it shouldn't be folded since
// we are calling from the MulOp.
EXPECT_TRUE(isa<TF::CastOp>(mul_op.getX().getDefiningOp()));
}
TEST_F(ConstantFoldingTest, FoldDepthwiseConvWeight) {
constexpr absl::string_view kModuleCode = R"mlir(
module {
func.func @test_fold_constant(%arg0: tensor<*xf32>) -> (tensor<?x?x?x3xf32>) {
%cst = "tf.Const"() {value = dense<2.000000e+00> : tensor<2x3x3x1xf32>} : () -> tensor<2x3x3x1xf32>
%cst_0 = "tf.Const"() {value = dense<0.400000e+00> : tensor<3xf32>} : () -> tensor<3xf32>
%cst_1 = "tf.Const"() {value = dense<0.500000e+00> : tensor<3xf32>} : () -> tensor<3xf32>
%cst_2 = "tf.Const"() {value = dense<3.0> : tensor<f32>} : () -> tensor<f32>
%w = "tf.Mul"(%cst, %cst_2) : (tensor<2x3x3x1xf32>, tensor<f32>) -> tensor<2x3x3x1xf32>
%0 = "tf.DepthwiseConv2dNative"(%arg0, %w) {data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<*xf32>, tensor<2x3x3x1xf32>) -> tensor<?x?x?x3xf32>
%1 = "tf.BiasAdd"(%0, %cst_0) {data_format = "NHWC"} : (tensor<?x?x?x3xf32>, tensor<3xf32>) -> tensor<?x?x?x3xf32>
%2 = "tf.Mul"(%1, %cst_1) : (tensor<?x?x?x3xf32>, tensor<3xf32>) -> tensor<?x?x?x3xf32>
func.return %2 : tensor<?x?x?x3xf32>
}
}
)mlir";
OwningOpRef<ModuleOp> module_op_ref = ParseModuleOpString(kModuleCode);
const auto test_func =
module_op_ref->lookupSymbol<func::FuncOp>("test_fold_constant");
ASSERT_THAT(test_func, NotNull());
RewritePatternSet patterns(ctx_.get());
patterns.add<ConstantFoldQuantizableOperands>(ctx_.get());
EXPECT_TRUE(succeeded(applyPatternsGreedily(test_func, std::move(patterns))));
auto depthwise_conv_op =
FindOperationOfType<TF::DepthwiseConv2dNativeOp>(test_func);
EXPECT_THAT(depthwise_conv_op, NotNull());
// The filter of the DepthwiseConv2dNativeOp is expected to be a constant.
EXPECT_TRUE(isa<TF::ConstOp>(depthwise_conv_op.getFilter().getDefiningOp()));
}
TEST_F(ConstantFoldingTest, DepthwiseConvWeightNotFoldable) {
constexpr absl::string_view kModuleCode = R"mlir(
module {
func.func @test_fold_constant(%arg0: tensor<*xf32>, %arg1: tensor<f32>) -> (tensor<?x?x?x3xf32>) {
%cst = "tf.Const"() {value = dense<2.000000e+00> : tensor<2x3x3x1xf32>} : () -> tensor<2x3x3x1xf32>
%cst_0 = "tf.Const"() {value = dense<0.400000e+00> : tensor<3xf32>} : () -> tensor<3xf32>
%cst_1 = "tf.Const"() {value = dense<0.500000e+00> : tensor<3xf32>} : () -> tensor<3xf32>
%w = "tf.Mul"(%cst, %arg1) : (tensor<2x3x3x1xf32>, tensor<f32>) -> tensor<2x3x3x1xf32>
%0 = "tf.DepthwiseConv2dNative"(%arg0, %w) {data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<*xf32>, tensor<2x3x3x1xf32>) -> tensor<?x?x?x3xf32>
%1 = "tf.BiasAdd"(%0, %cst_0) {data_format = "NHWC"} : (tensor<?x?x?x3xf32>, tensor<3xf32>) -> tensor<?x?x?x3xf32>
%2 = "tf.Mul"(%1, %cst_1) : (tensor<?x?x?x3xf32>, tensor<3xf32>) -> tensor<?x?x?x3xf32>
func.return %2 : tensor<?x?x?x3xf32>
}
}
)mlir";
OwningOpRef<ModuleOp> module_op_ref = ParseModuleOpString(kModuleCode);
const auto test_func =
module_op_ref->lookupSymbol<func::FuncOp>("test_fold_constant");
ASSERT_THAT(test_func, NotNull());
RewritePatternSet patterns(ctx_.get());
patterns.add<ConstantFoldQuantizableOperands>(ctx_.get());
EXPECT_TRUE(succeeded(applyPatternsGreedily(test_func, std::move(patterns))));
auto depthwise_conv_op =
FindOperationOfType<TF::DepthwiseConv2dNativeOp>(test_func);
EXPECT_THAT(depthwise_conv_op, NotNull());
// The filter of the DepthwiseConv2dNativeOp is not constant-foldable.
EXPECT_TRUE(isa<TF::MulOp>(depthwise_conv_op.getFilter().getDefiningOp()));
}
} // namespace
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/cc/convert_asset_args.cc
+142
View File
@@ -0,0 +1,142 @@
/* 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.
==============================================================================*/
#include "tensorflow/compiler/mlir/quantization/tensorflow/cc/convert_asset_args.h"
#include "llvm/ADT/SmallVector.h"
#include "mlir/Dialect/Func/IR/FuncOps.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/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/SymbolTable.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/quantization/common/func.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/compiler/mlir/tensorflow/translate/import_model.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
namespace mlir::quant {
namespace {
using ::mlir::tf_saved_model::AssetOp;
using ::mlir::tf_saved_model::kTfSavedModelIndexPathAttr;
using ::mlir::tf_saved_model::LookupBoundInputOfType;
using ::tensorflow::AssetFileDef;
// Given argument attributes `arg_attrs`, returns a new set of argument
// attributes where the "tf_saved_model.bound_input" attribute has been replaced
// with the "tf_saved_model.index_path" attribute. `index_path` is the element
// of the index path attribute.
SmallVector<NamedAttribute> ReplaceBoundInputAttrWithIndexPathAttr(
const ArrayRef<NamedAttribute> arg_attrs, const StringRef index_path,
Builder& builder) {
// Keep all other attributes except the tf_saved_model.bound_input attribute,
// as we are replacing it with tf_saved_model.index_path.
SmallVector<NamedAttribute> new_arg_attrs;
for (auto arg_attr : arg_attrs) {
if (arg_attr.getName() == "tf_saved_model.bound_input") continue;
new_arg_attrs.emplace_back(arg_attr);
}
const NamedAttribute index_path_attr(
builder.getStringAttr(kTfSavedModelIndexPathAttr),
builder.getStrArrayAttr({index_path}));
new_arg_attrs.emplace_back(index_path_attr);
return new_arg_attrs;
}
// Strips the "assets/" directory prefix, if `filename` begins with it. The
// SavedModel loader attaches the prefix for you during loading.
StringRef MaybeStripAssetDirectoryPrefix(const StringRef filename) {
if (filename.find("assets/") == 0) {
return filename.drop_front(7);
} else {
return filename;
}
}
AssetFileDef CreateAssetFileDef(const StringRef filename,
const StringRef tensor_name) {
AssetFileDef asset_file_def{};
asset_file_def.set_filename(MaybeStripAssetDirectoryPrefix(filename).str());
tensorflow::TensorInfo tensor_info{};
tensor_info.set_name(tensor_name.str());
*asset_file_def.mutable_tensor_info() = tensor_info;
return asset_file_def;
}
// Returns a list of "tf.entry_function" attribute's "inputs" comma-split
// values.
//
// Example: if `func_op` has attribute `tf.entry_function = {inputs =
// "arg0:0,arg1:0"}`, then this function returns `{"arg0:0", "arg1:0"}`.
SmallVector<StringRef> GetEntryFunctionInputs(func::FuncOp func_op) {
auto entry_function_attr =
func_op->getAttrOfType<DictionaryAttr>("tf.entry_function");
SmallVector<StringRef> inputs;
mlir::dyn_cast_or_null<StringAttr>(entry_function_attr.get("inputs"))
.strref()
.split(inputs, /*Separator=*/",");
return inputs;
}
void ConvertMainArgAttrs(func::FuncOp main_func_op, const int arg_idx,
const StringRef index_path) {
const ArrayRef<NamedAttribute> arg_attrs =
main_func_op.getArgAttrDict(arg_idx).getValue();
Builder builder(main_func_op.getContext());
SmallVector<NamedAttribute> new_arg_attrs =
ReplaceBoundInputAttrWithIndexPathAttr(arg_attrs, index_path, builder);
main_func_op.setArgAttrs(arg_idx, new_arg_attrs);
}
} // namespace
FailureOr<SmallVector<AssetFileDef>> ConvertAssetArgs(ModuleOp module_op) {
func::FuncOp main_func_op = FindMainFuncOp(module_op);
if (!main_func_op) return failure();
SmallVector<StringRef> input_names = GetEntryFunctionInputs(main_func_op);
SymbolTable symbol_table(module_op);
SmallVector<AssetFileDef> asset_file_defs;
for (BlockArgument argument : main_func_op.getArguments()) {
const int arg_idx = argument.getArgNumber();
auto asset_op =
LookupBoundInputOfType<AssetOp>(main_func_op, arg_idx, symbol_table);
if (!asset_op) continue;
const StringRef input_name = input_names[arg_idx];
ConvertMainArgAttrs(main_func_op, arg_idx, /*index_path=*/input_name);
// This assumes that the final tensor name in `GraphDef` is equal to
// `input_name`.
asset_file_defs.emplace_back(CreateAssetFileDef(
asset_op.getFilenameAttr(), /*tensor_name=*/input_name));
}
return asset_file_defs;
}
} // namespace mlir::quant
tensorflow/compiler/mlir/quantization/tensorflow/cc/convert_asset_args.h
+40
View File
@@ -0,0 +1,40 @@
/* 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_QUANTIZATION_TENSORFLOW_CC_CONVERT_ASSET_ARGS_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_CC_CONVERT_ASSET_ARGS_H_
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/core/protobuf/meta_graph.pb.h"
namespace mlir::quant {
// Converts arguments of the @main function that are bound to
// `tf_saved_model::AssetOp`s into regular tensor args. Returns `AsestFileDef`s
// that associates the arg with the asset.
//
// In detail, this function performs the following:
// * Replaces "tf_saved_model.bound_input" attributes to
// "tf_saved_model.index_path", if the bound input is attached to the
// `tf_saved_model::AssetOp`.
// * Strips the "assets/" prefix of the filename when setting it to
// `AssetFileDef`.
FailureOr<SmallVector<tensorflow::AssetFileDef>> ConvertAssetArgs(
ModuleOp module_op);
} // namespace mlir::quant
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_CC_CONVERT_ASSET_ARGS_H_
tensorflow/compiler/mlir/quantization/tensorflow/cc/convert_asset_args_test.cc
+169
View File
@@ -0,0 +1,169 @@
/* 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.
==============================================================================*/
#include "tensorflow/compiler/mlir/quantization/tensorflow/cc/convert_asset_args.h"
#include <gmock/gmock.h>
#include "absl/strings/string_view.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.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/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
namespace mlir::quant {
namespace {
using ::tensorflow::AssetFileDef;
using ::testing::Eq;
using ::testing::IsEmpty;
using ::testing::IsNull;
using ::testing::NotNull;
using ::testing::SizeIs;
class ConvertAssetArgsTest : public ::testing::Test {
protected:
ConvertAssetArgsTest() {
ctx_.loadDialect<func::FuncDialect, TF::TensorFlowDialect,
tf_saved_model::TensorFlowSavedModelDialect>();
}
// Parses `module_op_str` to create a `ModuleOp`. Checks whether the created
// module op is valid.
OwningOpRef<ModuleOp> ParseModuleOpString(
const absl::string_view module_op_str) {
auto module_op_ref = parseSourceString<ModuleOp>(module_op_str, &ctx_);
EXPECT_TRUE(module_op_ref);
return module_op_ref;
}
mlir::MLIRContext ctx_{};
};
func::FuncOp GetMainFuncOp(ModuleOp module_op) {
for (auto func_op : module_op.getOps<func::FuncOp>()) {
if (func_op.getSymName() == "main") {
return func_op;
}
}
return {};
}
TEST_F(ConvertAssetArgsTest, ConvertsSingleAssetArg) {
OwningOpRef<ModuleOp> module_op = ParseModuleOpString(R"mlir(
module {
"tf_saved_model.asset"() {filename = "assets/file_0.txt", sym_name = "__tf_saved_model_asset0"} : () -> ()
func.func @main(%arg_0: tensor<!tf_type.string> {tf_saved_model.bound_input = @__tf_saved_model_asset0}) -> () attributes {tf.entry_function = {inputs = "arg_0:0", outputs = ""}} {
return
}
}
)mlir");
FailureOr<SmallVector<AssetFileDef>> asset_file_defs =
ConvertAssetArgs(*module_op);
EXPECT_TRUE(succeeded(asset_file_defs));
EXPECT_THAT(*asset_file_defs, SizeIs(1));
const AssetFileDef& asset_file_def = *asset_file_defs->begin();
EXPECT_THAT(asset_file_def.filename(), Eq("file_0.txt"));
EXPECT_THAT(asset_file_def.tensor_info().name(), Eq("arg_0:0"));
func::FuncOp main_func_op = GetMainFuncOp(*module_op);
DictionaryAttr arg_attrs = main_func_op.getArgAttrDict(/*index=*/0);
EXPECT_THAT(arg_attrs.get("tf_saved_model.bound_input"), IsNull());
const ArrayRef<Attribute> index_path_attrs =
mlir::cast<ArrayAttr>(arg_attrs.get("tf_saved_model.index_path"))
.getValue();
EXPECT_THAT(index_path_attrs, SizeIs(1));
StringAttr index_path =
mlir::dyn_cast_or_null<StringAttr>(index_path_attrs[0]);
EXPECT_THAT(index_path, NotNull());
EXPECT_THAT(index_path, Eq("arg_0:0"));
}
TEST_F(ConvertAssetArgsTest, NonBoundedArgsNotModified) {
OwningOpRef<ModuleOp> module_op = ParseModuleOpString(R"mlir(
module {
func.func @main(%arg_0: tensor<!tf_type.string> {tf_saved_model.index_path = ["arg_0:0"]}) -> () attributes {tf.entry_function = {inputs = "arg_0:0", outputs = ""}} {
return
}
}
)mlir");
FailureOr<SmallVector<AssetFileDef>> asset_file_defs =
ConvertAssetArgs(*module_op);
EXPECT_TRUE(succeeded(asset_file_defs));
EXPECT_THAT(*asset_file_defs, IsEmpty());
func::FuncOp main_func_op = GetMainFuncOp(*module_op);
DictionaryAttr arg_attrs = main_func_op.getArgAttrDict(/*index=*/0);
EXPECT_THAT(arg_attrs.get("tf_saved_model.bound_input"), IsNull());
const ArrayRef<Attribute> index_path_attrs =
mlir::cast<ArrayAttr>(arg_attrs.get("tf_saved_model.index_path"))
.getValue();
EXPECT_THAT(index_path_attrs, SizeIs(1));
StringAttr index_path =
mlir::dyn_cast_or_null<StringAttr>(index_path_attrs[0]);
EXPECT_THAT(index_path, NotNull());
EXPECT_THAT(index_path, Eq("arg_0:0"));
}
TEST_F(ConvertAssetArgsTest, ArgsBoundedToGlobalTensorNotModified) {
// If the argument is not bound to AssetOp, it is not modified.
OwningOpRef<ModuleOp> module_op = ParseModuleOpString(R"mlir(
module {
"tf_saved_model.global_tensor"() {type = tensor<2xi32>, value = dense<2> : tensor<2xi32>, sym_name = "__tf_saved_model_x"} : () -> ()
func.func @main(%arg_0: tensor<!tf_type.resource<tensor<2xi32>>> {tf_saved_model.bound_input = @__tf_saved_model_x}) -> () attributes {tf.entry_function = {inputs = "arg_0:0", outputs = ""}} {
return
}
}
)mlir");
FailureOr<SmallVector<AssetFileDef>> asset_file_defs =
ConvertAssetArgs(*module_op);
EXPECT_TRUE(succeeded(asset_file_defs));
EXPECT_THAT(*asset_file_defs, IsEmpty());
func::FuncOp main_func_op = GetMainFuncOp(*module_op);
DictionaryAttr arg_attrs = main_func_op.getArgAttrDict(/*index=*/0);
EXPECT_THAT(arg_attrs.get("tf_saved_model.bound_input"), NotNull());
}
TEST_F(ConvertAssetArgsTest, FailsWhenNoMain) {
OwningOpRef<ModuleOp> module_op = ParseModuleOpString(R"mlir(module {})mlir");
FailureOr<SmallVector<AssetFileDef>> asset_file_defs =
ConvertAssetArgs(*module_op);
EXPECT_TRUE(failed(asset_file_defs));
}
} // namespace
} // namespace mlir::quant
tensorflow/compiler/mlir/quantization/tensorflow/cc/quantization_unit_loc.cc
+130
View File
@@ -0,0 +1,130 @@
/* 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.
==============================================================================*/
#include "tensorflow/compiler/mlir/quantization/tensorflow/cc/quantization_unit_loc.h"
#include <cstddef>
#include <optional>
#include <string>
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/Location.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
namespace mlir {
namespace quant {
namespace {
// Prefix and suffix to the QuantizationUnit string representation.
constexpr absl::string_view kQuantizationUnitPrefix = "QuantizationUnit(";
constexpr absl::string_view kQuantizationUnitSuffix = ")";
// Concatenates node name and func name with a "@" separator.
std::string ConcatNodeAndFuncName(absl::string_view node_name,
absl::string_view func_name) {
return absl::StrCat(node_name, "@", func_name);
}
// Generate a string to represent the QuantizationUnit.
std::string GenerateQuantizationUnitString(
const QuantizationUnitLoc::QuantizationUnit& unit) {
return absl::StrCat(kQuantizationUnitPrefix, unit.SerializeAsString(),
kQuantizationUnitSuffix);
}
std::optional<StringRef> CallerNameFromCallSiteLoc(CallSiteLoc callsite_loc) {
// loc(callsite("func" at "QuantizationUnit(...)"))
if (mlir::isa<NameLoc>(callsite_loc.getCaller())) {
return mlir::cast<NameLoc>(callsite_loc.getCaller()).getName().strref();
}
// loc(callsite("func" at callsite("QuantizationUnit(...)" at ...)))
if (mlir::isa<CallSiteLoc>(callsite_loc.getCaller())) {
CallSiteLoc caller_callsite_loc =
mlir::cast<CallSiteLoc>(callsite_loc.getCaller());
if (mlir::isa<NameLoc>(caller_callsite_loc.getCallee())) {
return mlir::cast<NameLoc>(caller_callsite_loc.getCallee())
.getName()
.strref();
}
}
return std::nullopt;
}
} // namespace
QuantizationUnitLoc::QuantizationUnitLoc(MLIRContext* context,
const QuantizationUnit& unit)
: CallSiteLoc(CallSiteLoc::get(
/*callee=*/NameLoc::get(
StringAttr::get(context, ConcatNodeAndFuncName(unit.node_name(),
unit.func_name())),
/*childLoc=*/NameLoc::get(
StringAttr::get(context, unit.op_type()))),
/*caller=*/NameLoc::get(StringAttr::get(
context, GenerateQuantizationUnitString(unit))))) {}
bool QuantizationUnitLoc::classof(Attribute attr) {
if (!llvm::isa<CallSiteLoc>(attr)) return false;
auto callsite_loc = llvm::dyn_cast<CallSiteLoc>(attr);
std::optional<StringRef> caller_name =
CallerNameFromCallSiteLoc(callsite_loc);
return caller_name && caller_name->starts_with(kQuantizationUnitPrefix) &&
caller_name->ends_with(kQuantizationUnitSuffix);
}
std::optional<QuantizationUnitLoc::QuantizationUnit>
FindQuantizationUnitFromLoc(Location loc) {
if (isa<QuantizationUnitLoc>(loc)) {
std::optional<StringRef> caller_name =
CallerNameFromCallSiteLoc(mlir::cast<CallSiteLoc>(loc));
if (!caller_name) {
return std::nullopt;
}
const size_t start_index = kQuantizationUnitPrefix.size();
const size_t end_index = caller_name->rfind(kQuantizationUnitSuffix);
std::string serialized_proto =
caller_name->substr(start_index, end_index - start_index).str();
QuantizationUnitLoc::QuantizationUnit quant_unit;
if (quant_unit.ParseFromString(serialized_proto)) {
return quant_unit;
}
} else if (isa<FusedLoc>(loc)) {
// If the op is rewritten, FusedLoc can be created.
for (Location child_loc : mlir::cast<FusedLoc>(loc).getLocations()) {
std::optional<QuantizationUnitLoc::QuantizationUnit> found_unit =
FindQuantizationUnitFromLoc(child_loc);
if (found_unit.has_value()) return found_unit;
}
} else if (isa<CallSiteLoc>(loc)) {
// If the graph is inlined, CallSiteLoc can be created.
return FindQuantizationUnitFromLoc(
mlir::cast<CallSiteLoc>(loc).getCallee());
}
return std::nullopt;
}
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/cc/quantization_unit_loc.h
+54
View File
@@ -0,0 +1,54 @@
/* 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_QUANTIZATION_TENSORFLOW_CC_QUANTIZATION_UNIT_LOC_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_CC_QUANTIZATION_UNIT_LOC_H_
#include <optional>
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Location.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
namespace mlir {
namespace quant {
// QuantizationUnitLoc uses CallSiteLoc as the base class so it can be printed
// with AsmPrinter and used to set the node name in MLIR to GraphDef exporter.
// The callee is named as `node_name@func_name` with child loc named as
// `op_type` while the caller is the quantization unit.
class QuantizationUnitLoc : public CallSiteLoc {
public:
using QuantizationUnit =
tensorflow::quantization::UnitWiseQuantizationSpec::QuantizationUnit;
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(QuantizationUnitLoc)
QuantizationUnitLoc(MLIRContext* context, const QuantizationUnit& unit);
// Checks if the given location is QuantizationUnitLoc. Users could call
// `isa<QuantizationUnitLoc>(loc)` to check if the type matches.
static bool classof(Attribute attr);
};
// Finds the QuantizationUnit from location info.
std::optional<QuantizationUnitLoc::QuantizationUnit>
FindQuantizationUnitFromLoc(Location loc);
} // namespace quant
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_CC_QUANTIZATION_UNIT_LOC_H_
tensorflow/compiler/mlir/quantization/tensorflow/cc/run_passes.cc
+52
View File
@@ -0,0 +1,52 @@
/* 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.
==============================================================================*/
#include "tensorflow/compiler/mlir/quantization/tensorflow/cc/run_passes.h"
#include <memory>
#include <optional>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/tensorflow/debugging/mlir_dump.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
#include "xla/tsl/platform/errors.h"
namespace tensorflow {
namespace quantization {
absl::Status RunPassesOnModuleOp(
std::optional<absl::string_view> mlir_dump_file_name,
mlir::PassManager& pass_manager, mlir::ModuleOp module_op) {
mlir::StatusScopedDiagnosticHandler statusHandler(module_op.getContext(),
/*propagate=*/true);
absl::StatusOr<std::unique_ptr<llvm::raw_ostream>> dump_file;
if (mlir_dump_file_name) {
TF_RETURN_IF_ERROR(tensorflow::quantization::MaybeEnableIrPrinting(
pass_manager, mlir_dump_file_name.value()));
}
if (failed(pass_manager.run(module_op))) {
return statusHandler.ConsumeStatus();
}
return absl::OkStatus();
}
} // namespace quantization
} // namespace tensorflow
tensorflow/compiler/mlir/quantization/tensorflow/cc/run_passes.h
+78
View File
@@ -0,0 +1,78 @@
/* 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_QUANTIZATION_TENSORFLOW_CC_RUN_PASSES_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_CC_RUN_PASSES_H_
#include <optional>
#include "absl/status/status.h"
#include "absl/strings/str_format.h"
#include "absl/strings/string_view.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/tensorflow/debugging/mlir_dump.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/statusor.h"
namespace tensorflow {
namespace quantization {
// Runs MLIR passes with `module_op`. The passes are added by calling
// `add_passes_func`, which is a callable receiving mlir::PassManager& as its
// only argument. `name` identifies the set of passes added by `add_passes_func`
// and is used for debugging. Changing the `name` does not modify the behavior
// of the passes.
//
// It will try to dump intermediate MLIRs if certain conditions are met. See the
// description from `MaybeEnableIrPrinting` for the details about the
// conditions.
//
// Returns a non-OK status when the pass run fails or it fails to create an MLIR
// dump file.
template <typename FuncT>
absl::Status RunPasses(const absl::string_view name, FuncT add_passes_func,
mlir::MLIRContext& ctx, mlir::ModuleOp module_op) {
mlir::PassManager pm{&ctx};
add_passes_func(pm);
mlir::StatusScopedDiagnosticHandler diagnostic_handler{&ctx};
TF_RETURN_IF_ERROR(MaybeEnableIrPrinting(pm, name));
if (failed(pm.run(module_op))) {
return absl::InternalError(
absl::StrFormat("Failed to run pass: %s. %s", name,
diagnostic_handler.ConsumeStatus().message()));
}
return absl::OkStatus();
}
// Runs MLIR passes with `module_op` on a `pass_manager`.
//
// It will try to dump intermediate MLIRs if certain conditions are met. See the
// description from `MaybeEnableIrPrinting` for the details about the
// conditions.
//
// Returns a non-OK status when the pass run fails or it fails to create an MLIR
// dump file.
absl::Status RunPassesOnModuleOp(
std::optional<absl::string_view> mlir_dump_file_name,
mlir::PassManager& pass_manager, mlir::ModuleOp module_op);
} // namespace quantization
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_CC_RUN_PASSES_H_
tensorflow/compiler/mlir/quantization/tensorflow/cc/save_variables.cc
+133
View File
@@ -0,0 +1,133 @@
/* 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/quantization/tensorflow/cc/save_variables.h"
#include <string>
#include <utility>
#include <vector>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "xla/tsl/platform/env.h"
#include "xla/tsl/platform/logging.h"
#include "xla/tsl/platform/status.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/ir/importexport/convert_tensor.h"
#include "tensorflow/core/util/tensor_bundle/tensor_bundle.h"
namespace tensorflow {
namespace quantization {
namespace {
using ::mlir::func::FuncOp;
using ::mlir::tf_saved_model::GetInitializerFunction;
using ::mlir::tf_saved_model::kTfSavedModelInitializerRestoreType;
// Adds the tensor that initializes the variable through the provided
// `assign_var_op` to the `bundle_writer` for saving to checkpoint. Returns the
// shared name of the variable if a variable is saved successfully. If the
// variable is not saved, returns an empty string.
absl::StatusOr<std::string> AddTensorToBundleWriter(
mlir::TF::AssignVariableOp assign_var_op, BundleWriter& bundle_writer) {
auto resource_operand = assign_var_op.getOperand(0);
auto var_handle_op =
llvm::dyn_cast<mlir::TF::VarHandleOp>(resource_operand.getDefiningOp());
if (!var_handle_op) {
assign_var_op->emitRemark(
"Operand idx 0 is not a tf.VarHandleOp. The initializing tensor is not "
"saved to checkpoint.");
return "";
}
auto assigned_value_operand = assign_var_op.getOperand(1);
auto const_op =
llvm::dyn_cast<mlir::TF::ConstOp>(assigned_value_operand.getDefiningOp());
if (!const_op) {
assign_var_op->emitRemark(
"Operand idx 1 is not a tf.ConstOp. The initializing tensor is not "
"saved to checkpoint.");
return "";
}
Tensor const_tensor{};
if (const absl::Status status = mlir::tfg::ConvertToTensor(
/*attr=*/const_op.getValue(), /*output_tensor=*/&const_tensor);
!status.ok()) {
return status;
}
if (!bundle_writer.Add(/*key=*/var_handle_op.getSharedName(), const_tensor)
.ok()) {
return bundle_writer.status();
}
return var_handle_op.getSharedName().str();
}
} // namespace
absl::StatusOr<std::vector<std::string>> SaveVariablesToCheckpoint(
const absl::string_view prefix, mlir::ModuleOp module_op) {
// Only the "tf.AssignVariableOp" patterns inside this initializer function
// will be searched.
FuncOp session_init_func_type_restore_op = GetInitializerFunction(
module_op, /*initializer_type=*/kTfSavedModelInitializerRestoreType);
if (!session_init_func_type_restore_op) {
LOG(INFO) << "No session initializer function with type 'restore_op'. No "
"variables are saved to checkpoint.";
return std::vector<std::string>{};
}
BundleWriter bundle_writer(Env::Default(), prefix);
if (!bundle_writer.status().ok()) {
return bundle_writer.status();
}
std::vector<std::string> saved_variable_shared_names;
for (auto assign_variable_op :
session_init_func_type_restore_op.getOps<mlir::TF::AssignVariableOp>()) {
if (const absl::StatusOr<std::string> variable_shared_name =
AddTensorToBundleWriter(assign_variable_op, bundle_writer);
!variable_shared_name.ok()) {
return variable_shared_name.status();
} else if (!variable_shared_name->empty()) {
// Empty string means the variable isn't applicable for saving.
saved_variable_shared_names.emplace_back(
std::move(*variable_shared_name));
VLOG(1) << "Saved a variable with shared_name: " << *variable_shared_name;
}
}
// Exit early if no variables are added.
if (saved_variable_shared_names.empty()) {
LOG(INFO) << "No variables are saved to checkpoint";
return saved_variable_shared_names;
}
if (!bundle_writer.Finish().ok()) {
return bundle_writer.status();
}
return saved_variable_shared_names;
}
} // namespace quantization
} // namespace tensorflow
tensorflow/compiler/mlir/quantization/tensorflow/cc/save_variables.h
+40
View File
@@ -0,0 +1,40 @@
/* 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_QUANTIZATION_TENSORFLOW_CC_SAVE_VARIABLES_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_CC_SAVE_VARIABLES_H_
#include <string>
#include <vector>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
namespace tensorflow {
namespace quantization {
// Saves variables in `module_op` to the checkpoint file inside `prefix`.
// It finds variables that are initialized with "tf.AssignVariableOp" inside the
// initializer function with type "restore_op". The "tf.Const"s used to
// initialize the variables are saved. This function does not modify the
// `module_op`. Returns a list of saved names of the saved variables.
absl::StatusOr<std::vector<std::string>> SaveVariablesToCheckpoint(
absl::string_view prefix, mlir::ModuleOp module_op);
} // namespace quantization
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_CC_SAVE_VARIABLES_H_
tensorflow/compiler/mlir/quantization/tensorflow/cc/save_variables_test.cc
+384
View File
@@ -0,0 +1,384 @@
/* 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/quantization/tensorflow/cc/save_variables.h"
#include <cstdint>
#include <string>
#include <vector>
#include <gmock/gmock.h>
#include "absl/cleanup/cleanup.h"
#include "absl/log/check.h"
#include "absl/status/status.h"
#include "absl/status/status_matchers.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "mlir/Dialect/Func/IR/FuncOps.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 "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/util/tensor_bundle/tensor_bundle.h"
namespace tensorflow {
namespace quantization {
namespace {
using ::tensorflow::test::AsTensor;
using ::tensorflow::test::ExpectEqual;
using ::testing::IsEmpty;
using ::testing::Not;
using ::testing::UnorderedElementsAre;
// This fixture simply wraps the Env and MLIRContext.
class SaveVariablesToCheckpointTest : public ::testing::Test {
protected:
SaveVariablesToCheckpointTest() : env_(Env::Default()) {
ctx_.loadDialect<mlir::func::FuncDialect, mlir::TF::TensorFlowDialect,
mlir::tf_saved_model::TensorFlowSavedModelDialect>();
}
absl::StatusOr<std::string> MakeTempDir() {
std::string tmp_dir{};
if (!env_->LocalTempFilename(&tmp_dir)) {
return absl::InternalError("Failed to create temp file.");
}
CHECK_OK(env_->CreateDir(tmp_dir));
return tmp_dir;
}
// Parses `module_op_str` to create a `ModuleOp`. Checks whether the created
// module op is valid.
mlir::OwningOpRef<mlir::ModuleOp> ParseModuleOpString(
const absl::string_view module_op_str) {
auto module_op_ref =
mlir::parseSourceString<mlir::ModuleOp>(module_op_str, &ctx_);
EXPECT_TRUE(module_op_ref);
return module_op_ref;
}
Env* env_{};
mlir::MLIRContext ctx_{};
};
TEST_F(SaveVariablesToCheckpointTest, VariableSavedToCheckpoint) {
constexpr absl::string_view kModuleCode = R"mlir(
module attributes {tf_saved_model.semantics} {
"tf_saved_model.session_initializer"() {initializers = [@init_func_restore_op]} : () -> ()
func.func @init_func_restore_op() -> () attributes {tf_saved_model.exported_names = ["restore"], tf_saved_model.initializer_type = "restore_op"} {
%cst = "tf.Const"() {device = "", value = dense<[1.0, 2.0]> : tensor<2xf32>} : () -> tensor<2xf32>
%0 = "tf.VarHandleOp"() {container = "", device = "/device:CPU:0", shared_name = "var_0"} : () -> tensor<!tf_type.resource<tensor<2xf32>>>
"tf.AssignVariableOp"(%0, %cst) : (tensor<!tf_type.resource<tensor<2xf32>>>, tensor<2xf32>) -> ()
return
}
}
)mlir";
mlir::OwningOpRef<mlir::ModuleOp> module_op_ref =
ParseModuleOpString(kModuleCode);
const absl::StatusOr<std::string> checkpoint_prefix = MakeTempDir();
EXPECT_TRUE(checkpoint_prefix.ok());
const absl::Cleanup checkpoint_prefix_cleanup = [this, &checkpoint_prefix]() {
int64_t undeleted_files, undeleted_dirs;
CHECK_OK(env_->DeleteRecursively(*checkpoint_prefix, &undeleted_files,
&undeleted_dirs));
};
const absl::StatusOr<std::vector<std::string>> variable_shared_names =
SaveVariablesToCheckpoint(*checkpoint_prefix, *module_op_ref);
EXPECT_TRUE(variable_shared_names.ok());
EXPECT_THAT(*variable_shared_names, UnorderedElementsAre("var_0"));
// Verify the saved variable.
BundleReader bundle_reader(env_, *checkpoint_prefix);
Tensor loaded_tensor{};
EXPECT_TRUE(bundle_reader.Lookup("var_0", &loaded_tensor).ok());
ExpectEqual(loaded_tensor, AsTensor<float>({1.0, 2.0}));
}
TEST_F(SaveVariablesToCheckpointTest, MultipleVariablesSavedToCheckpoint) {
// Module's session intializer contains two variables.
constexpr absl::string_view kModuleCode = R"mlir(
module attributes {tf_saved_model.semantics} {
"tf_saved_model.session_initializer"() {initializers = [@init_func_restore_op]} : () -> ()
func.func @init_func_restore_op() -> () attributes {tf_saved_model.exported_names = ["restore"], tf_saved_model.initializer_type = "restore_op"} {
%cst = "tf.Const"() {device = "", value = dense<[1.0, 2.0]> : tensor<2xf32>} : () -> tensor<2xf32>
%0 = "tf.VarHandleOp"() {container = "", device = "/device:CPU:0", shared_name = "var_0"} : () -> tensor<!tf_type.resource<tensor<2xf32>>>
"tf.AssignVariableOp"(%0, %cst) : (tensor<!tf_type.resource<tensor<2xf32>>>, tensor<2xf32>) -> ()
%cst_0 = "tf.Const"() {device = "", value = dense<[3, 4, 5, 6]> : tensor<4xi32>} : () -> tensor<4xi32>
%1 = "tf.VarHandleOp"() {container = "", device = "/device:CPU:0", shared_name = "var_1"} : () -> tensor<!tf_type.resource<tensor<4xi32>>>
"tf.AssignVariableOp"(%1, %cst_0) : (tensor<!tf_type.resource<tensor<4xi32>>>, tensor<4xi32>) -> ()
return
}
}
)mlir";
mlir::OwningOpRef<mlir::ModuleOp> module_op_ref =
ParseModuleOpString(kModuleCode);
const absl::StatusOr<std::string> checkpoint_prefix = MakeTempDir();
EXPECT_TRUE(checkpoint_prefix.ok());
const absl::Cleanup checkpoint_prefix_cleanup = [this, &checkpoint_prefix]() {
int64_t undeleted_files, undeleted_dirs;
CHECK_OK(env_->DeleteRecursively(*checkpoint_prefix, &undeleted_files,
&undeleted_dirs));
};
const absl::StatusOr<std::vector<std::string>> variable_shared_names =
SaveVariablesToCheckpoint(*checkpoint_prefix, *module_op_ref);
EXPECT_TRUE(variable_shared_names.ok());
EXPECT_THAT(*variable_shared_names, UnorderedElementsAre("var_0", "var_1"));
// Verify that both variables are saved correctly.
BundleReader bundle_reader(env_, *checkpoint_prefix);
Tensor loaded_var_0{};
EXPECT_TRUE(bundle_reader.Lookup("var_0", &loaded_var_0).ok());
ExpectEqual(loaded_var_0, AsTensor<float>({1.0, 2.0}));
Tensor loaded_var_1{};
EXPECT_TRUE(bundle_reader.Lookup("var_1", &loaded_var_1).ok());
ExpectEqual(loaded_var_1, AsTensor<int>({3, 4, 5, 6}));
}
TEST_F(SaveVariablesToCheckpointTest,
NoVariablesSavedWhenNoInitializerFunction) {
constexpr absl::string_view kModuleCode = R"mlir(
module attributes {tf_saved_model.semantics} {
"tf_saved_model.session_initializer"() {initializers = []} : () -> ()
}
)mlir";
mlir::OwningOpRef<mlir::ModuleOp> module_op_ref =
ParseModuleOpString(kModuleCode);
const absl::StatusOr<std::string> checkpoint_prefix = MakeTempDir();
EXPECT_TRUE(checkpoint_prefix.ok());
const absl::Cleanup checkpoint_prefix_cleanup = [this, &checkpoint_prefix]() {
int64_t undeleted_files, undeleted_dirs;
CHECK_OK(env_->DeleteRecursively(*checkpoint_prefix, &undeleted_files,
&undeleted_dirs));
};
const absl::StatusOr<std::vector<std::string>> variable_shared_names =
SaveVariablesToCheckpoint(*checkpoint_prefix, *module_op_ref);
EXPECT_TRUE(variable_shared_names.ok());
EXPECT_THAT(*variable_shared_names, IsEmpty());
// Verify that the checkpoint doesn't exist.
BundleReader bundle_reader(env_, *checkpoint_prefix);
EXPECT_THAT(bundle_reader.status(), Not(absl_testing::IsOk()));
}
TEST_F(SaveVariablesToCheckpointTest,
NoVariablesSavedWhenNoSessionInitializerOp) {
constexpr absl::string_view kModuleCode = R"mlir(
module {
func.func @my_func() -> () {
return
}
}
)mlir";
mlir::OwningOpRef<mlir::ModuleOp> module_op_ref =
ParseModuleOpString(kModuleCode);
const absl::StatusOr<std::string> checkpoint_prefix = MakeTempDir();
EXPECT_TRUE(checkpoint_prefix.ok());
const absl::Cleanup checkpoint_prefix_cleanup = [this, &checkpoint_prefix]() {
int64_t undeleted_files, undeleted_dirs;
CHECK_OK(env_->DeleteRecursively(*checkpoint_prefix, &undeleted_files,
&undeleted_dirs));
};
EXPECT_TRUE(
SaveVariablesToCheckpoint(*checkpoint_prefix, *module_op_ref).ok());
// Verify that the checkpoint doesn't exist.
BundleReader bundle_reader(env_, *checkpoint_prefix);
EXPECT_THAT(bundle_reader.status(), Not(absl_testing::IsOk()));
}
TEST_F(SaveVariablesToCheckpointTest,
NoVariablesSavedWhenNoSessionInitializerOpTypeRestoreOp) {
constexpr absl::string_view kModuleCode = R"mlir(
module attributes {tf_saved_model.semantics} {
"tf_saved_model.session_initializer"() {initializers = [@init_func_init_op]} : () -> ()
func.func @init_func_init_op() -> () attributes {tf_saved_model.exported_names = ["init"], tf_saved_model.initializer_type = "init_op"} {
%cst = "tf.Const"() {device = "", value = dense<[1.0, 2.0]> : tensor<2xf32>} : () -> tensor<2xf32>
%0 = "tf.VarHandleOp"() {container = "", device = "/device:CPU:0", shared_name = "var_0"} : () -> tensor<!tf_type.resource<tensor<2xf32>>>
"tf.AssignVariableOp"(%0, %cst) : (tensor<!tf_type.resource<tensor<2xf32>>>, tensor<2xf32>) -> ()
return
}
}
)mlir";
mlir::OwningOpRef<mlir::ModuleOp> module_op_ref =
ParseModuleOpString(kModuleCode);
const absl::StatusOr<std::string> checkpoint_prefix = MakeTempDir();
EXPECT_TRUE(checkpoint_prefix.ok());
const absl::Cleanup checkpoint_prefix_cleanup = [this, &checkpoint_prefix]() {
int64_t undeleted_files, undeleted_dirs;
CHECK_OK(env_->DeleteRecursively(*checkpoint_prefix, &undeleted_files,
&undeleted_dirs));
};
const absl::StatusOr<std::vector<std::string>> variable_shared_names =
SaveVariablesToCheckpoint(*checkpoint_prefix, *module_op_ref);
EXPECT_TRUE(variable_shared_names.ok());
EXPECT_THAT(*variable_shared_names, IsEmpty());
// Verify that the checkpoint doesn't exist.
BundleReader bundle_reader(env_, *checkpoint_prefix);
EXPECT_THAT(bundle_reader.status(), Not(absl_testing::IsOk()));
}
TEST_F(SaveVariablesToCheckpointTest, MutableVariablesNotSaved) {
// This function includes an AssignVariableOp that does not initialize the
// variable from a ConstOp. In this case, the variable is not saved to the
// checkpoint.
constexpr absl::string_view kModuleCode = R"mlir(
module attributes {tf_saved_model.semantics} {
"tf_saved_model.session_initializer"() {initializers = [@init_func_restore_op]} : () -> ()
func.func @init_func_restore_op() -> () attributes {tf_saved_model.exported_names = ["init"], tf_saved_model.initializer_type = "restore_op"} {
%cst = "tf.Const"() {device = "", value = dense<[1.0, 2.0]> : tensor<2xf32>} : () -> tensor<2xf32>
%add = "tf.AddV2"(%cst, %cst) : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>
%var_handle = "tf.VarHandleOp"() {container = "", device = "/device:CPU:0", shared_name = "var_0"} : () -> tensor<!tf_type.resource<tensor<2xf32>>>
"tf.AssignVariableOp"(%var_handle, %add) : (tensor<!tf_type.resource<tensor<2xf32>>>, tensor<2xf32>) -> ()
return
}
}
)mlir";
mlir::OwningOpRef<mlir::ModuleOp> module_op_ref =
ParseModuleOpString(kModuleCode);
const absl::StatusOr<std::string> checkpoint_prefix = MakeTempDir();
EXPECT_TRUE(checkpoint_prefix.ok());
const absl::Cleanup checkpoint_prefix_cleanup = [this, &checkpoint_prefix]() {
int64_t undeleted_files, undeleted_dirs;
CHECK_OK(env_->DeleteRecursively(*checkpoint_prefix, &undeleted_files,
&undeleted_dirs));
};
const absl::StatusOr<std::vector<std::string>> variable_shared_names =
SaveVariablesToCheckpoint(*checkpoint_prefix, *module_op_ref);
EXPECT_TRUE(variable_shared_names.ok());
EXPECT_THAT(*variable_shared_names, IsEmpty());
BundleReader bundle_reader(env_, *checkpoint_prefix);
EXPECT_THAT(bundle_reader.status(), Not(absl_testing::IsOk()));
}
TEST_F(SaveVariablesToCheckpointTest,
VariableNotSavedWhenNonVarHandleOpOperandForAssignVariableOp) {
constexpr absl::string_view kModuleCode = R"mlir(
module attributes {tf_saved_model.semantics} {
"tf_saved_model.session_initializer"() {initializers = [@init_func_restore_op]} : () -> ()
func.func @init_func_restore_op() -> () attributes {tf_saved_model.exported_names = ["init"], tf_saved_model.initializer_type = "restore_op"} {
%cst = "tf.Const"() {device = "", value = dense<[1.0, 2.0]> : tensor<2xf32>} : () -> tensor<2xf32>
%var_handle = "tf.VarHandleOp"() {container = "", device = "/device:CPU:0", shared_name = "var_0"} : () -> tensor<!tf_type.resource<tensor<2xf32>>>
%var_handle_cast = "tf.Cast"(%var_handle) : (tensor<!tf_type.resource<tensor<2xf32>>>) -> tensor<!tf_type.resource>
"tf.AssignVariableOp"(%var_handle_cast, %cst) : (tensor<!tf_type.resource>, tensor<2xf32>) -> ()
return
}
}
)mlir";
mlir::OwningOpRef<mlir::ModuleOp> module_op_ref =
ParseModuleOpString(kModuleCode);
const absl::StatusOr<std::string> checkpoint_prefix = MakeTempDir();
EXPECT_TRUE(checkpoint_prefix.ok());
const absl::Cleanup checkpoint_prefix_cleanup = [this, &checkpoint_prefix]() {
int64_t undeleted_files, undeleted_dirs;
CHECK_OK(env_->DeleteRecursively(*checkpoint_prefix, &undeleted_files,
&undeleted_dirs));
};
const absl::StatusOr<std::vector<std::string>> variable_shared_names =
SaveVariablesToCheckpoint(*checkpoint_prefix, *module_op_ref);
EXPECT_TRUE(variable_shared_names.ok());
EXPECT_THAT(*variable_shared_names, IsEmpty());
BundleReader bundle_reader(env_, *checkpoint_prefix);
EXPECT_THAT(bundle_reader.status(), Not(absl_testing::IsOk()));
}
TEST_F(SaveVariablesToCheckpointTest, FailsWhenDuplicateSharedName) {
// Saving variables fails when there are duplicate shared_names ("var_0").
constexpr absl::string_view kModuleCode = R"mlir(
module attributes {tf_saved_model.semantics} {
"tf_saved_model.session_initializer"() {initializers = [@init_func_restore_op]} : () -> ()
func.func @init_func_restore_op() -> () attributes {tf_saved_model.exported_names = ["restore"], tf_saved_model.initializer_type = "restore_op"} {
%cst = "tf.Const"() {device = "", value = dense<[1.0, 2.0]> : tensor<2xf32>} : () -> tensor<2xf32>
%0 = "tf.VarHandleOp"() {container = "", device = "/device:CPU:0", shared_name = "var_0"} : () -> tensor<!tf_type.resource<tensor<2xf32>>>
"tf.AssignVariableOp"(%0, %cst) : (tensor<!tf_type.resource<tensor<2xf32>>>, tensor<2xf32>) -> ()
%cst_0 = "tf.Const"() {device = "", value = dense<[3, 4, 5, 6]> : tensor<4xi32>} : () -> tensor<4xi32>
%1 = "tf.VarHandleOp"() {container = "", device = "/device:CPU:0", shared_name = "var_0"} : () -> tensor<!tf_type.resource<tensor<4xi32>>>
"tf.AssignVariableOp"(%1, %cst_0) : (tensor<!tf_type.resource<tensor<4xi32>>>, tensor<4xi32>) -> ()
return
}
}
)mlir";
mlir::OwningOpRef<mlir::ModuleOp> module_op_ref =
ParseModuleOpString(kModuleCode);
const absl::StatusOr<std::string> checkpoint_prefix = MakeTempDir();
EXPECT_TRUE(checkpoint_prefix.ok());
const absl::Cleanup checkpoint_prefix_cleanup = [this, &checkpoint_prefix]() {
int64_t undeleted_files, undeleted_dirs;
CHECK_OK(env_->DeleteRecursively(*checkpoint_prefix, &undeleted_files,
&undeleted_dirs));
};
EXPECT_FALSE(
SaveVariablesToCheckpoint(*checkpoint_prefix, *module_op_ref).ok());
}
} // namespace
} // namespace quantization
} // namespace tensorflow
tensorflow/compiler/mlir/quantization/tensorflow/debugging/BUILD
+82
View File
@@ -0,0 +1,82 @@
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load(
"//tensorflow:tensorflow.default.bzl",
"get_compatible_with_portable",
"tf_kernel_library",
)
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//tensorflow/compiler/mlir/quantization:__subpackages__",
],
licenses = ["notice"],
)
cc_library(
name = "mlir_dump",
srcs = ["mlir_dump.cc"],
hdrs = ["mlir_dump.h"],
compatible_with = get_compatible_with_portable(),
deps = [
"@com_google_absl//absl/log",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@tsl//tsl/platform:path",
"@tsl//tsl/platform:stringpiece",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:status",
"@xla//xla/tsl/platform:statusor",
],
)
tf_cc_test(
name = "mlir_dump_test",
srcs = ["mlir_dump_test.cc"],
compatible_with = get_compatible_with_portable(),
deps = [
":mlir_dump",
"@com_google_absl//absl/cleanup",
"@com_google_absl//absl/strings:string_view",
"@com_google_googletest//:gtest_main",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Parser",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:Transforms",
"@stablehlo//:stablehlo_ops",
"@tsl//tsl/platform:path",
"@xla//xla/tsl/lib/core:status_test_util",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:test",
"@xla//xla/tsl/platform:test_main",
],
)
tf_kernel_library(
name = "dump_tensor_op",
srcs = ["dump_tensor_op.cc"],
compatible_with = get_compatible_with_portable(),
deps = [
"//tensorflow/compiler/mlir/quantization/tensorflow:quantization_options_proto_cc",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/platform:path",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:errors",
],
)
tensorflow/compiler/mlir/quantization/tensorflow/debugging/dump_tensor_op.cc
+127
View File
@@ -0,0 +1,127 @@
/* 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.
==============================================================================*/
#include <memory>
#include <string>
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
#include "xla/tsl/platform/env.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/file_system.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/op_requires.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/lib/io/compression.h"
#include "tensorflow/core/lib/io/record_writer.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
absl::Status SaveSerializedProtoToFile(const absl::string_view serialized_proto,
const absl::string_view file_path,
tsl::Env* env) {
std::unique_ptr<tsl::WritableFile> file;
TF_RETURN_IF_ERROR(env->NewWritableFile(std::string(file_path), &file));
absl::Status append_result = file->Append(serialized_proto);
absl::Status close_result = file->Close();
return append_result.ok() ? close_result : append_result;
}
// `DumpTensor` op saves entire value of input to as a tensor proto into a
// specified directory and filename. When enabled is set to false, op is
// disabled and won't save any value. It also creates `QuantizationUnit` proto
// with `func_name` and `node_name` to identify the op.
REGISTER_OP("DumpTensor")
.Input("tensor_data: T")
.Attr("log_dir_path: string")
.Attr("file_name: string")
.Attr("T: type")
.Attr("enabled: bool")
.Attr("func_name: string")
.Attr("node_name: string")
.SetIsStateful();
class DumpTensorOp : public OpKernel {
public:
explicit DumpTensorOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
string log_dir_path;
string file_name;
string func_name;
string node_name;
OP_REQUIRES_OK(ctx, ctx->GetAttr("log_dir_path", &log_dir_path));
OP_REQUIRES_OK(ctx, ctx->GetAttr("enabled", &enabled_));
OP_REQUIRES_OK(ctx, ctx->GetAttr("file_name", &file_name));
OP_REQUIRES_OK(ctx, ctx->GetAttr("func_name", &func_name));
OP_REQUIRES_OK(ctx, ctx->GetAttr("node_name", &node_name));
OP_REQUIRES_OK(ctx, ctx->env()->RecursivelyCreateDir(log_dir_path));
std::string tensor_data_path = io::JoinPath(log_dir_path, file_name);
OP_REQUIRES_OK(
ctx, ctx->env()->NewWritableFile(tensor_data_path, &tensor_data_file_));
// Turn on Zlib compression.
io::RecordWriterOptions options =
io::RecordWriterOptions::CreateRecordWriterOptions(
io::compression::kZlib);
tensor_data_writer_ =
std::make_unique<io::RecordWriter>(tensor_data_file_.get(), options);
OP_REQUIRES(ctx, tensor_data_writer_ != nullptr,
absl::AbortedError("Could not create record writer"));
// Fetch func_name and node_name from attributes and save as proto.
quantization::UnitWiseQuantizationSpec::QuantizationUnit quant_unit_proto;
quant_unit_proto.set_func_name(func_name);
quant_unit_proto.set_node_name(node_name);
string quant_unit_path = io::JoinPath(log_dir_path, "quant_unit.pb");
OP_REQUIRES_OK(
ctx, SaveSerializedProtoToFile(quant_unit_proto.SerializeAsString(),
quant_unit_path, ctx->env()));
}
~DumpTensorOp() override {
(void)tensor_data_writer_->Flush();
(void)tensor_data_writer_->Close();
(void)tensor_data_file_->Close();
}
void Compute(OpKernelContext* ctx) override {
if (!enabled_) return;
const Tensor& tensor_data = ctx->input(0);
TensorProto tensor_proto;
tensor_data.AsProtoTensorContent(&tensor_proto);
OP_REQUIRES_OK(ctx, tensor_data_writer_->WriteRecord(
tensor_proto.SerializeAsString()));
}
private:
bool enabled_;
std::unique_ptr<tsl::WritableFile> tensor_data_file_;
std::unique_ptr<io::RecordWriter> tensor_data_writer_;
};
REGISTER_KERNEL_BUILDER(Name("DumpTensor").Device(DEVICE_CPU), DumpTensorOp);
} // namespace tensorflow
tensorflow/compiler/mlir/quantization/tensorflow/debugging/mlir_dump.cc
+242
View File
@@ -0,0 +1,242 @@
/* 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/quantization/tensorflow/debugging/mlir_dump.h"
#include <cstdint>
#include <cstdlib>
#include <memory>
#include <string>
#include <utility>
#include "absl/log/log.h"
#include "absl/memory/memory.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/match.h"
#include "absl/strings/string_view.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/Support/FormatVariadic.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/OperationSupport.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "xla/tsl/platform/env.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/file_system.h"
#include "xla/tsl/platform/statusor.h"
#include "tsl/platform/path.h"
#include "tsl/platform/stringpiece.h"
namespace tensorflow {
namespace quantization {
namespace {
// Retrieve the MLIR dump directory. The directory is read from the environment
// variable `TF_QUANT_MLIR_DUMP_PREFIX`. However, if a special value "sponge" is
// set to `TF_QUANT_MLIR_DUMP_PREFIX`, it uses the directory set in
// `TEST_UNDECLARED_OUTPUT_DIRS`. Returns `absl::FailedPreconditionError` if
// either:
// 1. `TF_QUANT_MLIR_DUMP_PREFIX` is not set (empty), or
// 2. `TEST_UNDECLARED_OUTPUT_DIRS` is not set (empty) when
// `TF_QUANT_MLIR_DUMP_PREFIX = "sponge"`.
absl::StatusOr<std::string> GetMlirDumpDir() {
auto dump_dir = std::string(
absl::NullSafeStringView(std::getenv("TF_QUANT_MLIR_DUMP_PREFIX")));
if (dump_dir.empty()) {
return absl::FailedPreconditionError(
"Environment variable not set: TF_QUANT_MLIR_DUMP_PREFIX, "
"IR dump file for TF quantization is not created.");
}
if (absl::EqualsIgnoreCase(dump_dir, "sponge")) {
if (!tsl::io::GetTestUndeclaredOutputsDir(&dump_dir)) {
return absl::FailedPreconditionError(
"Environment variable TF_QUANT_MLIR_DUMP_PREFIX=sponge but "
"TEST_UNDECLARED_OUTPUT_DIRS not set.");
}
}
return dump_dir;
}
// A simple wrapper of tsl::WritableFile so that mlir Pass infra can use it.
class WritableFileWrapper : public llvm::raw_ostream {
public:
~WritableFileWrapper() override { flush(); }
static absl::StatusOr<std::unique_ptr<WritableFileWrapper>> Create(
const std::string& filepath) {
std::unique_ptr<tsl::WritableFile> file;
TF_RETURN_IF_ERROR(tsl::Env::Default()->NewWritableFile(filepath, &file));
return absl::WrapUnique(new WritableFileWrapper(std::move(file)));
}
private:
explicit WritableFileWrapper(std::unique_ptr<tsl::WritableFile> file)
: file_(std::move(file)) {
SetBuffered();
}
uint64_t current_pos() const override {
int64_t position;
if (file_->Tell(&position).ok()) {
return position;
} else {
return -1;
}
}
void write_impl(const char* ptr, size_t size) override {
if (file_ && !file_->Append(absl::string_view(ptr, size)).ok()) {
file_ = nullptr;
}
}
std::unique_ptr<tsl::WritableFile> file_;
};
// Creates a new file to dump the intermediate MLIRs by prefixing the
// `dump_file_name` with the value of the TF_QUANT_MLIR_DUMP_PREFIX env
// variable. Returns absl::FailedPreconditionError if the env variable is not
// set or set to an empty string.
absl::StatusOr<std::unique_ptr<llvm::raw_ostream>> CreateMlirDumpFile(
const absl::string_view dump_file_name) {
const absl::StatusOr<std::string> dump_dir = GetMlirDumpDir();
if (!dump_dir.ok()) {
return dump_dir.status();
}
auto* env = tsl::Env::Default();
TF_RETURN_IF_ERROR(env->RecursivelyCreateDir(*dump_dir));
const std::string dump_file_path =
tsl::io::JoinPath(*dump_dir, dump_file_name);
TF_ASSIGN_OR_RETURN(std::unique_ptr<llvm::raw_ostream> file,
WritableFileWrapper::Create(dump_file_path));
LOG(INFO) << "IR dump file created: " << dump_file_path;
return file;
}
class PrinterConfig : public mlir::PassManager::IRPrinterConfig {
public:
explicit PrinterConfig(
absl::string_view dump_file_prefix, bool print_module_scope = false,
bool print_after_only_on_change = true,
mlir::OpPrintingFlags op_printing_flags = mlir::OpPrintingFlags())
: mlir::PassManager::IRPrinterConfig(
print_module_scope, print_after_only_on_change,
/*printAfterOnlyOnFailure=*/false, op_printing_flags),
mlir_pass_count_(1),
dump_file_prefix_(dump_file_prefix) {}
void printBeforeIfEnabled(mlir::Pass* pass, mlir::Operation* op,
PrintCallbackFn print_callback) override {
Dump(pass, print_callback, /*is_before=*/true);
}
void printAfterIfEnabled(mlir::Pass* pass, mlir::Operation* op,
PrintCallbackFn print_callback) override {
Dump(pass, print_callback, /*is_before=*/false);
}
private:
int64_t mlir_pass_count_;
absl::string_view dump_file_prefix_;
// Map from pass ptr to dump files and pass number.
//
// Each pass has unique and stable pointer, even for passes with the same
// name. E.g. a PassManager could have multiple Canonicalizer passes.
// We use this property to uniquely determine a Pass in a PassManager.
//
// If multiple consecutive func passes are applied to a Module. PassManager
// will iterate over the func in the outer loop and apply the passes in the
// inner loop. This may cause passes to run out-of-order. But the 1st runs of
// each pass are still in-order. So we use pass_to_number_map_ to keep track
// of the number for each pass.
llvm::DenseMap<mlir::Pass*, std::unique_ptr<llvm::raw_ostream>>
pass_to_dump_file_before_map_;
llvm::DenseMap<mlir::Pass*, std::unique_ptr<llvm::raw_ostream>>
pass_to_dump_file_after_map_;
llvm::DenseMap<mlir::Pass*, int64_t> pass_to_number_map_;
// Get the unique number for each pass.
int64_t GetPassNumber(mlir::Pass* pass) {
if (!pass_to_number_map_.contains(pass)) {
pass_to_number_map_[pass] = mlir_pass_count_++;
}
return pass_to_number_map_[pass];
}
void Dump(mlir::Pass* pass, PrintCallbackFn print_callback, bool is_before) {
auto& pass_to_dump_file_map = is_before ? pass_to_dump_file_before_map_
: pass_to_dump_file_after_map_;
if (!pass_to_dump_file_map.contains(pass)) {
std::string filename = llvm::formatv(
"{0}_{1,0+4}_{2}_{3}.mlir", dump_file_prefix_, GetPassNumber(pass),
pass->getName().str(), is_before ? "before" : "after");
absl::StatusOr<std::unique_ptr<llvm::raw_ostream>> dump_file =
CreateMlirDumpFile(filename);
if (!dump_file.ok()) {
LOG(WARNING) << "Failed to dump MLIR module to " << filename;
return;
}
pass_to_dump_file_map[pass] = std::move(*dump_file);
}
return print_callback(*(pass_to_dump_file_map[pass]));
}
};
} // namespace
void EnableIrPrinting(mlir::PassManager& pm,
absl::string_view file_name_prefix) {
mlir::OpPrintingFlags flag{};
flag.useLocalScope().elideLargeElementsAttrs().enableDebugInfo();
// IR printing requires multithreading disabled.
// Even if multithreading is already disabled, if we are executing within a
// pass-manager, disableMultithreading throws assertion fail. Below if
// statement ensures that disableMultithreading will not be executed if
// multithreading is already disabled.
if (pm.getContext()->isMultithreadingEnabled()) {
pm.getContext()->disableMultithreading();
}
// The configuration uses the default parameter values for
// `PassManager::enableIRPrinting`, except for the `printModuleScope`
// parameter, which is true by default. It is set to false to avoid the dump
// file size becoming too large when the passes are running on a large model.
pm.enableIRPrinting(std::make_unique<PrinterConfig>(
file_name_prefix, /*print_module_scope=*/false,
/*print_after_only_on_change=*/true, flag));
}
absl::Status MaybeEnableIrPrinting(mlir::PassManager& pm,
absl::string_view file_name_prefix) {
if (!VLOG_IS_ON(1)) {
LOG(INFO) << "Verbosity level too low to enable IR printing.";
return absl::OkStatus();
}
EnableIrPrinting(pm, file_name_prefix);
LOG(INFO) << "IR dump for TensorFlow quantization pipeline enabled.";
return absl::OkStatus();
}
} // namespace quantization
} // namespace tensorflow
tensorflow/compiler/mlir/quantization/tensorflow/debugging/mlir_dump.h
+45
View File
@@ -0,0 +1,45 @@
/* 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_QUANTIZATION_TENSORFLOW_DEBUGGING_MLIR_DUMP_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_DEBUGGING_MLIR_DUMP_H_
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "mlir/Pass/PassManager.h" // from @llvm-project
namespace tensorflow {
namespace quantization {
// Enables IR printing for `pm`. When the passes are run, each pass will dump to
// its own file with prefix `file_name_prefix`.
void EnableIrPrinting(mlir::PassManager &pm,
absl::string_view file_name_prefix);
// If verbosity level >= 1, this will dump intermediate IRs of passes to a file.
// The dumped mlir files with be under a directory determined by
// the TF_QUANT_MLIR_DUMP_PREFIX env variable. The PassManager will dump to a
// new file for each pass. The file name will have the format
// {file_name_prefix}_{pass_number}_{pass_name}_{before|after}.mlir.
// * `file_name_prefix` is from input.
// * `pass_number` increments from 1 for each pass.
// * `pass_name` is the name of the pass.
// * `before|after` indicates whether the dump occurs before or after the pass.
absl::Status MaybeEnableIrPrinting(mlir::PassManager &pm,
absl::string_view file_name_prefix);
} // namespace quantization
} // namespace tensorflow
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_DEBUGGING_MLIR_DUMP_H_
tensorflow/compiler/mlir/quantization/tensorflow/debugging/mlir_dump_test.cc
+193
View File
@@ -0,0 +1,193 @@
/* 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/quantization/tensorflow/debugging/mlir_dump.h"
#include <memory>
#include <string>
#include <vector>
#include "absl/strings/string_view.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/LogicalResult.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinDialect.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/Parser/Parser.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "mlir/Transforms/Passes.h" // from @llvm-project
#include "stablehlo/dialect/StablehloOps.h" // from @stablehlo
#include "xla/tsl/lib/core/status_test_util.h"
#include "xla/tsl/platform/env.h"
#include "xla/tsl/platform/test.h"
#include "tsl/platform/path.h"
namespace tensorflow {
namespace quantization {
namespace mlir_dump_test {
class NoOpPass
: public mlir::PassWrapper<NoOpPass, mlir::OperationPass<mlir::ModuleOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(NoOpPass)
NoOpPass() = default;
llvm::StringRef getArgument() const final { return "no-op-pass"; }
void runOnOperation() override {
// Noop pass does nothing on the operation.
}
};
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> CreateNoOpPass() {
return std::make_unique<NoOpPass>();
}
class ParentPass
: public mlir::PassWrapper<ParentPass,
mlir::OperationPass<mlir::ModuleOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ParentPass)
ParentPass() = default;
llvm::StringRef getArgument() const final { return "parent-pass"; }
void runOnOperation() override {
mlir::MLIRContext* ctx = &getContext();
mlir::ModuleOp module_op = getOperation();
mlir::PassManager pm(ctx);
pm.addPass(CreateNoOpPass());
EnableIrPrinting(pm, "dump2");
if (failed(pm.run(module_op))) {
signalPassFailure();
}
}
};
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> CreateParentPass() {
return std::make_unique<ParentPass>();
}
} // namespace mlir_dump_test
namespace {
using namespace tensorflow::quantization::mlir_dump_test;
class EnableIrPrintingTest : public ::testing::Test {
protected:
EnableIrPrintingTest() : env_(tsl::Env::Default()) {
if (!tsl::io::GetTestUndeclaredOutputsDir(&test_dir_)) {
test_dir_ = tsl::testing::TmpDir();
}
}
void SetUp() override {
tsl::setenv("TF_QUANT_MLIR_DUMP_PREFIX", test_dir_.c_str(), 1);
mlir::DialectRegistry dialects;
dialects.insert<mlir::BuiltinDialect, mlir::func::FuncDialect,
mlir::stablehlo::StablehloDialect>();
ctx_ = std::make_unique<mlir::MLIRContext>(dialects);
ctx_->loadAllAvailableDialects();
}
void TearDown() override {
// Delete files in the test directory.
std::vector<std::string> files;
TF_ASSERT_OK(
env_->GetMatchingPaths(tsl::io::JoinPath(test_dir_, "*"), &files));
for (const std::string& file : files) {
TF_ASSERT_OK(env_->DeleteFile(file));
}
}
tsl::Env* env_;
std::string test_dir_;
std::unique_ptr<mlir::MLIRContext> ctx_;
};
TEST_F(EnableIrPrintingTest, PassSuccessfullyRuns) {
mlir::PassManager pm = {ctx_.get()};
pm.addPass(CreateNoOpPass());
pm.addNestedPass<mlir::func::FuncOp>(mlir::createCanonicalizerPass());
pm.addNestedPass<mlir::func::FuncOp>(mlir::createCanonicalizerPass());
EnableIrPrinting(pm, "dump");
constexpr absl::string_view program = R"mlir(
module{
func.func @main(%arg0: tensor<10xf32>) -> tensor<10xf32> {
return %arg0 : tensor<10xf32>
}
func.func @func1(%arg0: tensor<10xf32>, %arg1: tensor<10xf32>) -> tensor<10xf32> {
%0 = stablehlo.add %arg0, %arg1 : tensor<10xf32>
%1 = stablehlo.add %arg0, %arg1 : tensor<10xf32>
return %0 : tensor<10xf32>
}
})mlir";
auto module_op = mlir::parseSourceString<mlir::ModuleOp>(program, ctx_.get());
const mlir::LogicalResult result = pm.run(module_op.get());
EXPECT_FALSE(failed(result));
TF_EXPECT_OK(tsl::Env::Default()->FileExists(
tsl::io::JoinPath(test_dir_,
"dump_0001_tensorflow::quantization::mlir_dump_test"
"::NoOpPass_before.mlir")));
TF_EXPECT_OK(tsl::Env::Default()->FileExists(
tsl::io::JoinPath(test_dir_, "dump_0002_CanonicalizerPass_before.mlir")));
TF_EXPECT_OK(tsl::Env::Default()->FileExists(
tsl::io::JoinPath(test_dir_, "dump_0002_CanonicalizerPass_after.mlir")));
TF_EXPECT_OK(tsl::Env::Default()->FileExists(
tsl::io::JoinPath(test_dir_, "dump_0003_CanonicalizerPass_before.mlir")));
}
TEST_F(EnableIrPrintingTest, NestedPassSuccessfullyRuns) {
mlir::MLIRContext ctx{};
mlir::PassManager pm = {&ctx};
pm.addPass(CreateParentPass());
EnableIrPrinting(pm, "dump");
mlir::OpBuilder builder(&ctx);
mlir::OwningOpRef<mlir::ModuleOp> module_op = mlir::ModuleOp::create(
builder, builder.getUnknownLoc()); /*ALLOW_MLIR_MODULE_OP_CREATE*/
const mlir::LogicalResult result = pm.run(*module_op);
EXPECT_FALSE(failed(result));
TF_EXPECT_OK(tsl::Env::Default()->FileExists(
tsl::io::JoinPath(test_dir_,
"dump_0001_tensorflow::quantization::mlir_dump_test"
"::ParentPass_before.mlir")));
TF_EXPECT_OK(tsl::Env::Default()->FileExists(
tsl::io::JoinPath(test_dir_,
"dump2_0001_tensorflow::quantization::mlir_dump_test"
"::NoOpPass_before.mlir")));
}
} // namespace
} // namespace quantization
} // namespace tensorflow
tensorflow/compiler/mlir/quantization/tensorflow/exported_model.proto
+66
View File
@@ -0,0 +1,66 @@
// 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 tensorflow.quantization;
import "tensorflow/core/framework/graph.proto";
import "tensorflow/core/protobuf/meta_graph.proto";
import "tensorflow/core/protobuf/saver.proto";
// Represents an exported TensorFlow model. It consists of a GraphDef and extra
// metadata required for building a SavedModel. This message is primarily used
// to "export" the model produced from various quantization passes in c++ to
// Python layer.
// Next ID: 11
message ExportedModel {
reserved 3, 4, 7, 9;
reserved "variable_shared_names";
reserved "restore_node_name";
reserved "save_node_name";
reserved "file_prefix_tensor_name";
GraphDef graph_def = 1;
// Name of the initialization node (TF Operation) used for initializing
// resources like hash tables upon loading.
string init_node_name = 2;
// Path to the directory where checkpoint files are saved. This directoy is
// not expected to be persistent (usually a temporary directory). When
// fetching the restore op (see `restore_node_name`), this value is provided
// to the "file_prefix" tensor to identify the checkpoint directory.
string checkpoint_dir = 5;
// Function name -> function alias mapping. This associates the quantized
// functions to the original functions' aliases. This information will be used
// to populate `MetaInfoDef`s `function_aliases` when the quantized model is
// exported to the saved model. This field is usually only populated for the
// TF2 models.
map<string, string> function_aliases = 6;
// Holds information about the asset files used for the model. It essentially
// associates asset file names with the tensors to which the asset file names
// should be fed.
repeated AssetFileDef asset_file_defs = 8;
// SaverDef including the information required for saving and restoring
// variables. This field is not set if there are no variables in the exported
// graph. If set, the fields `version`, `filename_tensor_name`,
// `restore_op_name` and `save_tensor_name` are populated.
SaverDef saver_def = 10;
}
tensorflow/compiler/mlir/quantization/tensorflow/gen_quantized_function_library.py
+232
View File
@@ -0,0 +1,232 @@
# 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.
# ==============================================================================
"""Generates the quantized function library contained header file."""
import ast
import re
import string
from typing import Sequence
from absl import app
from absl import flags
# TODO(b/263048929): Create a test for gen_quantized_function_library.
_OUTPUT_FILE = flags.DEFINE_string('output_file', None, 'output file location')
_SRCS = flags.DEFINE_string('src', None, 'source file locations')
_NAMESPACE = flags.DEFINE_string('namespace', 'mlir::quant',
'namespace in the generated file')
flags.mark_flags_as_required(['output_file', 'src'])
def _substitute_for_loop_template(module: str) -> str:
"""Substitutes the for loop templates in the given module."""
compiled_regex = re.compile(
r'^\s*for\s(.*?)\sin\s(\[.*?\])\s\{(.*?)\}\s//\send\sfor\n',
re.MULTILINE | re.DOTALL)
while True:
func_match = re.search(compiled_regex, module)
if func_match is None:
break
try:
arg_name = func_match.group(1)
arg_values = ast.literal_eval(func_match.group(2))
loop_template = string.Template(func_match.group(3))
except Exception as e: # pylint: disable=broad-except
raise ValueError('The loop template is in wrong format') from e
replacement_text = ''
for arg_value in arg_values:
arg_dict = {arg_name: arg_value}
replacement_text += '\\n'
replacement_text += _substitute_parameterization_template(
loop_template.safe_substitute(arg_dict))
module = re.sub(compiled_regex, replacement_text, module, count=1)
return module
def _substitute_parameterization_template(module: str) -> str:
"""Substitutes all the function templates in the given module."""
compiled_regex = re.compile(
r'^\s*parameters(\[.*?\])\n?(^\s*(?:func\.)+func.*?\{.*?(?:func\.)+return.*?\}\n)',
re.MULTILINE | re.DOTALL)
while True:
func_match = re.search(compiled_regex, module)
if func_match is None:
break
try:
value_list = ast.literal_eval(func_match.group(1))
# Escapes template $-based substitutions for attributes containing $.
# $$ is replaced with a single $.
func_template = string.Template(
func_match.group(2).replace('tfdtype$DT_', 'tfdtype$$DT_'))
except Exception as e: # pylint: disable=broad-except
raise ValueError('The function template is in wrong format') from e
replacement_text = ''
for value_dict in value_list:
for key, value in value_dict.items():
# Replace single quote to double quote since single quote around a
# string are not valid in the MLIR representation.
value_dict[key] = str(value).replace("'", '"')
replacement_text += '\\n'
replacement_text += func_template.substitute(value_dict)
module = re.sub(compiled_regex, replacement_text, module, count=1)
return module
def _format_snake_case_op_name(s):
"""Formats the op name to snake case."""
s = s.replace('2D', '2d').replace('3D', '3d')
snake_case = ''.join(['_' + i.lower() if i.isupper() else i for i in s
]).lstrip('_')
return snake_case.replace('mat_mul', 'matmul').replace('bias_add', 'bias')
def _substitute_impl_function_name_template(module: str) -> str:
"""Generates the op-specific implementation function name."""
compiled_regex = re.compile(r'GenerateImplFunctionName\(([\w\s]+)\)')
while True:
func_match = re.search(compiled_regex, module)
if func_match is None:
break
text = func_match.group(1)
function_name = 'internal_{}_fn'.format(_format_snake_case_op_name(text))
module = re.sub(compiled_regex, function_name, module, count=1)
return module
def _substitute_quantized_function_name_template(module: str) -> str:
"""Generates the quantized function name."""
compiled_regex = re.compile(
r'GenerateQuantizedFunctionName(\([\w\s\'\"\[\],]+\))')
while True:
func_match = re.search(compiled_regex, module)
if func_match is None:
break
# Make sure the string ends with ",)" so the parsed value is a tuple.
argument_string = func_match.group(1)
if not argument_string.endswith(',)'):
argument_string = argument_string[:-1] + ',)'
arguments = ast.literal_eval(argument_string)
if len(arguments) < 1 or len(arguments) > 2:
raise ValueError(
'Wrong number of arguments to GenerateQuantizedFunctionName')
quantized_ops = arguments[0]
if not quantized_ops:
raise ValueError('The quantized_ops list must not be empty')
# Add op names to the function name.
function_name = 'quantized_{}'.format(
_format_snake_case_op_name(quantized_ops[0]))
if len(quantized_ops) > 1:
function_name += '_with_{}'.format(
_format_snake_case_op_name(quantized_ops[1]))
if len(quantized_ops) > 1:
for quantized_op in quantized_ops[2:]:
function_name += '_and_{}'.format(
_format_snake_case_op_name(quantized_op))
# Add suffix based on output type.
suffix = '_fn'
if len(arguments) > 1 and arguments[1] == 'f32':
suffix = '_float_output_fn'
function_name += suffix
module = re.sub(compiled_regex, function_name, module, count=1)
return module
def main(_: Sequence[str]) -> None:
namespaces = _NAMESPACE.value.split('::')
src_files = _SRCS.value.split(' ')
file_prefix = 'quantized_function_library'
module_prefix = 'kQuantizedFunctionLibraryInMLIR'
modules = []
for src_file in src_files:
with open(src_file, 'r') as f:
content = f.read()
# Skip the copyright in the source file.
module_match = re.search(r'(^module\s\{)(.*)(^\})', content,
re.MULTILINE | re.DOTALL)
if module_match is None:
raise ValueError("Couldn't find module in the function library")
module = module_match.group()
# Substitute all the function templates.
out = re.split(file_prefix, src_file)
if len(out) != 2:
raise ValueError('The file name must start with {}'.format(file_prefix))
tag = out[1][:-5] # the last five values = ".mlir"
module = _substitute_for_loop_template(module)
module = _substitute_parameterization_template(module)
module = _substitute_quantized_function_name_template(module)
module = _substitute_impl_function_name_template(module)
modules.append((tag, module))
with open(_OUTPUT_FILE.value, 'w') as f:
f.write("""/* 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_QUANTIZATION_TENSORFLOW_PASSES_QUANTIZED_FUNCTION_LIBRARY_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_PASSES_QUANTIZED_FUNCTION_LIBRARY_H_
""")
for namespace in namespaces:
f.write('namespace {0} {{\n'.format(namespace))
for tag, module in modules:
f.write('constexpr char {0}[] ='.format(module_prefix + tag.upper()))
for line in module.splitlines():
f.write('\n "')
f.write(line.rstrip().replace('"', r'\"'))
f.write('\\n"')
f.write(';\n')
for namespace in reversed(namespaces):
f.write('}} // namespace {0}\n'.format(namespace))
f.write(
'#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_PASSES_QUANTIZED_FUNCTION_LIBRARY_H_'
)
if __name__ == '__main__':
app.run(main)
tensorflow/compiler/mlir/quantization/tensorflow/internal_visibility_allowlist.bzl
+5
View File
@@ -0,0 +1,5 @@
"""Internal visibility rules."""
def internal_visibility_allowlist():
"""Returns a list of the packages that can depend on internal targets."""
return []
tensorflow/compiler/mlir/quantization/tensorflow/ops/BUILD
+93
View File
@@ -0,0 +1,93 @@
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//tensorflow/compiler/mlir/quantization/tensorflow:internal_visibility_allowlist_package",
],
licenses = ["notice"],
)
cc_library(
name = "tf_op_quant_spec",
srcs = [
"tf_op_quant_spec.cc",
],
hdrs = ["tf_op_quant_spec.h"],
compatible_with = get_compatible_with_portable(),
deps = [
"//tensorflow/compiler/mlir/quantization/common/quantization_lib",
"//tensorflow/compiler/mlir/quantization/tensorflow:quantization_options_proto_cc",
"//tensorflow/compiler/mlir/tensorflow",
"@com_google_absl//absl/container:flat_hash_set",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
)
tf_cc_test(
name = "tf_op_quant_spec_test",
srcs = ["tf_op_quant_spec_test.cc"],
deps = [
":tf_op_quant_spec",
"//tensorflow/compiler/mlir/quantization/tensorflow:quantization_options_proto_cc",
"@com_google_googletest//:gtest_main",
],
)
cc_library(
name = "tf_quantize_op",
srcs = [
"tf_quantize_op.cc",
],
hdrs = ["tf_quantize_op.h"],
compatible_with = get_compatible_with_portable(),
deps = [
"//tensorflow/compiler/mlir/quantization/common/quantization_lib",
"//tensorflow/compiler/mlir/quantization/tensorflow:quantization_options_proto_cc",
"//tensorflow/compiler/mlir/quantization/tensorflow/utils:tf_quantize_op_utils",
"//tensorflow/compiler/mlir/tensorflow",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:optional",
"@llvm-project//mlir:Dialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:QuantOps",
"@llvm-project//mlir:Support",
],
)
tf_cc_test(
name = "tf_quantize_op_test",
srcs = ["tf_quantize_op_test.cc"],
deps = [
":tf_quantize_op",
"//tensorflow/compiler/mlir/quantization/common:attrs_and_constraints",
"//tensorflow/compiler/mlir/quantization/tensorflow:quantization_options_proto_cc",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:QuantOps",
"@llvm-project//mlir:Support",
],
)
cc_library(
name = "uniform_op_quant_spec",
srcs = [
"uniform_op_quant_spec.cc",
],
hdrs = ["uniform_op_quant_spec.h"],
compatible_with = get_compatible_with_portable(),
deps = [
"//tensorflow/compiler/mlir/quantization/common/quantization_lib",
"//tensorflow/compiler/mlir/tensorflow",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
)
tensorflow/compiler/mlir/quantization/tensorflow/ops/tf_op_quant_spec.cc
+162
View File
@@ -0,0 +1,162 @@
/* 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/quantization/tensorflow/ops/tf_op_quant_spec.h"
#include <memory>
#include <optional>
#include <vector>
#include "absl/container/flat_hash_set.h"
#include "llvm/Support/Casting.h"
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinTypeInterfaces.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/quantization/common/quantization_lib/quantization_utils.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
namespace mlir {
namespace quant {
// TODO - b/296503614: [Converter Component][TF-Quantizer] Reflect custom traits
// from TF-Quantizer to stableHLO quantization
bool IsOpWithDataMovementTrait(Operation* op) {
// Supported data movement ops. These ops do not perform any computations and
// has one result operand.
return isa<TF::IdentityOp, TF::CastOp, TF::ReshapeOp, TF::XlaShardingOp,
TF::GatherOp, TF::GatherV2Op, TF::XlaGatherOp, TF::ExpandDimsOp,
TF::SqueezeOp, TF::TransposeOp>(op);
}
bool IsOpWithQuantizableTrait(Operation* op) {
// Supported quantizable ops.
return isa<TF::XlaConvV2Op, TF::XlaDotV2Op, TF::MatMulOp, TF::Conv2DOp,
TF::GatherOp, TF::GatherV2Op, TF::XlaGatherOp,
TF::ResourceGatherOp, TF::DepthwiseConv2dNativeOp, TF::Conv3DOp,
TF::BatchMatMulV2Op, TF::EinsumOp>(op);
}
bool IsOpWithInt8TypeOperand(Operation* op) {
return (isa<TF::XlaConvV2Op, TF::XlaDotV2Op, TF::XlaGatherOp, TF::GatherOp,
TF::GatherV2Op>(op));
}
bool IsValueWithQuantizablePrecision(Value val) {
auto type = mlir::dyn_cast<ShapedType>(val.getType());
if (!type) return false;
// Supported original tensor data types.
if (type.getElementType().isF32() || type.getElementType().isBF16())
return true;
return false;
}
std::optional<tensorflow::quantization::QuantizationComponentSpec>
GetWeightComponentSpec(
const tensorflow::quantization::QuantizationOptions& quantization_options) {
for (auto& cur_spec : quantization_options.quantization_method()
.quantization_component_specs()) {
if (cur_spec.quantization_component() ==
tensorflow::quantization::QuantizationComponentSpec::COMPONENT_WEIGHT)
return cur_spec;
}
return std::nullopt;
}
// TODO(b/228928859): Improve the getter function to match attributes rather
// than function name.
std::unique_ptr<OpQuantSpec> GetTFOpQuantSpec(Operation* op) {
auto spec = std::make_unique<OpQuantSpec>();
if (auto call_op = dyn_cast<TF::PartitionedCallOp>(op)) {
StringRef function_name =
mlir::cast<FlatSymbolRefAttr>(call_op.getFAttr()).getValue();
if (!function_name.starts_with("composite_")) {
return spec;
}
if (function_name.contains("depthwise_conv2d")) {
spec->coeff_op_quant_dim[1] = 3;
if (function_name.contains("with_bias")) {
spec->biases_params[2] = {{0, 1}, GetUniformQuantizedTypeForBias};
}
} else if (function_name.contains("conv2d")) {
spec->coeff_op_quant_dim[1] = 3;
if (function_name.contains("with_bias")) {
spec->biases_params[2] = {{0, 1}, GetUniformQuantizedTypeForBias};
}
} else if (function_name.contains("matmul")) {
spec->coeff_op_quant_dim[1] = -1;
if (function_name.contains("with_bias") ||
function_name.contains("and_bias")) {
spec->biases_params[2] = {{0, 1}, GetUniformQuantizedTypeForBias};
}
} else if (function_name.contains("einsum")) {
spec->coeff_op_quant_dim[1] = -1;
if (function_name.contains("with_bias")) {
spec->biases_params[2] = {{0, 1}, GetUniformQuantizedTypeForBias};
}
} else if (function_name.contains("conv3d")) {
spec->coeff_op_quant_dim[1] = 4;
if (function_name.contains("with_bias")) {
spec->biases_params[2] = {{0, 1}, GetUniformQuantizedTypeForBias};
}
} else if (function_name.contains("batch_matmul")) {
spec->coeff_op_quant_dim[1] = -1;
if (function_name.contains("with_bias")) {
spec->biases_params[2] = {{0, 1}, GetUniformQuantizedTypeForBias};
}
} else if (function_name.contains("gather")) {
// Note that gather has axis attribute that specifies channel axis.
spec->coeff_op_quant_dim[0] = -1;
}
for (auto quantizable_operand : spec->coeff_op_quant_dim) {
spec->quantizable_operands.insert(quantizable_operand.first);
}
}
return spec;
}
std::unique_ptr<OpQuantScaleSpec> GetTfQuantScaleSpec(Operation* op) {
auto scale_spec = std::make_unique<OpQuantScaleSpec>();
if (llvm::isa<
// clang-format off
// go/keep-sorted start
TF::AvgPoolOp,
TF::ConcatOp,
TF::ConcatV2Op,
TF::ExpandDimsOp,
TF::IdentityNOp,
TF::IdentityOp,
TF::MaxPoolOp,
TF::PadV2Op,
TF::RankOp,
TF::ReshapeOp,
TF::SelectOp,
TF::SelectV2Op,
TF::ShapeNOp,
TF::ShapeOp,
TF::SizeOp,
TF::SqueezeOp,
TF::TransposeOp
// go/keep-sorted end
// clang-format on
>(op)) {
scale_spec->has_same_scale_requirement = true;
}
return scale_spec;
}
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/ops/tf_op_quant_spec.h
+61
View File
@@ -0,0 +1,61 @@
/* 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.
==============================================================================*/
// Functions for quantization specifications of TensorFlow ops.
#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_OPS_TF_OP_QUANT_SPEC_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_OPS_TF_OP_QUANT_SPEC_H_
#include <memory>
#include <optional>
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
namespace mlir {
namespace quant {
// Check if the op has data movement trait. Ops with this trait do not perform
// any computations but just move data and has one result operand.
bool IsOpWithDataMovementTrait(Operation* op);
// Check if the op is quantizable. Currently, the scope of quantizable op is
// limited to compute intense operations and the ops that supports integer
// operands.
bool IsOpWithQuantizableTrait(Operation* op);
// Check if the op's operand accepts int8 type.
bool IsOpWithInt8TypeOperand(Operation* op);
// Check if the data is in quantizable precision. Currently, a value in f32 or
// bf16 is quantizable.
bool IsValueWithQuantizablePrecision(Value val);
std::optional<tensorflow::quantization::QuantizationComponentSpec>
GetWeightComponentSpec(
const tensorflow::quantization::QuantizationOptions& quantization_options);
// Returns the spec for the given operation that can be used for both of
// dynamic and static range quantization.
std::unique_ptr<OpQuantSpec> GetTFOpQuantSpec(Operation* op);
// Returns quantization scale specs (fixed output, same scale) for a TF op.
std::unique_ptr<OpQuantScaleSpec> GetTfQuantScaleSpec(Operation* op);
} // namespace quant
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_OPS_TF_OP_QUANT_SPEC_H_
tensorflow/compiler/mlir/quantization/tensorflow/ops/tf_op_quant_spec_test.cc
+47
View File
@@ -0,0 +1,47 @@
/* 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.
==============================================================================*/
#include "tensorflow/compiler/mlir/quantization/tensorflow/ops/tf_op_quant_spec.h"
#include <gtest/gtest.h>
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
namespace mlir::quant {
namespace {
using QuantizationOptions = tensorflow::quantization::QuantizationOptions;
using QuantizationComponentSpec =
tensorflow::quantization::QuantizationComponentSpec;
TEST(TfOpQuantSpecTest, WeightComponentSpecExist) {
QuantizationOptions quant_options;
QuantizationComponentSpec quant_spec;
quant_spec.set_quantization_component(
QuantizationComponentSpec::COMPONENT_WEIGHT);
quant_spec.set_tensor_type(QuantizationComponentSpec::TENSORTYPE_INT_8);
auto mutable_quant_method = quant_options.mutable_quantization_method();
*mutable_quant_method->add_quantization_component_specs() = quant_spec;
auto output = GetWeightComponentSpec(quant_options);
EXPECT_TRUE(output.has_value());
}
TEST(TfOpQuantSpecTest, WeightComponentSpecDoNotExist) {
QuantizationOptions quant_options;
auto output = GetWeightComponentSpec(quant_options);
EXPECT_FALSE(output.has_value());
}
} // namespace
} // namespace mlir::quant
tensorflow/compiler/mlir/quantization/tensorflow/ops/tf_quantize_op.cc
+263
View File
@@ -0,0 +1,263 @@
/* 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.
==============================================================================*/
#include "tensorflow/compiler/mlir/quantization/tensorflow/ops/tf_quantize_op.h"
#include <functional>
#include <optional>
#include "absl/strings/str_cat.h"
#include "absl/strings/str_join.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/Quant/IR/QuantTypes.h" // from @llvm-project
#include "mlir/IR/Block.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/BuiltinTypeInterfaces.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Diagnostics.h" // from @llvm-project
#include "mlir/IR/Matchers.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/IR/SymbolTable.h" // from @llvm-project
#include "mlir/IR/Types.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/utils/tf_quantize_op_utils.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
namespace mlir {
namespace quant {
namespace {
constexpr StringRef kDequantizeFunctionName = "composite_dequantize";
constexpr StringRef kUniformQuantizationFunctionName = "uniform";
// Pre-actions before adding quantization logics. It creates a function with the
// func_name where input_val is an input and result_type is a result.
func::FuncOp PrepareFunctionRegister(PatternRewriter& rewriter, Value input_val,
ShapedType result_type,
StringRef func_name,
Value& func_input_arg) {
Operation* input_op = input_val.getDefiningOp();
Operation* insertion_point = input_op->getParentOfType<func::FuncOp>();
if (!insertion_point) insertion_point = input_op->getParentOfType<ModuleOp>();
rewriter.setInsertionPointAfter(insertion_point);
UnrankedTensorType create_unknown_input_shape =
CreateUnknownShapeFromElementType(input_val.getType());
UnrankedTensorType create_unknown_output_shape =
CreateUnknownShapeFromElementType(result_type);
FunctionType func_type =
FunctionType::get(rewriter.getContext(), {create_unknown_input_shape},
{create_unknown_output_shape});
func::FuncOp quantization_func =
func::FuncOp::create(rewriter, input_op->getLoc(), func_name, func_type);
OpBuilder::InsertionGuard guard = OpBuilder::InsertionGuard(rewriter);
ArrayRef<Type> inputs = quantization_func.getFunctionType().getInputs();
Block* block = rewriter.createBlock(
&quantization_func.getBody(), quantization_func.begin(), inputs,
SmallVector<Location>(inputs.size(), quantization_func.getLoc()));
func_input_arg = block->getArgument(0);
return quantization_func;
}
// Post-actions after adding quantization logics. Post-actions include
// 1) Adding the created function in the symbol table
// 2) Creating a PartitionedCallOp in the main graph that calls the created
// function.
TF::PartitionedCallOp FinalizeFunctionRegister(
PatternRewriter& rewriter, Value input, Value output,
func::FuncOp& quantization_func, Operation* quantized_op,
StringRef func_name, IRRewriter::InsertPoint original_point,
Type quantize_result_type) {
func::ReturnOp::create(rewriter, input.getLoc(), ArrayRef<Value>({output}));
quantization_func.setVisibility(func::FuncOp::Visibility::Private);
SymbolTable symbol_table(quantized_op->getParentOfType<ModuleOp>());
symbol_table.insert(quantization_func);
FlatSymbolRefAttr func_name_attr =
FlatSymbolRefAttr::get(rewriter.getStringAttr(func_name));
rewriter.restoreInsertionPoint(original_point);
auto quantize_call = TF::PartitionedCallOp::create(
rewriter, quantized_op->getLoc(), quantize_result_type, input,
/*args_attrs=*/nullptr, /*res_attrs=*/nullptr, func_name_attr,
/*config=*/"", /*config_proto=*/"", /*executor_type=*/"");
return quantize_call;
}
// Acts as a register of a function where the body has a sequence of operations
// required to execute certain quantization scheme's quant/dequantization
// logics.
std::optional<TF::PartitionedCallOp> RegisterOperationsInFuncOp(
StringRef func_name, PatternRewriter& rewriter, QuantizedType quant_type,
Value input_val, ShapedType result_type,
std::function<Operation*(PatternRewriter&, Operation*, Value, ShapedType,
QuantizedType)>
quantization_operations_func) {
Operation* input_op = input_val.getDefiningOp();
auto original_point = rewriter.saveInsertionPoint();
auto unique_func_name = func_name.str();
SymbolTable symbol_table(input_op->getParentOfType<ModuleOp>());
while (symbol_table.lookup(unique_func_name)) {
absl::StrAppend(&unique_func_name, "_");
}
Value func_input_arg;
// Creates a function.
func::FuncOp func_op = PrepareFunctionRegister(
rewriter, input_val, result_type, unique_func_name, func_input_arg);
// Fills the body.
Operation* last_op_in_func =
quantization_operations_func(rewriter, func_op.getOperation(),
func_input_arg, result_type, quant_type);
// Connect the function in the existing graph.
auto end_call_op = FinalizeFunctionRegister(
rewriter, input_val, last_op_in_func->getResult(0), func_op, input_op,
unique_func_name, original_point, result_type);
return end_call_op;
}
QuantizedType CalculateUniformQuantParams(
PatternRewriter& rewriter, TF::ConstOp op,
tensorflow::quantization::QuantizationComponentSpec& weight_spec) {
// TODO - b/278949920: Enable Per-Channel Quantization for XLA Opset
// Currently, support symmetric, per-tensor, signed int8
const bool kIsNarrowRange = true;
const bool kIsSigned = true;
const int kBitWidth = 8;
DenseFPElementsAttr attr;
if (!matchPattern(op->getResult(0), m_Constant(&attr))) return nullptr;
QuantizedType quant_type =
mlir::dyn_cast<quant::QuantizedType>(GetUniformQuantizedTypeForWeight(
attr, /*symmetric=*/kIsNarrowRange && kIsSigned, kBitWidth, kIsSigned,
kIsNarrowRange, /*is_legacy_float*/ false));
return quant_type;
}
// Add uniform quantization's quantization logic.
std::optional<Value> AddUniformQuantizeOps(PatternRewriter& rewriter,
TF::ConstOp op,
QuantizedType quant_type) {
DenseFPElementsAttr attr;
if (!matchPattern(op->getResult(0), m_Constant(&attr))) {
return nullptr;
}
Type expressed_type = op.getResult().getType();
Type quantized_type = quant_type.castFromExpressedType(expressed_type);
ShapedType shaped_quantized_type = mlir::cast<ShapedType>(quantized_type);
DenseElementsAttr tensor_proto_attr =
mlir::dyn_cast<DenseElementsAttr>(Quantize(attr, shaped_quantized_type));
if (!tensor_proto_attr) {
return nullptr;
}
Type storage_type =
mlir::cast<QuantizedType>(shaped_quantized_type.getElementType())
.getStorageType();
ShapedType new_type = shaped_quantized_type.clone(storage_type);
rewriter.setInsertionPointAfter(op);
auto const_op =
TF::ConstOp::create(rewriter, op.getLoc(), new_type, tensor_proto_attr);
auto new_identity_op = TF::IdentityOp::create(rewriter, op->getLoc(),
const_op.getType(), const_op);
return new_identity_op.getResult();
}
Operation* LogicsForUniformDequanization(PatternRewriter& rewriter,
Operation* func_op, Value input_val,
ShapedType original_input_tensor_type,
QuantizedType quant_type) {
auto loc = input_val.getLoc();
rewriter.setInsertionPointToStart(
&(cast<func::FuncOp>(func_op)).getBody().front());
UnrankedTensorType create_unknown_input_shape =
CreateUnknownShapeFromElementType(original_input_tensor_type);
auto new_cast_op =
TF::CastOp::create(rewriter, loc, create_unknown_input_shape, input_val);
// TODO - b/278949920: Enable Per-Channel Quantization for XLA Opset
auto qtype = mlir::dyn_cast<UniformQuantizedType>(quant_type);
TensorType scale_type = RankedTensorType::get({}, rewriter.getF32Type());
Value scale_op = TF::ConstOp::create(
rewriter, loc, scale_type,
DenseFPElementsAttr::get(scale_type,
{static_cast<float>(qtype.getScale())}));
if (original_input_tensor_type.getElementType().isBF16()) {
// Add bf16 cast op after scale to match with the next op's data
// type.
scale_op = TF::CastOp::create(
rewriter, loc, UnrankedTensorType::get(rewriter.getBF16Type()),
scale_op);
}
auto mul_op = TF::MulOp::create(rewriter, loc, new_cast_op.getType(),
scale_op, new_cast_op);
return mul_op;
}
// Add uniform quantization's dequantization logic.
std::optional<TF::PartitionedCallOp> AddUniformDequantizeOps(
PatternRewriter& rewriter, QuantizedType quant_type,
Value val_to_dequantize, ShapedType result_type) {
auto func_name = absl::StrJoin(
{kDequantizeFunctionName, kUniformQuantizationFunctionName}, "_");
std::optional<TF::PartitionedCallOp> dequant_op = RegisterOperationsInFuncOp(
func_name, rewriter, quant_type, val_to_dequantize, result_type,
LogicsForUniformDequanization);
return dequant_op;
}
} // namespace
// Generate quantize and dequantize functions with uniform quantization.
std::optional<TF::PartitionedCallOp> ApplyUniformQuantization(
PatternRewriter& rewriter, TF::ConstOp op,
tensorflow::quantization::QuantizationComponentSpec& weight_spec) {
QuantizedType quant_type =
CalculateUniformQuantParams(rewriter, op, weight_spec);
if (!quant_type) return nullptr;
std::optional<Value> quantized_val =
AddUniformQuantizeOps(rewriter, op, quant_type);
if (!quantized_val.has_value()) return std::nullopt;
std::optional<TF::PartitionedCallOp> dequantized_val =
AddUniformDequantizeOps(rewriter, quant_type, quantized_val.value(),
mlir::cast<ShapedType>(op.getType()));
return dequantized_val;
}
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/ops/tf_quantize_op.h
+45
View File
@@ -0,0 +1,45 @@
/* 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.
==============================================================================*/
// This file provides a list of supported quantization algorithms in the format
// of "apply<Name of the Quantization Algorithm>Quantization".
// After applying the function, a quantize/dequantize functions are created
// where the body of each function contains a specific quantization algorithm.
// The input of the quantize function has one operand of
// IsValueWithQuantizablePrecision and the output is a tensor with supported
// quantized precision (like int8). For dequantize function, it is the other way
// around.
#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_OPS_TF_QUANTIZE_OP_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_OPS_TF_QUANTIZE_OP_H_
#include <optional>
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/Traits.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
namespace mlir {
namespace quant {
std::optional<TF::PartitionedCallOp> ApplyUniformQuantization(
PatternRewriter& rewriter, TF::ConstOp op,
tensorflow::quantization::QuantizationComponentSpec& weight_spec);
} // namespace quant
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_OPS_TF_QUANTIZE_OP_H_
tensorflow/compiler/mlir/quantization/tensorflow/ops/tf_quantize_op_test.cc
+71
View File
@@ -0,0 +1,71 @@
/* 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.
==============================================================================*/
#include "tensorflow/compiler/mlir/quantization/tensorflow/ops/tf_quantize_op.h"
#include <optional>
#include <gtest/gtest.h>
#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/IR/PatternMatch.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
namespace mlir::quant {
namespace {
using QuantizationComponentSpec =
tensorflow::quantization::QuantizationComponentSpec;
class EmptyPatternRewriter : public mlir::PatternRewriter {
public:
explicit EmptyPatternRewriter(const OpBuilder& other_builder)
: mlir::PatternRewriter(other_builder) {}
~EmptyPatternRewriter() override = default;
};
TEST(TfQuantOpTest, applyUniformQuantization) {
MLIRContext context;
OwningOpRef<ModuleOp> module(ModuleOp::create(UnknownLoc::get(&context)));
OpBuilder builder(&module->getBodyRegion());
context.loadDialect<TF::TensorFlowDialect, quant::QuantDialect,
func::FuncDialect>();
EmptyPatternRewriter pattern_rewriter(builder);
Value value = CreateConstValue<float>(builder, module->getLoc(), {1024, 2},
SmallVector<float>(2048, 0));
QuantizationComponentSpec quant_spec;
quant_spec.set_quantization_component(
QuantizationComponentSpec::COMPONENT_WEIGHT);
quant_spec.set_tensor_type(QuantizationComponentSpec::TENSORTYPE_INT_8);
std::optional<TF::PartitionedCallOp> dequantize_op = ApplyUniformQuantization(
pattern_rewriter, cast<TF::ConstOp>(value.getDefiningOp()), quant_spec);
EXPECT_TRUE(dequantize_op.has_value());
EXPECT_EQ(dequantize_op.value().func().getName().str(),
"composite_dequantize_uniform");
}
} // namespace
} // namespace mlir::quant
tensorflow/compiler/mlir/quantization/tensorflow/ops/uniform_op_quant_spec.cc
+41
View File
@@ -0,0 +1,41 @@
/* 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/quantization/tensorflow/ops/uniform_op_quant_spec.h"
#include <memory>
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
namespace mlir::quant {
std::unique_ptr<OpQuantSpec> GetUniformOpQuantSpec(Operation* op) {
auto spec = std::make_unique<OpQuantSpec>();
if (isa<TF::UniformQuantizedConvolutionHybridOp>(op) ||
isa<TF::UniformQuantizedConvolutionOp>(op)) {
spec->coeff_op_quant_dim[1] = 3;
} else if (isa<TF::UniformQuantizedDotHybridOp>(op)) {
spec->coeff_op_quant_dim[1] = -1;
}
for (auto quantizable_operand : spec->coeff_op_quant_dim) {
spec->quantizable_operands.insert(quantizable_operand.first);
}
return spec;
}
} // namespace mlir::quant
tensorflow/compiler/mlir/quantization/tensorflow/ops/uniform_op_quant_spec.h
+35
View File
@@ -0,0 +1,35 @@
/* 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.
==============================================================================*/
// Functions for quantization specifications of Uniform Quantized ops.
#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_OPS_UNIFORM_OP_QUANT_SPEC_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_OPS_UNIFORM_OP_QUANT_SPEC_H_
#include <memory>
#include "mlir/IR/Operation.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"
namespace mlir {
namespace quant {
// Returns the spec for the given operation that can be used for both of
// dynamic and static range quantization.
std::unique_ptr<OpQuantSpec> GetUniformOpQuantSpec(Operation* op);
} // namespace quant
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_OPS_UNIFORM_OP_QUANT_SPEC_H_
tensorflow/compiler/mlir/quantization/tensorflow/passes/add_dump_tensor_op.cc
+322
View File
@@ -0,0 +1,322 @@
/* 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.
==============================================================================*/
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/CommandLine.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/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/IR/SymbolTable.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/LogicalResult.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.h"
#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"
#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/cc/quantization_unit_loc.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/passes.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/tf_quant_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/xla_call_module_attrs.h"
#include "tensorflow/core/platform/path.h"
namespace mlir {
namespace quant {
namespace {
using ::stablehlo::quantization::DebuggerConfig;
using DebuggerType = DebuggerConfig::DebuggerType;
constexpr StringRef kOriginalEntryFuncAttrName = "_original_entry_function";
constexpr StringRef kCompositeFuncPrefix = "composite_";
constexpr StringRef kEmptyNodeName = "_empty_node";
// Returns a pair: `func_name` and `node_name` for the lifted function. In TF
// quantizer, both are filled. For StableHLO quantizer, the func_name is only
// filled and node_name is always set to "_empty_node".
std::pair<std::string, std::string> GetFuncNameAndNodeName(
TF::PartitionedCallOp call_op, const FlatSymbolRefAttr &f_attr) {
std::optional<QuantizationUnitLoc::QuantizationUnit> quant_unit =
FindQuantizationUnitFromLoc(call_op->getLoc());
return std::make_pair(quant_unit->func_name(), quant_unit->node_name());
}
std::pair<std::string, std::string> GetFuncNameAndNodeName(
TF::XlaCallModuleOp call_op, const FlatSymbolRefAttr &f_attr) {
return std::make_pair(f_attr.getValue().str(), kEmptyNodeName.str());
}
Operation *DuplicateOp(TF::PartitionedCallOp call_op, PatternRewriter &rewriter,
const StringAttr &new_ref_func_name) {
// Create PartitionedCallOp to the copied composite function. This
// PartitionedCallOp does not have kQuantTraitAttrName, and therefore won't
// get quantized.
auto new_call_op = TF::PartitionedCallOp::create(
rewriter, call_op.getLoc(), call_op.getResultTypes(),
call_op.getOperands(), call_op.getArgAttrsAttr(),
call_op.getResAttrsAttr(), FlatSymbolRefAttr::get(new_ref_func_name));
return new_call_op;
}
Operation *DuplicateOp(TF::XlaCallModuleOp call_op, PatternRewriter &rewriter,
const StringAttr &new_ref_func_name) {
// Create XlaCallModuleOp to the copied composite function. This
// XlaCallModuleOp does not have kQuantTraitAttrName, and therefore won't get
// quantized.
auto new_call_op = TF::XlaCallModuleOp::create(
rewriter, call_op.getLoc(), call_op.getResultTypes(),
call_op.getOperands(), call_op.getVersionAttr(), call_op.getModuleAttr(),
call_op.getSoutAttr());
new_call_op->setAttr(TF::kStablehloEntryFunctionAttrName,
rewriter.getStringAttr(new_ref_func_name.getValue()));
new_call_op->setAttrs(call_op->getAttrs());
new_call_op->setAttr(TF::kStablehloVersionAttrName,
call_op->getAttr(TF::kStablehloVersionAttrName));
new_call_op->removeAttr(rewriter.getStringAttr(kQuantTraitAttrName));
FlatSymbolRefAttr new_func_name_attr =
FlatSymbolRefAttr::get(rewriter.getContext(), new_ref_func_name);
new_call_op->setAttr(TF::kStablehloEntryFunctionAttrName, new_func_name_attr);
new_call_op->setAttr(kOriginalEntryFuncAttrName, new_ref_func_name);
return new_call_op;
}
// AddDumpTensorOp pass adds DumpTensorOp - which saves entire value of its
// input into a file - to quantizable layer's output.
class AddDumpTensorOpPass
: public PassWrapper<AddDumpTensorOpPass, OperationPass<ModuleOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AddDumpTensorOpPass)
explicit AddDumpTensorOpPass() = default;
explicit AddDumpTensorOpPass(DebuggerType debugger_type,
std::string log_dir_path)
: log_dir_path_(std::move(log_dir_path)) {
debugger_type_ = debugger_type;
}
AddDumpTensorOpPass(const AddDumpTensorOpPass &other) {
debugger_type_ = other.debugger_type_;
log_dir_path_ = other.log_dir_path_;
}
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-add-dump-tensor-op";
}
StringRef getDescription() const final {
// This is a brief description of the pass.
return "Add DumpTensor ops after quantizable ops";
}
void getDependentDialects(DialectRegistry &registry) const override {
registry.insert<TF::TensorFlowDialect>();
registry.insert<quant::QuantDialect>();
registry.insert<mlir::quant::ir::TFQuantDialect>();
}
private:
void runOnOperation() override;
Option<DebuggerType> debugger_type_{
*this, "debugger_type",
llvm::cl::init(DebuggerConfig::DEBUGGER_TYPE_UNSPECIFIED),
llvm::cl::values(
clEnumValN(DebuggerConfig::DEBUGGER_TYPE_WHOLE_MODEL, "whole_model",
"Whole model verify"),
clEnumValN(DebuggerConfig::DEBUGGER_TYPE_INT_PER_LAYER,
"int_per_layer", "Int Per-layer verify"),
clEnumValN(DebuggerConfig::DEBUGGER_TYPE_FLOAT_PER_LAYER,
"float_per_layer", "Float Per-layer verify"))};
std::string log_dir_path_ = "/tmp/dumps";
};
template <typename LiftedOpT>
class AddDumpTensorOp : public OpRewritePattern<LiftedOpT> {
public:
// Does not take ownership of context, which must refer to a valid value that
// outlives this object.
explicit AddDumpTensorOp(MLIRContext *context, DebuggerType debugger_type,
std::string log_dir_path)
: OpRewritePattern<LiftedOpT>(context),
debugger_type_(debugger_type),
log_dir_path_(std::move(log_dir_path)) {}
LogicalResult matchAndRewrite(LiftedOpT op,
PatternRewriter &rewriter) const override {
if (match(op).failed()) {
return failure();
}
rewrite(op, rewriter);
return success();
}
private:
SmallVector<NamedAttribute> CreateDumpAttributes(
PatternRewriter &rewriter, const StringRef folder_name,
const StringRef file_name, const bool enabled, const StringRef func_name,
const StringRef node_name) const {
SmallVector<NamedAttribute> dump_attributes{
rewriter.getNamedAttr("log_dir_path",
rewriter.getStringAttr(folder_name)),
rewriter.getNamedAttr("file_name", rewriter.getStringAttr(file_name)),
// The op is disabled by default. Otherwise, values will be saved
// during calibration.
rewriter.getNamedAttr("enabled", rewriter.getBoolAttr(enabled)),
rewriter.getNamedAttr("func_name", rewriter.getStringAttr(func_name)),
rewriter.getNamedAttr("node_name", rewriter.getStringAttr(node_name)),
};
return dump_attributes;
}
StringAttr DuplicateFunction(Operation *op,
const FlatSymbolRefAttr &f_attr) const {
ModuleOp module = op->getParentOfType<ModuleOp>();
SymbolTable symbol_table(module);
const func::FuncOp ref_func =
dyn_cast_or_null<func::FuncOp>(symbol_table.lookup(f_attr.getValue()));
func::FuncOp new_ref_func = dyn_cast<func::FuncOp>(ref_func->clone());
return symbol_table.insert(new_ref_func);
}
LogicalResult match(LiftedOpT op) const {
if (!op->hasAttr(kQuantTraitAttrName) || op->getNumResults() != 1) {
return failure();
}
Value result = op->getResult(0);
for (auto user : result.getUsers()) {
if (dyn_cast_or_null<TF::DumpTensorOp>(user)) return failure();
}
const FlatSymbolRefAttr f_attr = GetFuncAttr(op);
if (!f_attr.getValue().starts_with(kCompositeFuncPrefix)) return failure();
return success();
}
void rewrite(LiftedOpT op, PatternRewriter &rewriter) const {
// Only support ops with 1 results
Value result = op->getResult(0);
rewriter.setInsertionPointAfterValue(result);
// In Whole model, we first need to set file_name as
// unquantized_tensor_data.pb as it is used by unquantized dump model.
// After saving unquantized dump model, the file name will be changed to
// quantized_tensor_data.pb.
// Since this process doesn't happen for per layer, we need to set file_name
// as quantized_tensor_data.pb here.
// TODO: b/296933893 - Refactor the debugger code when no quantize option
// is added
std::string file_name =
debugger_type_ == DebuggerConfig::DEBUGGER_TYPE_WHOLE_MODEL
? "unquantized_tensor_data.pb"
: "quantized_tensor_data.pb";
const FlatSymbolRefAttr f_attr = GetFuncAttr(op);
// In TF::PartitionedCallOp case, func_name and node_name are filled.
// But in TF::XlaCallModuleOp case, node_name is `kEmptyNodeName` since
// debugging and selective quantization of StableHLO Quantizer only uses
// func_name for op matching.
auto [func_name, node_name] = GetFuncNameAndNodeName(op, f_attr);
std::string folder_name =
tensorflow::io::JoinPath(log_dir_path_, f_attr.getValue());
// Attach DumpTensorOp to its output layer.
SmallVector<NamedAttribute> dump_attributes =
CreateDumpAttributes(rewriter, folder_name, file_name,
/*enabled=*/true, func_name, node_name);
TF::DumpTensorOp::create(rewriter, op->getLoc(), TypeRange{}, result,
dump_attributes);
// Per-layer mode.
if (debugger_type_ == DebuggerConfig::DEBUGGER_TYPE_INT_PER_LAYER ||
debugger_type_ == DebuggerConfig::DEBUGGER_TYPE_FLOAT_PER_LAYER) {
// Duplicate composite function and op of quantizable layer for creating
// unquantized layer.
StringAttr new_ref_func_name = DuplicateFunction(op, f_attr);
Operation *new_op = DuplicateOp(op, rewriter, new_ref_func_name);
// Attach second DumpTensorOp to its output unquantized layer.
SmallVector<NamedAttribute> dump_attributes = CreateDumpAttributes(
rewriter, folder_name, /*file_name=*/"unquantized_tensor_data.pb",
/*enabled=*/true, func_name, node_name);
TF::DumpTensorOp::create(rewriter, op.getLoc(), TypeRange{},
new_op->getResult(0), dump_attributes);
if (debugger_type_ == DebuggerConfig::DEBUGGER_TYPE_FLOAT_PER_LAYER) {
// Swap all uses between call_op and ref_call_op, except for the
// particular use that owns DumpTensor.
rewriter.replaceUsesWithIf(
op.getResult(0), new_op->getResult(0), [](OpOperand &use) -> bool {
return !isa<TF::DumpTensorOp>(use.getOwner());
});
}
}
}
DebuggerType debugger_type_;
std::string log_dir_path_;
};
static PassRegistration<AddDumpTensorOpPass> pass;
void AddDumpTensorOpPass::runOnOperation() {
MLIRContext *ctx = &getContext();
RewritePatternSet patterns(ctx);
ModuleOp module = getOperation();
patterns.add<AddDumpTensorOp<TF::PartitionedCallOp>,
AddDumpTensorOp<TF::XlaCallModuleOp>>(ctx, debugger_type_,
log_dir_path_);
if (failed(applyPatternsGreedily(module, std::move(patterns)))) {
module.emitError() << "quant-add-dump-tensor-op failed.";
signalPassFailure();
}
}
} // namespace
std::unique_ptr<OperationPass<ModuleOp>> CreateAddDumpTensorOpPass(
DebuggerType debugger_type, std::string log_dir_path) {
return std::make_unique<AddDumpTensorOpPass>(debugger_type,
std::move(log_dir_path));
}
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/passes/add_quantization_unit_loc.cc
+201
View File
@@ -0,0 +1,201 @@
/* 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.
==============================================================================*/
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include "absl/algorithm/container.h"
#include "absl/strings/match.h"
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/Location.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/PatternMatch.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/LogicalResult.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/tensorflow/cc/quantization_unit_loc.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/ops/tf_op_quant_spec.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/passes.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
namespace mlir {
namespace quant {
namespace {
using QuantizationUnit =
tensorflow::quantization::UnitWiseQuantizationSpec::QuantizationUnit;
// Adds QuantizationUnitLoc to quantizable layers.
class AddQuantizationUnitLocPass
: public PassWrapper<AddQuantizationUnitLocPass,
OperationPass<func::FuncOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AddQuantizationUnitLocPass)
explicit AddQuantizationUnitLocPass() = default;
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-add-quantization-unit-loc";
}
StringRef getDescription() const final {
return "Add QuantizationUnitLoc to quantizable layers.";
}
private:
void runOnOperation() override;
};
// TF graph nodes are imported with one of following location patterns:
// FusedLoc[NameLoc(op_type:), ..., NameLoc(node_name@func_name)] or
// FusedLoc[NameLoc(op_type:), ..., CallSiteLoc(node_name@func_name)]. See
// tensorflow/compiler/mlir/tensorflow/translate/import_model.cc for more
// details.
bool IsImportLocPattern(FusedLoc loc) {
ArrayRef<Location> locations = mlir::cast<FusedLoc>(loc).getLocations();
if (locations.size() < 2 || !isa<NameLoc>(locations.front())) return false;
StringRef op_type_with_suffix =
mlir::cast<NameLoc>(locations.front()).getName().strref();
if (!op_type_with_suffix.ends_with(":")) return false;
return absl::c_all_of(locations, [](Location loc) {
return isa<NameLoc>(loc) ||
(isa<CallSiteLoc>(loc) &&
isa<NameLoc>(mlir::cast<CallSiteLoc>(loc).getCallee()));
});
}
// Finds the pattern of the location created by `ImporterBase::GetLocation`
// in `tensorflow/compiler/mlir/tensorflow/translate/import_model.cc`.
void FindQuantizationUnitsRecursively(Location loc,
SmallVector<QuantizationUnit>& units) {
if (!isa<FusedLoc>(loc)) return;
auto set_node_and_func_name = [](QuantizationUnit& new_unit,
StringRef name_loc_id) {
if (name_loc_id.contains("@")) {
new_unit.set_node_name(name_loc_id.split('@').first.str());
new_unit.set_func_name(name_loc_id.split('@').second.str());
} else {
new_unit.set_node_name(name_loc_id.str());
}
};
ArrayRef<Location> locations = mlir::cast<FusedLoc>(loc).getLocations();
if (IsImportLocPattern(mlir::cast<FusedLoc>(loc))) {
QuantizationUnit new_unit;
// Op type is a NameLoc with the ":" suffix.
StringRef op_type_with_suffix =
mlir::cast<NameLoc>(locations.front()).getName().strref();
StringRef op_type =
op_type_with_suffix.substr(0, op_type_with_suffix.size() - 1);
new_unit.set_op_type(op_type.str());
if (isa<NameLoc>(locations.back())) {
StringRef name_loc_id =
mlir::cast<NameLoc>(locations.back()).getName().strref();
set_node_and_func_name(new_unit, name_loc_id);
} else {
Location callee = mlir::cast<CallSiteLoc>(locations.back()).getCallee();
StringRef name_loc_id = mlir::cast<NameLoc>(callee).getName().strref();
set_node_and_func_name(new_unit, name_loc_id);
}
units.push_back(new_unit);
} else {
for (Location child_loc : locations) {
FindQuantizationUnitsRecursively(child_loc, units);
}
}
}
// Finds the QuantizationUnit from location.
std::optional<QuantizationUnit> FindQuantizationUnit(Operation* op) {
SmallVector<QuantizationUnit> quant_units;
FindQuantizationUnitsRecursively(op->getLoc(), quant_units);
if (quant_units.size() == 1) {
return *quant_units.begin();
}
// Among units, return the one with the same type as given op.
StringRef given_op_type = op->getName().getStringRef();
for (const QuantizationUnit& quant_unit : quant_units) {
if (absl::StrContains(given_op_type.lower(),
StringRef(quant_unit.op_type()).lower())) {
return quant_unit;
}
}
return std::nullopt;
}
class AddQuantizationUnitLoc : public RewritePattern {
public:
explicit AddQuantizationUnitLoc(MLIRContext* context)
: RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}
private:
LogicalResult matchAndRewrite(Operation* op,
PatternRewriter& rewriter) const override {
if (!IsOpWithQuantizableTrait(op) ||
FindQuantizationUnitFromLoc(op->getLoc()).has_value()) {
return failure();
}
std::optional<QuantizationUnit> quantization_unit =
FindQuantizationUnit(op);
if (!quantization_unit.has_value()) return failure();
if (quantization_unit->func_name().empty()) {
std::string func_name =
op->getParentOfType<func::FuncOp>().getSymNameAttr().str();
quantization_unit->set_func_name(func_name);
}
QuantizationUnitLoc unit_loc(getContext(), quantization_unit.value());
op->setLoc(unit_loc);
return success();
}
};
void AddQuantizationUnitLocPass::runOnOperation() {
MLIRContext* ctx = &getContext();
RewritePatternSet patterns(ctx);
func::FuncOp func = getOperation();
patterns.add<AddQuantizationUnitLoc>(ctx);
if (failed(applyPatternsGreedily(func, std::move(patterns)))) {
func.emitError() << "quant-add-quantization-unit-loc pattern "
"conversion did not converge.";
signalPassFailure();
}
}
} // namespace
// Creates an instance of `AddQuantizationUnitLocPass`.
std::unique_ptr<OperationPass<func::FuncOp>>
CreateAddQuantizationUnitLocPass() {
return std::make_unique<AddQuantizationUnitLocPass>();
}
static PassRegistration<AddQuantizationUnitLocPass> pass;
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/passes/cast_bf16_ops_to_f32.cc
+151
View File
@@ -0,0 +1,151 @@
/* 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.
==============================================================================*/
#include <memory>
#include <utility>
#include "mlir/IR/BuiltinOps.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/PatternMatch.h" // from @llvm-project
#include "mlir/IR/TypeUtilities.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 "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
namespace mlir {
namespace quant {
namespace {
class CastBf16OpsToF32Pass
: public PassWrapper<CastBf16OpsToF32Pass, OperationPass<ModuleOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(CastBf16OpsToF32Pass)
explicit CastBf16OpsToF32Pass() = default;
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-cast-bf16-ops-to-f32";
}
StringRef getDescription() const final {
return "Cast BF16 operations to F32.";
}
void runOnOperation() override;
};
class CastBf16OpsToF32 : public RewritePattern {
public:
explicit CastBf16OpsToF32(MLIRContext* context)
: RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}
LogicalResult matchAndRewrite(Operation* op,
PatternRewriter& rewriter) const override {
if (match(op).failed()) {
return failure();
}
rewrite(op, rewriter);
return success();
}
private:
LogicalResult match(Operation* op) const {
if (isa<TF::CastOp, TF::ConstOp>(op) ||
op->getName().hasTrait<OpTrait::ZeroOperands>()) {
return failure();
}
for (Value input : op->getOperands()) {
if (getElementTypeOrSelf(input).isBF16()) {
return success();
}
}
for (Value value : op->getResults()) {
if (getElementTypeOrSelf(value).isBF16()) {
return success();
}
}
return failure();
}
void rewrite(Operation* op, PatternRewriter& rewriter) const {
// Casts inputs of the operation.
for (int i = 0; i < op->getNumOperands(); i++) {
Value input = op->getOperand(i);
if (getElementTypeOrSelf(input).isBF16()) {
Value f32_cast = TF::CastOp::create(
rewriter, op->getLoc(),
CloneTypeWithNewElementType(input.getType(), rewriter.getF32Type()),
input);
op->setOperand(i, f32_cast);
}
}
// Casts BF16 outputs of the operation.
for (Value value : op->getResults()) {
if (getElementTypeOrSelf(value).isBF16()) {
value.setType(CloneTypeWithNewElementType(value.getType(),
rewriter.getF32Type()));
rewriter.setInsertionPointAfterValue(value);
for (Operation* user : op->getUsers()) {
for (int i = 0; i < user->getNumOperands(); i++) {
if (user->getOperand(i) == value) {
Value bf16_cast = TF::CastOp::create(
rewriter, user->getLoc(),
CloneTypeWithNewElementType(value.getType(),
rewriter.getBF16Type()),
value);
user->setOperand(i, bf16_cast);
}
}
}
}
}
}
};
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/cast_bf16_ops_to_f32.inc"
void CastBf16OpsToF32Pass::runOnOperation() {
MLIRContext* ctx = &getContext();
RewritePatternSet patterns(ctx);
auto module_op = getOperation();
patterns.add<CastBf16OpsToF32>(ctx);
populateWithGenerated(patterns);
if (failed(applyPatternsGreedily(module_op, std::move(patterns)))) {
module_op.emitError() << "quant-cast-bf16-ops-to-f32 failed.";
signalPassFailure();
}
}
} // namespace
// Creates an instance of the Cast BF16 ops to F32 pass.
std::unique_ptr<OperationPass<ModuleOp>> CreateCastBf16OpsToF32Pass() {
return std::make_unique<CastBf16OpsToF32Pass>();
}
static PassRegistration<CastBf16OpsToF32Pass> pass;
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/passes/cast_bf16_ops_to_f32.td
+33
View File
@@ -0,0 +1,33 @@
/* 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.
==============================================================================*/
include "mlir/IR/OpBase.td"
include "mlir/IR/PatternBase.td"
include "mlir/Dialect/Func/IR/FuncOps.td"
include "mlir/Dialect/Arith/IR/ArithOps.td"
include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.td"
include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"
//===----------------------------------------------------------------------===//
// Pattern rules for converting bfloat16 operations to fp32 conversions.
//===----------------------------------------------------------------------===//
// Remove unneeded redundant cast ops like (f32 -> bf16 -> f32).
def RemoveUnneededCastOps : Pat<
(TF_CastOp:$output
(TF_CastOp
$input, $truncate_0), $truncate_1),
(replaceWithValue $input),
[(AreTheSameElementType $input, $output)]>;
tensorflow/compiler/mlir/quantization/tensorflow/passes/constants.h
+42
View File
@@ -0,0 +1,42 @@
/* 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_QUANTIZATION_TENSORFLOW_PASSES_CONSTANTS_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_PASSES_CONSTANTS_H_
#include "llvm/ADT/StringRef.h"
#include "mlir/Support/LLVM.h" // from @llvm-project
namespace mlir {
namespace quant {
// Name of the save function. The "tf_quant__" prefix is for avoiding conflict
// with existing function's name.
inline constexpr StringRef kTfQuantSaveFuncName = "tf_quant__save";
// Name of the TensorFlow Operation to be fetched to save the variables to
// checkpoint. This save op follows the SavedModel's load semantics, so it
// should return the file prefix of the checkpoint as a string tensor.
inline constexpr StringRef kTfQuantSaveOpName = "tf_quant__save_op";
// Name the file prefix string tensor. The tensor is used to identify the prefix
// to the checkpoint where the variables are saved / loaded. This may be present
// in a function argument's "tf_saved_model.index_path" attribute to identify
// the file prefix function argument.
inline constexpr StringRef kTfFilePrefix = "__tf_file_prefix";
} // namespace quant
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_PASSES_CONSTANTS_H_
tensorflow/compiler/mlir/quantization/tensorflow/passes/convert_custom_aggregation_op_to_quant_stats.cc
+129
View File
@@ -0,0 +1,129 @@
/* 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 <algorithm>
#include <memory>
#include <string>
#include <tuple>
#include <utility>
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/SourceMgr.h"
#include "mlir/Dialect/Quant/IR/Quant.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.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/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/passes.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/tf_quant_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
namespace mlir {
namespace quant {
namespace {
class ConvertCustomAggregationOpToQuantStatsPass
: public PassWrapper<ConvertCustomAggregationOpToQuantStatsPass,
OperationPass<func::FuncOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
ConvertCustomAggregationOpToQuantStatsPass)
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-convert-tf-custom-aggregator-op-to-quant-stats";
}
StringRef getDescription() const final {
// This is a brief description of the pass.
return "Convert tf.CustomAggregator op to quant.Stats";
}
void getDependentDialects(DialectRegistry &registry) const override {
registry.insert<TF::TensorFlowDialect>();
registry.insert<quant::QuantDialect>();
registry.insert<mlir::quant::ir::TFQuantDialect>();
}
void runOnOperation() override;
};
class ConvertCustomAggregationOpToQuantStats
: public OpRewritePattern<TF::CustomAggregatorOp> {
public:
// Does not take ownership of context, which must refer to a valid value that
// outlives this object.
explicit ConvertCustomAggregationOpToQuantStats(MLIRContext *context)
: OpRewritePattern<TF::CustomAggregatorOp>(context) {}
LogicalResult matchAndRewrite(TF::CustomAggregatorOp op,
PatternRewriter &rewriter) const override {
FloatAttr min = mlir::dyn_cast_or_null<FloatAttr>(op->getAttr("min"));
FloatAttr max = mlir::dyn_cast_or_null<FloatAttr>(op->getAttr("max"));
// When there are no min and max attributes, remove op.
if (min == nullptr || max == nullptr) {
op.getOutput().replaceAllUsesWith(op.getInput());
rewriter.eraseOp(op);
return success();
}
// The layer stats contain only the first min/max pairs.
ElementsAttr layer_stats = DenseFPElementsAttr::get(
RankedTensorType::get({2}, rewriter.getF32Type()),
{static_cast<float>(min.getValueAsDouble()),
static_cast<float>(max.getValueAsDouble())});
ElementsAttr axis_stats;
IntegerAttr axis;
mlir::quant::ir::StatisticsOp stats_op =
mlir::quant::ir::StatisticsOp::create(rewriter, op->getLoc(),
op.getInput(), layer_stats,
axis_stats, axis);
op.getOutput().replaceAllUsesWith(stats_op.getResult());
return success();
}
};
static PassRegistration<ConvertCustomAggregationOpToQuantStatsPass> pass;
void ConvertCustomAggregationOpToQuantStatsPass::runOnOperation() {
MLIRContext *ctx = &getContext();
RewritePatternSet patterns(ctx);
func::FuncOp func = getOperation();
patterns.add<ConvertCustomAggregationOpToQuantStats>(ctx);
if (failed(applyPatternsGreedily(func, std::move(patterns)))) {
func.emitError()
<< "quant-convert-tf-custom-aggregator-op-to-quant-stats failed.";
signalPassFailure();
}
}
} // namespace
std::unique_ptr<OperationPass<func::FuncOp>>
CreateConvertCustomAggregationOpToQuantStatsPass() {
return std::make_unique<ConvertCustomAggregationOpToQuantStatsPass>();
}
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/passes/convert_fake_quant_to_qdq.cc
+89
View File
@@ -0,0 +1,89 @@
/* 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/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/Quant/IR/Quant.h" // from @llvm-project
#include "mlir/IR/DialectRegistry.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/PatternMatch.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 "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project // IWYU pragma: keep, for applyPatternsGreedily
#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/passes.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/utils/fake_quant_utils.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
namespace mlir {
namespace quant {
namespace {
class ConvertFakeQuantToQdqPass
: public PassWrapper<ConvertFakeQuantToQdqPass,
OperationPass<func::FuncOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ConvertFakeQuantToQdqPass)
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-convert-fake-quant-to-qdq";
}
StringRef getDescription() const final {
// This is a brief description of the pass.
return "Convert Fake Quant op to quant.qcast and quant.dcast pairs";
}
void getDependentDialects(DialectRegistry& registry) const override {
registry.insert<TF::TensorFlowDialect>();
registry.insert<quant::QuantDialect>();
registry.insert<mlir::quant::ir::TFQuantDialect>();
}
void runOnOperation() override;
};
static PassRegistration<ConvertFakeQuantToQdqPass> pass;
void ConvertFakeQuantToQdqPass::runOnOperation() {
MLIRContext* ctx = &getContext();
func::FuncOp func = getOperation();
if (failed(
ConvertFakeQuantOps(func, ctx, /*use_fake_quant_num_bits=*/false))) {
func.emitError() << "quant-convert-fake-quant-to-qdq pass failed.";
signalPassFailure();
}
// For removing dead FakeQuant* ops
RewritePatternSet patterns(ctx);
if (failed(applyPatternsGreedily(func, std::move(patterns)))) {
signalPassFailure();
}
}
} // namespace
std::unique_ptr<OperationPass<func::FuncOp>> CreateConvertFakeQuantToQdqPass() {
return std::make_unique<ConvertFakeQuantToQdqPass>();
}
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/passes/convert_tf_xla_op_to_tf_op.cc
+342
View File
@@ -0,0 +1,342 @@
/* 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.
==============================================================================*/
#include <cstdint>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/container/flat_hash_set.h"
#include "absl/strings/match.h"
#include "absl/strings/str_cat.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 "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/MLIRContext.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/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/LogicalResult.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/passes.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "xla/xla_data.pb.h"
namespace mlir {
namespace quant {
namespace {
class ConvertTfXlaOpToTfOpPass
: public PassWrapper<ConvertTfXlaOpToTfOpPass,
OperationPass<func::FuncOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ConvertTfXlaOpToTfOpPass)
ConvertTfXlaOpToTfOpPass() = default;
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-convert-tf-xla-op-to-tf-op";
}
StringRef getDescription() const final {
// This is a brief description of the pass.
return "Apply converting Tensorflow Xla ops to non-xla ops.";
}
void getDependentDialects(DialectRegistry& registry) const override {
registry.insert<TF::TensorFlowDialect, arith::ArithDialect>();
}
void runOnOperation() override;
};
// Generate an einsum equation from the given DotDimensionNumber.
std::string CreateEinsumEquation(
const xla::DotDimensionNumbers& dot_dimension_numbers, const int lhs_rank,
const int rhs_rank) {
// Prepare necessary indices.
absl::flat_hash_set<int64_t> lhs_batch_idx, rhs_batch_idx;
absl::flat_hash_set<int64_t> lhs_contract_idx, rhs_contract_idx;
lhs_batch_idx.insert(dot_dimension_numbers.lhs_batch_dimensions().begin(),
dot_dimension_numbers.lhs_batch_dimensions().end());
lhs_contract_idx.insert(
dot_dimension_numbers.lhs_contracting_dimensions().begin(),
dot_dimension_numbers.lhs_contracting_dimensions().end());
rhs_batch_idx.insert(dot_dimension_numbers.rhs_batch_dimensions().begin(),
dot_dimension_numbers.rhs_batch_dimensions().end());
rhs_contract_idx.insert(
dot_dimension_numbers.rhs_contracting_dimensions().begin(),
dot_dimension_numbers.rhs_contracting_dimensions().end());
// Generate equation.
std::string lhs_eq = "";
std::string rhs_eq = "";
std::string out_eq = "";
char c = 'a';
std::vector<char> lhs_batch_dims;
std::vector<char> lhs_contract_dims;
for (int i = 0; i < lhs_rank; i++) {
absl::StrAppend(&lhs_eq, std::string(1, c));
if (lhs_batch_idx.contains(i)) {
lhs_batch_dims.push_back(c);
} else if (lhs_contract_idx.contains(i)) {
lhs_contract_dims.push_back(c);
}
c++;
}
int batch_trace_idx = 0;
int contract_trace_idx = 0;
const bool rhs_only_batch = lhs_batch_dims.empty();
for (int i = 0; i < rhs_rank; i++) {
if (rhs_batch_idx.contains(i)) {
if (rhs_only_batch) {
rhs_eq.push_back(c);
lhs_batch_dims.push_back(c);
c++;
} else {
rhs_eq.push_back(lhs_batch_dims[batch_trace_idx]);
batch_trace_idx++;
}
} else if (rhs_contract_idx.contains(i)) {
absl::StrAppend(&rhs_eq,
std::string(1, lhs_contract_dims[contract_trace_idx]));
contract_trace_idx++;
} else {
rhs_eq += c;
c++;
}
}
// Create out_eq by merging lhs and rhs.
// In XlaDotv2 style - batch dim - leftover from lhs - leftover from rhs.
for (const char c : lhs_batch_dims) {
absl::StrAppend(&out_eq, std::string(1, c));
}
for (const char c : lhs_eq) {
if (!absl::StrContains(out_eq, c) && !absl::StrContains(rhs_eq, c)) {
absl::StrAppend(&out_eq, std::string(1, c));
}
}
for (const char c : rhs_eq) {
if (!absl::StrContains(out_eq, c) && !absl::StrContains(lhs_eq, c)) {
absl::StrAppend(&out_eq, std::string(1, c));
}
}
return absl::StrCat(lhs_eq, ",", rhs_eq, "->", out_eq);
}
Value CreateEinsumOpFromXlaDotV2Op(OpBuilder& builder, const Location loc,
Value lhs, Value rhs, Value output,
StringAttr dot_dimension_numbers_str) {
xla::DotDimensionNumbers dot_dimension_numbers;
dot_dimension_numbers.ParseFromString(dot_dimension_numbers_str.str());
SmallVector<Value> input_arguments = {lhs, rhs};
const int lhs_rank = mlir::cast<ShapedType>(lhs.getType()).getShape().size();
const int rhs_rank = mlir::cast<ShapedType>(rhs.getType()).getShape().size();
const std::string einsum_equation =
CreateEinsumEquation(dot_dimension_numbers, lhs_rank, rhs_rank);
return TF::EinsumOp::create(builder, loc, output.getType(), input_arguments,
builder.getStringAttr(einsum_equation));
}
// Restores the collapsed dimensions to the `tensor_type`. `collapsed_dims`
// designate the dimension indices that were collapsed to produce `tensor_type`.
// The restored dimensions' sizes are 1, according to the semantics of
// `XlaGatherOp (https://www.tensorflow.org/xla/operation_semantics#gather). The
// resulting type's shape has `tensor_type.size() + collapsed_dims.size()`
// dimensions.
RankedTensorType RestoreCollapsedDimensions(
const RankedTensorType tensor_type,
const absl::flat_hash_set<int64_t>& collapsed_dims) {
ArrayRef<int64_t> original_tensor_shape = tensor_type.getShape();
const int output_tensor_rank =
original_tensor_shape.size() + collapsed_dims.size();
auto shape_itr = tensor_type.getShape().begin();
// Populate the dimensions of the output shape, including the restored
// dimensions.
SmallVector<int64_t> output_shape(output_tensor_rank);
for (int i = 0; i < output_tensor_rank; i++) {
if (collapsed_dims.contains(i)) {
// The collapsed dimension's size should have been 1, so it restores the
// dimension with size 1.
output_shape[i] = 1;
} else {
output_shape[i] = *shape_itr;
shape_itr++;
}
}
return RankedTensorType::get(output_shape, tensor_type.getElementType());
}
// Determines the output type of the `SliceOp` when it is being inserted in
// place of a `XlaGatherOp`. When the dimensions of `xla_gather_op_output_type`
// is known, the `collapsed_dims` are restored. `xla_gather_op_output_type` is
// the result of collapsing the `collapsed_dims`, but the `SliceOp`'s output
// should not have the dimensions collapsed already. Returns
// `xla_gather_op_output_type` unchanged if the rank is unknown.
//
// Examples:
// * If `xla_gather_op_output_type` == tensor<*xf32>, then it returns:
// tensor<*xf32>.
// * If `xla_gather_op_output_type` == tensor<3x5xi32> and `collapsed_dims` ==
// {0}, then it returns: tensor<1x3x5xi32>.
// * If `xla_gather_op_output_type` == tensor<3x5xf32> and `collapsed_dims` ==
// {1, 3}, then it returns: tensor<3x1x5x1xf32>.
Type GetSliceOpOutputType(Type xla_gather_op_output_type,
const absl::flat_hash_set<int64_t>& collapsed_dims) {
if (auto ranked_output_type =
mlir::dyn_cast<RankedTensorType>(xla_gather_op_output_type);
ranked_output_type) {
return RestoreCollapsedDimensions(ranked_output_type, collapsed_dims);
}
return xla_gather_op_output_type;
}
// TODO (b/275225582): Supports Xla Gather op in general case.
bool IsXlaGatherWithoutBatch(Value operand, Value start_indices) {
auto operand_type = mlir::dyn_cast_or_null<ShapedType>(operand.getType());
auto start_indices_type =
mlir::dyn_cast_or_null<ShapedType>(start_indices.getType());
if (start_indices_type == nullptr || operand_type == nullptr) return false;
return start_indices_type.getShape().size() == 1;
}
Value CreateSliceAndReshapeOpFromXlaGatherOpWithoutBatch(
OpBuilder& builder, const Location loc, Value operand, Value start_indices,
Value slice_sizes, Value output, StringAttr dimension_numbers_str) {
// Reads dimension numbers.
xla::GatherDimensionNumbers dimension_numbers;
dimension_numbers.ParseFromString(dimension_numbers_str.str());
// Construct full start_indices with given start_indices and
// start_index_map.
const ArrayRef<int64_t> operand_shape =
mlir::cast<ShapedType>(operand.getType()).getShape();
const int64_t operand_rank = operand_shape.size();
// Fills zeros if start_index is not given in start_indices.
Value empty_start_indices = TF::FillOp::create(
builder, loc, RankedTensorType::get({operand_rank}, builder.getI64Type()),
/*shape=*/Create1DConstValue<int64_t>(builder, loc, {operand_rank}),
/*value=*/CreateScalarConstValue<int64_t>(builder, loc, 0));
// Converts start_index_map proto to tensor.
const int64_t index_map_size = dimension_numbers.start_index_map().size();
SmallVector<int64_t> indices(index_map_size);
for (int64_t i = 0; i < index_map_size; i++) {
indices[i] = dimension_numbers.start_index_map()[i];
}
// Fill elements from start_indices with start_index_map
Value scattered_start_indices = TF::TensorScatterUpdateOp::create(
builder, loc, empty_start_indices,
/*indices=*/
TF::ReshapeOp::create(
builder, loc,
RankedTensorType::get({index_map_size, 1}, builder.getI64Type()),
Create1DConstValue<int64_t>(builder, loc, indices),
Create1DConstValue<int64_t>(builder, loc, {index_map_size, 1})),
/*value=*/
TF::CastOp::create(
builder, loc,
RankedTensorType::get(
mlir::cast<ShapedType>(start_indices.getType()).getShape(),
builder.getI64Type()),
start_indices));
absl::flat_hash_set<int64_t> collapsed_dims;
collapsed_dims.insert(dimension_numbers.collapsed_slice_dims().begin(),
dimension_numbers.collapsed_slice_dims().end());
// Slice operand by constructed start_indices and slice_sizes.
auto slice_op = TF::SliceOp::create(
builder, loc, GetSliceOpOutputType(output.getType(), collapsed_dims),
operand,
/*start_indices=*/scattered_start_indices,
/*slice_sizes=*/
TF::CastOp::create(
builder, loc,
RankedTensorType::get(
mlir::cast<ShapedType>(slice_sizes.getType()).getShape(),
builder.getI64Type()),
slice_sizes));
// Collapses dimensions by reshaping.
SmallVector<int64_t> new_shape(operand_rank - collapsed_dims.size());
for (int64_t i = 0, j = 0; i < operand_rank; i++) {
if (!collapsed_dims.contains(i)) {
new_shape[j++] = operand_shape[i];
}
}
if (!new_shape.empty()) new_shape[0] = -1;
return TF::ReshapeOp::create(builder, loc, output.getType(), slice_op,
Create1DConstValue(builder, loc, new_shape));
}
bool IsPrecisionEmpty(StringAttr prec_str) {
xla::PrecisionConfig prec;
prec.ParseFromString(prec_str.str());
return !prec.operand_precision_size();
}
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/convert_tf_xla_op_to_tf_op.inc"
void ConvertTfXlaOpToTfOpPass::runOnOperation() {
MLIRContext* ctx = &getContext();
auto func = getOperation();
// The pattern includes
// - Converting XlaDotV2Op to EinsumOp
// - Converting XlaGatherOp to SliceOp
RewritePatternSet patterns(ctx);
populateWithGenerated(patterns);
if (failed(applyPatternsGreedily(func, std::move(patterns)))) {
func.emitError() << "quant-converting-tf-xla-op-to-tf-op failed.";
signalPassFailure();
}
}
} // namespace
std::unique_ptr<OperationPass<func::FuncOp>> CreateConvertTfXlaOpToTfOpPass() {
return std::make_unique<ConvertTfXlaOpToTfOpPass>();
}
static PassRegistration<ConvertTfXlaOpToTfOpPass> pass;
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/passes/convert_tf_xla_op_to_tf_op.td
+51
View File
@@ -0,0 +1,51 @@
/* 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.
==============================================================================*/
include "mlir/IR/OpBase.td"
include "mlir/IR/PatternBase.td"
include "mlir/Dialect/Func/IR/FuncOps.td"
include "mlir/Dialect/Arith/IR/ArithOps.td"
include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.td"
include "tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td"
// Only handles the case where precision config is default.
def IsPrecisionEmpty :
Constraint<CPred<"IsPrecisionEmpty($0)">>;
// Creates Einsum Op from XlaDotV2 Op by generating equation.
def CreateEinsumOpFromXlaDotV2Op : NativeCodeCall<
"CreateEinsumOpFromXlaDotV2Op($_builder, $_loc, $0...)">;
// Convert XlaDotV2 Op to Einsum Op with above two functions.
def ConvertXlaDotV2OpToEinsumOp : Pat<
(TF_XlaDotV2Op:$dot $lhs, $rhs, $dot_dimension_numbers, $precision_config),
(CreateEinsumOpFromXlaDotV2Op $lhs, $rhs, $dot, $dot_dimension_numbers),
[(IsPrecisionEmpty $precision_config)]>;
// Only handles the case where batch_dimension is empty.
def IsXlaGatherWithoutBatch :
Constraint<CPred<"IsXlaGatherWithoutBatch($0, $1)">>;
// Create Slice op from XlaGather op without batch dimension.
def CreateSliceAndReshapeOpFromXlaGatherOpWithoutBatch : NativeCodeCall<
"CreateSliceAndReshapeOpFromXlaGatherOpWithoutBatch($_builder, $_loc, $0...)">;
// Convert XlaGather op without batch to Slice op with above two functions.
def ConvertXlaGatherOpWithoutBatch : Pat<
(TF_XlaGatherOp:$gather $operand,
$start_indices, $slice_sizes, $dimension_numbers, $indices_are_sorted),
(CreateSliceAndReshapeOpFromXlaGatherOpWithoutBatch $operand,
$start_indices, $slice_sizes, $gather, $dimension_numbers),
[(IsXlaGatherWithoutBatch $operand, $start_indices)]>;
tensorflow/compiler/mlir/quantization/tensorflow/passes/convert_tpu_model_to_cpu.cc
+155
View File
@@ -0,0 +1,155 @@
/* 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.
==============================================================================*/
#include <memory>
#include <utility>
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/IR/SymbolTable.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/LogicalResult.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/passes.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/remove_identity_op_pattern.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/core/tpu/tpu_defs.h"
namespace mlir {
namespace quant {
namespace {
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/convert_tpu_model_to_cpu.inc"
// Convert a TPU model to be compatible on CPU by rewriting/removing TPU ops.
class ConvertTpuModelToCpuPass
: public PassWrapper<ConvertTpuModelToCpuPass, OperationPass<ModuleOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ConvertTpuModelToCpuPass)
explicit ConvertTpuModelToCpuPass() = default;
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-convert-tpu-model-to-cpu";
}
StringRef getDescription() const final {
return "Convert TPU models to CPU by rewriting TPU related operations.";
}
void runOnOperation() override;
};
class RemoveTpuOp : public RewritePattern {
public:
explicit RemoveTpuOp(MLIRContext* context)
: RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}
private:
LogicalResult matchAndRewrite(Operation* op,
PatternRewriter& rewriter) const override {
// Remove `_tpu_replicate` attributes on each operation first.
if (op->hasAttr(tensorflow::kTPUReplicateAttr)) {
op->removeAttr(tensorflow::kTPUReplicateAttr);
return success();
}
// Remove TPU operations.
if (isa<TF::TPUReplicateMetadataOp, TF::TPUCompilationResultOp,
TF::TPUOrdinalSelectorOp>(op)) {
op->erase();
} else if (auto replicated_input_op =
dyn_cast_or_null<TF::TPUReplicatedInputOp>(op)) {
// TODO(b/267700110): Handle multiple input/output cases.
rewriter.replaceOp(replicated_input_op, replicated_input_op.getInputs());
} else if (auto replicated_output_op =
dyn_cast_or_null<TF::TPUReplicatedOutputOp>(op)) {
// TODO(b/267700110): Handle multiple input/output cases.
rewriter.replaceOp(replicated_output_op, replicated_output_op.getInput());
} else {
return failure();
}
return success();
}
};
class ReplaceTpuPartitionedCallOpWithPartitionedCallOp
: public OpRewritePattern<TF::TPUPartitionedCallOp> {
public:
using OpRewritePattern<TF::TPUPartitionedCallOp>::OpRewritePattern;
private:
LogicalResult matchAndRewrite(TF::TPUPartitionedCallOp call_op,
PatternRewriter& rewriter) const override {
auto f_attr = mlir::dyn_cast<FlatSymbolRefAttr>(call_op.getFAttr());
auto module_op = call_op->getParentOfType<ModuleOp>();
SymbolTable symbol_table(module_op);
auto f_name = f_attr.getValue();
func::FuncOp float_func =
dyn_cast<func::FuncOp>(symbol_table.lookup(f_name));
if (!float_func) {
return failure();
}
rewriter.setInsertionPointAfter(call_op);
// The TPUPartitionedCall has a TPUOrdinalSelectorOp for its last argument
// which should be removed. So the replaced PartitionedCall op should keep
// its original arguments except for the last element.
SmallVector<Value> args = call_op.getOperands().drop_back();
rewriter.replaceOpWithNewOp<TF::PartitionedCallOp>(
call_op, float_func.getResultTypes(), args, call_op.getArgAttrsAttr(),
call_op.getResAttrsAttr(), f_attr);
return success();
}
};
void ConvertTpuModelToCpuPass::runOnOperation() {
MLIRContext* ctx = &getContext();
RewritePatternSet patterns(ctx);
ModuleOp module_op = getOperation();
patterns.add<ReplaceTpuPartitionedCallOpWithPartitionedCallOp,
ReplaceBatchFunctionOpToPartitionedCallOp>(ctx);
patterns.add<RemoveTpuOp>(ctx);
patterns.add<RemoveIdentity>(ctx);
if (failed(applyPatternsGreedily(module_op, std::move(patterns)))) {
module_op.emitError() << "quant-convert-tpu-model-to-cpu pattern "
"conversion did not converge.";
signalPassFailure();
return;
}
}
} // namespace
// Creates an instance of `ConvertTpuModelToCpuPass`.
std::unique_ptr<OperationPass<ModuleOp>> CreateConvertTpuModelToCpuPass() {
return std::make_unique<ConvertTpuModelToCpuPass>();
}
static PassRegistration<ConvertTpuModelToCpuPass> pass;
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/passes/convert_tpu_model_to_cpu.td
+39
View File
@@ -0,0 +1,39 @@
/* 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.
==============================================================================*/
include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"
include "mlir/IR/OpBase.td"
include "mlir/IR/PatternBase.td"
include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.td"
// Combines the two variadic arguments ($in_tensors and $captured_tensors).
def GetBatchFunctionOpArgOperands:
NativeCodeCall<"cast<TF::BatchFunctionOp>($0[0].getDefiningOp()).getArgOperands()">;
def CreateEmptyDictAttr : NativeCodeCall<"$_builder.getArrayAttr({})">;
// Replaces `TF_BatchFunctionOp` into `TF_PartitionedCallOp` that calls the
// same $f. This may be required, for example, when inlining is desired,
// because `TF_BatchFunctionOp` doesn't have the `CallOpInterface` trait.
def ReplaceBatchFunctionOpToPartitionedCallOp : Pat<
(TF_BatchFunctionOp:$src_op_res
$_, $_, $f, $_, $_, $_, $_, $_, $_, $_, $_, $_, $_, $_, $_, $_, $_, $_, $_, $_, $_, $_),
(TF_PartitionedCallOp
(GetBatchFunctionOpArgOperands $src_op_res),
/*arg_attrs=*/(CreateEmptyDictAttr),
/*res_attrs=*/(CreateEmptyDictAttr),
$f,
/*config=*/(CreateStringAttr<"">),
/*config_proto=*/(CreateStringAttr<"">),
/*executor_type=*/(CreateStringAttr<"">))>;
tensorflow/compiler/mlir/quantization/tensorflow/passes/duplicate_shape_determining_constants.cc
+374
View File
@@ -0,0 +1,374 @@
/* 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 <array>
#include <iterator>
#include <memory>
#include "absl/algorithm/container.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/Debug.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Builders.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/quantization/tensorflow/passes/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
// Required to use LLVM_DEBUG macro.
#define DEBUG_TYPE "quant-duplicate-shape-determining-constants"
namespace mlir {
namespace quant {
namespace {
// This pass duplicates constants that affect or determine the shape of a tensor
// after being used in a computation for some op. Some specific operands of TF
// ops (like the `dim` argument for `TF::ExpandDimsOp`) determine the shape of
// the resulting tensor. If these operands are constants, they are duplicated
// and replace the shape-determining operands. Each duplicated constant will
// only be used as the shape-determining operand; it will not replace other
// usages of the original constant. If the operands are not constants (i.e.
// results of some other computation), then the pass recursively traverses the
// call tree upwards and duplicates all constants found in the subtree in a
// similar manner.
//
// This pass may be used to avoid placing shape-determining constants in the CPU
// graph and pass them as arguments to the TPU graph (via `TPUPartitionedCall`).
// If this happens, the XLA compiler cannot recognize such arguments as
// constants and may result in an error.
//
// A set of predefined ops and operand indices is used to determine whether an
// operand is a target for constant duplication.
class DuplicateShapeDeterminingConstantsPass
: public PassWrapper<DuplicateShapeDeterminingConstantsPass,
OperationPass<func::FuncOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
DuplicateShapeDeterminingConstantsPass)
StringRef getArgument() const final {
return "quant-duplicate-shape-determining-constants";
}
StringRef getDescription() const final {
return "Duplicates shape-determining constants. A shape-determining "
"constant is a constant that are transitively used to change or "
"determine the shape of a tensor. For example, the second argument "
"'dim' to TF::ExpandDimsOp specifies the dimension index to expand.";
}
void runOnOperation() override;
};
// Returns True iff the otuput value of `op` is either a compile time constant
// or bounded from the XLA compiler's perspective, even if it is not a
// `ConstOp`.
bool IsOutputCompileTimeConstantOrBounded(Operation* op) {
return llvm::isa_and_nonnull<TF::ShapeOp, TF::ShapeNOp, TF::RankOp,
TF::SizeOp, TF::TensorArraySizeV3Op,
TF::XlaSetBoundOp>(op);
}
// Recursively duplicate constants for `op_operands` upward.
void RecursivelyDuplicateConstantsForOperands(
llvm::ArrayRef<OpOperand*> op_operands) {
// Target operands to duplicate if it is a ConstOp.
llvm::SmallVector<OpOperand*, 4> duplication_targets{op_operands.begin(),
op_operands.end()};
int target_idx = 0;
while (target_idx < duplication_targets.size()) {
OpOperand* curr_operand = duplication_targets[target_idx];
target_idx++;
Operation* owning_op = curr_operand->getOwner();
Operation* defining_op = curr_operand->get().getDefiningOp();
if (llvm::isa_and_nonnull<TF::ConstOp>(defining_op)) {
// No need to clone if this is the only use.
if (defining_op->hasOneUse()) {
LLVM_DEBUG(llvm::dbgs()
<< "Not duplicating constant operand since it has only one "
"usage. Op: "
<< curr_operand->getOperandNumber()
<< ", operand idx: " << curr_operand->getOperandNumber()
<< ", loc: " << owning_op->getLoc() << "\n");
continue;
}
mlir::OpBuilder builder{owning_op->getContext()};
builder.setInsertionPointAfter(defining_op);
auto const_op_cloned = builder.clone(*defining_op);
// Replace the operand with the duplicated op.
owning_op->setOperand(curr_operand->getOperandNumber(),
const_op_cloned->getResult(0));
LLVM_DEBUG(llvm::dbgs()
<< "Duplicated constant operand from: "
<< owning_op->getName().getStringRef()
<< ", operand idx: " << curr_operand->getOperandNumber()
<< ", loc: " << const_op_cloned->getLoc() << "\n");
} else if (IsOutputCompileTimeConstantOrBounded(defining_op)) {
// Stop the recursion early when the output of the defining op is
// considered compile-time constant from the XLA compiler's perspective.
continue;
} else if (!defining_op) {
// One example for this case is when `curr_operand` is a function
// argument.
owning_op->emitWarning()
<< "Operand idx (zero-based): " << curr_operand->getOperandNumber()
<< " does not have a defining op and cannot be duplicated.";
} else {
// If the operand's defining is not a ConstOp, recursively traverse
// "upwards" to find ConstOps that transitively produces the current
// operand and duplicate them.
auto op_operands = defining_op->getOpOperands();
absl::c_transform(
op_operands, std::back_inserter(duplication_targets),
[](OpOperand& op_operand) -> OpOperand* { return &op_operand; });
}
}
}
// Evaluate `operand_idx` w.r.t. `op`'s operands. If `operand_idx` is a positive
// number or a zero, it is returned as it is. If it is a negative number, it
// means it is counting backwards and will return the zero-based operand index
// for `op`.
//
// `operand_idx` should be within the range: [-num_operands, num_operands - 1].
int EvaluateOperandIdx(const int operand_idx, Operation& op) {
if (operand_idx < 0) {
// Calculate the actual index if a negative value is provided for
// `operand_idx`.
return op.getNumOperands() + operand_idx;
}
return operand_idx;
}
// Returns the pointers to operands at `operand_indices` of `op`.
llvm::SmallVector<OpOperand*> GetOperands(Operation& op,
llvm::ArrayRef<int> operand_indices) {
llvm::SmallVector<OpOperand*> operands{};
for (const int operand_idx : operand_indices) {
const int evaluated_operand_idx = EvaluateOperandIdx(operand_idx, op);
operands.emplace_back(&op.getOpOperand(evaluated_operand_idx));
}
return operands;
}
// Represents an op type and its operand indices that should be "compile time
// constant" from the XLA compiler's point of view.
template <typename OpT, int... OperandIdx>
struct CompileTimeConstantOperand {
static_assert(
sizeof...(OperandIdx) > 0,
"CompileTimeConstantOperand should have at least one operand index.");
using OpType = OpT;
// Returns the indices of operands that should be compile time constants.
static constexpr std::array<int, sizeof...(OperandIdx)> OperandIndices() {
return {OperandIdx...};
}
};
// Finds all op of type `T::OpType` `func_op` and recursively duplicates
// constants used at the op's operands at `T::OperandIndices()`. It sequentially
// does the same thing for `Ts`.
template <typename T, typename... Ts>
void DuplicateShapeDeterminingConstants(func::FuncOp func_op) {
for (auto op : func_op.getOps<typename T::OpType>()) {
RecursivelyDuplicateConstantsForOperands(
GetOperands(*op, T::OperandIndices()));
}
// Do the same thing for the rest of `Ts`.
if constexpr (sizeof...(Ts) != 0) {
DuplicateShapeDeterminingConstants<Ts...>(func_op);
}
}
void DuplicateShapeDeterminingConstantsPass::runOnOperation() {
func::FuncOp func_op = getOperation();
DuplicateShapeDeterminingConstants<
// go/keep-sorted start
CompileTimeConstantOperand<TF::AllToAllOp, 1>, // $group_assignment
CompileTimeConstantOperand<TF::ArgMaxOp, 1>, // $dimension
CompileTimeConstantOperand<TF::ArgMinOp, 1>, // $dimension
// $orig_input_shape
CompileTimeConstantOperand<TF::AvgPool3DGradOp, 0>,
// $orig_input_shape
CompileTimeConstantOperand<TF::AvgPoolGradOp, 0>,
// $block_shape, $crops
CompileTimeConstantOperand<TF::BatchToSpaceNDOp, 1, 2>,
CompileTimeConstantOperand<TF::BatchToSpaceOp, 1>, // $crops
CompileTimeConstantOperand<TF::BincountOp, 1>, // $size
CompileTimeConstantOperand<TF::BroadcastArgsOp, 0, 1>, // $s0, $s1
// $s0, $s1
CompileTimeConstantOperand<TF::BroadcastGradientArgsOp, 0, 1>,
CompileTimeConstantOperand<TF::BroadcastToOp, 1>, // $shape
/// $group_assignment
CompileTimeConstantOperand<TF::CollectiveAssignGroupV2Op, 0>,
// $source_target_pairs
CompileTimeConstantOperand<TF::CollectivePermuteOp, 1>,
// $group_size, $group_key
CompileTimeConstantOperand<TF::CollectiveReduceV2Op, 1, 2>,
CompileTimeConstantOperand<TF::ConcatV2Op, -1>, // (variadic) $axis
// $filter_sizes
CompileTimeConstantOperand<TF::Conv2DBackpropFilterOp, 1>,
CompileTimeConstantOperand<TF::Conv2DBackpropInputOp, 0>, // $input_sizes
// $filter_sizes
CompileTimeConstantOperand<TF::Conv3DBackpropFilterV2Op, 1>,
// $input_sizes
CompileTimeConstantOperand<TF::Conv3DBackpropInputV2Op, 0>,
// $group_assignment
CompileTimeConstantOperand<TF::CrossReplicaSumOp, 1>,
CompileTimeConstantOperand<TF::CumprodOp, 1>, // $axis
CompileTimeConstantOperand<TF::CumsumOp, 1>, // $axis
CompileTimeConstantOperand<TF::CumulativeLogsumexpOp, 1>, // $axis
// $filter_sizes
CompileTimeConstantOperand<TF::DepthwiseConv2dNativeBackpropFilterOp, 1>,
// $input_sizes
CompileTimeConstantOperand<TF::DepthwiseConv2dNativeBackpropInputOp, 0>,
CompileTimeConstantOperand<TF::EmptyOp, 0>, // $shape
// $element_shape, $max_num_elements
CompileTimeConstantOperand<TF::EmptyTensorListOp, 0, 1>,
CompileTimeConstantOperand<TF::ExpandDimsOp, 1>, // $dim
CompileTimeConstantOperand<TF::FillOp, 0>, // $dims
CompileTimeConstantOperand<TF::GatherV2Op, 2>, // $axis
CompileTimeConstantOperand<TF::IRFFT2DOp, 1>, // $fft_length
CompileTimeConstantOperand<TF::IRFFT3DOp, 1>, // $fft_length
CompileTimeConstantOperand<TF::IRFFTOp, 1>, // $fft_length
CompileTimeConstantOperand<TF::InTopKV2Op, 2>, // $k
CompileTimeConstantOperand<TF::LinSpaceOp, 2>, // $num
CompileTimeConstantOperand<TF::ListDiffOp, 0, 1>, // $x, $y
// $k, $padding_value
CompileTimeConstantOperand<TF::MatrixDiagPartV3Op, 1, 2>,
// $k, $num_rows, $num_cols, $padding_value
CompileTimeConstantOperand<TF::MatrixDiagV2Op, 1, 2, 3, 4>,
// $k, $num_rows, $num_cols, $padding_value
CompileTimeConstantOperand<TF::MatrixDiagV3Op, 1, 2, 3, 4>,
CompileTimeConstantOperand<TF::MatrixSetDiagV2Op, 2>, // $k
CompileTimeConstantOperand<TF::MatrixSetDiagV3Op, 2>, // $k
CompileTimeConstantOperand<TF::MaxOp, 1>, // $reduction_indices
// $ksize, $strides
CompileTimeConstantOperand<TF::MaxPoolGradGradV2Op, 3, 4>,
// $ksize, $strides
CompileTimeConstantOperand<TF::MaxPoolGradV2Op, 2, 3>,
CompileTimeConstantOperand<TF::MaxPoolV2Op, 1, 2>, // $ksize, $strides
CompileTimeConstantOperand<TF::MeanOp, 1>, // $reduction_indices
CompileTimeConstantOperand<TF::MirrorPadGradOp, 1>, // $paddings
CompileTimeConstantOperand<TF::MirrorPadOp, 1>, // $paddings
CompileTimeConstantOperand<TF::MultinomialOp, 1>, // $num_samples
// $max_output_size
CompileTimeConstantOperand<TF::NonMaxSuppressionV3Op, 2>,
// $max_output_size
CompileTimeConstantOperand<TF::NonMaxSuppressionV4Op, 2>,
CompileTimeConstantOperand<TF::OneHotOp, 1>, // $depth
CompileTimeConstantOperand<TF::PadOp, 1>, // $paddings
CompileTimeConstantOperand<TF::PadV2Op, 1>, // $paddings
// $shape
CompileTimeConstantOperand<TF::ParameterizedTruncatedNormalOp, 0>,
CompileTimeConstantOperand<TF::RFFT2DOp, 1>, // $fft_length
CompileTimeConstantOperand<TF::RFFT3DOp, 1>, // $fft_length
CompileTimeConstantOperand<TF::RFFTOp, 1>, // $fft_length
CompileTimeConstantOperand<TF::RandomStandardNormalOp, 0>, // $shape
CompileTimeConstantOperand<TF::RandomUniformIntOp, 0>, // $shape
CompileTimeConstantOperand<TF::RandomUniformOp, 0>, // $shape
// $start, $limit, $delta
CompileTimeConstantOperand<TF::RangeOp, 0, 1, 2>,
CompileTimeConstantOperand<TF::ReshapeOp, 1>, // $shape
CompileTimeConstantOperand<TF::ResizeBilinearOp, 1>, // $size
CompileTimeConstantOperand<TF::ResizeNearestNeighborOp, 1>, // $size
// $begin, $end, $strides
CompileTimeConstantOperand<TF::ResourceStridedSliceAssignOp, 1, 2, 3>,
CompileTimeConstantOperand<TF::ReverseOp, 1>, // $dims
CompileTimeConstantOperand<TF::ReverseV2Op, 1>, // $axis
CompileTimeConstantOperand<TF::ScatterNdOp, 2>, // $shape
CompileTimeConstantOperand<TF::SegmentSumV2Op, 2>, // $num_segments
CompileTimeConstantOperand<TF::SliceOp, 1, 2>, // $begin, $size
CompileTimeConstantOperand<TF::SparseToDenseOp, 1>, // $output_shape
CompileTimeConstantOperand<TF::SplitOp, 0>, // $split_dim
// $size_splits, $split_dim
CompileTimeConstantOperand<TF::SplitVOp, 1, 2>,
CompileTimeConstantOperand<TF::StackV2Op, 0>, // $max_size
// $num_samples
CompileTimeConstantOperand<TF::StatelessMultinomialOp, 1>,
// $shape, $begin, $end, $strides
CompileTimeConstantOperand<TF::StridedSliceGradOp, 0, 1, 2, 3>,
// $begin, $end, $strides
CompileTimeConstantOperand<TF::StridedSliceOp, 1, 2, 3>,
CompileTimeConstantOperand<TF::SumOp, 1>, // $reduction_indices
CompileTimeConstantOperand<TF::TensorArraySplitV3Op, 2>, // $lengths
CompileTimeConstantOperand<TF::TensorArrayV3Op, 0>, // $size
// $element_shape
CompileTimeConstantOperand<TF::TensorListFromTensorOp, 1>,
// $element_shape, $num_elements
CompileTimeConstantOperand<TF::TensorListReserveOp, 0, 1>,
// $begin, $end, $strides
CompileTimeConstantOperand<TF::TensorStridedSliceUpdateOp, 1, 2, 3>,
CompileTimeConstantOperand<TF::TileOp, 1>, // $multiples
CompileTimeConstantOperand<TF::TopKV2Op, 1>, // $k
CompileTimeConstantOperand<TF::TransposeOp, 1>, // $perm
CompileTimeConstantOperand<TF::TruncatedNormalOp, 0>, // $shape
CompileTimeConstantOperand<TF::UnsortedSegmentMaxOp, 2>, // $num_segments
CompileTimeConstantOperand<TF::UnsortedSegmentMinOp, 2>, // $num_segments
CompileTimeConstantOperand<TF::UnsortedSegmentSumOp, 2>, // $num_segments
// $broadcast_dims
CompileTimeConstantOperand<TF::XlaBroadcastHelperOp, 2>,
// $window_strides, $padding, $lhs_dilation, $rhs_dilation,
// $feature_group_count
CompileTimeConstantOperand<TF::XlaConvOp, 2, 3, 4, 5, 6>,
// $window_strides, $padding, $lhs_dilation, $rhs_dilation,
// $feature_group_count
CompileTimeConstantOperand<TF::XlaConvV2Op, 2, 3, 4, 5, 6>,
CompileTimeConstantOperand<TF::XlaDynamicSliceOp, 2>, // $slice_indices
CompileTimeConstantOperand<TF::XlaGatherOp, 2>, // $slice_sizes
// $padding_low, $padding_high, $padding_interior
CompileTimeConstantOperand<TF::XlaPadOp, 2, 3, 4>,
// $window_dimensions, $window_strides, $base_dilations,
// $window_dilations, $padding
CompileTimeConstantOperand<TF::XlaReduceWindowOp, 2, 3, 4, 5, 6>,
// $dim_index
CompileTimeConstantOperand<TF::XlaRemoveDynamicDimensionSizeOp, 1>,
// $window_dimensions, $window_strides, $padding
CompileTimeConstantOperand<TF::XlaSelectAndScatterOp, 1, 2, 3>,
CompileTimeConstantOperand<TF::XlaSetBoundOp, 1>, // $bound
// $dim_index
CompileTimeConstantOperand<TF::XlaSetDynamicDimensionSizeOp, 1>
// go/keep-sorted end
>(func_op);
}
static PassRegistration<DuplicateShapeDeterminingConstantsPass> pass{};
} // namespace
std::unique_ptr<OperationPass<func::FuncOp>>
CreateDuplicateShapeDeterminingConstantsPass() {
return std::make_unique<DuplicateShapeDeterminingConstantsPass>();
}
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/passes/insert_custom_aggregation_ops.cc
+370
View File
@@ -0,0 +1,370 @@
/* 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 <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include "absl/status/statusor.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/Twine.h"
#include "llvm/Support/CommandLine.h"
#include "mlir/Dialect/Func/IR/FuncOps.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/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/IR/TypeUtilities.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/LogicalResult.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/common/lift_as_function_call.h"
#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/calibration_parameters.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/passes.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/tf_quant_ops.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
namespace mlir {
namespace quant {
namespace {
using ::stablehlo::quantization::CalibrationOptions;
using ::stablehlo::quantization::Method;
constexpr StringRef kQuantTraitAttrName = "_tfl_quant_trait";
// Whether the op is a call op to lifted composite function.
bool IsCallToQuantizableLiftedFunction(Operation *op) {
if (!op) return false;
if (auto xla_call_module_op = dyn_cast_or_null<TF::XlaCallModuleOp>(op);
xla_call_module_op != nullptr) {
absl::StatusOr<Method> method = GetQuantizationMethod(xla_call_module_op);
if (method.ok() && method->has_static_range_ptq()) return true;
}
TF::PartitionedCallOp call_op = dyn_cast_or_null<TF::PartitionedCallOp>(op);
return call_op && call_op->hasAttrOfType<StringAttr>(kQuantTraitAttrName) &&
call_op->getAttrOfType<StringAttr>(kQuantTraitAttrName).getValue() ==
llvm::StringRef(
QuantTraitValues[QuantizationTrait::FullyQuantizable]);
}
// Returns the composite function name.
std::optional<StringRef> GetCompsiteFunctionName(Operation *op) {
if (!IsCallToQuantizableLiftedFunction(op)) return std::nullopt;
if (auto xla_call_module_op = dyn_cast_or_null<TF::XlaCallModuleOp>(op);
xla_call_module_op != nullptr) {
auto entry_function_attr = xla_call_module_op->getAttrOfType<StringAttr>(
kOriginalStablehloEntryFunctionAttrName);
if (!entry_function_attr) return std::nullopt;
return entry_function_attr.getValue();
} else {
TF::PartitionedCallOp call_op = dyn_cast_or_null<TF::PartitionedCallOp>(op);
const auto f_attr = mlir::dyn_cast<FlatSymbolRefAttr>(call_op.getFAttr());
if (!f_attr) return std::nullopt;
return f_attr.getValue();
}
}
class InsertCustomAggregationOpsPass
: public PassWrapper<InsertCustomAggregationOpsPass,
OperationPass<func::FuncOp>> {
public:
explicit InsertCustomAggregationOpsPass() : test_mode_(true) {
initializeForTest();
}
explicit InsertCustomAggregationOpsPass(const CalibrationOptions &calib_opts)
: test_mode_(false), calib_opts_(calib_opts) {}
InsertCustomAggregationOpsPass(const InsertCustomAggregationOpsPass &other) {
test_mode_ = other.test_mode_;
test_case_ = other.test_case_;
calib_opts_ = other.calib_opts_;
initializeForTest();
}
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(InsertCustomAggregationOpsPass)
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-insert-custom-aggregation-ops";
}
StringRef getDescription() const final {
// This is a brief description of the pass.
return "Insert custom aggregation ops for the calibration procedure";
}
void getDependentDialects(DialectRegistry &registry) const override {
registry.insert<TF::TensorFlowDialect>();
}
void runOnOperation() override;
private:
enum TestCase {
TEST_CASE_MIN_MAX,
TEST_CASE_AVERAGE_MIN_MAX,
TEST_CASE_HISTOGRAM_PERCENTILE,
TEST_CASE_HISTOGRAM_MSE_BRUTEFORCE,
TEST_CASE_HISTOGRAM_MSE_MAX_FREQUENCY,
TEST_CASE_HISTOGRAM_MSE_SYMMETRIC,
};
bool test_mode_;
CalibrationOptions calib_opts_;
Option<TestCase> test_case_{
*this, "test-case",
llvm::cl::desc(
"Select a the test case for testing various calibration methods. It "
"sets the value of calib_opts_ when test_mode_ is true."),
llvm::cl::init(TEST_CASE_MIN_MAX),
llvm::cl::values(
clEnumValN(TEST_CASE_MIN_MAX, "MIN_MAX",
"Uses MIN_MAX calibration method"),
clEnumValN(TEST_CASE_AVERAGE_MIN_MAX, "AVERAGE_MIN_MAX",
"Uses AVERAGE_MIN_MAX calibration method"),
clEnumValN(TEST_CASE_HISTOGRAM_PERCENTILE, "HISTOGRAM_PERCENTILE",
"Uses HISTOGRAM_PERCENTILE calibration method"),
clEnumValN(TEST_CASE_HISTOGRAM_MSE_BRUTEFORCE,
"HISTOGRAM_MSE_BRUTEFORCE",
"Uses HISTOGRAM_MSE_BRUTEFORCE calibration method"),
clEnumValN(TEST_CASE_HISTOGRAM_MSE_MAX_FREQUENCY,
"HISTOGRAM_MSE_MAX_FREQUENCY",
"Uses HISTOGRAM_MSE_MAX_FREQUENCY calibration "
"method"),
clEnumValN(TEST_CASE_HISTOGRAM_MSE_SYMMETRIC,
"HISTOGRAM_MSE_SYMMETRIC",
"Uses HISTOGRAM_MSE_SYMMETRIC calibration "
"method"))};
// Initialize for tests.
void initializeForTest() {
if (!test_mode_) return;
switch (test_case_.getValue()) {
case TEST_CASE_MIN_MAX:
calib_opts_.set_calibration_method(
CalibrationOptions::CALIBRATION_METHOD_MIN_MAX);
break;
case TEST_CASE_AVERAGE_MIN_MAX:
calib_opts_.set_calibration_method(
CalibrationOptions::CALIBRATION_METHOD_AVERAGE_MIN_MAX);
break;
case TEST_CASE_HISTOGRAM_PERCENTILE: {
calib_opts_.set_calibration_method(
CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_PERCENTILE);
auto calibration_parameters =
CalibrationOptions::CalibrationParameters();
calibration_parameters.set_num_bins(512);
calibration_parameters.set_min_percentile(0.001);
calibration_parameters.set_max_percentile(99.999);
calib_opts_.mutable_calibration_parameters()->CopyFrom(
calibration_parameters);
break;
}
case TEST_CASE_HISTOGRAM_MSE_BRUTEFORCE: {
calib_opts_.set_calibration_method(
CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_MSE_BRUTEFORCE);
auto calibration_parameters =
CalibrationOptions::CalibrationParameters();
calibration_parameters.set_num_bins(512);
calib_opts_.mutable_calibration_parameters()->CopyFrom(
calibration_parameters);
break;
}
case TEST_CASE_HISTOGRAM_MSE_MAX_FREQUENCY: {
calib_opts_.set_calibration_method(
CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_MSE_MAX_FREQUENCY);
auto calibration_parameters =
CalibrationOptions::CalibrationParameters();
calibration_parameters.set_num_bins(512);
calib_opts_.mutable_calibration_parameters()->CopyFrom(
calibration_parameters);
break;
}
case TEST_CASE_HISTOGRAM_MSE_SYMMETRIC: {
calib_opts_.set_calibration_method(
CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_MSE_SYMMETRIC);
auto calibration_parameters =
CalibrationOptions::CalibrationParameters();
calibration_parameters.set_num_bins(512);
calib_opts_.mutable_calibration_parameters()->CopyFrom(
calibration_parameters);
break;
}
}
}
};
static PassRegistration<InsertCustomAggregationOpsPass> pass;
class AddCustomAggregationOp : public RewritePattern {
public:
// Does not take ownership of context, which must refer to a valid value that
// outlives this object.
explicit AddCustomAggregationOp(MLIRContext *context,
const CalibrationOptions &calib_opts)
: RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context),
calib_opts_(calib_opts) {}
LogicalResult matchAndRewrite(Operation *op,
PatternRewriter &rewriter) const override {
// Return early if the given operator is the custom aggregator op.
if (dyn_cast_or_null<TF::CustomAggregatorOp>(op)) return failure();
// The CustomAggregatorOp is only added after quantizable values.
SmallVector<Value> quantizable_values;
SmallVector<std::string> aggregator_ids;
if (IsCallToQuantizableLiftedFunction(op)) {
std::optional<StringRef> composite_function_name =
GetCompsiteFunctionName(op);
if (!composite_function_name.has_value()) return failure();
// Quantize inputs of quantizable composite functions.
for (OpOperand &input : op->getOpOperands()) {
Type element_type = getElementTypeOrSelf(input.get().getType());
// Non-float cases won't be calibrated.
if (!element_type.isF32()) {
continue;
}
// Skip when there is any already existing CustomAggregatorOp found.
Operation *defining_op = input.get().getDefiningOp();
if (dyn_cast_or_null<TF::CustomAggregatorOp>(defining_op)) {
continue;
}
// Skip calibration when the given operand comes from a constant.
if (defining_op != nullptr &&
defining_op->hasTrait<OpTrait::ConstantLike>()) {
continue;
}
quantizable_values.push_back(input.get());
aggregator_ids.push_back(
(llvm::Twine(composite_function_name.value()) + "_arg_" +
llvm::Twine(input.getOperandNumber()) + "_calibration_method_" +
llvm::Twine(calib_opts_.calibration_method()))
.str());
}
} else {
// Quantize output of fully quantizable composite functions.
for (Value input : op->getOperands()) {
auto defining_op = input.getDefiningOp();
std::optional<StringRef> composite_function_name =
GetCompsiteFunctionName(defining_op);
if (!composite_function_name.has_value()) continue;
// Do not add CustomAggregatorOp after Gather since it is a weight-only
// quantizable op.
if (auto call_op =
dyn_cast_or_null<TF::PartitionedCallOp>(defining_op)) {
StringRef function_name =
mlir::cast<FlatSymbolRefAttr>(call_op.getFAttr()).getValue();
if (function_name.contains("gather")) continue;
}
quantizable_values.push_back(input);
// All composite functions have a single result at the moment.
aggregator_ids.push_back((llvm::Twine(composite_function_name.value()) +
"_calibration_method_" +
llvm::Twine(calib_opts_.calibration_method()))
.str());
}
}
if (quantizable_values.empty()) return failure();
int32_t effective_num_bins = GetNumBins(calib_opts_);
for (auto [value, aggregator_id] :
llvm::zip_equal(quantizable_values, aggregator_ids)) {
// ID attribute will have empty value for now.
SmallVector<NamedAttribute, 5> attributes{
rewriter.getNamedAttr("id", rewriter.getStringAttr(aggregator_id)),
rewriter.getNamedAttr(
"calibration_method",
rewriter.getI32IntegerAttr(calib_opts_.calibration_method())),
rewriter.getNamedAttr("num_bins",
rewriter.getI32IntegerAttr(effective_num_bins)),
rewriter.getNamedAttr(
"min_percentile",
rewriter.getF32FloatAttr(
calib_opts_.calibration_parameters().min_percentile())),
rewriter.getNamedAttr(
"max_percentile",
rewriter.getF32FloatAttr(
calib_opts_.calibration_parameters().max_percentile())),
};
SmallVector<Type, 4> output_types{
value.getType(),
RankedTensorType::get({}, rewriter.getF32Type()),
RankedTensorType::get({}, rewriter.getF32Type()),
RankedTensorType::get({effective_num_bins}, rewriter.getI64Type()),
};
// Insert custom aggregation op between operand and operator.
rewriter.setInsertionPointAfterValue(value);
Operation *aggregator_op = TF::CustomAggregatorOp::create(
rewriter, op->getLoc(), output_types, value, attributes);
Value aggregator_op_result = aggregator_op->getOpResult(0);
value.replaceAllUsesExcept(aggregator_op_result, aggregator_op);
}
return success();
}
private:
CalibrationOptions calib_opts_;
};
void InsertCustomAggregationOpsPass::runOnOperation() {
MLIRContext *ctx = &getContext();
RewritePatternSet patterns(ctx);
func::FuncOp func = getOperation();
patterns.add<AddCustomAggregationOp>(ctx, calib_opts_);
if (failed(applyPatternsGreedily(func, std::move(patterns)))) {
func.emitError() << "quant-insert-custom-aggregation-ops failed.";
signalPassFailure();
}
}
} // namespace
std::unique_ptr<OperationPass<func::FuncOp>>
CreateInsertCustomAggregationOpsPass(const CalibrationOptions &calib_opts) {
return std::make_unique<InsertCustomAggregationOpsPass>(calib_opts);
}
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/passes/insert_main_function.cc
+446
View File
@@ -0,0 +1,446 @@
/* 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 <algorithm>
#include <cstdint>
#include <iterator>
#include <memory>
#include <string>
#include <vector>
#include "absl/algorithm/container.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_join.h"
#include "absl/strings/str_split.h"
#include "absl/strings/string_view.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/StringMap.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/StringSet.h" // IWYU pragma: keep
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#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/OperationSupport.h" // from @llvm-project
#include "mlir/IR/SymbolTable.h" // from @llvm-project
#include "mlir/IR/TypeRange.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/Interfaces/FunctionInterfaces.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/quantization/common/attrs_and_constraints.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/compiler/mlir/tensorflow/translate/import_model.h"
namespace mlir {
namespace quant {
namespace {
using ::mlir::tf_saved_model::kTfSavedModelExportedNamesAttr;
using ::mlir::tf_saved_model::kTfSavedModelIndexPathAttr;
using ::tensorflow::kImportModelDefaultGraphFuncName;
constexpr StringRef kEntryFunctionAttr = "tf.entry_function";
// The ConvertMlirToGraphdef requires the provided input module to have a main
// function, which might not exist in case of multi-signature graphs. In that
// case, this pass will create a new main function, which calls signature
// functions.
//
// An already existing @main function will be renamed by attaching a numeric
// suffix like `@main_0` to avoid conflict with the newly created main function.
class InsertMainFunctionPass
: public PassWrapper<InsertMainFunctionPass, OperationPass<ModuleOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(InsertMainFunctionPass)
explicit InsertMainFunctionPass() = default;
StringRef getArgument() const override {
return "quant-insert-main-function";
}
StringRef getDescription() const override {
return "Inserts the main function to the module.";
}
void runOnOperation() override;
};
// Checks if a FuncOp is exported.
bool IsExported(func::FuncOp op) {
auto exported_names =
op->getAttrOfType<ArrayAttr>(kTfSavedModelExportedNamesAttr);
return exported_names && !exported_names.empty();
}
// Check if a function is an entry function.
bool IsEntryFunction(func::FuncOp op) {
return op->hasAttr(kEntryFunctionAttr);
}
// Returns true iff the provided FuncOp is qualified to be included in the main
// function.
bool ShouldIncludeInMainFunction(func::FuncOp func_op) {
return !func_op.isPrivate() && IsExported(func_op) &&
IsEntryFunction(func_op);
}
// Sets a function to be private so it can be referred internally.
void SetFunctionPrivate(func::FuncOp func) {
func.setVisibility(SymbolTable::Visibility::Private);
// The `tf_saved_model` attributes can only be applied to public functions.
for (auto& attr : func->getAttrs()) {
StringRef attr_name = attr.getName().getValue();
if (attr_name.starts_with("tf_saved_model.")) {
func->removeAttr(attr_name);
}
}
auto iface = cast<FunctionOpInterface>(func.getOperation());
for (int i = 0; i < func.getNumArguments(); ++i) {
for (auto& attr : iface.getArgAttrs(i)) {
const StringAttr& attr_name = attr.getName();
if (attr_name.getValue().starts_with("tf_saved_model.")) {
func.removeArgAttr(i, attr_name);
}
}
}
for (int i = 0; i < func.getNumResults(); ++i) {
for (auto& attr : iface.getResultAttrs(i)) {
const StringAttr& attr_name = attr.getName();
if (attr_name.getValue().starts_with("tf_saved_model.")) {
func.removeResultAttr(i, attr_name);
}
}
}
}
// Information to identify an output in its node and in the model output list.
// Ex: If the model output list is ["add:0", "topk:0": "topk:1"], then the
// output corresponding to "topk:1" will have output_index=2 and tensor_index=1.
struct OutputInfo {
// The index of this output in the model output list.
int32_t output_index;
// The index of this output in its node.
int32_t tensor_index;
// The output value.
Value value;
};
// Makes input/output names across entry functions unique if necessary. If a
// duplicated name is found, this function will add signature prefix for all the
// input/output names.
void GetUniqueInputOutputNodeNames(ModuleOp module_op,
std::vector<std::string>& input_name_vec,
std::vector<std::string>& output_name_vec) {
bool need_prefix_for_input_name = false;
bool need_prefix_for_output_name = false;
std::vector<StringRef> fn_input_name_vec, fn_output_name_vec;
StringSet<> input_name_set, output_name_set;
for (auto func_op : module_op.getOps<func::FuncOp>()) {
if (!ShouldIncludeInMainFunction(func_op)) continue;
if (auto tf_attrs =
func_op->getAttrOfType<DictionaryAttr>(kEntryFunctionAttr)) {
StringRef function_name = func_op.getSymName();
if (auto inputs_attr = tf_attrs.get("inputs")) {
const std::string inputs_attr_str =
mlir::cast<StringAttr>(inputs_attr).getValue().str();
std::vector<std::string> fn_input_names =
absl::StrSplit(inputs_attr_str, ',', absl::SkipEmpty());
for (StringRef input_name : fn_input_names) {
if (input_name_set.contains(input_name)) {
// Found a duplicated name, all input names will be prefixed by
// their corresponding function names.
need_prefix_for_input_name = true;
}
input_name_set.insert(input_name);
fn_input_name_vec.push_back(function_name);
}
input_name_vec.insert(input_name_vec.end(),
std::make_move_iterator(fn_input_names.begin()),
std::make_move_iterator(fn_input_names.end()));
}
if (auto outputs_attr = tf_attrs.get("outputs")) {
const std::string outputs_attr_str =
mlir::cast<StringAttr>(outputs_attr).getValue().str();
std::vector<std::string> fn_output_names =
absl::StrSplit(outputs_attr_str, ',', absl::SkipEmpty());
for (StringRef output_name : fn_output_names) {
if (output_name_set.contains(output_name)) {
// Found a duplicated name, all output names will be prefixed by
// their corresponding function names.
need_prefix_for_output_name = true;
}
output_name_set.insert(output_name);
fn_output_name_vec.push_back(function_name);
}
output_name_vec.insert(output_name_vec.end(),
std::make_move_iterator(fn_output_names.begin()),
std::make_move_iterator(fn_output_names.end()));
}
}
}
if (need_prefix_for_input_name) {
absl::c_transform(
input_name_vec, fn_input_name_vec, input_name_vec.begin(),
[](const std::string& input_name, const StringRef fn_name) {
return absl::StrCat(fn_name.str(), "_", input_name);
});
}
if (need_prefix_for_output_name) {
absl::c_transform(
output_name_vec, fn_output_name_vec, output_name_vec.begin(),
[](const std::string& output_name, const StringRef fn_name) {
return absl::StrCat(fn_name.str(), "_", output_name);
});
}
}
// Creates a main function which calls other exported functions.
bool CreateMainFunction(ModuleOp module_op) {
MLIRContext* context = module_op.getContext();
OpBuilder builder(context);
std::vector<std::string> input_names, output_names;
GetUniqueInputOutputNodeNames(module_op, input_names, output_names);
// Collects argument and result types.
llvm::SmallVector<Location> arg_locs;
llvm::SmallVector<Type> arg_types, result_types;
for (auto func_op : module_op.getOps<func::FuncOp>()) {
if (!ShouldIncludeInMainFunction(func_op)) continue;
arg_types.append(func_op.getArgumentTypes().begin(),
func_op.getArgumentTypes().end());
auto& return_op = func_op.getBody().getBlocks().front().back();
result_types.append(return_op.getOperandTypes().begin(),
return_op.getOperandTypes().end());
for (const auto& arg : func_op.getArguments()) {
arg_locs.push_back(arg.getLoc());
}
}
// Creates a new main function.
auto func_type = FunctionType::get(context, arg_types, result_types);
auto main_func = func::FuncOp::create(
builder, module_op.getLoc(), kImportModelDefaultGraphFuncName, func_type);
builder.createBlock(&main_func.getBody(), main_func.begin(), arg_types,
arg_locs);
SmallVector<NamedAttribute> func_attrs;
func_attrs.push_back(
{StringAttr::get(context, "inputs"),
StringAttr::get(context, absl::StrJoin(input_names, ","))});
func_attrs.push_back(
{StringAttr::get(context, "outputs"),
StringAttr::get(context, absl::StrJoin(output_names, ","))});
auto dictAttr = DictionaryAttr::get(context, func_attrs);
main_func->setAttr(StringAttr::get(context, kEntryFunctionAttr), dictAttr);
main_func->setAttr(
kTfSavedModelExportedNamesAttr,
builder.getStrArrayAttr({kImportModelDefaultGraphFuncName}));
if (input_names.size() != main_func.getNumArguments() ||
output_names.size() != main_func.getNumResults()) {
module_op.emitError()
<< "Number of inputs and outputs in the tf.entry_function attribute "
"mismatched. [Input] Expected: "
<< input_names.size() << ", got: " << main_func.getNumArguments()
<< ". [Output] Expected: " << output_names.size()
<< ", got: " << main_func.getNumResults();
return false;
}
const int num_args = main_func.getNumArguments();
for (int i = 0; i < num_args; ++i) {
main_func.setArgAttr(
i, kTfSavedModelIndexPathAttr,
ArrayAttr::get(context, {StringAttr::get(context, input_names[i])}));
}
const int num_results = main_func.getNumResults();
for (int i = 0; i < num_results; ++i) {
main_func.setResultAttr(
i, kTfSavedModelIndexPathAttr,
ArrayAttr::get(context, {StringAttr::get(context, output_names[i])}));
}
// Creates PartitionedCall ops to call exported functions.
auto guard = OpBuilder::InsertionGuard(builder);
int arg_idx = 0;
int result_idx = 0;
llvm::SmallVector<Value> call_op_returns;
for (auto func_op : module_op.getOps<func::FuncOp>()) {
if (!ShouldIncludeInMainFunction(func_op)) continue;
llvm::ArrayRef<BlockArgument> new_args = llvm::ArrayRef(
main_func.getArguments().begin() + arg_idx, func_op.getNumArguments());
arg_idx += func_op.getNumArguments();
llvm::ArrayRef<Type> new_types = llvm::ArrayRef(
result_types.begin() + result_idx, func_op.getNumResults());
result_idx += func_op.getNumResults();
auto call_op = TF::PartitionedCallOp::create(
builder, module_op.getLoc(), new_types, new_args,
/*args_attrs=*/nullptr,
/*res_attrs=*/nullptr,
SymbolRefAttr::get(context, func_op.getSymName()),
/*config=*/builder.getStringAttr(""),
/*config_proto=*/builder.getStringAttr(""),
/*executor_type=*/builder.getStringAttr(""));
call_op_returns.append(call_op.getResults().begin(),
call_op.getResults().end());
SetFunctionPrivate(func_op);
}
// Creates Identity/IdentityN ops for returing values. This allows us to
// restore the same output tensor names in python.
int32_t output_count = 0;
// Map from node name to the list of the OutputInfos of its outputs that are
// used as the model outputs.
llvm::StringMap<llvm::SmallVector<OutputInfo>> node_to_output_map;
for (auto [output_name, call_op_return] :
llvm::zip(output_names, call_op_returns)) {
std::vector<std::string> name_and_index =
absl::StrSplit(output_name, ':', absl::SkipEmpty());
llvm::StringRef node_name = name_and_index.front();
int32_t tensor_index = 0;
if (name_and_index.size() > 1) {
tensor_index = std::stoi(name_and_index.back());
}
node_to_output_map[node_name].push_back(
{output_count++, tensor_index, call_op_return});
}
Value scalar_one =
CreateScalarConstValue<float>(builder, builder.getUnknownLoc(), 1.0);
llvm::SmallVector<Value> returning_values(output_count, Value());
for (const auto& node_name : node_to_output_map.keys()) {
auto node_output_tensors = node_to_output_map[node_name];
NameLoc new_loc = NameLoc::get(builder.getStringAttr(node_name));
int32_t max_tensor_index = 0;
absl::c_for_each(node_output_tensors,
[&max_tensor_index](const OutputInfo& output_info) {
max_tensor_index =
std::max(max_tensor_index, output_info.tensor_index);
});
// Create IdentityOp or IdentityNOp based on the number of outputs.
Operation* identity_op;
if (max_tensor_index == 0) {
Value output_value = node_output_tensors.front().value;
identity_op = TF::IdentityOp::create(
builder, new_loc, output_value.getType(), output_value);
} else {
llvm::SmallVector<Value> input_values(node_output_tensors.size(),
scalar_one);
for (const auto& [output_index, tensor_index, tensor_value] :
node_output_tensors) {
input_values[tensor_index] = tensor_value;
}
identity_op = TF::IdentityNOp::create(
builder, new_loc, TypeRange(ValueRange(input_values)), input_values);
}
for (const auto& [output_index, tensor_index, tensor_value] :
node_output_tensors) {
returning_values[output_index] = identity_op->getResult(tensor_index);
}
}
func::ReturnOp::create(builder, main_func.getBody().getLoc(),
returning_values);
// Adds the new function to symbol table.
SymbolTable symbol_table(module_op);
symbol_table.insert(main_func);
return true;
}
// Creates a new function name by attaching a number suffix
// (`main_func_name_{i}`) and incrementing it until there are no conflicts.
std::string CreateNewFuncName(const StringRef main_func_name,
SymbolTable& symbol_table) {
int suffix_id = 0;
std::string new_func_name =
absl::StrCat(main_func_name.str(), "_", suffix_id);
while (symbol_table.lookup(new_func_name)) {
suffix_id++;
new_func_name = absl::StrCat(main_func_name.str(), "_", suffix_id);
}
return new_func_name;
}
// Renames the existing @main function to avoid conflict with the newly
// created main function. When it is renamed, its usages will also be replaced.
// It will be renamed by attaching a number suffix like `@main_{i}`, until there
// are no conflicts. This function is a no-op when no function called @main
// exists.
LogicalResult RenameExistingMainFunction(ModuleOp module_op) {
SymbolTable symbol_table(module_op);
auto main_func_op =
symbol_table.lookup<func::FuncOp>(kImportModelDefaultGraphFuncName);
if (!main_func_op) {
return success();
}
const std::string new_func_name =
CreateNewFuncName(main_func_op.getSymName(), symbol_table);
main_func_op.setSymName(new_func_name);
return symbol_table.replaceAllSymbolUses(
main_func_op, StringAttr::get(module_op.getContext(), new_func_name),
module_op);
}
void InsertMainFunctionPass::runOnOperation() {
ModuleOp module_op = getOperation();
if (failed(RenameExistingMainFunction(module_op))) {
module_op->emitError("Failed to rename existing function `@main`.");
signalPassFailure();
}
if (!CreateMainFunction(module_op)) {
signalPassFailure();
}
}
} // namespace
std::unique_ptr<OperationPass<ModuleOp>> CreateInsertMainFunctionPass() {
return std::make_unique<InsertMainFunctionPass>();
}
static PassRegistration<InsertMainFunctionPass> pass([] {
return CreateInsertMainFunctionPass();
});
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/passes/insert_quantized_functions.cc
+223
View File
@@ -0,0 +1,223 @@
/* 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 "absl/container/flat_hash_map.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/SourceMgr.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/UB/IR/UBOps.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Diagnostics.h" // from @llvm-project
#include "mlir/IR/DialectRegistry.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "mlir/IR/SymbolTable.h" // from @llvm-project
#include "mlir/Parser/Parser.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassManager.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 "mlir/Transforms/Passes.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/passes.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/quantized_function_library.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
namespace mlir {
namespace quant {
namespace {
using QuantMethod = tensorflow::quantization::QuantizationMethod::PresetMethod;
using ::tensorflow::quantization::OpSet;
class InsertQuantizedFunctionsPass
: public PassWrapper<InsertQuantizedFunctionsPass,
OperationPass<ModuleOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(InsertQuantizedFunctionsPass)
explicit InsertQuantizedFunctionsPass() = default;
explicit InsertQuantizedFunctionsPass(QuantMethod quantization_method,
OpSet op_set) {
quantization_method_ = quantization_method;
op_set_ = op_set;
}
InsertQuantizedFunctionsPass(const InsertQuantizedFunctionsPass& other) {
quantization_method_ = other.quantization_method_;
op_set_ = other.op_set_;
}
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-insert-quantized-functions";
}
StringRef getDescription() const final {
// This is a brief description of the pass.
return "Insert quantized functions into the module";
}
void getDependentDialects(DialectRegistry& registry) const override {
registry.insert<TF::TensorFlowDialect, func::FuncDialect, ub::UBDialect>();
}
private:
void runOnOperation() override;
// Returns the function library for the given quantization method and opset
// pair.
llvm::StringRef GetFunctionLibrary(QuantMethod quantization_method,
OpSet op_set);
Option<QuantMethod> quantization_method_{
*this, "quantization-method",
llvm::cl::init(tensorflow::quantization::QuantizationMethod::
METHOD_STATIC_RANGE_INT8),
llvm::cl::desc("Choose quantization method."),
llvm::cl::values(
clEnumValN(tensorflow::quantization::QuantizationMethod::
METHOD_STATIC_RANGE_INT8,
"ptq", "Post-training static-range quantization"),
clEnumValN(tensorflow::quantization::QuantizationMethod::
METHOD_DYNAMIC_RANGE_INT8,
"drq", "Post-training dynamic-range quantizaiton"),
clEnumValN(tensorflow::quantization::QuantizationMethod::
METHOD_STATIC_RANGE_WEIGHT_ONLY_INT8,
"weight_only", "Post-training weight_only quantizaiton"))};
Option<OpSet> op_set_{
*this, "target-opset", llvm::cl::init(OpSet::TF),
llvm::cl::desc("Choose target opset."),
llvm::cl::values(
clEnumValN(OpSet::TF, "TF",
"Uses TF ops that mimic quantization behavior"),
clEnumValN(OpSet::XLA, "XLA", "Uses TF XLA ops"),
clEnumValN(OpSet::UNIFORM_QUANTIZED, "UNIFORM_QUANTIZED",
"Uses TF Uniform Quantized ops"))};
};
llvm::StringRef InsertQuantizedFunctionsPass::GetFunctionLibrary(
QuantMethod quantization_method, OpSet op_set) {
absl::flat_hash_map<OpSet, llvm::StringRef> function_library_map;
if (quantization_method ==
tensorflow::quantization::QuantizationMethod::METHOD_DYNAMIC_RANGE_INT8) {
function_library_map = {
{OpSet::TF, kQuantizedFunctionLibraryInMLIR_TF_DRQ},
{OpSet::UNIFORM_QUANTIZED,
kQuantizedFunctionLibraryInMLIR_UNIFORM_QUANTIZED_DRQ},
{OpSet::XLA, kQuantizedFunctionLibraryInMLIR_TF_DRQ}};
} else if (quantization_method ==
tensorflow::quantization::QuantizationMethod::
METHOD_STATIC_RANGE_WEIGHT_ONLY_INT8) {
// Uniform quantized opset is not supported for weight-only as inputs for
// weight quantization are floats. And only dequantize_i8 is used from the
// quantized function library.
function_library_map = {
{OpSet::TF, kQuantizedFunctionLibraryInMLIR},
{OpSet::XLA, kQuantizedFunctionLibraryInMLIR_XLA_WEIGHT_ONLY}};
} else {
function_library_map = {{OpSet::TF, kQuantizedFunctionLibraryInMLIR},
{OpSet::UNIFORM_QUANTIZED,
kQuantizedFunctionLibraryInMLIR_UNIFORM_QUANTIZED},
{OpSet::XLA, kQuantizedFunctionLibraryInMLIR}};
}
auto it = function_library_map.find(op_set);
if (it != function_library_map.end()) {
return it->second;
}
return llvm::StringRef();
}
static PassRegistration<InsertQuantizedFunctionsPass> pass;
void InsertQuantizedFunctionsPass::runOnOperation() {
ModuleOp module = getOperation();
SymbolTable symbol_table(module);
std::unique_ptr<llvm::MemoryBuffer> mem_buffer;
llvm::StringRef quantized_function_library =
GetFunctionLibrary(quantization_method_, op_set_);
if (quantized_function_library.empty()) {
emitError(module.getLoc())
<< "Failed to get function library for the opset.";
signalPassFailure();
return;
}
mem_buffer =
llvm::MemoryBuffer::getMemBuffer(quantized_function_library,
/*BufferName=*/"",
/*RequiresNullTerminator=*/false);
llvm::SourceMgr source_mgr;
source_mgr.AddNewSourceBuffer(std::move(mem_buffer), llvm::SMLoc());
OwningOpRef<ModuleOp> module_ref =
parseSourceFile<ModuleOp>(source_mgr, module.getContext());
// Inline and optimize loaded functions.
MLIRContext* context = &getContext();
PassManager pm(context);
pm.addPass(createInlinerPass());
pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
pm.addNestedPass<func::FuncOp>(createCSEPass());
StatusScopedDiagnosticHandler diagnostic_handler(context);
if (failed(pm.run(*module_ref))) {
emitError(module.getLoc()) << "failed to apply the optimization: "
<< diagnostic_handler.ConsumeStatus().message();
signalPassFailure();
return;
}
// Copy all functions used by this signature to the final MLIR module.
for (func::FuncOp func : module_ref->getOps<func::FuncOp>()) {
// Do nothing if the function already exists.
if (symbol_table.lookup(func.getSymName()) != nullptr) continue;
// Set the function to private and insert to the module.
func::FuncOp new_func = func.clone();
new_func.setPrivate();
symbol_table.insert(new_func);
// For consistency, we require all quantized composite function to have
// the "tf_quant.quantized_ops" attribute.
if (!new_func.getSymName().starts_with("quantized_")) continue;
if (!new_func->hasAttrOfType<ArrayAttr>("tf_quant.quantized_ops")) {
new_func->emitError() << "Missing \"tf_quant.quantized_ops\" "
"attribute in the quantized composite function.";
signalPassFailure();
}
}
}
} // namespace
// Creates an instance of the pass for inserting quantized functions.
std::unique_ptr<OperationPass<ModuleOp>> CreateInsertQuantizedFunctionsPass(
QuantMethod quantization_method, OpSet target_opset) {
return std::make_unique<InsertQuantizedFunctionsPass>(quantization_method,
target_opset);
}
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/passes/insert_restore_op.cc
+227
View File
@@ -0,0 +1,227 @@
/* 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 <cstdint>
#include <memory>
#include <string>
#include <vector>
#include "absl/log/log.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/StringRef.h"
#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 "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Location.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/quantization/tensorflow/passes/constants.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
namespace mlir {
namespace quant {
namespace {
using ::mlir::tf_saved_model::GetInitializerFunction;
using ::mlir::tf_saved_model::kTfSavedModelIndexPathAttr;
using ::mlir::tf_saved_model::kTfSavedModelInitializerRestoreType;
// This pass creates a RestoreV2 op in the initializer function with
// type "restore_op" that initializes variables from checkpoint. It finds
// tf.AssignVariableOp(tf.VarHandleOp, tf.Const) patterns in the initializer
// function and replaces tf.Consts with the results of RestoreV2.
class InsertRestoreOpPass
: public PassWrapper<InsertRestoreOpPass, OperationPass<ModuleOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(InsertRestoreOpPass)
explicit InsertRestoreOpPass() = default;
// The argument used to refer to the pass in the textual format (e.g. on the
// commandline).
StringRef getArgument() const final { return "quant-insert-restore-op"; }
StringRef getDescription() const final {
return "Creates RestoreV2 op to initialize the variables in the "
"initializer function (`tf_saved_model.initializer_type == "
"'restore_op'`). Replaces each occurrence of "
"`tf.AssignVariableOp(tf.VarHandleOp, tf.Const)` patterns with "
"`tf.AssignVariableOp(tf.VarHandleOp, restore_op_output#N)`, where "
"`restore_op_output#N` is the Nth output of the newly created "
"RestoreV2Op.";
}
void runOnOperation() override;
};
// Finds `tf.AssignVariableOp(tf.VarHandleOp, tf.Const)` patterns and returns
// the `tf.VarHandleOp`s that are initialized by these `tf.AssignVariableOp`s.
std::vector<TF::VarHandleOp> CollectVariableOps(
func::FuncOp session_init_func) {
std::vector<TF::VarHandleOp> var_handle_ops{};
for (auto assign_variable_op : llvm::make_early_inc_range(
session_init_func.getOps<TF::AssignVariableOp>())) {
Value resource_operand = assign_variable_op.getOperand(0);
Value assigned_value_operand = assign_variable_op.getOperand(1);
if (auto var_handle_op =
dyn_cast<TF::VarHandleOp>(resource_operand.getDefiningOp());
var_handle_op &&
isa<TF::ConstOp>(assigned_value_operand.getDefiningOp())) {
var_handle_ops.emplace_back(var_handle_op);
}
}
return var_handle_ops;
}
// Creates a `ConstOp` of 1-dimensional TF::StringType out of `str_values`.
TF::ConstOp Create1DStringConst(const ArrayRef<std::string> str_values,
const Location loc, OpBuilder& builder) {
const auto tensor_type =
RankedTensorType::get(/*shape=*/{static_cast<int64_t>(str_values.size())},
/*elementType=*/builder.getType<TF::StringType>());
return TF::ConstOp::create(
builder, loc,
DenseStringElementsAttr::get(
tensor_type,
SmallVector<StringRef>(str_values.begin(), str_values.end())));
}
// Creates a new argument for `func_op` that accepts a string tensor containing
// the checkpoint file's prefix.
BlockArgument InsertFilePrefixArgument(func::FuncOp func_op,
OpBuilder& builder) {
const auto filename_op_type = RankedTensorType::get(
/*shape=*/{}, /*elementType=*/builder.getType<TF::StringType>());
const auto file_prefix_attr = builder.getStringAttr(kTfFilePrefix);
const auto arg_attrs = builder.getDictionaryAttr({builder.getNamedAttr(
kTfSavedModelIndexPathAttr, builder.getArrayAttr({file_prefix_attr}))});
const int insert_idx = func_op.getNumArguments();
(void)func_op.insertArgument(insert_idx, /*argType=*/filename_op_type,
arg_attrs, NameLoc::get(file_prefix_attr));
return func_op.getArgument(insert_idx);
}
// Creates a 1D string array constant for "tensor_names" input of `RestoreV2`
// op. The `ConstOp` will be created at `builder`'s current insertion point.
TF::ConstOp CreateTensorNamesConst(const ArrayRef<std::string> tensor_names,
OpBuilder& builder) {
const auto loc = NameLoc::get(builder.getStringAttr("tensor_names"));
return Create1DStringConst(tensor_names, loc, builder);
}
// Creates a 1D string array constant for "shape_and_slices" input of
// `RestoreV2` op. The `ConstOp` will be created at `builder`'s current
// insertion point. It will be filled with `size` empty strings.
TF::ConstOp CreateShapeAndSlicesConst(const int size, OpBuilder& builder) {
const SmallVector<std::string> shape_and_slices_values(size, /*Value=*/"");
const auto loc = NameLoc::get(builder.getStringAttr("shape_and_slices"));
return Create1DStringConst(shape_and_slices_values, loc, builder);
}
// Creates a `tf.RestoreV2Op` that loads the variable values from the checkpoint
// file. The loaded tensors will be used to initialize `tf.VarHandleOp`s via
// `tf.AssignVariableOp`s.
void CreateRestoreV2Op(std::vector<TF::VarHandleOp>& target_var_handle_ops,
func::FuncOp session_init_func) {
SmallVector<Type> tensor_types{};
SmallVector<std::string> tensor_names{};
for (auto var_handle_op : target_var_handle_ops) {
tensor_names.emplace_back(var_handle_op.getSharedName().str());
// Location must be set to the same name as the shared name. The Location is
// later tranlated to the op's name when exported to `GraphDef`. This is
// required to find the correct variable name to restore when it is
// imported back to MLIR. When importing the graph to MLIR, the name of the
// op is used to retrieve the tensor values of each variable. See
// `InitializeVariablesInSessionInitializer` for further details.
const auto loc = NameLoc::get(StringAttr::get(
var_handle_op.getContext(), var_handle_op.getSharedName()));
var_handle_op->setLoc(loc);
// Ex) If VarHandleOp's type is tensor<!tf_type.resource<tensor<1xf32>>>,
// then tensor<1xf32> is the subtype.
tensor_types.emplace_back(var_handle_op.resource_subtype());
}
auto builder =
OpBuilder::atBlockTerminator(&session_init_func.getBody().front());
const BlockArgument filename_arg =
InsertFilePrefixArgument(session_init_func, builder);
TF::ConstOp tensor_names_const =
CreateTensorNamesConst(tensor_names, builder);
TF::ConstOp shape_and_slices_const =
CreateShapeAndSlicesConst(tensor_names.size(), builder);
auto restore_op = TF::RestoreV2Op::create(
builder, session_init_func.getLoc(),
/*tensors=*/tensor_types,
/*prefix=*/filename_arg, tensor_names_const, shape_and_slices_const);
for (auto [idx, restore_result] : llvm::enumerate(restore_op.getResults())) {
TF::AssignVariableOp::create(builder, restore_op.getLoc(),
target_var_handle_ops[idx], restore_result);
}
}
// TODO(b/261813194): Do not create a new RestoreV2 op when a RestoreV2 op
// already exists.
void InsertRestoreOpPass::runOnOperation() {
ModuleOp module_op = getOperation();
func::FuncOp session_init_func = GetInitializerFunction(
module_op, /*initializer_type=*/kTfSavedModelInitializerRestoreType);
if (!session_init_func) {
LOG(INFO) << "No session initializer function with type 'restore_op'. "
"RestoreV2 op will not be created.";
return;
}
std::vector<TF::VarHandleOp> target_var_handle_ops =
CollectVariableOps(session_init_func);
if (target_var_handle_ops.empty()) {
LOG(INFO) << "There are no VarHandleOps to restore. RestoreV2 op will not "
"be created.";
return;
}
CreateRestoreV2Op(target_var_handle_ops, session_init_func);
}
static PassRegistration<InsertRestoreOpPass> pass{};
} // namespace
std::unique_ptr<OperationPass<ModuleOp>> CreateInsertRestoreOpPass() {
return std::make_unique<InsertRestoreOpPass>();
}
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/passes/insert_save_op.cc
+255
View File
@@ -0,0 +1,255 @@
/* 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.
==============================================================================*/
#include <cstdint>
#include <memory>
#include <string>
#include "absl/log/log.h"
#include "llvm/ADT/STLExtras.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Block.h" // from @llvm-project
#include "mlir/IR/Builders.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/IRMapping.h" // from @llvm-project
#include "mlir/IR/Location.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/quantization/tensorflow/passes/constants.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
namespace mlir {
namespace quant {
namespace {
using ::mlir::tf_saved_model::GetInitializerFunction;
using ::mlir::tf_saved_model::kTfSavedModelInitializerRestoreType;
constexpr StringRef kTfQuantSaveV2OpName = "tf_quant__save_save_v2";
constexpr StringRef kTfQuantSaveReturnOpName = "tf_quant__save_return";
// A pass that creates a new function that wraps the newly created SaveV2 op.
// The new function's name is "tf_quant__save". The function accepts a single
// string tensor as argument, which specifies the path to the checkpoint to
// which the variable's tensor values are saved. It finds
// `tf.AssignVariableOp(tf.VarHandleOp, tf.Const)` pattern in the initializer
// function of type "restore_op" to identify the VarHandleOps that should be
// saved using the SaveV2 op.
class InsertSaveOpPass
: public PassWrapper<InsertSaveOpPass, OperationPass<ModuleOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(InsertSaveOpPass)
explicit InsertSaveOpPass() = default;
// The argument used to refer to the pass in the textual format (e.g. on the
// commandline).
StringRef getArgument() const final { return "quant-insert-save-op"; }
StringRef getDescription() const final {
return "Inserts a new function that wraps a SaveV2 op. The SaveV2 op saves "
"the values of the VarHandleOps that are found in the initializer "
"function of 'restore_op' type.";
}
void runOnOperation() override;
};
// Finds `tf.AssignVariableOp(tf.VarHandleOp, tf.Const)` patterns and removes
// `tf.AssignVariableOp`s and `tf.Const`s. Collects and returns the
// `tf.VarHandleOp`s that are initialized by these `tf.AssignVariableOp`s.
SmallVector<TF::VarHandleOp> CollectVariableOps(
func::FuncOp session_init_func) {
SmallVector<TF::VarHandleOp> var_handle_ops{};
for (auto assign_variable_op : llvm::make_early_inc_range(
session_init_func.getOps<TF::AssignVariableOp>())) {
Value resource_operand = assign_variable_op.getOperand(0);
auto var_handle_op =
dyn_cast<TF::VarHandleOp>(resource_operand.getDefiningOp());
if (!var_handle_op) continue;
Value assigned_value_operand = assign_variable_op.getOperand(1);
auto const_op =
dyn_cast<TF::ConstOp>(assigned_value_operand.getDefiningOp());
if (!const_op) continue;
var_handle_ops.emplace_back(var_handle_op);
}
return var_handle_ops;
}
// Creates a `ConstOp` of 1-dimensional TF::StringType out of `str_values`.
TF::ConstOp Create1DStringConst(const ArrayRef<std::string> str_values,
const Location loc, OpBuilder& builder) {
const auto tensor_type =
RankedTensorType::get(/*shape=*/{static_cast<int64_t>(str_values.size())},
/*elementType=*/builder.getType<TF::StringType>());
return TF::ConstOp::create(
builder, loc,
DenseStringElementsAttr::get(
tensor_type,
SmallVector<StringRef>(str_values.begin(), str_values.end())));
}
// Creates a 1D string array constant for "tensor_names" input of `RestoreV2`
// op. The `ConstOp` will be created at `builder`'s current insertion point.
TF::ConstOp CreateTensorNamesConst(const ArrayRef<std::string> tensor_names,
OpBuilder& builder) {
const auto loc = NameLoc::get(builder.getStringAttr("tensor_names"));
return Create1DStringConst(tensor_names, loc, builder);
}
// Creates a 1D string array constant for "shape_and_slices" input of
// `RestoreV2` op. The `ConstOp` will be created at `builder`'s current
// insertion point. It will be filled with `size` empty strings.
TF::ConstOp CreateShapeAndSlicesConst(const int size, OpBuilder& builder) {
const SmallVector<std::string> shape_and_slices_values(size, /*Value=*/"");
const auto loc = NameLoc::get(builder.getStringAttr("shape_and_slices"));
return Create1DStringConst(shape_and_slices_values, loc, builder);
}
// Returns cloned `VarHandleOp`s. Assumes `save_func`'s body is empty.
SmallVector<TF::VarHandleOp> CloneVarHandleOpsIntoSaveFunc(
func::FuncOp save_func, const ArrayRef<TF::VarHandleOp> var_handle_ops) {
Block& save_op_block = save_func.getBody().front();
IRMapping mapper{};
SmallVector<TF::VarHandleOp> cloned_var_handle_ops = {};
for (auto var_handle_op : var_handle_ops) {
Operation* cloned_var_handle_op = var_handle_op->clone(mapper);
save_op_block.push_back(cloned_var_handle_op);
cloned_var_handle_ops.push_back(
cast<TF::VarHandleOp>(cloned_var_handle_op));
}
return cloned_var_handle_ops;
}
// Creates and returns a `TF::SaveV2Op` for the `var_handle_ops`. For each
// VarHandleOp in `var_handle_ops` the tensor value is read via
// `TF::ReadVariableOp` and provided as arguments to the newly created SaveV2
// op.
TF::SaveV2Op CreateSaveV2Op(func::FuncOp save_func,
const ArrayRef<TF::VarHandleOp> var_handle_ops) {
auto builder = OpBuilder::atBlockEnd(&save_func.getBody().front());
SmallVector<std::string> tensor_names = {};
SmallVector<Value> tensor_values = {};
for (auto var_handle_op : var_handle_ops) {
tensor_names.emplace_back(var_handle_op.getSharedName().str());
auto read_var_op = TF::ReadVariableOp::create(
builder, var_handle_op.getLoc(), var_handle_op.resource_subtype(),
var_handle_op);
tensor_values.emplace_back(read_var_op.getResult());
}
TF::ConstOp tensor_names_const =
CreateTensorNamesConst(tensor_names, builder);
TF::ConstOp shape_and_slices_const =
CreateShapeAndSlicesConst(tensor_names.size(), builder);
BlockArgument filename_arg = save_func.getArgument(0);
return TF::SaveV2Op::create(
builder, NameLoc::get(builder.getStringAttr(kTfQuantSaveV2OpName)),
/*prefix=*/filename_arg, tensor_names_const, shape_and_slices_const,
/*tensors=*/tensor_values);
}
// Creates and returns a new `FuncOp` named "tf_quant__save". The resulting
// `FuncOp`'s body has no ops.
func::FuncOp CreateEmptySaveFunc(ModuleOp module_op) {
OpBuilder builder(module_op);
builder.setInsertionPointToEnd(&module_op.getBodyRegion().front());
auto filename_input_type = RankedTensorType::get(
/*shape=*/{}, /*elementType=*/builder.getType<TF::StringType>());
FunctionType func_type = builder.getFunctionType(
/*inputs=*/{filename_input_type}, /*results=*/{});
auto save_func = func::FuncOp::create(
builder, NameLoc::get(builder.getStringAttr(kTfQuantSaveFuncName)),
/*sym_name=*/kTfQuantSaveFuncName, func_type);
save_func.addEntryBlock();
save_func.setPrivate();
return save_func;
}
// Creates a save function that contains the `TF::SaveV2Op` for the variables in
// `var_handle_ops`. The `var_handle_ops` are cloned into the new function and
// provides the tensor values to be saved. The new function is a private
// function and has one argument for the file prefix (the directory to the
// checkpoint).
void CreateSaveFunc(ModuleOp module_op,
const ArrayRef<TF::VarHandleOp> var_handle_ops) {
func::FuncOp save_func = CreateEmptySaveFunc(module_op);
const SmallVector<TF::VarHandleOp> cloned_var_handle_ops =
CloneVarHandleOpsIntoSaveFunc(save_func, var_handle_ops);
CreateSaveV2Op(save_func, cloned_var_handle_ops);
// Create a "func.return".
auto builder = OpBuilder::atBlockEnd(&save_func.getBody().front());
func::ReturnOp::create(
builder, NameLoc::get(builder.getStringAttr(kTfQuantSaveReturnOpName)));
}
void InsertSaveOpPass::runOnOperation() {
ModuleOp module_op = getOperation();
func::FuncOp session_init_func = GetInitializerFunction(
module_op, /*initializer_type=*/kTfSavedModelInitializerRestoreType);
if (!session_init_func) {
LOG(INFO) << "No session initializer function with type 'restore_op'. "
"SaveV2 op will not be created.";
return;
}
SmallVector<TF::VarHandleOp> target_var_handle_ops =
CollectVariableOps(session_init_func);
if (target_var_handle_ops.empty()) {
LOG(INFO) << "There are no VarHandleOps to save. SaveV2 op will not "
"be created.";
return;
}
CreateSaveFunc(module_op, target_var_handle_ops);
}
static PassRegistration<InsertSaveOpPass> pass{};
} // namespace
std::unique_ptr<OperationPass<ModuleOp>> CreateInsertSaveOpPass() {
return std::make_unique<InsertSaveOpPass>();
}
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/passes/lift_hashtable_ops_as_args.cc
+225
View File
@@ -0,0 +1,225 @@
/* 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.
==============================================================================*/
#include <memory>
#include <utility>
#include "absl/strings/str_cat.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/StringMap.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Block.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/IRMapping.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/SymbolTable.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/LogicalResult.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/constants.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/manipulate_model_attr.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/compiler/mlir/tensorflow/translate/import_model.h"
namespace mlir {
namespace quant {
namespace {
constexpr StringRef kSharedNameAttr = "shared_name";
class LiftHashTableOpsAsArgsPass
: public PassWrapper<LiftHashTableOpsAsArgsPass, OperationPass<ModuleOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LiftHashTableOpsAsArgsPass)
explicit LiftHashTableOpsAsArgsPass() = default;
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-lift-hashtable-ops-as-args";
}
StringRef getDescription() const final {
return "Lifts HashTable ops as function arguments.";
}
void runOnOperation() override;
};
// Checks if the given op is a Hashtable op.
bool IsHashTableOp(Operation* op) {
return llvm::isa<TF::HashTableOp, TF::HashTableV2Op,
TF::MutableHashTableV2Op>(op);
}
// Checks if the function is the main or initializer function.
bool IsMainOrInitializerFunction(ModuleOp module, func::FuncOp func) {
if (func.getSymName() ==
llvm::StringRef(tensorflow::kImportModelDefaultGraphFuncName) ||
func.getSymName() == kTfQuantSaveFuncName) {
return true;
}
for (func::FuncOp init_func :
tf_saved_model::GetInitializerFunctions(module)) {
if (func.getSymName() == init_func.getSymName()) {
return true;
}
}
return false;
}
// Checks if the function is only used by supported ops. Returns false when the
// function has no uses. Currently, only PartitionedCall is supported.
// TODO(b/284222309): Support lifting for functions called by control flow.
bool UsedBySupportedOps(ModuleOp module, func::FuncOp func) {
auto function_uses =
SymbolTable::getSymbolUses(func, &module.getBodyRegion());
if (!function_uses.has_value()) return false;
for (auto& function_use : function_uses.value()) {
if (!llvm::isa<TF::PartitionedCallOp, TF::StatefulPartitionedCallOp>(
function_use.getUser())) {
return false;
}
}
return true;
}
// Returns the `shared_name` attribute value if exists. If not, returns an
// empty string.
StringRef GetSharedName(Operation* op) {
if (!op->hasAttrOfType<StringAttr>(kSharedNameAttr)) return "";
return op->getAttrOfType<StringAttr>(kSharedNameAttr).getValue();
}
// Checks if the HashTable is initialized. This function assumes that the
// HashTable is initialized if it appears in the initializer since it can't
// check the actual value.
bool IsResourceInitialized(ModuleOp module_op, Operation* hash_table) {
StringRef shared_name = GetSharedName(hash_table);
if (shared_name.empty()) return false;
for (func::FuncOp init_func_op :
tf_saved_model::GetInitializerFunctions(module_op)) {
for (Operation& op : init_func_op.getBody().getOps()) {
StringRef other_shared_name = GetSharedName(&op);
if (IsHashTableOp(&op) && other_shared_name == shared_name) {
return true;
}
}
}
return false;
}
// Lifts HashTable ops in the target function as function arguments and returns
// the lifted ops. These ops will then be added to the caller function and
// passed to the target function.
LogicalResult LiftHashTableOpsToArguments(ModuleOp module_op,
func::FuncOp target_func) {
if (!llvm::hasSingleElement(target_func)) return success();
if (!UsedBySupportedOps(module_op, target_func)) return success();
if (IsMainOrInitializerFunction(module_op, target_func)) return success();
llvm::StringMap<int> shared_name_to_arg_idx;
llvm::SmallVector<std::pair<Operation*, int>> lifted_op_and_arg_idx;
Block& block = target_func.front();
auto func_type = target_func.getFunctionType();
for (Operation& op : block.without_terminator()) {
StringRef shared_name = GetSharedName(&op);
if (shared_name.empty() || !IsHashTableOp(&op)) continue;
if (!IsResourceInitialized(module_op, &op)) continue;
auto it =
shared_name_to_arg_idx.insert({shared_name, block.getNumArguments()});
if (it.second) {
auto resource_type = op.getResult(0).getType();
op.getResult(0).replaceAllUsesWith(
block.addArgument(resource_type, op.getLoc()));
AddEntryFunctionInput(
absl::StrCat("hash_table_", it.first->getValue(), ":0"), target_func);
// Avoid deleting the op here, clone it to the caller function first.
lifted_op_and_arg_idx.emplace_back(&op, it.first->getValue());
} else {
op.getResult(0).replaceAllUsesWith(
block.getArgument(it.first->getValue()));
op.erase();
}
}
if (lifted_op_and_arg_idx.empty()) return success();
// Update the function signature as well as its uses.
target_func.setType(FunctionType::get(target_func.getContext(),
block.getArgumentTypes(),
func_type.getResults()));
IRMapping mapping;
OpBuilder builder(module_op);
OpBuilder::InsertionGuard g(builder);
// The function has been checked to have at least one use.
auto function_uses =
SymbolTable::getSymbolUses(target_func, &module_op.getBodyRegion());
for (auto& function_use : function_uses.value()) {
auto call_op = function_use.getUser();
auto caller_func = call_op->getParentOfType<func::FuncOp>();
if (!caller_func) return failure();
builder.setInsertionPoint(call_op);
for (auto [lifted_op, arg_idx] : lifted_op_and_arg_idx) {
auto new_op = builder.clone(*lifted_op, mapping);
call_op->insertOperands(arg_idx, new_op->getResult(0));
}
// Try to lift recursively until the main function.
if (failed(LiftHashTableOpsToArguments(module_op, caller_func))) {
return failure();
}
}
// Erase the lifted operations explicitly.
for (auto [lifted_op, arg_idx] : lifted_op_and_arg_idx) {
lifted_op->erase();
}
return success();
}
void LiftHashTableOpsAsArgsPass::runOnOperation() {
auto module_op = getOperation();
for (auto func_op : module_op.getOps<func::FuncOp>()) {
if (failed(LiftHashTableOpsToArguments(module_op, func_op))) {
signalPassFailure();
return;
}
}
}
static PassRegistration<LiftHashTableOpsAsArgsPass> pass;
} // namespace
std::unique_ptr<OperationPass<ModuleOp>> CreateLiftHashTableOpsAsArgsPass() {
return std::make_unique<LiftHashTableOpsAsArgsPass>();
}
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/passes/lift_quantizable_spots_as_functions.cc
+417
View File
@@ -0,0 +1,417 @@
/* 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 <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/raw_ostream.h"
#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 "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/BuiltinTypeInterfaces.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/IR/SymbolTable.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassRegistry.h" // from @llvm-project
#include "mlir/Rewrite/FrozenRewritePatternSet.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
#include "re2/re2.h"
#include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.h"
#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"
#include "tensorflow/compiler/mlir/quantization/common/lift_as_function_call.h" // IWYU pragma: keep
#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/cc/quantization_unit_loc.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/ops/tf_op_quant_spec.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
namespace mlir {
namespace quant {
namespace {
using QuantizationUnit =
::tensorflow::quantization::UnitWiseQuantizationSpec::QuantizationUnit;
using ::tensorflow::quantization::OpSet;
using ::tensorflow::quantization::QuantizationComponentSpec;
using ::tensorflow::quantization::QuantizationMethod;
using ::tensorflow::quantization::QuantizationOptions;
using ::tensorflow::quantization::UnitWiseQuantizationSpec;
class LiftQuantizableSpotsAsFunctionsPass
: public PassWrapper<LiftQuantizableSpotsAsFunctionsPass,
OperationPass<ModuleOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
LiftQuantizableSpotsAsFunctionsPass)
LiftQuantizableSpotsAsFunctionsPass() : test_mode_(true) {
initializeForTest();
}
explicit LiftQuantizableSpotsAsFunctionsPass(
const QuantizationOptions& quant_options)
: quant_options_(quant_options), test_mode_(false) {}
LiftQuantizableSpotsAsFunctionsPass(
const LiftQuantizableSpotsAsFunctionsPass& other) {
quant_options_ = other.quant_options_;
test_mode_ = other.test_mode_;
op_set_ = other.op_set_;
initializeForTest();
}
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-lift-quantizable-spots-as-functions";
}
StringRef getDescription() const final {
// This is a brief description of the pass.
return "Replace quantization candidates with composite functions into the "
"module";
}
void getDependentDialects(DialectRegistry& registry) const override {
registry.insert<TF::TensorFlowDialect>();
}
void runOnOperation() override;
private:
QuantizationOptions quant_options_;
bool test_mode_;
Option<OpSet> op_set_{
*this, "target-opset", llvm::cl::init(OpSet::TF),
llvm::cl::desc("Choose target opset."),
llvm::cl::values(
clEnumValN(OpSet::TF, "TF",
"Uses TF ops that mimic quantization behavior"),
clEnumValN(OpSet::XLA, "XLA", "Uses TF XLA ops"),
clEnumValN(OpSet::UNIFORM_QUANTIZED, "UNIFORM_QUANTIZED",
"Uses TF Uniform Quantized ops"))};
// Initialize for tests.
void initializeForTest() {
if (!test_mode_) return;
op_set_.setCallback([this](const OpSet& new_op_set) {
quant_options_.set_op_set(new_op_set);
});
// Set the test quantization method to static-range.
if (quant_options_.quantization_method().preset_method() ==
QuantizationMethod::METHOD_UNSPECIFIED) {
quant_options_.mutable_quantization_method()->set_preset_method(
QuantizationMethod::METHOD_STATIC_RANGE_INT8);
}
if (quant_options_.quantization_method()
.quantization_component_specs()
.empty()) {
auto add_new_spec =
[this](QuantizationComponentSpec::QuantizationComponent component,
QuantizationComponentSpec::TensorType type) {
QuantizationComponentSpec* new_spec =
quant_options_.mutable_quantization_method()
->add_quantization_component_specs();
new_spec->set_quantization_component(component);
new_spec->set_tensor_type(type);
};
add_new_spec(QuantizationComponentSpec::COMPONENT_ACTIVATION,
QuantizationComponentSpec::TENSORTYPE_INT_8);
add_new_spec(QuantizationComponentSpec::COMPONENT_WEIGHT,
QuantizationComponentSpec::TENSORTYPE_INT_8);
add_new_spec(QuantizationComponentSpec::COMPONENT_BIAS,
QuantizationComponentSpec::TENSORTYPE_INT_32);
}
if (quant_options_.unit_wise_quantization_specs().empty()) {
// Opt-out a node named `test_opt_out`.
UnitWiseQuantizationSpec* new_spec =
quant_options_.add_unit_wise_quantization_specs();
QuantizationUnit* new_unit = new_spec->add_unit();
new_unit->set_node_name("test_opt_out");
new_spec->mutable_quantization_method()->set_preset_method(
QuantizationMethod::METHOD_NO_QUANTIZE);
}
}
};
class CheckQuantizableOps
: public mlir::OpRewritePattern<TF::PartitionedCallOp> {
public:
explicit CheckQuantizableOps(MLIRContext* context,
const QuantizationOptions& quant_options)
: OpRewritePattern<TF::PartitionedCallOp>(context),
quant_options_(quant_options) {}
private:
LogicalResult matchAndRewrite(TF::PartitionedCallOp call_op,
PatternRewriter& rewriter) const override {
StringRef function_name =
mlir::cast<FlatSymbolRefAttr>(call_op.getFAttr()).getValue();
if (!function_name.starts_with("composite_") ||
!call_op->hasAttr(kQuantTraitAttrName)) {
return failure();
}
absl::Status check_status;
// TODO(b/270906404): Support weight-only gather for uniform quantized opset
// in PTQ mode
if (quant_options_.op_set() == OpSet::UNIFORM_QUANTIZED &&
function_name.contains("gather")) {
check_status.Update(absl::InternalError("Weight-only op is skipped."));
}
if (quant_options_.op_set() == OpSet::XLA) {
check_status.Update(checkQuantizableOpsForXla(call_op, function_name));
}
// Only the composite functions with f32 inputs are quantizable.
if (call_op.getResults().size() == 1 &&
!mlir::cast<ShapedType>(call_op->getResult(0).getType())
.getElementType()
.isF32()) {
check_status.Update(absl::InternalError(
"Composite functions for quantization should be f32 type."));
}
// The OK status means this op is quantizable. Return failure since the
// pattern doesn't rewrite anything yet.
if (check_status.ok()) return failure();
call_op->removeAttr(kQuantTraitAttrName);
removeAttrMapAttribute(call_op, function_name, check_status.message());
return success();
}
// Get the quantization method to apply to this composite function. If set,
// the unit-wise quantization method overrides the default one.
std::optional<QuantizationMethod> getUnitWiseQuantizationMethod(
TF::PartitionedCallOp call_op) const {
// If unit-wise quantization config is found, overwrite the default config.
auto quantization_unit = FindQuantizationUnitFromLoc(call_op.getLoc());
if (!quantization_unit.has_value()) return std::nullopt;
for (const auto& unit_config :
quant_options_.unit_wise_quantization_specs()) {
for (const auto& unit : unit_config.unit()) {
if (!unit.op_type().empty() &&
quantization_unit.value().op_type() != unit.op_type()) {
continue;
}
if (!unit.node_name().empty()) {
const RE2 node_name_regex(unit.node_name());
if (!RE2::FullMatch(quantization_unit.value().node_name(),
node_name_regex)) {
continue;
}
}
if (!unit.func_name().empty()) {
const RE2 func_name_regex(unit.func_name());
if (!RE2::FullMatch(quantization_unit.value().func_name(),
func_name_regex)) {
continue;
}
}
// Overrides the default quantization method.
return unit_config.quantization_method();
}
}
return std::nullopt;
}
absl::Status checkQuantizableOpsForXla(TF::PartitionedCallOp call_op,
StringRef function_name) const {
// Disable quantization for the DepthwiseConv since it has no benefits in
// the XLA opset.
if (function_name.contains("depthwise_conv2d")) {
return absl::InternalError(
"DepthwiseConv2D doesn't get any benefit of quantization in XLA.");
} else if (function_name.contains("conv2d")) {
// For Conv2D, the channel dimension must be static to calculate the
// feature group count.
if (!HasStaticShapeAtDims(call_op->getOperand(0), /*dims=*/3)) {
return absl::InternalError(
"The channel dimension of Conv2D is required to be static.");
}
} else if (function_name.contains("conv3d")) {
// For Conv3D, the channel dimension must be static to calculate the
// feature group count.
if (!HasStaticShapeAtDims(call_op->getOperand(0), /*dims=*/4)) {
return absl::InternalError(
"The channel dimension of Conv3D is required to be static.");
}
} else if (function_name.contains("batch_matmul")) {
// For BatchMatMul, the input must be ranked to determine the batch
// dimensions.
ShapedType shaped_type =
mlir::dyn_cast<ShapedType>(call_op->getOperand(0).getType());
if (!shaped_type || !shaped_type.hasRank()) {
return absl::InternalError("The input of BatchMatMul must have rank.");
}
} else if (function_name.contains("gather")) {
// This op is guaranteed to be a constant as ODS checks IsConstTensor.
// Check if the number of elements meets the requirement.
int64_t num_elements =
mlir::cast<ShapedType>(call_op.getOperand(0).getType())
.getNumElements();
if (num_elements < quant_options_.min_num_elements_for_weights()) {
return absl::InternalError(
"The params of Gather have fewer number of elements than "
"the `min_num_elements_for_weights`.");
}
}
// Disable quantization if the quantization method is NO_QUANTIZE.
QuantizationMethod quantization_method =
quant_options_.quantization_method();
if (quantization_method.quantization_component_specs().empty()) {
return absl::InternalError(
"The quantization method has been set to METHOD_NO_QUANTIZE.");
}
// The unit-wise quantization config should override the loser-grained
// quantization config, such as `enable_two_input_tensors`.
bool is_unitwise_quantization_enabled = false;
std::optional<QuantizationMethod> unit_wise_quantization_method =
getUnitWiseQuantizationMethod(call_op);
if (unit_wise_quantization_method.has_value()) {
if (unit_wise_quantization_method.value()
.quantization_component_specs()
.empty()) {
return absl::InternalError(
"The unit-wise quantization method has been set to "
"METHOD_NO_QUANTIZE.");
}
is_unitwise_quantization_enabled = true;
}
std::unique_ptr<OpQuantSpec> spec = GetTFOpQuantSpec(call_op);
for (auto iter : spec->coeff_op_quant_dim) {
Operation* preceding_op = call_op.getOperand(iter.first).getDefiningOp();
// The XLA opset only supports constant filter/weight at the moment.
bool is_weight_constant =
preceding_op && preceding_op->hasTrait<OpTrait::ConstantLike>();
// There might be q/dq ops after the filter/weight.
if (auto dq_op =
llvm::dyn_cast_or_null<mlir::quant::ir::DequantizeCastOp>(
preceding_op)) {
if (auto q_op = llvm::dyn_cast_or_null<mlir::quant::ir::QuantizeCastOp>(
dq_op.getArg().getDefiningOp())) {
Operation* q_op_input = q_op.getArg().getDefiningOp();
is_weight_constant =
q_op_input && q_op_input->hasTrait<OpTrait::ConstantLike>();
}
}
if (!is_weight_constant) {
if (!function_name.contains("matmul") &&
!function_name.contains("einsum")) {
return absl::InternalError(
"Non-constant weights are not supported at the moment,"
" except matmul and einsum.");
} else if (!quant_options_.enable_two_input_tensors() &&
!is_unitwise_quantization_enabled) {
return absl::InternalError(
"Quantization is disabled for this op due to the non-constant "
"weight. You can enable it by setting `enable_two_input_tensors` "
"to true or using unit-wise quantization config.");
}
}
}
return absl::OkStatus();
}
void removeAttrMapAttribute(TF::PartitionedCallOp call_op,
StringRef function_name,
StringRef error_message) const {
ModuleOp module = call_op->getParentOfType<ModuleOp>();
SymbolTable symbol_table(module);
mlir::func::FuncOp composite_func =
dyn_cast<func::FuncOp>(symbol_table.lookup(function_name));
if (!composite_func) return;
composite_func.walk([&](Operation* op) {
if (op->hasAttr(kAttrMapAttribute)) {
op->removeAttr(kAttrMapAttribute);
std::string log_message;
llvm::raw_string_ostream log_stream(log_message);
op->getLoc().print(log_stream);
log_stream << ": Quantization disabled on this op: ";
log_stream << error_message << "\n";
log_stream << "See the current operation:\n";
op->print(log_stream);
VLOG(2) << log_message;
}
});
}
const QuantizationOptions& quant_options_;
};
static PassRegistration<LiftQuantizableSpotsAsFunctionsPass> pass;
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/lift_quantizable_spots_as_functions.inc"
void LiftQuantizableSpotsAsFunctionsPass::runOnOperation() {
MLIRContext* ctx = &getContext();
RewritePatternSet patterns(ctx);
ModuleOp module = getOperation();
populateWithGenerated(patterns);
patterns.add<CheckQuantizableOps>(ctx, quant_options_);
FrozenRewritePatternSet frozen_patterns(std::move(patterns));
// Iterate over the sorted list of functions to keep the order deterministic.
for (func::FuncOp func : GetSortedFunctions(module)) {
if (failed(applyPatternsGreedily(func, frozen_patterns))) {
func.emitError() << "quant-lift-quantizable-spots-as-functions failed.";
signalPassFailure();
}
}
}
} // namespace
std::unique_ptr<OperationPass<ModuleOp>>
CreateLiftQuantizableSpotsAsFunctionsPass(
const QuantizationOptions& quant_options) {
return std::make_unique<LiftQuantizableSpotsAsFunctionsPass>(quant_options);
}
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/passes/lift_quantizable_spots_as_functions.td
+390
View File
@@ -0,0 +1,390 @@
/* 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 "mlir/IR/OpBase.td"
include "mlir/Dialect/Func/IR/FuncOps.td"
include "mlir/Dialect/Arith/IR/ArithOps.td"
include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"
include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.td"
include "tensorflow/compiler/mlir/quantization/common/lift_as_function_call.td"
//===----------------------------------------------------------------------===//
// Helper functions.
//===----------------------------------------------------------------------===//
class IsFusedOpEndsWith<string OpName> : AttrConstraint<
CPred<"!llvm::cast<ArrayAttr>($_self).empty() && "
"llvm::cast<ArrayAttr>($_self)[llvm::cast<ArrayAttr>($_self).size() - 1]."
"cast<::mlir::StringAttr>().str() == \"" # OpName # "\"">,
"Matching fused '" # OpName # "' op at the end">;
//===----------------------------------------------------------------------===//
// Pattern rules for lifting ops as functions
//===----------------------------------------------------------------------===//
def LiftConv : Pat<
(TF_Conv2DOp:$res $input, $filter, $strides, $use_cudnn_on_gpu, $padding,
$explicit_paddings, IsDataFormatNHWC:$data_format, $dilations),
(LiftAsTFPartitionedCall<"composite_conv2d_fn">
(ArgumentList $input, $filter),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"strides"> $strides),
(NamedAttr<"use_cudnn_on_gpu"> $use_cudnn_on_gpu),
(NamedAttr<"padding"> $padding),
(NamedAttr<"explicit_paddings"> $explicit_paddings),
(NamedAttr<"dilations"> $dilations))),
[(IsNotInLiftedFunc $res)], [], (addBenefit 1)>;
def LiftDepthwiseConv : Pat<
(TF_DepthwiseConv2dNativeOp:$res $input, $filter, $strides, $padding,
$explicit_paddings, IsDataFormatNHWC:$data_format, $dilations),
(LiftAsTFPartitionedCall<"composite_depthwise_conv2d_fn">
(ArgumentList $input, $filter),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"strides"> $strides),
(NamedAttr<"padding"> $padding),
(NamedAttr<"explicit_paddings"> $explicit_paddings),
(NamedAttr<"dilations"> $dilations))),
[(IsNotInLiftedFunc $res)], [], (addBenefit 1)>;
def LiftMatMul : Pat<
(TF_MatMulOp:$res $a, $b, $transpose_a, $transpose_b, $grad_a, $grad_b),
(LiftAsTFPartitionedCall<"composite_matmul_fn">
(ArgumentList $a, $b),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"transpose_a"> $transpose_a),
(NamedAttr<"transpose_b"> $transpose_b))),
[(IsNotInLiftedFunc $res)], [], (addBenefit 1)>;
def LiftConv3D : Pat<
(TF_Conv3DOp:$res $input, $filter, $strides, $padding,
IsDataFormatNDHWC:$data_format, $dilations),
(LiftAsTFPartitionedCall<"composite_conv3d_fn">
(ArgumentList $input, $filter),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"strides"> $strides),
(NamedAttr<"padding"> $padding),
(NamedAttr<"dilations"> $dilations))),
[(IsNotInLiftedFunc $res)], [], (addBenefit 1)>;
def LiftBatchMatMul : Pat<
(TF_BatchMatMulV2Op:$res $x, $y, $adj_x, $adj_y, $grad_x, $grad_y),
(LiftAsTFPartitionedCall<"composite_batch_matmul_fn">
(ArgumentList $x, $y),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"adj_x"> $adj_x),
(NamedAttr<"adj_y"> $adj_y))),
[(IsNotInLiftedFunc $res)], [], (addBenefit 1)>;
def LiftEinsum : Pat<
(TF_EinsumOp:$res $input, $equation),
(LiftAsTFPartitionedCall<"composite_einsum_fn">
(ArgumentList $input),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"equation"> $equation))),
[(IsNotInLiftedFunc $res),
(IsEinsumSupportedByXlaDotV2 $equation)
], [], (addBenefit 1)>;
//===----------------------------------------------------------------------===//
// Pattern rules for lifting ops with bias as functions
//===----------------------------------------------------------------------===//
def LiftDepthwiseConv2dNativeWithBias : Pat<
(TF_BiasAddOp:$res
(TF_DepthwiseConv2dNativeOp $input, $filter, $strides, $padding,
$explicit_paddings, IsDataFormatNHWC:$data_format, $dilations),
$bias, IsDataFormatNHWC:$bias_data_format),
(LiftAsTFPartitionedCall<"composite_depthwise_conv2d_with_bias_fn">
(ArgumentList $input, $filter, $bias),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"strides"> $strides),
(NamedAttr<"padding"> $padding),
(NamedAttr<"explicit_paddings"> $explicit_paddings),
(NamedAttr<"dilations"> $dilations))),
[(IsNotInLiftedFunc $res)], [], (addBenefit 5)>;
def LiftConv2dWithBias : Pat<
(TF_BiasAddOp:$res
(TF_Conv2DOp $input, $filter, $strides, $use_cudnn_on_gpu, $padding,
$explicit_paddings, IsDataFormatNHWC:$data_format, $dilations),
$bias, IsDataFormatNHWC:$bias_data_format),
(LiftAsTFPartitionedCall<"composite_conv2d_with_bias_fn">
(ArgumentList $input, $filter, $bias),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"strides"> $strides),
(NamedAttr<"use_cudnn_on_gpu"> $use_cudnn_on_gpu),
(NamedAttr<"padding"> $padding),
(NamedAttr<"explicit_paddings"> $explicit_paddings),
(NamedAttr<"dilations"> $dilations))),
[(IsNotInLiftedFunc $res)], [], (addBenefit 5)>;
def LiftMatmulWithBias : Pat<
(TF_BiasAddOp:$res
(TF_MatMulOp $a, $b, $transpose_a, $transpose_b, $grad_a, $grad_b),
$bias, IsDataFormatNHWC:$bias_data_format),
(LiftAsTFPartitionedCall<"composite_matmul_with_bias_fn">
(ArgumentList $a, $b, $bias),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"transpose_a"> $transpose_a),
(NamedAttr<"transpose_b"> $transpose_b))),
[(IsNotInLiftedFunc $res)], [], (addBenefit 5)>;
// TODO(b/278493977): Create generic implementation of lifting any fused op
// with any reshaping op
def LiftMatmulWithReshapeAndBias : Pat<
(TF_BiasAddOp:$res
(TF_ReshapeOp:$out
(TF_MatMulOp $a, $b, $transpose_a, $transpose_b, $grad_a, $grad_b),
$shape),
$bias, IsDataFormatNHWC:$bias_data_format),
(LiftAsTFPartitionedCall<"composite_matmul_with_reshape_and_bias_fn">
(ArgumentList $a, $b, $bias, $shape),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"transpose_a"> $transpose_a),
(NamedAttr<"transpose_b"> $transpose_b))),
[(IsNotInLiftedFunc $res)], [], (addBenefit 5)>;
def LiftConv3dWithBias : Pat<
(TF_BiasAddOp:$res
(TF_Conv3DOp $input, $filter, $strides, $padding,
IsDataFormatNDHWC:$data_format, $dilations),
$bias, $bias_data_format),
(LiftAsTFPartitionedCall<"composite_conv3d_with_bias_fn">
(ArgumentList $input, $filter, $bias),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"strides"> $strides),
(NamedAttr<"padding"> $padding),
(NamedAttr<"dilations"> $dilations))),
[(IsNotInLiftedFunc $res)], [], (addBenefit 5)>;
def LiftBatchMatMulWithBias : Pat<
(TF_BiasAddOp:$res
(TF_BatchMatMulV2Op $x, $y, $adj_x, $adj_y, $grad_x, $grad_y),
$bias, IsDataFormatNHWC:$bias_data_format),
(LiftAsTFPartitionedCall<"composite_batch_matmul_with_bias_fn">
(ArgumentList $x, $y, $bias),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"adj_x"> $adj_x),
(NamedAttr<"adj_y"> $adj_y))),
[(IsNotInLiftedFunc $res)], [], (addBenefit 5)>;
def LiftEinsumWithBias : Pat<
(TF_BiasAddOp:$res
(TF_EinsumOp $input, $equation),
$bias, IsDataFormatNHWC:$bias_data_format),
(LiftAsTFPartitionedCall<"composite_einsum_with_bias_fn">
(AppendToVector (ArgumentList $input), $bias),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"equation"> $equation))),
[(IsNotInLiftedFunc $res),
(IsEinsumSupportedByXlaDotV2 $equation)],
[], (addBenefit 5)>;
//===----------------------------------------------------------------------===//
// Pattern rules for lifting ops with bias and activation as functions
//===----------------------------------------------------------------------===//
multiclass LiftCompositeOpsWithActivation<Op ActivationOp, string ActivationName> {
def LiftConvWith#ActivationOp : Pat<
(ActivationOp:$res
(TF_Conv2DOp $input, $filter, $strides, $use_cudnn_on_gpu, $padding,
$explicit_paddings, IsDataFormatNHWC:$data_format, $dilations)),
(LiftAsTFPartitionedCall<"composite_conv2d_with_"# ActivationName #"_fn">
(ArgumentList $input, $filter),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"strides"> $strides),
(NamedAttr<"use_cudnn_on_gpu"> $use_cudnn_on_gpu),
(NamedAttr<"padding"> $padding),
(NamedAttr<"explicit_paddings"> $explicit_paddings),
(NamedAttr<"dilations"> $dilations))),
[(IsNotInLiftedFunc $res)], [], (addBenefit 10)>;
def LiftConv2dWithBiasAnd#LastFusedOp : Pat<
(ActivationOp:$res
(TF_BiasAddOp
(TF_Conv2DOp $input, $filter, $strides, $use_cudnn_on_gpu, $padding,
$explicit_paddings, IsDataFormatNHWC:$data_format, $dilations),
$bias, IsDataFormatNHWC:$bias_data_format)),
(LiftAsTFPartitionedCall<"composite_conv2d_with_bias_and_"# ActivationName #"_fn">
(ArgumentList $input, $filter, $bias),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"strides"> $strides),
(NamedAttr<"use_cudnn_on_gpu"> $use_cudnn_on_gpu),
(NamedAttr<"padding"> $padding),
(NamedAttr<"explicit_paddings"> $explicit_paddings),
(NamedAttr<"dilations"> $dilations))),
[(IsNotInLiftedFunc $res)], [], (addBenefit 10)>;
def LiftDepthwiseConv2dNativeWith#ActivationOp : Pat<
(ActivationOp:$res
(TF_DepthwiseConv2dNativeOp $input, $filter, $strides, $padding,
$explicit_paddings, IsDataFormatNHWC:$data_format, $dilations)),
(LiftAsTFPartitionedCall<"composite_depthwise_conv2d_with_"# ActivationName #"_fn">
(ArgumentList $input, $filter),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"strides"> $strides),
(NamedAttr<"padding"> $padding),
(NamedAttr<"explicit_paddings"> $explicit_paddings),
(NamedAttr<"dilations"> $dilations))),
[(IsNotInLiftedFunc $res)], [], (addBenefit 10)>;
def LiftDepthwiseConv2dNativeWithBiasAnd#LastFusedOp : Pat<
(ActivationOp:$res
(TF_BiasAddOp
(TF_DepthwiseConv2dNativeOp $input, $filter, $strides, $padding,
$explicit_paddings, IsDataFormatNHWC:$data_format, $dilations),
$bias, IsDataFormatNHWC:$bias_data_format)),
(LiftAsTFPartitionedCall<"composite_depthwise_conv2d_with_bias_and_"# ActivationName #"_fn">
(ArgumentList $input, $filter, $bias),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"strides"> $strides),
(NamedAttr<"padding"> $padding),
(NamedAttr<"explicit_paddings"> $explicit_paddings),
(NamedAttr<"dilations"> $dilations))),
[(IsNotInLiftedFunc $res)], [], (addBenefit 10)>;
def LiftMatmulWith#ActivationOp : Pat<
(ActivationOp:$res
(TF_MatMulOp $a, $b, $transpose_a, $transpose_b, $grad_a, $grad_b)),
(LiftAsTFPartitionedCall<"composite_matmul_with_"# ActivationName #"_fn">
(ArgumentList $a, $b),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"transpose_a"> $transpose_a),
(NamedAttr<"transpose_b"> $transpose_b))),
[(IsNotInLiftedFunc $res)], [], (addBenefit 10)>;
def LiftMatmulWithBiasAnd#LastFusedOp : Pat<
(ActivationOp:$res
(TF_BiasAddOp
(TF_MatMulOp $a, $b, $transpose_a, $transpose_b, $grad_a, $grad_b),
$bias, IsDataFormatNHWC:$bias_data_format)),
(LiftAsTFPartitionedCall<"composite_matmul_with_bias_and_"# ActivationName #"_fn">
(ArgumentList $a, $b, $bias),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"transpose_a"> $transpose_a),
(NamedAttr<"transpose_b"> $transpose_b))),
[(IsNotInLiftedFunc $res)], [], (addBenefit 10)>;
def LiftConv3dWith#ActivationOp : Pat<
(ActivationOp:$res
(TF_Conv3DOp $input, $filter, $strides, $padding,
IsDataFormatNDHWC:$data_format, $dilations)),
(LiftAsTFPartitionedCall<"composite_conv3d_with_"# ActivationName #"_fn">
(ArgumentList $input, $filter),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"strides"> $strides),
(NamedAttr<"padding"> $padding),
(NamedAttr<"dilations"> $dilations))),
[(IsNotInLiftedFunc $res)], [], (addBenefit 10)>;
def LiftConv3dWithBiasAnd#LastFusedOp : Pat<
(ActivationOp:$res
(TF_BiasAddOp
(TF_Conv3DOp $input, $filter, $strides, $padding,
IsDataFormatNDHWC:$data_format, $dilations),
$bias, $bias_data_format)),
(LiftAsTFPartitionedCall<"composite_conv3d_with_bias_and_"# ActivationName #"_fn">
(ArgumentList $input, $filter, $bias),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"strides"> $strides),
(NamedAttr<"padding"> $padding),
(NamedAttr<"dilations"> $dilations))),
[(IsNotInLiftedFunc $res)], [], (addBenefit 10)>;
def LiftBatchMatMulWith#ActivationOp : Pat<
(ActivationOp:$res
(TF_BatchMatMulV2Op $x, $y, $adj_x, $adj_y, $grad_x, $grad_y)),
(LiftAsTFPartitionedCall<"composite_batch_matmul_with_"# ActivationName #"_fn">
(ArgumentList $x, $y),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"adj_x"> $adj_x),
(NamedAttr<"adj_y"> $adj_y))),
[(IsNotInLiftedFunc $res)], [], (addBenefit 10)>;
def LiftBatchMatMulWithBiasAnd#LastFusedOp : Pat<
(ActivationOp:$res
(TF_BiasAddOp
(TF_BatchMatMulV2Op $x, $y, $adj_x, $adj_y, $grad_x, $grad_y),
$bias, IsDataFormatNHWC:$bias_data_format)),
(LiftAsTFPartitionedCall<"composite_batch_matmul_with_bias_and_"# ActivationName #"_fn">
(ArgumentList $x, $y, $bias),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"adj_x"> $adj_x),
(NamedAttr<"adj_y"> $adj_y))),
[(IsNotInLiftedFunc $res)], [], (addBenefit 10)>;
def LiftEinsumWith#ActivationOp : Pat<
(ActivationOp:$res
(TF_EinsumOp $input, $equation)),
(LiftAsTFPartitionedCall<"composite_einsum_with_"# ActivationName #"_fn">
(ArgumentList $input),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"equation"> $equation))),
[(IsNotInLiftedFunc $res),
(IsEinsumSupportedByXlaDotV2 $equation)],
[], (addBenefit 10)>;
def LiftEinsumWithBiasAnd#LastFusedOp : Pat<
(ActivationOp:$res
(TF_BiasAddOp
(TF_EinsumOp $input, $equation),
$bias, IsDataFormatNHWC:$bias_data_format)),
(LiftAsTFPartitionedCall<"composite_einsum_with_bias_and_"# ActivationName #"_fn">
(AppendToVector (ArgumentList $input), $bias),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"equation"> $equation))),
[(IsNotInLiftedFunc $res),
(IsEinsumSupportedByXlaDotV2 $equation)],
[], (addBenefit 10)>;
}
defm : LiftCompositeOpsWithActivation<TF_ReluOp, "relu">;
defm : LiftCompositeOpsWithActivation<TF_Relu6Op, "relu6">;
def LiftGather : Pat<
(TF_GatherV2Op:$res $params, $indices, $axis, $batch_dims),
(LiftAsTFPartitionedCall<"composite_gather_fn">
(ArgumentList $params, $indices, $axis),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"batch_dims"> $batch_dims))),
[(IsNotInLiftedFunc $res), (IsConstTensor $params)], [], (addBenefit 1)>;
tensorflow/compiler/mlir/quantization/tensorflow/passes/lift_quantizable_spots_as_functions_drq.cc
+214
View File
@@ -0,0 +1,214 @@
/* 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 <cstdint>
#include <memory>
#include <utility>
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/CommandLine.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/Diagnostics.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassRegistry.h" // from @llvm-project
#include "mlir/Rewrite/FrozenRewritePatternSet.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/common/lift_as_function_call.h" // IWYU pragma: keep
#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/ops/tf_op_quant_spec.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/passes.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
namespace mlir {
namespace quant {
namespace {
using QuantMethod =
::tensorflow::quantization::QuantizationMethod::PresetMethod;
using ::tensorflow::quantization::OpSet;
class LiftQuantizableSpotsAsFunctionsDRQPass
: public PassWrapper<LiftQuantizableSpotsAsFunctionsDRQPass,
OperationPass<ModuleOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
LiftQuantizableSpotsAsFunctionsDRQPass)
// Constructor used by the PassRegistration. This is only used by test.
explicit LiftQuantizableSpotsAsFunctionsDRQPass() = default;
// Constructor used by manually creating the pass.
explicit LiftQuantizableSpotsAsFunctionsDRQPass(
const QuantMethod quantization_method, const OpSet target_opset,
const int min_num_elements_for_weights) {
quantization_method_ = quantization_method;
target_opset_ = target_opset;
min_num_elements_for_weights_ = min_num_elements_for_weights;
}
LiftQuantizableSpotsAsFunctionsDRQPass(
const LiftQuantizableSpotsAsFunctionsDRQPass& other) {
quantization_method_ = other.quantization_method_;
target_opset_ = other.target_opset_;
min_num_elements_for_weights_ = other.min_num_elements_for_weights_;
}
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-lift-quantizable-spots-as-functions-drq";
}
StringRef getDescription() const final {
// This is a brief description of the pass.
return "Replace quantization candidates with composite functions into the "
"module for post-training dynamic range case";
}
void getDependentDialects(DialectRegistry& registry) const override {
registry.insert<TF::TensorFlowDialect>();
}
void runOnOperation() override;
private:
Option<OpSet> target_opset_{
*this, "target-opset", llvm::cl::init(OpSet::TF),
llvm::cl::desc("Choose target opset."),
llvm::cl::values(
clEnumValN(OpSet::TF, "TF",
"Uses TF ops that mimic quantization behavior"),
clEnumValN(OpSet::XLA, "XLA", "Uses TF XLA ops"),
clEnumValN(OpSet::UNIFORM_QUANTIZED, "UNIFORM_QUANTIZED",
"Uses TF Uniform Quantized ops"))};
Option<int64_t> min_num_elements_for_weights_{
*this, "min-num-elements-for-weights", llvm::cl::init(0),
llvm::cl::desc("The minimum required number of elements in a weight "
"array to apply quantization.")};
Option<QuantMethod> quantization_method_{
*this, "quantization-method",
llvm::cl::init(tensorflow::quantization::QuantizationMethod::
METHOD_DYNAMIC_RANGE_INT8),
llvm::cl::desc("Choose quantization method."),
llvm::cl::values(
clEnumValN(tensorflow::quantization::QuantizationMethod::
METHOD_DYNAMIC_RANGE_INT8,
"drq", "Post-training dynamic-range quantizaiton"),
clEnumValN(tensorflow::quantization::QuantizationMethod::
METHOD_STATIC_RANGE_WEIGHT_ONLY_INT8,
"weight_only", "Post-training weight_only quantizaiton"))};
};
class CheckQuantizableOps
: public mlir::OpRewritePattern<TF::PartitionedCallOp> {
public:
explicit CheckQuantizableOps(MLIRContext* context,
const QuantMethod quantization_method,
const OpSet target_opset,
const int min_num_elements_for_weights)
: OpRewritePattern<TF::PartitionedCallOp>(context),
quantization_method_(quantization_method),
target_opset_(target_opset),
min_num_elements_for_weights_(min_num_elements_for_weights) {}
private:
LogicalResult matchAndRewrite(TF::PartitionedCallOp call_op,
PatternRewriter& rewriter) const override {
std::unique_ptr<OpQuantSpec> spec = GetTFOpQuantSpec(call_op);
if (spec->quantizable_operands.empty()) return failure();
for (auto idx : spec->quantizable_operands) {
// This op is guaranteed to be a constant as ODS checks IsConstTensor.
// Check if the number of elements meets the requirement.
int current_num_elements =
mlir::cast<ShapedType>(call_op.getOperand(idx).getType())
.getNumElements();
if (current_num_elements < min_num_elements_for_weights_) {
call_op.emitRemark("Quantization is skipped for ")
<< call_op->getName().getStringRef().str() << " because it has "
<< current_num_elements
<< " elements which is fewer than the threshold("
<< min_num_elements_for_weights_ << " elements).";
call_op->removeAttr(kQuantTraitAttrName);
}
}
StringRef function_name =
mlir::cast<FlatSymbolRefAttr>(call_op.getFAttr()).getValue();
if ((quantization_method_ == tensorflow::quantization::QuantizationMethod::
METHOD_DYNAMIC_RANGE_INT8) &&
(function_name.contains("batch_matmul") ||
function_name.contains("conv3d"))) {
call_op->removeAttr(kQuantTraitAttrName);
}
// TODO(b/270906404): Support weight-only gather for uniform quantized opset
// in PTQ mode
if (target_opset_ == OpSet::UNIFORM_QUANTIZED &&
function_name.contains("gather")) {
call_op->removeAttr(kQuantTraitAttrName);
}
return failure();
}
QuantMethod quantization_method_;
OpSet target_opset_;
int min_num_elements_for_weights_;
};
static PassRegistration<LiftQuantizableSpotsAsFunctionsDRQPass> pass;
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/lift_quantizable_spots_as_functions_drq.inc"
void LiftQuantizableSpotsAsFunctionsDRQPass::runOnOperation() {
MLIRContext* ctx = &getContext();
RewritePatternSet patterns(ctx);
ModuleOp module = getOperation();
populateWithGenerated(patterns);
patterns.add<CheckQuantizableOps>(ctx, quantization_method_, target_opset_,
min_num_elements_for_weights_);
FrozenRewritePatternSet frozen_patterns(std::move(patterns));
for (auto func : module.getOps<func::FuncOp>()) {
if (failed(applyPatternsGreedily(func, frozen_patterns))) {
func.emitError()
<< "quant-lift-quantizable-spots-as-functions-drq failed.";
signalPassFailure();
}
}
}
} // namespace
std::unique_ptr<OperationPass<ModuleOp>>
CreateLiftQuantizableSpotsAsFunctionsDRQPass(
const QuantMethod quantization_method, const OpSet target_opset,
const int min_num_elements_for_weights) {
return std::make_unique<LiftQuantizableSpotsAsFunctionsDRQPass>(
quantization_method, target_opset, min_num_elements_for_weights);
}
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/passes/lift_quantizable_spots_as_functions_drq.td
+93
View File
@@ -0,0 +1,93 @@
/* 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 "mlir/IR/OpBase.td"
include "mlir/Dialect/Func/IR/FuncOps.td"
include "mlir/Dialect/Arith/IR/ArithOps.td"
include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"
include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.td"
include "tensorflow/compiler/mlir/quantization/common/lift_as_function_call.td"
//===----------------------------------------------------------------------===//
// Pattern rules for lifting ops as functions
//===----------------------------------------------------------------------===//
def LiftConv : Pat<
(TF_Conv2DOp:$res $input, $filter, $strides, $use_cudnn_on_gpu, $padding,
$explicit_paddings, IsDataFormatNHWC:$data_format, $dilations),
(LiftAsTFPartitionedCall<"composite_conv2d_fn">
(ArgumentList $input, $filter),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"strides"> $strides),
(NamedAttr<"use_cudnn_on_gpu"> $use_cudnn_on_gpu),
(NamedAttr<"padding"> $padding),
(NamedAttr<"explicit_paddings"> $explicit_paddings),
(NamedAttr<"dilations"> $dilations))),
[(IsNotInLiftedFunc $res), (IsConstTensor $filter)], [], (addBenefit 1)>;
def LiftDepthwiseConv : Pat<
(TF_DepthwiseConv2dNativeOp:$res $input, $filter, $strides, $padding,
$explicit_paddings, IsDataFormatNHWC:$data_format, $dilations),
(LiftAsTFPartitionedCall<"composite_depthwise_conv2d_fn">
(ArgumentList $input, $filter),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"strides"> $strides),
(NamedAttr<"padding"> $padding),
(NamedAttr<"explicit_paddings"> $explicit_paddings),
(NamedAttr<"dilations"> $dilations))),
[(IsNotInLiftedFunc $res), (IsConstTensor $filter)], [], (addBenefit 1)>;
def LiftMatMul : Pat<
(TF_MatMulOp:$res $a, $b, $transpose_a, $transpose_b, $grad_a, $grad_b),
(LiftAsTFPartitionedCall<"composite_matmul_fn">
(ArgumentList $a, $b),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"transpose_a"> $transpose_a),
(NamedAttr<"transpose_b"> $transpose_b))),
[(IsNotInLiftedFunc $res), (IsConstTensor $b)], [], (addBenefit 1)>;
def LiftGather : Pat<
(TF_GatherV2Op:$res $params, $indices, $axis, $batch_dims),
(LiftAsTFPartitionedCall<"composite_gather_fn">
(ArgumentList $params, $indices, $axis),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"batch_dims"> $batch_dims))),
[(IsNotInLiftedFunc $res), (IsConstTensor $params)], [], (addBenefit 1)>;
def LiftConv3D : Pat<
(TF_Conv3DOp:$res $input, $filter, $strides, $padding,
IsDataFormatNDHWC:$data_format, $dilations),
(LiftAsTFPartitionedCall<"composite_conv3d_fn">
(ArgumentList $input, $filter),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"strides"> $strides),
(NamedAttr<"padding"> $padding),
(NamedAttr<"dilations"> $dilations))),
[(IsNotInLiftedFunc $res), (IsConstTensor $filter)], [], (addBenefit 1)>;
def LiftBatchMatMul : Pat<
(TF_BatchMatMulV2Op:$res $x, $y, $adj_x, $adj_y, $grad_x, $grad_y),
(LiftAsTFPartitionedCall<"composite_batch_matmul_fn">
(ArgumentList $x, $y),
(ResultList $res),
(NamedAttributeList
(NamedAttr<"adj_x"> $adj_x),
(NamedAttr<"adj_y"> $adj_y))),
[(IsNotInLiftedFunc $res), (IsConstTensor $y)], [], (addBenefit 1)>;
tensorflow/compiler/mlir/quantization/tensorflow/passes/manipulate_model_attr.cc
+60
View File
@@ -0,0 +1,60 @@
/* 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.
==============================================================================*/
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/manipulate_model_attr.h"
#include <string>
#include <utility>
#include "llvm/ADT/StringExtras.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
namespace mlir {
namespace quant {
constexpr StringRef kTfEntryFunctionAttr = "tf.entry_function";
void AddEntryFunctionInput(StringRef input_name, func::FuncOp func_op) {
auto entry_func_attr =
func_op->getAttrOfType<DictionaryAttr>(kTfEntryFunctionAttr);
if (!entry_func_attr) return;
auto entry_func_attrs = SmallVector<NamedAttribute>(entry_func_attr.begin(),
entry_func_attr.end());
MLIRContext* ctx = func_op.getContext();
for (auto& named_attr : entry_func_attrs) {
if (named_attr.getName() != "inputs") continue;
// Splits the "inputs" field to retrieve individual input names. Ignores
// empty strings.
SmallVector<StringRef> inputs_attrs{};
cast<StringAttr>(named_attr.getValue())
.strref()
.split(inputs_attrs, /*Separator=*/',', /*MaxSplit=*/-1,
/*KeepEmpty=*/false);
inputs_attrs.emplace_back(input_name);
const std::string new_inputs_attr_str =
llvm::join(std::move(inputs_attrs), /*Separator=*/",");
named_attr.setValue(StringAttr::get(ctx, new_inputs_attr_str));
}
func_op->setAttr(kTfEntryFunctionAttr,
DictionaryAttr::get(ctx, entry_func_attrs));
}
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/passes/manipulate_model_attr.h
+32
View File
@@ -0,0 +1,32 @@
/* 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_QUANTIZATION_TENSORFLOW_PASSES_MANIPULATE_MODEL_ATTR_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_PASSES_MANIPULATE_MODEL_ATTR_H_
#include "llvm/ADT/StringRef.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
namespace mlir {
namespace quant {
// Adds a new input name to the `inputs` field of the `tf.entry_function`
// attribute if the attribute exist in the given function. Otherwise, no
// attribute is modified.
void AddEntryFunctionInput(StringRef input_name, func::FuncOp func_op);
} // namespace quant
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_PASSES_MANIPULATE_MODEL_ATTR_H_
tensorflow/compiler/mlir/quantization/tensorflow/passes/mark_functions_noinline.cc
+125
View File
@@ -0,0 +1,125 @@
/* 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 <string>
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/StringSet.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Builders.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/quantization/tensorflow/passes/passes.h"
// Required when using LLVM_DEBUG macro.
#define DEBUG_TYPE "mark-functions-noinline"
namespace mlir {
namespace quant {
namespace {
// Name of the boolean attribute indicating whether the function can be
// inlined or not.
constexpr StringRef kTfNoinlineAttr = "tf._noinline";
// This pass marks functions with the attribute `tf._noinline = true` so that
// they aren't inlined by the `InlinerPass`. The names of the functions to be
// marked noinline should be specified by the `noinline-functions` option.
class MarkFunctionsNoinlinePass
: public PassWrapper<MarkFunctionsNoinlinePass,
OperationPass<func::FuncOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MarkFunctionsNoinlinePass)
explicit MarkFunctionsNoinlinePass()
: MarkFunctionsNoinlinePass(
/*noinline_functions=*/ArrayRef<std::string>{}) {}
// `noinline_functions` is a list of function names to be marked noinline.
explicit MarkFunctionsNoinlinePass(
const ArrayRef<std::string> noinline_functions)
: noinline_functions_(CreateNoinlineFunctionsOption(noinline_functions)) {
}
MarkFunctionsNoinlinePass(const MarkFunctionsNoinlinePass& other)
: MarkFunctionsNoinlinePass() {
noinline_functions_ = other.noinline_functions_;
}
StringRef getArgument() const final { return "mark-functions-noinline"; }
StringRef getDescription() const final {
return "Marks a function whose name is in `noinline-functions` option with "
"the attribute `tf._noinline = true`. This attributes the function "
"from being inlined by the `InlinerPass`.";
}
void runOnOperation() override;
private:
ListOption<std::string> CreateNoinlineFunctionsOption(
const ArrayRef<std::string> noinline_functions) {
return {*this, "noinline-functions",
llvm::cl::desc(
"Name of the functions that should be marked "
"tf._noinline = true to prevent inlining. The name of the "
"function should exactly match to be marked noinline."),
llvm::cl::list_init<std::string>(noinline_functions),
llvm::cl::ZeroOrMore};
}
// Gets a set of function names from `noinline_functions_`.
StringSet<> GetNoinlineFunctionsSet() {
StringSet<> noinline_functions;
noinline_functions.insert(noinline_functions_.begin(),
noinline_functions_.end());
return noinline_functions;
}
// Names of the functions to be marked noinline.
ListOption<std::string> noinline_functions_;
};
void MarkFunctionsNoinlinePass::runOnOperation() {
const StringSet<> noinline_functions = GetNoinlineFunctionsSet();
func::FuncOp func_op = getOperation();
Builder builder(&getContext());
// Adds the `tf._noinline = true` attribute to the function if the name
// matches.
if (noinline_functions.contains(func_op.getSymName())) {
func_op->setAttr(kTfNoinlineAttr, builder.getBoolAttr(true));
LLVM_DEBUG(llvm::dbgs()
<< "Marked tf._noinline = true: " << func_op.getSymName());
}
}
static PassRegistration<MarkFunctionsNoinlinePass> pass{};
} // namespace
std::unique_ptr<OperationPass<func::FuncOp>> CreateMarkFunctionsNoinlinePass(
const ArrayRef<std::string> noinline_functions) {
return std::make_unique<MarkFunctionsNoinlinePass>(noinline_functions);
}
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/passes/merge_duplicate_resource_ops.cc
+149
View File
@@ -0,0 +1,149 @@
/* 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.
==============================================================================*/
#include <memory>
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringMap.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/Operation.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/quantization/tensorflow/passes/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
namespace mlir {
namespace quant {
namespace {
using ::mlir::tf_executor::GraphOp;
using ::mlir::tf_executor::IslandOp;
constexpr StringRef kSharedNameAttr = "shared_name";
class MergeDuplicateResourceOpsPass
: public PassWrapper<MergeDuplicateResourceOpsPass,
OperationPass<func::FuncOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MergeDuplicateResourceOpsPass)
StringRef getArgument() const final {
return "quant-merge-duplicate-resource-ops";
}
StringRef getDescription() const final {
return "Merge resource ops that have the same shared name.";
}
void runOnOperation() override;
};
// Checks if the island op contains a resource op like Variable or Hashtable
// and returns that resource op. Otherwise, returns null.
Operation* GetResourceOp(Operation* op) {
// Check if the island has only one block thats contain two ops, including
// one resource op and one Yield op.
auto island_op = llvm::dyn_cast_or_null<IslandOp>(op);
if (!island_op || !island_op.getBody().hasOneBlock()) return nullptr;
auto& island_block = island_op.getBody().front();
if (++island_block.begin() != --island_block.end()) return nullptr;
Operation* resource_op = &island_block.front();
if (llvm::isa<TF::VarHandleOp, TF::HashTableOp, TF::HashTableV2Op,
TF::MutableHashTableV2Op>(resource_op)) {
return resource_op;
}
return nullptr;
}
// Returns the `shared_name` attribute value if exists. If not, returns an
// empty string.
StringRef GetSharedName(Operation* op) {
if (!op->hasAttrOfType<StringAttr>(kSharedNameAttr)) return "";
return op->getAttrOfType<StringAttr>(kSharedNameAttr).getValue();
}
// Gets the GraphOp from the function op. Returns an empty op iff it doesn't
// exist.
// TODO(b/284222084): Move executor dialect utilities to a new library.
GraphOp GetGraphOpFromFuncOp(func::FuncOp func_op) {
if (func_op->getNumRegions() == 0 || func_op.getBody().empty()) return {};
auto graph_op_range = func_op.front().without_terminator();
if (llvm::hasSingleElement(graph_op_range)) {
// The pass runs on a valid tf_executor dialect, so the op should be the
// GraphOp.
return cast<GraphOp>(graph_op_range.begin());
}
return {};
}
void MergeDuplicateResourceOpsPass::runOnOperation() {
func::FuncOp func_op = getOperation();
GraphOp graph_op = GetGraphOpFromFuncOp(func_op);
if (!graph_op) return;
llvm::StringMap<Operation*> shared_name_to_resource;
llvm::SmallVector<Operation*> ops_to_remove;
for (Operation& op : graph_op.GetBody().without_terminator()) {
Operation* resource_op = GetResourceOp(&op);
if (!resource_op) continue;
StringRef shared_name = GetSharedName(resource_op);
if (shared_name.empty()) continue;
if (!shared_name_to_resource.contains(shared_name)) {
shared_name_to_resource[shared_name] = resource_op;
continue;
}
auto existing_resource = shared_name_to_resource[shared_name];
if (resource_op->getName().getStringRef() !=
existing_resource->getName().getStringRef() ||
resource_op->getResult(0).getType() !=
existing_resource->getResult(0).getType()) {
resource_op->emitOpError(
"This op has the same `shared_name` but different type with another "
"resource op in the function");
signalPassFailure();
return;
}
op.replaceAllUsesWith(existing_resource->getParentOp()->getResults());
ops_to_remove.push_back(&op);
}
// Remove op after the loop to avoid crash.
for (Operation* op : ops_to_remove) {
op->erase();
}
}
static PassRegistration<MergeDuplicateResourceOpsPass> pass{};
} // namespace
std::unique_ptr<OperationPass<func::FuncOp>>
CreateMergeDuplicateResourceOpsPass() {
return std::make_unique<MergeDuplicateResourceOpsPass>();
}
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/passes/merge_initializer_function_ops_to_main.cc
+403
View File
@@ -0,0 +1,403 @@
/* 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 <array>
#include <memory>
#include <string>
#include <utility>
#include "absl/cleanup/cleanup.h"
#include "absl/container/flat_hash_map.h"
#include "absl/log/log.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_format.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/raw_ostream.h"
#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 "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/DialectRegistry.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/SymbolTable.h" // from @llvm-project
#include "mlir/IR/TypeRange.h" // from @llvm-project
#include "mlir/IR/Types.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/LogicalResult.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/common/func.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/manipulate_model_attr.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/passes.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_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
namespace mlir {
namespace quant {
namespace {
using ::mlir::tf_executor::FetchOp;
using ::mlir::tf_executor::GraphOp;
using ::mlir::tf_executor::IslandOp;
using ::mlir::tf_saved_model::GetInitializerFunctions;
using ::mlir::tf_saved_model::GetSessionInitializerOp;
using ::mlir::tf_saved_model::kTfSavedModelInitializerInitType;
using ::mlir::tf_saved_model::kTfSavedModelInitializerRestoreType;
using ::mlir::tf_saved_model::kTfSavedModelInitializerTypeAttr;
using ::mlir::tf_saved_model::SessionInitializerOp;
// Array of initializer functions' types. The corresponding initializer
// functions should be merged in this order. This is because:
// 1) Variable restoration usually happens before initialization of other
// resources when a SavedModel is loaded. This ordering follows this semantic.
// 2) The `tf_saved_model` dialect requires that the arguments with
// `tf_saved_model.index_path` attributes should precede those with
// `tf_saved_model.bound_input` attributes. The init function of type
// `kTfSavedModelInitializerRestoreType` usually has an argument with
// `tf_saved_model.index_path`, whereas the init function of type
// `kTfSavedModelInitializerInitType` may have arguments with
// `tf_saved_model.bound_input`. This ordering avoids breaking the argument
// ordering constraint.
constexpr std::array<StringRef, 2> kInitializerTypesByMergeOrder = {
kTfSavedModelInitializerRestoreType, kTfSavedModelInitializerInitType};
// This pass moves all ops from initializer functions to the main function. A
// new `tf.NoOp` that has control dependency to the initializer function for
// non-variable resources will be created. The control output of the new
// `tf.NoOp` will be merged into the main function's `FetchOp`.
class MergeInitializerFunctionOpsToMainPass
: public PassWrapper<MergeInitializerFunctionOpsToMainPass,
OperationPass<ModuleOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
MergeInitializerFunctionOpsToMainPass)
explicit MergeInitializerFunctionOpsToMainPass() = default;
StringRef getArgument() const override {
return "quant-merge-initializer-function-ops-to-main";
}
StringRef getDescription() const override {
return "Moves all ops from the initializer functions to the main function. "
"A new `tf.NoOp` that has a control dependency to the initializer "
"function for non-variable resources will be created. Its control "
"output will be merged into the main function's `FetchOp`. The "
"initializer functions will be removed after this pass.";
}
void runOnOperation() override;
private:
void getDependentDialects(DialectRegistry& registry) const override {
registry
.insert<TF::TensorFlowDialect, tf_executor::TensorFlowExecutorDialect,
tf_saved_model::TensorFlowSavedModelDialect>();
}
};
// Returns true iff func_op has either no Region or the body has no Blocks.
bool IsFuncOpEmpty(func::FuncOp func_op) {
return func_op->getNumRegions() == 0 || func_op.getBody().empty();
}
// Gets the GraphOp from the function op. Returns an empty op iff it doesn't
// exist.
GraphOp GetGraphOpFromFuncOp(func::FuncOp func_op) {
if (IsFuncOpEmpty(func_op)) return {};
auto graph_op_range = func_op.front().without_terminator();
if (llvm::hasSingleElement(graph_op_range)) {
// The pass runs on a valid tf_executor dialect, so the op should be the
// GraphOp.
return cast<GraphOp>(graph_op_range.begin());
}
return {};
}
// Gets the string representation of the type name.
std::string GetTypeName(const Type type) {
std::string type_name{};
auto os = llvm::raw_string_ostream{type_name};
os << type;
return type_name;
}
// Retrieves the value of `tf_saved_model.initializer_type` attribute from the
// initializer function. Assumes that there exists such an attribute.
std::string GetInitializerType(func::FuncOp init_func_op) {
return init_func_op
->getAttrOfType<StringAttr>(kTfSavedModelInitializerTypeAttr)
.str();
}
// An initializer function should satisfy the follwing conditions:
// * Its GraphOp should only have control outputs.
// * "tf_saved_model.initializer_type" attribute must exist.
LogicalResult ValidateInitFunc(func::FuncOp init_func_op) {
GraphOp graph_op = GetGraphOpFromFuncOp(init_func_op);
if (!graph_op) return success(); // Consider empty FuncOp valid.
FetchOp fetch_op = graph_op.GetFetch();
for (const Value fetch : fetch_op.getFetches()) {
if (!mlir::isa<tf_executor::ControlType>(fetch.getType())) {
fetch_op.emitError(absl::StrFormat(
"Validation failed for the initializer function: %s. "
"All initializer function's fetches should be "
"tf_executor::ControlType. Got: %s.",
init_func_op.getName().str(), GetTypeName(fetch.getType())));
return failure();
}
}
if (const auto init_type_attr = init_func_op->getAttrOfType<StringAttr>(
kTfSavedModelInitializerTypeAttr);
!init_type_attr) {
return init_func_op->emitError() << "Initializer func op does not have "
"tf_saved_model.initializer_type "
"attribute. Func op: "
<< init_func_op.getSymName();
}
return success();
}
// Returns initializer_type -> init_func_op mapping from the session_init_op's
// initializers. The initializer functions are validated for whether it can be
// moved to the main function. Returns failure() iff validation fails.
FailureOr<absl::flat_hash_map<std::string, func::FuncOp>> GetInitFuncOps(
ModuleOp module_op) {
absl::flat_hash_map<std::string, func::FuncOp> init_func_ops;
for (func::FuncOp init_func_op : GetInitializerFunctions(module_op)) {
if (failed(ValidateInitFunc(init_func_op))) {
return failure();
}
init_func_ops[GetInitializerType(init_func_op)] = init_func_op;
}
return init_func_ops;
}
// Creates new arguments to the main function that corresponds to the source
// function's arguments. Returns the `IRMapping` that contains the
// relationship.
IRMapping CloneSrcFuncArgumentsToMainFunc(func::FuncOp src_func_op,
func::FuncOp main_func_op) {
IRMapping mapper{};
for (auto [src_arg_idx, src_arg] :
llvm::enumerate(src_func_op.getArguments())) {
// No need to create a mapping when there is no usage - it will not affect
// the cloning.
if (src_arg.use_empty()) continue;
const unsigned main_arg_idx = main_func_op.getNumArguments();
const DictionaryAttr main_arg_attr =
src_func_op.getArgAttrDict(src_arg_idx);
(void)main_func_op.insertArgument(main_arg_idx, src_arg.getType(),
main_arg_attr, src_arg.getLoc());
const std::string new_input_name =
absl::StrCat(GetInitializerType(src_func_op), "_", src_arg_idx, ":0");
AddEntryFunctionInput(new_input_name, main_func_op);
// During cloning, let it know that the source function's argument
// corresponds to the main function's newly created argument when cloning
// ops from src -> main.
BlockArgument main_arg = main_func_op.getArgument(main_arg_idx);
mapper.map(src_arg, main_arg);
}
return mapper;
}
// Copies ops from `src_func_op` to `main_body` except for the FetchOps. Returns
// the fetch values in the main GraphOp corresponding to the original fetch
// values from `src_func_op`. Returns an empty vector when `src_func_op` is
// empty. `main_func_op` must have a GraphOp.
SmallVector<Value> CopyOpsToMainFunction(func::FuncOp src_func_op,
func::FuncOp main_func_op) {
GraphOp src_graph_op = GetGraphOpFromFuncOp(src_func_op);
if (!src_graph_op) {
VLOG(1) << "Function " << src_func_op.getName().str()
<< " does not have a tf_executor::GraphOp. No ops are copied to "
"the main function.";
return {};
}
GraphOp main_graph_op = GetGraphOpFromFuncOp(main_func_op);
FetchOp main_fetch_op = main_graph_op.GetFetch();
const absl::Cleanup erase_main_fetch_op = [main_fetch_op]() mutable {
main_fetch_op.erase();
};
// TODO(b/245473863): Handle when assets are actually used in the body.
IRMapping mapper = CloneSrcFuncArgumentsToMainFunc(src_func_op, main_func_op);
// Clones each op from src to main_body.
Block& main_body = main_graph_op.GetBody();
Block& src_body = src_graph_op.GetBody();
for (Operation& op : src_body.without_terminator()) {
main_body.push_back(op.clone(mapper));
}
// Relocate the main function's FetchOp at the last.
main_body.push_back(main_fetch_op->clone(mapper));
// Clone the source's FetchOp, but do not push to the main function's body.
// The clone is only needed to identify the fetch operands.
auto cloned_fetch_op = cast<FetchOp>(src_graph_op.GetFetch()->clone(mapper));
const absl::Cleanup erase_cloned_fetch_op = [cloned_fetch_op]() mutable {
cloned_fetch_op.erase();
};
return llvm::to_vector(cloned_fetch_op.getFetches());
}
// Creates a new `IslandOp` that wraps a `TF::NoOp`. The `IslandOp` has control
// dependencies to the values provided.
IslandOp CreateNoOpWithControlDependencies(
const Location loc, GraphOp main_graph_op,
const ArrayRef<Value> control_dependencies) {
auto builder = OpBuilder::atBlockTerminator(&main_graph_op.GetBody());
auto wrapper_island_op = IslandOp::create(
builder, loc, /*outputs=*/TypeRange{},
/*control=*/tf_executor::ControlType::get(builder.getContext()),
/*controlInputs=*/control_dependencies);
wrapper_island_op.getBody().emplaceBlock();
// Create a NoOp inside the IslandOp.
auto guard = OpBuilder::InsertionGuard(builder);
builder.setInsertionPointToStart(&wrapper_island_op.GetBody());
TF::NoOp::create(builder, loc);
tf_executor::YieldOp::create(builder, loc);
return wrapper_island_op;
}
// Adds a new fetch operand for the main function's GraphOp.
void AddFetchOperandToMain(GraphOp main_graph_op, const Value fetch_operand) {
FetchOp old_fetch = main_graph_op.GetFetch();
const absl::Cleanup erase_old_fetch = [old_fetch]() mutable {
old_fetch.erase();
};
auto fetches = llvm::to_vector(old_fetch.getFetches());
fetches.emplace_back(fetch_operand);
auto builder = OpBuilder::atBlockTerminator(&main_graph_op.GetBody());
FetchOp::create(builder, main_graph_op.getLoc(), std::move(fetches));
}
// Creates a new Location for the initializer function. This creates a loc by
// attaching a to the initializer function's type so that it is identifiable.
Location CreateInitOpLoc(MLIRContext* ctx, func::FuncOp init_func_ops) {
const std::string init_type = GetInitializerType(init_func_ops);
const std::string name =
absl::StrCat(init_type, "_", init_func_ops.getName().str());
return NameLoc::get(StringAttr::get(ctx, name));
}
void MergeInitializerFunctionOpsToMainPass::runOnOperation() {
ModuleOp module_op = getOperation();
MLIRContext* ctx = module_op.getContext();
func::FuncOp main_func_op = FindMainFuncOp(module_op);
if (!main_func_op) {
module_op.emitError("Main function op not found.");
return signalPassFailure();
}
GraphOp main_graph_op = GetGraphOpFromFuncOp(main_func_op);
if (!main_graph_op) return;
SessionInitializerOp session_init_op = GetSessionInitializerOp(module_op);
if (!session_init_op) return;
// initializer_type -> init_func_op mapping.
SymbolTable symbol_table{module_op};
FailureOr<absl::flat_hash_map<std::string, func::FuncOp>> init_func_ops =
GetInitFuncOps(module_op);
if (failed(init_func_ops)) {
module_op->emitError("Validation on initializer functions failed.");
return signalPassFailure();
} else if (init_func_ops->empty()) {
VLOG(1) << "No initializer functions found.";
return;
}
// Find the initializer functions and clone their ops to @main.
for (const StringRef init_type : kInitializerTypesByMergeOrder) {
const auto it = init_func_ops->find(init_type);
if (it == init_func_ops->end()) continue;
func::FuncOp init_func_op = it->second;
const SmallVector<Value> init_op_fetches =
CopyOpsToMainFunction(init_func_op, main_func_op);
if (init_op_fetches.empty()) {
VLOG(1) << "No fetch values exist from initializer functions.";
return;
}
// Creates a NoOp that has control dependency to the initializer function
// for non-variables.
const Location init_op_loc = CreateInitOpLoc(ctx, init_func_op);
IslandOp noop_wrapper_island_op = CreateNoOpWithControlDependencies(
init_op_loc, main_graph_op,
/*control_dependencies=*/init_op_fetches);
AddFetchOperandToMain(
main_graph_op,
/*fetch_operand=*/noop_wrapper_island_op.getControl());
symbol_table.erase(init_func_op);
}
// Empties the "initializers" attribute from the `SessionInitializerOp` since
// all ops of the initializer ops are cloned into @main.
session_init_op.setInitializersAttr(ArrayAttr::get(ctx, {}));
}
} // namespace
std::unique_ptr<OperationPass<ModuleOp>>
CreateMergeInitializerFunctionOpsToMainPass() {
return std::make_unique<MergeInitializerFunctionOpsToMainPass>();
}
// Registers MergeInitializerFunctionOpsToMainPass.
static PassRegistration<MergeInitializerFunctionOpsToMainPass> pass([] {
return CreateMergeInitializerFunctionOpsToMainPass();
});
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/passes/merge_save_function_ops_to_main.cc
+299
View File
@@ -0,0 +1,299 @@
/* 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.
==============================================================================*/
#include <memory>
#include <utility>
#include "absl/algorithm/container.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#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 "mlir/IR/BuiltinOps.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/Value.h" // from @llvm-project
#include "mlir/IR/ValueRange.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/quantization/tensorflow/passes/constants.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/manipulate_model_attr.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
#include "tensorflow/compiler/mlir/tensorflow/translate/import_model.h"
namespace mlir {
namespace quant {
namespace {
using ::mlir::tf_executor::FetchOp;
using ::mlir::tf_executor::GraphOp;
using ::mlir::tf_executor::IslandOp;
using ::mlir::tf_saved_model::kTfSavedModelIndexPathAttr;
using ::tensorflow::kImportModelDefaultGraphFuncName;
class MergeSaveFunctionOpsToMainPass
: public PassWrapper<MergeSaveFunctionOpsToMainPass,
OperationPass<ModuleOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MergeSaveFunctionOpsToMainPass)
explicit MergeSaveFunctionOpsToMainPass() = default;
StringRef getArgument() const override {
return "quant-merge-save-function-ops-to-main";
}
StringRef getDescription() const override {
return "Merge the save function's ops to the main function. The save "
"function will be removed after the pass.";
}
void runOnOperation() override;
};
// Returns true iff func_op has either no Region or the body has no Blocks.
bool IsFuncOpEmpty(func::FuncOp func_op) {
return func_op->getNumRegions() == 0 || func_op.getBody().empty();
}
// Gets the GraphOp from the function op. Returns an empty op iff it doesn't
// exist.
GraphOp GetGraphOpFromFuncOp(func::FuncOp func_op) {
if (IsFuncOpEmpty(func_op)) return {};
auto graph_op_range = func_op.front().without_terminator();
if (llvm::hasSingleElement(graph_op_range)) {
// The pass runs on a valid tf_executor dialect, so the op should be the
// GraphOp.
return cast<GraphOp>(graph_op_range.begin());
}
return {};
}
// Gets the "main" function from the module. Returns an empty op iff it doesn't
// exist.
func::FuncOp GetMainFunction(ModuleOp module_op) {
const auto main_func_id =
StringAttr::get(module_op.getContext(), kImportModelDefaultGraphFuncName);
auto func_ops = module_op.getOps<func::FuncOp>();
auto main_func_itr = absl::c_find_if(func_ops, [&main_func_id](auto func_op) {
return func_op.getName() == main_func_id;
});
if (main_func_itr == func_ops.end()) return {};
return *main_func_itr;
}
func::FuncOp GetSaveFuncOp(ModuleOp module_op) {
for (auto func_op : module_op.getOps<func::FuncOp>()) {
if (func_op.getSymName() == kTfQuantSaveFuncName) return func_op;
}
return nullptr;
}
// Adds the file prefix argument to `main_func_op`. The file prefix argument
// is the argument whose "tf_saved_model.index_path" attribute has
// "__tf_file_prefix". Its type is `tensor<!tf_type.string>`. Also, the value
// "__tf_file_prefix:0" is appended to the "tf.entry_function" attribute's
// "inputs" key.
BlockArgument CreateFilePrefixArg(func::FuncOp main_func_op) {
Builder builder(main_func_op);
// Add a new argument of type `tensor<!tf_type.string>` and update the
// function type.
auto file_prefix_arg_type =
RankedTensorType::get(/*shape=*/{}, builder.getType<TF::StringType>());
BlockArgument new_file_prefix_arg =
main_func_op.getBody().front().addArgument(
file_prefix_arg_type,
NameLoc::get(builder.getStringAttr(kTfFilePrefix)));
SmallVector<Type> input_types(main_func_op.getArgumentTypes());
input_types.emplace_back(file_prefix_arg_type);
main_func_op.setType(
builder.getFunctionType(input_types, main_func_op.getResultTypes()));
// Add "__tf_file_prefix" to the "tf_saved_model.index_path" attribute for the
// newly created argument.
main_func_op.setArgAttr(new_file_prefix_arg.getArgNumber(),
/*name=*/kTfSavedModelIndexPathAttr,
/*value=*/builder.getStrArrayAttr({kTfFilePrefix}));
// Append the "__tf_file_prefix:0" to the "tf.entry_function" attribute's
// item keyed by "inputs".
AddEntryFunctionInput(Twine(kTfFilePrefix).concat(":0").str(), main_func_op);
return new_file_prefix_arg;
}
// Finds the file prefix argument from `main_func_op`. The file prefix argument
// is the argument whose "tf_saved_model.index_path" attribute has
// "__tf_file_prefix". If such an argument doesn't exist, returns a null value.
BlockArgument GetFilePrefixArg(func::FuncOp main_func_op) {
for (int i = 0; i < main_func_op.getNumArguments(); i++) {
auto index_path_attr =
main_func_op.getArgAttrOfType<ArrayAttr>(i, kTfSavedModelIndexPathAttr);
if (index_path_attr && !index_path_attr.empty() &&
mlir::cast<StringAttr>(index_path_attr[0]) == kTfFilePrefix) {
return main_func_op.getArgument(i);
}
}
return {};
}
// Returns the existing file prefix argument from the `main_func_op`. The file
// prefix argument is the argument whose "tf_saved_model.index_path" attribute
// has "__tf_file_prefix". If such an argument doesn't exist, creates a new file
// prefix argument and returns it.
BlockArgument GetOrCreateFilePrefixArg(func::FuncOp main_func_op) {
if (BlockArgument main_file_prefix_arg = GetFilePrefixArg(main_func_op);
main_file_prefix_arg) {
return main_file_prefix_arg;
} else {
return CreateFilePrefixArg(main_func_op);
}
}
// Clones ops from `src_graph_op` to `dst_graph_op`. The `dst_graph_op`'s
// `FetchOp` will be used without modified. Returns the fetch operands from the
// `scr_graph_op`.
Value CloneGraphOps(GraphOp src_graph_op, GraphOp dst_graph_op,
IRMapping& mapper) {
Block& main_body = dst_graph_op.GetBody();
// Take the reference of the main graph's FetchOp to later move to the end.
FetchOp main_fetch_op = dst_graph_op.GetFetch();
Block& save_func_body = src_graph_op.GetBody();
for (Operation& op : save_func_body.without_terminator()) {
main_body.push_back(op.clone(mapper));
}
// Relocate the main function's FetchOp to the last.
main_body.push_back(main_fetch_op->clone(mapper));
main_fetch_op.erase();
auto cloned_fetch_op = cast<FetchOp>(src_graph_op.GetFetch()->clone(mapper));
Value control_fetch = *cloned_fetch_op.getFetches().begin();
cloned_fetch_op.erase();
return control_fetch;
}
// Creates a new `IdentityOp` wrapped by an `IslandOp`. The identity op returns
// the `main_file_prefix_arg` and has control dependencies to `control_inputs`.
IslandOp CreateFilePrefixIdentityOp(const BlockArgument main_file_prefix_arg,
const ArrayRef<Value> control_inputs,
GraphOp main_graph_op) {
MLIRContext& ctx = *main_graph_op.getContext();
const auto name_loc = NameLoc::get(StringAttr::get(&ctx, kTfQuantSaveOpName));
auto builder = OpBuilder::atBlockTerminator(&main_graph_op.GetBody());
// Create an IslandOp that will wrap the IdentityOp. Add a control dependency
// for the newly copied save function.
auto wrapper_island_op = IslandOp::create(
builder, name_loc, TypeRange{main_file_prefix_arg.getType()},
tf_executor::ControlType::get(&ctx), ValueRange(control_inputs));
wrapper_island_op.getBody().emplaceBlock();
builder.setInsertionPointToStart(&wrapper_island_op.GetBody());
auto identity_op = TF::IdentityOp::create(
builder, name_loc, /*result_types=*/main_file_prefix_arg.getType(),
/*input=*/main_file_prefix_arg);
tf_executor::YieldOp::create(builder, name_loc, identity_op.getResult());
return wrapper_island_op;
}
// Appends `value` to the arguments of the `FetchOp` of `graph_op`.
void AppendValueToFetch(GraphOp graph_op, Value value) {
FetchOp old_main_fetch = graph_op.GetFetch();
auto fetches = llvm::to_vector(old_main_fetch.getFetches());
fetches.emplace_back(value);
auto builder = OpBuilder::atBlockTerminator(&graph_op.GetBody());
FetchOp::create(builder, old_main_fetch.getLoc(), std::move(fetches));
old_main_fetch.erase();
}
void MergeSaveFunctionOpsToMain(func::FuncOp save_func_op,
func::FuncOp main_func_op) {
GraphOp main_graph_op = GetGraphOpFromFuncOp(main_func_op);
if (!main_graph_op) return;
GraphOp save_func_graph_op = GetGraphOpFromFuncOp(save_func_op);
if (!save_func_graph_op) return;
IRMapping mapper{};
BlockArgument main_file_prefix_arg = GetOrCreateFilePrefixArg(main_func_op);
// TODO(b/268452435): This part assumes that the save function is always valid
// and has the argument. Add a validation function to filter out any invalid
// inputs.
mapper.map(save_func_op.getArgument(0), main_file_prefix_arg);
Value save_control_fetch =
CloneGraphOps(save_func_graph_op, main_graph_op, mapper);
IslandOp file_prefix_identity_wrapper = CreateFilePrefixIdentityOp(
main_file_prefix_arg, /*control_inputs=*/{save_control_fetch},
main_graph_op);
// Adds the newly created identity op's control output to the main's fetches.
AppendValueToFetch(main_graph_op, file_prefix_identity_wrapper.getControl());
}
} // namespace
void MergeSaveFunctionOpsToMainPass::runOnOperation() {
ModuleOp module_op = getOperation();
func::FuncOp main_func_op = GetMainFunction(module_op);
if (!main_func_op) {
module_op.emitError("Main function op not found.");
return signalPassFailure();
}
func::FuncOp save_func_op = GetSaveFuncOp(module_op);
if (!save_func_op) return;
MergeSaveFunctionOpsToMain(save_func_op, main_func_op);
// Erase the save function when all ops are successfully cloned.
save_func_op.erase();
}
std::unique_ptr<OperationPass<ModuleOp>>
CreateMergeSaveFunctionOpsToMainPass() {
return std::make_unique<MergeSaveFunctionOpsToMainPass>();
}
static PassRegistration<MergeSaveFunctionOpsToMainPass> pass([] {
return CreateMergeSaveFunctionOpsToMainPass();
});
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/passes/optimize.cc
+70
View File
@@ -0,0 +1,70 @@
/* 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/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/PatternMatch.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/LogicalResult.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.h" // IWYU pragma: keep - required to use `IsSplatValueEqual`.
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/passes.h"
namespace mlir::quant {
namespace {
// Applies optimization after quantization.
class OptimizePass
: public PassWrapper<OptimizePass, OperationPass<func::FuncOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(OptimizePass)
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-optimize";
}
StringRef getDescription() const final {
// This is a brief description of the pass.
return "Applies optimization after quantization";
}
void runOnOperation() override;
};
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/optimize.inc"
void OptimizePass::runOnOperation() {
RewritePatternSet patterns(&getContext());
populateWithGenerated(patterns);
auto func = getOperation();
if (failed(applyPatternsGreedily(func, std::move(patterns)))) {
signalPassFailure();
}
}
} // namespace
std::unique_ptr<OperationPass<func::FuncOp>> CreateOptimizePass() {
return std::make_unique<OptimizePass>();
}
static PassRegistration<OptimizePass> pass;
} // namespace mlir::quant
tensorflow/compiler/mlir/quantization/tensorflow/passes/optimize.td
+61
View File
@@ -0,0 +1,61 @@
/* 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 "mlir/IR/OpBase.td"
include "mlir/IR/PatternBase.td"
include "mlir/Dialect/Func/IR/FuncOps.td"
include "mlir/Dialect/Arith/IR/ArithOps.td"
include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.td"
include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"
// Remove redundant `CastOp` to int8 if the input is properly clipped.
def RemoveRedundantCastOps : Pat<
(TF_CastOp:$root_cast
(TF_CastOp:$i8_cast
(TF_ClipByValueOp:$clip $input, $min_value, $max_value),
ConstBoolAttrFalse:$truncate2),
ConstBoolAttrFalse:$truncate1),
(TF_CastOp $clip, ConstBoolAttrFalse),
[(TensorOf<[I8]> $i8_cast),
(TensorOf<[I32]> $clip),
(IsIntSplatValueEqual<"int32_t", "-128"> $min_value),
(IsIntSplatValueEqual<"int32_t", "127"> $max_value)]>;
// This pattern optimizes:
// (x + cst1) + cst2 -> x + cst
// (x - cst1) - cst2 -> x - cst
// Where: cst = cst1 + cst2
foreach BinaryOp = [TF_AddV2Op, TF_SubOp] in {
def OptimizeConsecutive#BinaryOp : Pat<
(BinaryOp
(BinaryOp $x, (TF_ConstOp:$cst1 $cst1_value)),
(TF_ConstOp:$cst2 $cst2_value)),
(BinaryOp
$x, (TF_AddV2Op $cst1, $cst2))>;
}
// This pattern optimizes:
// (x + cst1) - cst2 -> x - cst
// (x - cst1) + cst2 -> x + cst
// Where: cst = cst2 - cst1
foreach BinaryOpPair = [[TF_AddV2Op, TF_SubOp],
[TF_SubOp, TF_AddV2Op]] in {
def OptimizeConsecutive#BinaryOpPair[0]#BinaryOpPair[1] : Pat<
(BinaryOpPair[0]
(BinaryOpPair[1] $x, (TF_ConstOp:$cst1 $cst1_value)),
(TF_ConstOp:$cst2 $cst2_value)),
(BinaryOpPair[0]
$x, (TF_SubOp $cst2, $cst1))>;
}
tensorflow/compiler/mlir/quantization/tensorflow/passes/passes.h
+250
View File
@@ -0,0 +1,250 @@
/* 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_QUANTIZATION_TENSORFLOW_PASSES_PASSES_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_PASSES_PASSES_H_
#include <memory>
#include <optional>
#include <string>
#include "absl/strings/string_view.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_config.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
namespace mlir {
namespace quant {
// Creates a main function if it doesn't exist in the module. This is a
// workaround to make ConvertMlirToGraphdef work for multi-signatures graphs.
// TODO(b/204265523): Removes this pass after the exporting MLIR to SavedModel
// path is available.
std::unique_ptr<OperationPass<ModuleOp>> CreateInsertMainFunctionPass();
// Converts FakeQuant ops to quant.qcast and quant.dcast (QDQ) pairs.
std::unique_ptr<OperationPass<func::FuncOp>> CreateConvertFakeQuantToQdqPass();
// Lifts the quantizable spots as composite functions.
std::unique_ptr<OperationPass<ModuleOp>>
CreateLiftQuantizableSpotsAsFunctionsPass(
const tensorflow::quantization::QuantizationOptions& quant_options);
// Apply graph optimizations such as fusing and constant folding to prepare
// lifting.
std::unique_ptr<OperationPass<func::FuncOp>> CreatePrepareLiftingPass(
tensorflow::quantization::OpSet target_opset);
// Lifts the dynamic range quantizable spots as composite functions.
std::unique_ptr<OperationPass<ModuleOp>>
CreateLiftQuantizableSpotsAsFunctionsDRQPass(
tensorflow::quantization::QuantizationMethod::PresetMethod
quantization_method,
tensorflow::quantization::OpSet op_set, int min_num_elements_for_weights);
// Replaces tf.CustomAggregator ops with quant.Stats ops for finalizing the
// calibration procedure.
std::unique_ptr<OperationPass<func::FuncOp>>
CreateConvertCustomAggregationOpToQuantStatsPass();
// Inserts quantized function library.
std::unique_ptr<OperationPass<ModuleOp>> CreateInsertQuantizedFunctionsPass(
tensorflow::quantization::QuantizationMethod::PresetMethod
quantization_method,
tensorflow::quantization::OpSet target_opset);
// Inserts custom aggregation operators for the calibration procedure.
std::unique_ptr<OperationPass<func::FuncOp>>
CreateInsertCustomAggregationOpsPass(
const ::stablehlo::quantization::CalibrationOptions& calib_opts);
// Replaces composite functions with quantized composite functions. After this
// pass runs, functions in the given graph will be replaced with their quantized
// versions. By doing so, the quantization will be applied to the given input.
// mlir_dump_file_prefix is an optional field that is used for debugging to save
// mlir dump files.
std::unique_ptr<OperationPass<ModuleOp>> CreateQuantizeCompositeFunctionsPass(
tensorflow::quantization::QuantizationMethod::PresetMethod
quantization_method,
tensorflow::quantization::OpSet target_opset,
bool enable_per_channel_quantization, int min_num_elements_for_weights,
bool enable_legacy_weight_only = false,
std::optional<const absl::string_view> mlir_dump_file_prefix =
std::nullopt);
// Converts dequantize-(quantizable) call-quantize pattern to a single call op
// that has quantized input and output types. It is expected for this pass to
// emit illegal IR with unsupported quantized input and output types. The
// pass following immediately after this one will be responsible for legalizing
// input and output types by unwrapping quantization parameters.
std::unique_ptr<OperationPass<func::FuncOp>> CreateQuantizePass();
// Overloading of CreateQuantizePass which takes QuantizationSpecs.
std::unique_ptr<OperationPass<func::FuncOp>> CreateQuantizePass(
QuantizationSpecs quant_specs,
tensorflow::quantization::OpSet target_opset);
// Creates an instance of the PrepareQuantize pass, which will perform similar
// transformations as TFL::PrepareQuantizePass.
std::unique_ptr<OperationPass<func::FuncOp>> CreatePrepareQuantizePass(
const QuantizationSpecs& quant_specs,
tensorflow::quantization::QuantizationMethod::PresetMethod
quantization_method);
// Creates an instance of the PrepareQuantizeDRQ pass, which will
// perform similar transformations as TFL::PrepareQuantizeDynamicRangePass.
std::unique_ptr<OperationPass<ModuleOp>> CreatePrepareQuantizeDRQPass(
const QuantizationSpecs& quant_specs,
tensorflow::quantization::OpSet op_set);
// Creates an instance of the PreprocessOp pass, which will perform op
// preprocessing to allow multi-axis quantization, prior to quantization.
std::unique_ptr<OperationPass<ModuleOp>> CreatePreprocessOpPass(
tensorflow::quantization::OpSet op_set,
tensorflow::quantization::QuantizationMethod::PresetMethod
quantization_method,
bool enable_per_channel_quantization);
// Creates an instance of the PostQuantize pass, which will remove unnecessary
// ops from the final quantized graph.
std::unique_ptr<OperationPass<func::FuncOp>> CreatePostQuantizePass();
// Applies optimization patterns after quantization.
std::unique_ptr<OperationPass<mlir::func::FuncOp>> CreateOptimizePass();
// Creates an instance of the ReplaceCastHacksWithTFXLAOpsPass, which will
// replace mixed-type convolution and matmul cast hacks by XLA Conv2DOp and
// MatmulOp.
std::unique_ptr<OperationPass<func::FuncOp>>
CreateReplaceCastHacksWithTFXLAOpsPass();
// Creates a pass that moves & merges initializer function's ops into the @main
// function. This pass should be run on a valid tf_executor dialect. The control
// output of the initializer function for non-variable resource initialization
// will be passed on as a dependency to a new `tf.NoOp`, whose control output
// will be merged into the main function's FetchOp. The initializer functions
// will be removed.
//
// Running this pass essentially has the effect of inlining the initializer
// functions into the main graph. This is beneficial when we wish to find and
// fetch the node that restores resources, after the ModuleOp has been exported
// as GraphDef.
std::unique_ptr<OperationPass<ModuleOp>>
CreateMergeInitializerFunctionOpsToMainPass();
// Creates a pass that moves & merges the "@tf_quant__save" function to "@main"
// function. A new `IdentityOp` will be created. It will have control dependency
// to the save function and returns the file_prefix argument (typed
// `tensor<!tf_type.string>`). The file_prefix argument, which can be identified
// if the "tf_saved_model.index_path" attribute has "__tf_file_prefix", will be
// reused if it already exist in @main. Otherwise a new file prefix argument
// will be created. @tf_quant__save function will be erased.
//
// Running this pass essentially has the effect of inlining the @tf_quant__save
// into the main graph. This is beneficial when we wish to find and fetch
// the node that saves the variables, after the ModuleOp has been exported as
// GraphDef.
std::unique_ptr<OperationPass<ModuleOp>> CreateMergeSaveFunctionOpsToMainPass();
// Creates a pass that "unfreezes" ConstOps into variables. Each ConstOp's use
// will be replaced by a VarHandleOp -> ReadVariableOp pattern. The newly
// created variables will be initialized in the session initializer function via
// AssignVariableOps.
std::unique_ptr<OperationPass<ModuleOp>> CreateUnfreezeConstantsPass();
// Creates a pass that duplicates constants that affect the shape of a tensor
// after some computation.
std::unique_ptr<OperationPass<func::FuncOp>>
CreateDuplicateShapeDeterminingConstantsPass();
// Creates a pass that creates a RestoreV2 op in the initializer function with
// type "restore_op" that initializes variables from the checkpoint. It finds
// tf.AssignVariableOp(tf.VarHandleOp, tf.Const) patterns in the initializer
// function and replaces tf.Consts with the results of RestoreV2.
std::unique_ptr<OperationPass<ModuleOp>> CreateInsertRestoreOpPass();
// Creates a pass that creates a new function that wraps the newly created
// SaveV2 op. The new function's name is "tf_quant__save". The function accepts
// a single string tensor as argument, which specifies the path to the
// checkpoint to which the variable's tensor values are saved. It finds
// `tf.AssignVariableOp(tf.VarHandleOp, tf.Const)` pattern in the initializer
// function of type "restore_op" to identify the VarHandleOps that should be
// saved using the SaveV2 op.
std::unique_ptr<OperationPass<ModuleOp>> CreateInsertSaveOpPass();
// Creates a pass that marks functions with the attribute `tf._noinline = true`
// to avoid being inlined by the `InlinerPass`. `noinline_functions` is the name
// of the functions to mark.
std::unique_ptr<OperationPass<func::FuncOp>> CreateMarkFunctionsNoinlinePass(
ArrayRef<std::string> noinline_functions);
// Removes `tf.AssignVariableOp(tf.VarHandleOp, tf.Const)` patterns from the
// initializer function (type = "restore_op").
// Note: initializing values (`tf.Const`s) will be removed and this may result
// in an information loss and uninitialized variables eventually. Make sure that
// this effect is desired (e.g. there is a `tf.RestoreV2Op` that restores the
// variables instead).
std::unique_ptr<OperationPass<ModuleOp>>
CreateRemoveVariableInitializationByConstPass();
// Creates a pass that converts Tensorflow Xla ops to non-Xla ops.
std::unique_ptr<OperationPass<func::FuncOp>> CreateConvertTfXlaOpToTfOpPass();
// Creates a pass that converts TPU models for CPU by removing TPU related ops
// such as TPUPartitionedCall, TPUReplicatedOp, etc. The TF quantizer does not
// work with models specifically designed for TPU, so this pass makes the input
// TPU model compatible with the TF quantizer by rewriting the TPU ops. The
// output model of this pass is expected to be ready for the TF quantizer.
std::unique_ptr<OperationPass<ModuleOp>> CreateConvertTpuModelToCpuPass();
// Creates a pass that casts BFloat16 operations to Float32 operations. This
// pass is a part of the ConvertTpuModelToCpu pass to support BF16 optimized TPU
// model quantization.
std::unique_ptr<OperationPass<ModuleOp>> CreateCastBf16OpsToF32Pass();
// Creates a pass that lifts HashTable ops as function arguments. In the graph
// execution mode, resource ops with the same `shared_name` attribute point to
// the same underlying resource. This is not true in the eager execution mode.
// Lifting resource ops as arguments will help unifying them across functions.
std::unique_ptr<OperationPass<ModuleOp>> CreateLiftHashTableOpsAsArgsPass();
// Creates a pass that merges duplicate resource ops in each function. Two
// resource ops are considered duplicated if they have the same `shared_name`.
std::unique_ptr<OperationPass<func::FuncOp>>
CreateMergeDuplicateResourceOpsPass();
// Apply quantization to weights based on the provided schemes.
std::unique_ptr<OperationPass<ModuleOp>> CreateQuantizeWeightsPass(
const tensorflow::quantization::QuantizationOptions& quant_options);
// Propagate quantized type through allowed ops.
std::unique_ptr<OperationPass<ModuleOp>> CreatePropagateQuantizeTypePass();
// Create a pass that inserts dump tensor to quantizable layer's output.
std::unique_ptr<OperationPass<ModuleOp>> CreateAddDumpTensorOpPass(
::stablehlo::quantization::DebuggerConfig::DebuggerType debugger_type,
std::string log_dir_path);
// Creates a pass that add QuantizationUnitLoc to quantizable layers.
std::unique_ptr<OperationPass<func::FuncOp>> CreateAddQuantizationUnitLocPass();
} // namespace quant
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_PASSES_PASSES_H_
tensorflow/compiler/mlir/quantization/tensorflow/passes/post_quantize.cc
+160
View File
@@ -0,0 +1,160 @@
/* 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 transformation pass applies some clean up steps after quantization.
#include <memory>
#include <string>
#include <utility>
#include "llvm/Support/Casting.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/Quant/IR/QuantTypes.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "mlir/IR/PatternMatch.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/LogicalResult.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"
#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h" // IWYU pragma: keep
//===----------------------------------------------------------------------===//
// The post-quantize Passes.
//
namespace mlir {
namespace quant {
namespace {
// Applies all the clean up steps after quantization.
class PostQuantizePass
: public PassWrapper<PostQuantizePass, OperationPass<func::FuncOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PostQuantizePass)
// Constructor used by the PassRegistration. This will remove the adaptor ops.
explicit PostQuantizePass() = default;
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-post-quantize";
}
StringRef getDescription() const final {
// This is a brief description of the pass.
return "Apply post quantization clean up after quantization";
}
void runOnOperation() override;
};
enum RemoveVolatileOpsType {
// Remove all volatile quant-dequant ops.
kPreserveNone,
// Preserve volatile quant-dequants for input and output ops.
kPreserveInputsAndOutputs,
};
// Remove the back-to-back quantize and dequantize ops with volatile attribute.
template <RemoveVolatileOpsType remove_volatile_ops_type>
struct RemoveVolatileOps
: public OpRewritePattern<mlir::quant::ir::DequantizeCastOp> {
explicit RemoveVolatileOps(MLIRContext* context)
: OpRewritePattern<mlir::quant::ir::DequantizeCastOp>(context, 1) {}
LogicalResult matchAndRewrite(mlir::quant::ir::DequantizeCastOp op,
PatternRewriter& rewriter) const override {
auto input_op = op.getArg().getDefiningOp();
if (auto q =
llvm::dyn_cast_or_null<mlir::quant::ir::QuantizeCastOp>(input_op)) {
if (!q->getAttr(kVolatileOpAttrName)) return failure();
if (remove_volatile_ops_type == kPreserveInputsAndOutputs) {
// Don't remove leading and trailing QDQ for PTQ workflow, so the io
// modifying lib can work correctly.
if (!q.getArg().getDefiningOp()) return failure();
if (op->hasOneUse() &&
op->user_begin()->hasTrait<OpTrait::IsTerminator>())
return failure();
}
// If the quantize op is a requantize op, it is being used in other scale
// adjustments and should be kept. Instead, moving dequantize op before
// the requantize op to remove the unnecessary requantize op.
if (auto qtype =
QuantizedType::getQuantizedElementType(q.getArg().getType())) {
rewriter.setInsertionPoint(op);
rewriter.replaceOpWithNewOp<mlir::quant::ir::DequantizeCastOp>(
op, op.getResult().getType(), q.getArg());
return success();
}
op.replaceAllUsesWith(q.getArg());
return success();
}
return failure();
}
};
// The StorageCastOp is used to cast from a quantized type to its storage type
// or the opposite. If none of its input and output is quantized, the op has
// no effect and should be removed.
class RemoveRedundantScast
: public mlir::OpRewritePattern<mlir::quant::ir::StorageCastOp> {
public:
explicit RemoveRedundantScast(MLIRContext* context)
: OpRewritePattern<mlir::quant::ir::StorageCastOp>(context) {}
private:
LogicalResult matchAndRewrite(mlir::quant::ir::StorageCastOp scast_op,
PatternRewriter& rewriter) const override {
if (QuantizedType::getQuantizedElementType(scast_op.getArg().getType()) ||
QuantizedType::getQuantizedElementType(scast_op.getType())) {
return failure();
}
scast_op.replaceAllUsesWith(scast_op.getArg());
return success();
}
};
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/post_quantize.inc"
void PostQuantizePass::runOnOperation() {
RewritePatternSet patterns(&getContext());
auto func = getOperation();
auto* ctx = func.getContext();
patterns.add<FoldTrivalRequantizeOp<mlir::quant::ir::QuantizeCastOp>,
RemoveVolatileOps<kPreserveNone>, RemoveRedundantScast>(ctx);
populateWithGenerated(patterns);
if (failed(applyPatternsGreedily(func, std::move(patterns)))) {
signalPassFailure();
}
}
} // namespace
// Creates an instance of the TensorFlow dialect PostQuantize pass.
std::unique_ptr<OperationPass<func::FuncOp>> CreatePostQuantizePass() {
return std::make_unique<PostQuantizePass>();
}
static PassRegistration<PostQuantizePass> pass;
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/passes/post_quantize.td
+35
View File
@@ -0,0 +1,35 @@
/* 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 "mlir/IR/OpBase.td"
include "mlir/IR/PatternBase.td"
include "mlir/Dialect/Func/IR/FuncOps.td"
include "mlir/Dialect/Arith/IR/ArithOps.td"
include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"
include "mlir/Dialect/Arith/IR/ArithOps.td"
include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.td"
include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.td"
// Re-orders the Identity op following a quantized composite function. This
// allows the QuantizeCompositeFunctionsPass to merge the DequantizeCast with
// the quantized composite function to optimize the requantization part.
def ReorderIdentityFollowingQuantizedFunction : Pat<
(Quantization_DequantizeCastOp:$output
(Quantization_StorageCastOp
(TF_IdentityOp
(Quantization_StorageCastOp $value)))),
(TF_IdentityOp
(Quantization_DequantizeCastOp
$value, (returnType (GetValueType $output))))>;
tensorflow/compiler/mlir/quantization/tensorflow/passes/prepare_lifting.cc
+356
View File
@@ -0,0 +1,356 @@
/* 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 <cstddef>
#include <cstdint>
#include <iterator>
#include <memory>
#include <utility>
#include <vector>
#include "absl/algorithm/container.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/Sequence.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/CommandLine.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/BuiltinTypeInterfaces.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/DialectRegistry.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Matchers.h" // from @llvm-project
#include "mlir/IR/PatternMatch.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/LogicalResult.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.h"
#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/cc/constant_fold.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/passes.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/remove_identity_op_pattern.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/einsum.h"
namespace mlir {
namespace quant {
namespace {
using ::tensorflow::quantization::OpSet;
class PrepareLiftingPass
: public PassWrapper<PrepareLiftingPass, OperationPass<func::FuncOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PrepareLiftingPass)
PrepareLiftingPass() = default;
explicit PrepareLiftingPass(OpSet op_set) { op_set_ = op_set; }
PrepareLiftingPass(const PrepareLiftingPass& other) {
op_set_ = other.op_set_;
}
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-prepare-lifting";
}
StringRef getDescription() const final {
// This is a brief description of the pass.
return "Apply graph optimizations such as fusing and constant folding to "
"prepare lifting.";
}
void getDependentDialects(DialectRegistry& registry) const override {
registry.insert<TF::TensorFlowDialect, arith::ArithDialect>();
}
void runOnOperation() override;
private:
Option<OpSet> op_set_{
*this, "target-opset", llvm::cl::init(OpSet::TF),
llvm::cl::desc("Choose target opset."),
llvm::cl::values(
clEnumValN(OpSet::TF, "TF",
"Uses TF ops that mimic quantization behavior"),
clEnumValN(OpSet::XLA, "XLA", "Uses TF XLA ops"),
clEnumValN(OpSet::UNIFORM_QUANTIZED, "UNIFORM_QUANTIZED",
"Uses TF Uniform Quantized ops"))};
};
// Check if given indices in `val1` has same number of elements as given
// indices in `val2`.
bool HasEqualElementSize(Value val1, Value val2, ArrayRef<int> val1_indices,
ArrayRef<int> val2_indices) {
ShapedType val1_shape = mlir::cast<ShapedType>(val1.getType());
ShapedType val2_shape = mlir::cast<ShapedType>(val2.getType());
if (!val1_shape.hasRank() || !val2_shape.hasRank()) return false;
int val1_result = 1;
int val2_result = 1;
for (auto idx : val1_indices) {
if (idx < 0) idx = idx + val1_shape.getRank();
if (idx >= val1_shape.getRank() || val1_shape.isDynamicDim(idx)) {
return false;
}
val1_result *= val1_shape.getDimSize(idx);
}
for (auto idx : val2_indices) {
if (idx < 0) idx = idx + val2_shape.getRank();
if (idx >= val2_shape.getRank() || val2_shape.isDynamicDim(idx)) {
return false;
}
val2_result *= val2_shape.getDimSize(idx);
}
return val1_result == val2_result;
}
// Checks if a shape has dim sizes of all ones except the right most dim.
bool ReshapableTo1DTensor(ShapedType rhs_shape) {
for (auto rank = 0; rank < rhs_shape.getRank() - 1; rank++) {
if (rhs_shape.getDimSize(rank) != 1) {
return false;
}
}
return true;
}
Value ReshapeTo1DTensor(OpBuilder& builder, Location loc, Value value) {
auto shape = mlir::cast<ShapedType>(value.getType());
if (shape.getRank() != 1) {
SmallVector<int64_t> new_shape;
new_shape.push_back(shape.getNumElements());
value = TF::ReshapeOp::create(builder, loc, value,
Create1DConstValue(builder, loc, new_shape));
}
return ConstantFoldOpIfPossible(value.getDefiningOp()).front();
}
// Matches convolution op with "NHWC" data format or matmul op with false adj_y.
// The list of supported ops in this function is:
// - Conv2DOp
// - Conv3DOp
// - DepthwiseConv2dNativeOp
// - MatMulOp
// - BatchMatMulV2Op
LogicalResult MatchSupportedAffineOp(Operation* op, Value& binding_output,
Value& binding_input,
Value& binding_weight) {
bool is_supported_affine_op = false;
if (llvm::isa<TF::Conv2DOp, TF::Conv3DOp, TF::DepthwiseConv2dNativeOp>(op)) {
if (const auto data_format = op->getAttrOfType<StringAttr>("data_format")) {
is_supported_affine_op =
data_format.getValue() == "NHWC" || data_format.getValue() == "NDHWC";
}
} else if (llvm::isa<TF::BatchMatMulV2Op>(op)) {
if (const auto adj_y = op->getAttrOfType<BoolAttr>("adj_y")) {
is_supported_affine_op = !adj_y.getValue();
}
} else if (llvm::isa<TF::MatMulOp>(op)) {
if (const auto adj_y = op->getAttrOfType<BoolAttr>("transpose_b")) {
is_supported_affine_op = !adj_y.getValue();
}
}
if (!is_supported_affine_op) return failure();
// Bind input, output and weight to the given values.
binding_output = op->getResult(0);
binding_input = op->getOperand(0);
binding_weight = op->getOperand(1);
return success();
}
// Makes the 1D value broadcastable with the `rhs_shape`.
Value MakeOneDimValueBroadcastable(OpBuilder& builder, Location loc,
Value value, ShapedType rhs_shape) {
ShapedType value_shape = mlir::dyn_cast_or_null<ShapedType>(value.getType());
if (!value_shape || value_shape.getRank() != 1 ||
!value_shape.hasStaticShape() || !rhs_shape.hasStaticShape()) {
return {};
}
int64_t num_elements = value_shape.getNumElements();
SmallVector<int64_t> new_shape;
for (auto idx : llvm::reverse(llvm::seq<int32_t>(0, rhs_shape.getRank()))) {
const int64_t rhs_dim = rhs_shape.getDimSize(idx);
if (num_elements % rhs_dim != 0) {
return {};
}
new_shape.push_back(rhs_dim);
num_elements = num_elements / rhs_dim;
if (num_elements == 1) break;
}
absl::c_reverse(new_shape);
auto reshape_op = TF::ReshapeOp::create(
builder, loc, value, Create1DConstValue(builder, loc, new_shape));
return ConstantFoldOpIfPossible(reshape_op).front();
}
// Checks if a value can be symmetrically quantized.
bool CanBeSymmetricallyQuantized(Value weight) {
auto dq_op = weight.getDefiningOp<mlir::quant::ir::DequantizeCastOp>();
if (!dq_op) return true;
auto qtype =
mlir::cast<TensorType>(dq_op.getArg().getType()).getElementType();
if (auto uniform_type = llvm::dyn_cast_or_null<UniformQuantizedType>(qtype)) {
return uniform_type.getZeroPoint() == 0;
} else if (auto per_axis_type =
llvm::dyn_cast_or_null<UniformQuantizedPerAxisType>(qtype)) {
return absl::c_all_of(per_axis_type.getZeroPoints(),
[](int64_t x) { return x == 0; });
}
return false;
}
// Multiplies two 1D arrays with broadcasting support.
template <typename T>
SmallVector<T> MultiplyTwoArrays(ArrayRef<T> a, ArrayRef<T> b) {
auto get_value_at = [](ArrayRef<T> v, size_t i) -> T {
if (v.size() == 1) return v.front();
return v[i];
};
size_t max_size = std::max(a.size(), b.size());
SmallVector<T> result(max_size);
for (size_t i : llvm::seq<size_t>(0, max_size)) {
result[i] = get_value_at(a, i) * get_value_at(b, i);
}
return result;
}
// Multiplies the value followed by a FakeQuant op and adjusts the quantization
// params. This function only supports symmetrically quantized values.
Value MultiplyFakeQuantValue(OpBuilder& builder, Location loc, Value value,
Value multiplier) {
auto dq_op = value.getDefiningOp<mlir::quant::ir::DequantizeCastOp>();
if (!dq_op) {
auto mul_op = TF::MulOp::create(builder, loc, value, multiplier);
return mul_op.getResult();
}
auto q_op = dq_op.getArg().getDefiningOp<mlir::quant::ir::QuantizeCastOp>();
if (!q_op) return {};
Value float_value = q_op.getArg();
Value new_value = TF::MulOp::create(builder, loc, float_value, multiplier);
auto new_value_type = mlir::cast<TensorType>(new_value.getType());
// Get multiplier value in double.
DenseFPElementsAttr multiplier_attr;
if (!matchPattern(multiplier, m_Constant(&multiplier_attr)) ||
mlir::cast<ShapedType>(multiplier_attr.getType()).getRank() > 1) {
return {};
}
std::vector<double> multiplier_values;
absl::c_transform(multiplier_attr, std::back_inserter(multiplier_values),
[](auto v) { return FloatAttr::getValueAsDouble(v); });
ArrayRef<double> multiplier_array(multiplier_values.data(),
multiplier_values.size());
// Multiply the quantization parameters by the multiplier.
QuantizedType new_qtype;
auto element_type = mlir::cast<TensorType>(q_op.getType()).getElementType();
if (auto uniform_type = llvm::dyn_cast<UniformQuantizedType>(element_type)) {
if (multiplier_attr.isSplat()) {
double new_scale = multiplier_array.front() * uniform_type.getScale();
new_qtype = UniformQuantizedType::get(
uniform_type.getFlags(), uniform_type.getStorageType(),
uniform_type.getExpressedType(), new_scale,
uniform_type.getZeroPoint(), uniform_type.getStorageTypeMin(),
uniform_type.getStorageTypeMax());
} else {
auto new_scales =
MultiplyTwoArrays(multiplier_array, {uniform_type.getScale()});
int32_t quantized_dim = new_value_type.getRank() - 1;
auto new_zero_points =
SmallVector<int64_t>(new_scales.size(), uniform_type.getZeroPoint());
new_qtype = UniformQuantizedPerAxisType::get(
uniform_type.getFlags(), uniform_type.getStorageType(),
uniform_type.getExpressedType(), new_scales, new_zero_points,
quantized_dim, uniform_type.getStorageTypeMin(),
uniform_type.getStorageTypeMax());
}
} else if (auto per_axis_type =
llvm::dyn_cast_or_null<UniformQuantizedPerAxisType>(
element_type)) {
auto new_scales =
MultiplyTwoArrays(multiplier_array, per_axis_type.getScales());
new_qtype = UniformQuantizedPerAxisType::get(
per_axis_type.getFlags(), per_axis_type.getStorageType(),
per_axis_type.getExpressedType(), new_scales,
per_axis_type.getZeroPoints(), per_axis_type.getQuantizedDimension(),
per_axis_type.getStorageTypeMin(), per_axis_type.getStorageTypeMax());
}
auto quantize = mlir::quant::ir::QuantizeCastOp::create(
builder, q_op.getLoc(), new_value_type.clone(new_qtype), new_value);
auto dequantize = mlir::quant::ir::DequantizeCastOp::create(
builder, dq_op.getLoc(), new_value_type, quantize.getResult());
return dequantize.getResult();
}
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/prepare_lifting.inc"
void PrepareLiftingPass::runOnOperation() {
MLIRContext* ctx = &getContext();
auto func = getOperation();
// The pattern includes decomposing batch normalization ops, fusing add/mul
// with a constant operand to a preceding affine operation.
RewritePatternSet patterns(ctx);
populateWithGenerated(patterns);
patterns.add<RemoveIdentity, ConstantFoldQuantizableOperands>(ctx);
if (op_set_ != OpSet::XLA) {
// Convert Einsum into BatchMatMul for non-XLA opsets.
// For the uniform opset, it is requested to maintain the BatchMatmul logic.
// For the TF opset, since we need to test the effect we remain it as a
// future work.
patterns.add<TF::ConvertTFEinsumOp>(ctx);
}
if (failed(applyPatternsGreedily(func, std::move(patterns)))) {
func.emitError() << "quant-prepare-lifting failed.";
signalPassFailure();
}
}
} // namespace
std::unique_ptr<OperationPass<func::FuncOp>> CreatePrepareLiftingPass(
const OpSet target_opset) {
return std::make_unique<PrepareLiftingPass>(target_opset);
}
static PassRegistration<PrepareLiftingPass> pass;
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/passes/prepare_lifting.td
+209
View File
@@ -0,0 +1,209 @@
/* 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 "mlir/IR/OpBase.td"
include "mlir/IR/PatternBase.td"
include "mlir/Dialect/Func/IR/FuncOps.td"
include "mlir/Dialect/Arith/IR/ArithOps.td"
include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.td"
include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"
include "mlir/Dialect/Arith/IR/ArithOps.td"
// Converts arith.constant ops from freezing passes back to tf.Const ops.
def ConvertArithConstToTfConst : Pat<
(Arith_ConstantOp:$res DenseElementsAttr:$value),
(TF_ConstOp $value),
[(AnyStaticShapeTensor $res)]>;
// Remove CheckNumerics op
def RemoveCheckNumerics : Pat<
(TF_CheckNumericsOp $arg, $msg),
(replaceWithValue $arg)>;
// Remove StopGradient op
def RemoveStopGradient : Pat<
(TF_StopGradientOp $arg),
(replaceWithValue $arg)>;
// Converts tf.FusedBatchNormV3 into a sequence of more primitive arithmetic
// operations. Specifically, performs the following calculation:
//
// (x - mean) * scale / sqrt(variance + epsilon) + offset
//
// Let multiplier = scale / sqrt(variance + epsilon),
// to compute
// (x - mean) * scale / sqrt(variance + epsilon) + offset,
// is then to compute
// (x * multiplier) + (offset - mean * multiplier).
//
// TODO(b/228916181): There is a known issue with this DDR rule that it doesn't
// take into account broadcasting conditions. If the issue needs to be handled,
// see tensorflow/compiler/mlir/lite/transforms/prepare_tf.cc
def FoldFusedBatchNormV3: Pattern<
(TF_FusedBatchNormV3Op:$root
$x, $scale, $offset, $mean, $variance,
F32Attr:$epsilon, $exponential_avg_factor,
$data_format, IsFalseBoolAttr:$is_training),
[(TF_AddV2Op
(TF_MulOp
$x,
(TF_MulOp:$multiplier
$scale,
(TF_RsqrtOp
(TF_AddV2Op $variance,
(TF_ConstOp $epsilon))))),
(TF_SubOp $offset, (TF_MulOp $mean, $multiplier))),
// We already guaranteed that the last five results have no use so it does
// not matter what value we provide here for replacement.
/*batch_mean=*/(replaceWithValue $x),
/*batch_variance=*/(replaceWithValue $x),
/*reserve_space_1=*/(replaceWithValue $x),
/*reserve_space_2=*/(replaceWithValue $x),
/*reserve_space_3=*/(replaceWithValue $x)],
[(HasNoUseOf:$root__1), (HasNoUseOf:$root__2),
(HasNoUseOf:$root__3), (HasNoUseOf:$root__4),
(HasNoUseOf:$root__5)]>;
class HasEqualElementSize<list<int> shape_1, list<int> shape_2> : Constraint<
CPred<"HasEqualElementSize($0, $1,"
"llvm::ArrayRef<int>({" # !interleave(shape_1, ", ") # "}),"
"llvm::ArrayRef<int>({" # !interleave(shape_2, ", ") # "}))">,
"Checks if the given dimensions contain the same number of elements.">;
def ReshapableTo1DTensor : Constraint<
CPred<"ReshapableTo1DTensor(llvm::cast<ShapedType>($0.getType()))">,
"Checks if the value dims are all ones except the right most dim">;
def ReshapeTo1DTensor : NativeCodeCall<
"ReshapeTo1DTensor($_builder, $_loc, $0)">;
def HasEqualShape : Constraint<CPred<
"llvm::cast<ShapedType>($0.getType()).hasRank() && "
"llvm::cast<ShapedType>($1.getType()).hasRank() && "
"llvm::cast<ShapedType>($0.getType()).getShape() == llvm::cast<ShapedType>($1.getType()).getShape()">,
"Checks if the shapes of tensors are same.">;
// Make the 1D value $0 broadcastable with the shape of $1.
def MakeOneDimValueBroadcastable : NativeCodeCall<
"MakeOneDimValueBroadcastable($_builder, $_loc, $0, llvm::cast<ShapedType>($1.getType()))">;
// Match convolution op with "NHWC" data format or matmul op.
def SupportedAffineOpMatcher : NativeCodeCall<
"MatchSupportedAffineOp($_self, $0, $1, $2)">;
// Checks if a value can be symetrically quantized.
def CanBeSymmetricallyQuantized : Constraint<CPred<"CanBeSymmetricallyQuantized($0)">>;
// Multiplies the value followed by a FakeQuant op and adjusts its params.
def MultiplyFakeQuantValue : NativeCodeCall<
"MultiplyFakeQuantValue($_builder, $_loc, $0...)">;
// Convert AddV2Op following an AffineOp to BiasAddOp.
// For Conv3D, even though the Conv3D op has "NDHWC" data format, the BiasAdd
// will still has the data format of "NHWC".
def ConvertAddToBiasAdd : Pat<
(TF_AddV2Op
(SupportedAffineOpMatcher $conv_out, $input, $weight),
(TF_ConstOp:$add_rhs IsFloatElementsAttr:$add_rhs_value)),
(TF_BiasAddOp $conv_out, $add_rhs, (CreateStringAttr<"NHWC">)),
[(HasRankOf<1> $add_rhs_value),
(HasEqualElementSize<[-1], [0]> $conv_out, $add_rhs)], [], (addBenefit -1)>;
// Convert conv+sub+mul pattern to conv+mul+add.
// (conv - sub) * mul -> conv * mul + (-sub) * mul
//
// This is needed to support Conv+BatchNorm pattern from Jax models converted
// using jax2tf w/o native serialization. Note that Jax2tf patterns always
// extend bias shapes to a rank of 4, e.g. 1x1x1x5.
def ConvertSubMulToMulAdd : Pat<
(TF_MulOp
(TF_SubOp
(SupportedAffineOpMatcher $conv_out, $input, $weight),
(TF_ConstOp:$sub_rhs IsFloatElementsAttr:$sub_rhs_value)),
(TF_ConstOp:$mul_rhs IsFloatElementsAttr:$mul_rhs_value)),
(TF_AddV2Op
(TF_MulOp $conv_out, (ReshapeTo1DTensor $mul_rhs)),
(TF_MulOp
(TF_NegOp (ReshapeTo1DTensor $sub_rhs)),
(ReshapeTo1DTensor $mul_rhs))),
[(ReshapableTo1DTensor $mul_rhs),
(ReshapableTo1DTensor $sub_rhs),
(HasEqualElementSize<[-1], [-1]> $conv_out, $mul_rhs),
(HasEqualElementSize<[-1], [-1]> $conv_out, $sub_rhs)]>;
// TODO(b/278493977): Create generic implementation of lifting any fused op
// with any reshaping op
def ConvertAddWithReshapeToBiasAddWithReshape : Pat<
(TF_AddV2Op
(TF_ReshapeOp:$reshape_out
(SupportedAffineOpMatcher $_, $_, $_),
$_
),
(TF_ConstOp:$add_rhs IsFloatElementsAttr:$add_rhs_value)),
(TF_BiasAddOp $reshape_out, $add_rhs, (CreateStringAttr<"NHWC">)),
[(HasRankOf<1> $add_rhs_value),
(HasEqualElementSize<[-1], [0]> $reshape_out, $add_rhs)]>;
// Fuse consecutive BiasAddOp and an AddV2Op.
// We also handle the case where add_rhs has rank 4.
def FuseBiasAndAddV2 : Pat<
(TF_AddV2Op
(TF_BiasAddOp:$bias_add
$conv_out,
(TF_ConstOp:$bias IsFloatElementsAttr:$bias_value), $data_format),
(TF_ConstOp:$add_rhs IsFloatElementsAttr:$add_rhs_value)),
(TF_BiasAddOp
$conv_out, (TF_AddV2Op $bias, (ReshapeTo1DTensor $add_rhs)), $data_format),
[(HasOneUse $bias_add),
(ReshapableTo1DTensor $add_rhs),
(HasEqualElementSize<[-1], [-1]> $bias, $add_rhs)]>;
// Fuse AffineOp followed by an MulOp patterns.
def FuseAffineOpAndMul : Pat<
(TF_MulOp
(SupportedAffineOpMatcher $conv_out, $input, $weight),
(TF_ConstOp:$mul_rhs IsFloatElementsAttr:$mul_rhs_value)),
(CloneOpWithReplacedOperands
(GetDefiningOp $conv_out),
$input,
(MultiplyFakeQuantValue $weight,
(MakeOneDimValueBroadcastable $mul_rhs, $weight))),
[(HasOneUse $conv_out),
(HasRankOf<1> $mul_rhs_value),
(HasStaticShapeConstraint $weight),
(CanBeSymmetricallyQuantized $weight),
(HasEqualElementSize<[-1], [0]> $conv_out, $mul_rhs)]>;
// Fuse AffineOp followed by an BiasAddOp and an MulOp patterns.
def FuseAffineOpWithBiasAddAndMul : Pat<
(TF_MulOp
(TF_BiasAddOp:$bias_add
(SupportedAffineOpMatcher $conv_out, $input, $weight),
$bias, $data_format),
(TF_ConstOp:$mul_rhs IsFloatElementsAttr:$mul_rhs_value)),
(TF_BiasAddOp
(CloneOpWithReplacedOperands
(GetDefiningOp $conv_out),
$input,
(MultiplyFakeQuantValue $weight,
(MakeOneDimValueBroadcastable $mul_rhs, $weight))),
(MultiplyFakeQuantValue $bias, $mul_rhs), $data_format),
[(HasOneUse $conv_out),
(HasOneUse $bias_add),
(HasRankOf<1> $mul_rhs_value),
(HasStaticShapeConstraint $weight),
(CanBeSymmetricallyQuantized $weight),
(CanBeSymmetricallyQuantized $bias),
(HasEqualShape $bias, $mul_rhs_value)]>;
tensorflow/compiler/mlir/quantization/tensorflow/passes/prepare_quantize.cc
+439
View File
@@ -0,0 +1,439 @@
/* 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.
==============================================================================*/
// Copied and modified from
// //third_party/tensorflow/compiler/mlir/lite/transforms/prepare_quantize.cc
// This transformation pass applies quantization propagation on TF dialect.
#include <iterator>
#include <memory>
#include <optional>
#include <utility>
#include "llvm/Support/Casting.h"
#include "llvm/Support/CommandLine.h"
#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/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinTypeInterfaces.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Operation.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/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 "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"
#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_config.h"
#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_driver.h"
#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/ops/tf_op_quant_spec.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h" // IWYU pragma: keep
//===----------------------------------------------------------------------===//
// The prepare-quantize Pass.
//
namespace mlir {
namespace quant {
namespace {
using ::mlir::quant::ir::DequantizeCastOp;
using ::mlir::quant::ir::QuantizeCastOp;
using QuantMethod = tensorflow::quantization::QuantizationMethod::PresetMethod;
// Applies prepare quantization on the model in TF dialect. This pass runs
// before the quantization pass and propagate the quantization parameters
// across ops. This step is necessary for post-training quantization and also
// making the quantization rule for some operations in the quantization-aware
// training quantization simpler.
class PrepareQuantizePass
: public PassWrapper<PrepareQuantizePass, OperationPass<func::FuncOp>> {
void getDependentDialects(DialectRegistry& registry) const override {
registry.insert<TF::TensorFlowDialect, ::mlir::quant::QuantDialect,
::mlir::quant::ir::TFQuantDialect>();
}
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PrepareQuantizePass)
// Constructor used by the PassRegistration and enforce uint8 quantization.
// This is only used by test.
explicit PrepareQuantizePass() {
quant_specs_.inference_type = tensorflow::DT_QINT8;
}
// Constructor used by manually creating the pass.
explicit PrepareQuantizePass(const QuantizationSpecs& quant_specs,
QuantMethod quantization_method)
: quant_specs_(quant_specs) {
quant_specs_.inference_type = tensorflow::DT_QINT8;
enable_per_channel_quantization_ = !quant_specs_.disable_per_channel;
enable_post_training_quantize_ =
(quantization_method == tensorflow::quantization::QuantizationMethod::
METHOD_STATIC_RANGE_INT8);
}
PrepareQuantizePass(const PrepareQuantizePass& other) {
quant_specs_ = other.quant_specs_;
enable_post_training_quantize_ = other.enable_post_training_quantize_;
enable_per_channel_quantization_ = !quant_specs_.disable_per_channel;
}
explicit PrepareQuantizePass(const QuantizationSpecs& quant_specs)
: quant_specs_(quant_specs) {
enable_post_training_quantize_ = quant_specs.post_training_quantization;
}
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-prepare-quantize";
}
StringRef getDescription() const final {
// This is a brief description of the pass.
return "Prepare TF dialect for quantization";
}
void runOnOperation() override;
private:
// Set the quantization parameters of the input nodes. These parameters are
// converted from the user specified input value ranges. The input nodes with
// non-float tensor types will be skipped because they are not quantizable.
// Return true if number of input nodes doesn't equal to that of the input
// ranges.
bool SetInputNodesQuantizationParams(func::FuncOp func);
// 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 RemoveRedundantStats(func::FuncOp func);
// Verify the quantization specification is expected for quantizing the
// current function.
bool IsLegalQuantSpecs(func::FuncOp func) {
if (func.getName() == quant_specs_.target_func) {
return func.getNumArguments() == quant_specs_.input_ranges.size();
}
return true;
}
// Get the min and max values from the quantization specification for the
// current function and argument index. Uses default values if the function
// is specified in the `quantize_allowlist`.
std::pair<std::optional<double>, std::optional<double>>
GetMinMaxValuesForArgument(llvm::StringRef func_name, int index) {
if (func_name == quant_specs_.target_func) {
return quant_specs_.input_ranges[index];
} else {
return {0.0, 255.0};
}
}
// Apply some sanity check and report some warnings for those who don't follow
// the best quantization practice. This also fixes some simple violations.
void SanityCheckAndAdjustment(func::FuncOp func);
// Whether the func contains Quantize ops. This is used to determine whether
// to use the quantization parameters from the fixed output range property.
bool ContainsQuantizeOps(func::FuncOp func);
QuantizationSpecs quant_specs_;
Option<bool> enable_post_training_quantize_{
*this, "post-training-quantize", llvm::cl::init(false),
llvm::cl::desc("Enable post training quantization. Only used in tests.")};
// A local flag is needed for testing conditions in
// prepare_quantize_ptq_per_channel.mlir.
Option<bool> enable_per_channel_quantization_{
*this, "enable-per-channel-quantization", llvm::cl::init(false),
llvm::cl::desc("Whether enable per-channel quantized weights.")};
};
bool PrepareQuantizePass::SetInputNodesQuantizationParams(func::FuncOp func) {
StringRef func_name = func.getName();
auto has_quantize_op = [&](const Value arg) {
return (arg.hasOneUse() && llvm::isa<QuantizeCastOp>(*arg.user_begin()));
};
bool need_to_set_input_nodes_quantization_params = false;
for (const BlockArgument arg : func.getArguments()) {
auto shaped = mlir::dyn_cast<ShapedType>(arg.getType());
if (shaped && mlir::isa<FloatType>(shaped.getElementType()) &&
!has_quantize_op(arg)) {
need_to_set_input_nodes_quantization_params = true;
break;
}
}
if (!need_to_set_input_nodes_quantization_params) {
return false;
}
// If the validation fails, the pass should stop immediately.
if (!IsLegalQuantSpecs(func)) {
return true;
}
OpBuilder builder(func);
bool is_signed = quant_specs_.IsSignedInferenceType();
IntegerAttr num_bits =
builder.getI32IntegerAttr(quant_specs_.GetQuantizationTypeWidth());
BoolAttr narrow_range = builder.getBoolAttr(false);
auto add_quantize_op = [&](Location loc, Type input_type, Block* block,
Block::iterator insertion_point, Value arg,
int i) {
if (auto shaped = mlir::dyn_cast<ShapedType>(input_type)) {
if (mlir::isa<FloatType>(shaped.getElementType())) {
// If there are existing quantize ops, they are from training and we
// should respect them.
if (has_quantize_op(arg)) {
return;
}
auto min_max = GetMinMaxValuesForArgument(func_name, i);
// The input min/max or mean/std are not specified, then skip.
if (!min_max.first.has_value() || !min_max.second.has_value()) return;
TypeAttr params = quant::GetQuantizedTypeAttr(
builder, input_type, builder.getF64FloatAttr(min_max.first.value()),
builder.getF64FloatAttr(min_max.second.value()),
/*quant_dim=*/-1, num_bits, narrow_range, is_signed);
builder.setInsertionPoint(block, insertion_point);
auto q_op =
QuantizeCastOp::create(builder, loc, params.getValue(), arg);
auto dq_op = DequantizeCastOp::create(builder, loc, input_type,
q_op.getResult());
arg.replaceAllUsesWith(dq_op.getResult());
q_op.setOperand(arg);
}
}
};
for (int i = 0, e = func.getNumArguments(); i != e; ++i) {
BlockArgument arg = func.getArgument(i);
auto* arg_block = arg.getOwner();
add_quantize_op(arg.getLoc(), arg.getType(), arg_block,
std::next(arg_block->begin(), i), arg, i);
}
return false;
}
bool PrepareQuantizePass::RemoveRedundantStats(func::FuncOp func) {
return RemoveRedundantStatsOps(func, GetTFOpQuantSpec, GetTfQuantScaleSpec);
}
static Value Quantized(Operation* user) {
if (auto q = llvm::dyn_cast_or_null<QuantizeCastOp>(user)) {
if (auto dq = llvm::dyn_cast_or_null<DequantizeCastOp>(
*q.getResult().user_begin())) {
return dq.getResult();
}
}
return {};
}
void PrepareQuantizePass::SanityCheckAndAdjustment(func::FuncOp func) {
// If an op output has two users: one of them is a quantize op and another
// one is returned directly, we decide to return the quantized result instead,
// so this op can be quantized. This is only applied on the returned result
// because the error will not be accumulated.
func.walk([&](func::ReturnOp ret) {
int i = 0;
for (Value returned : ret.getOperands()) {
llvm::SmallVector<Value, 4> quantized;
for (auto user : returned.getUsers()) {
if (auto q = Quantized(user)) {
quantized.push_back(q);
}
}
if (quantized.size() == 1) {
ret.setOperand(i, quantized.front());
}
i++;
}
});
// Check for (Quant (Dequant $in), $qA) "qdq" pairs that couldn't be
// eliminated at this point. This only occurs for the pattern
// (Quant (Dequant (Quant $in, $qB)), $qA) $qB != $qA
// where the qdq pair denotes a non-trivial requantization of an
// already quantized value. Since this makes little sense (directly quantizing
// (Quant $in, $qA) would introduce less quantization noise) the likely cause
// is an minor error in constructing the original network model that
// introduced back-to-back Fake Quantization operations. Hence: emit a
// warning. N.b. at this point we're (teporarility) in the quantization
// dialect (presumably enable re-use in xla etc) quantfork::*QuantizeCastOp
// we're matching here.
//
func.walk([&](QuantizeCastOp q_op) {
// If up with end up with
auto dq_op =
dyn_cast_or_null<DequantizeCastOp>(q_op.getOperand().getDefiningOp());
if (!dq_op) {
return;
}
auto dq_arg = dq_op.getOperand();
if (!dq_arg.hasOneUse()) {
// The initial quantization is used someplace else ... so it might be
// reasonable for it to requantized for another purpose.
// Ideally would want to still check whether requantization narrows
// rather than widens the representation.
return;
}
// Invariant:
// isa<quantfork::QuantizeCastOp>(dq_arg.getDefiningOp()) -->
// getdq_arg.getType() != q_op.getResult().getType()
//
// as otherwise qdq pair would have been optimized away.
auto qd_arg_def_q_op =
dyn_cast_or_null<QuantizeCastOp>(dq_arg.getDefiningOp());
if (!qd_arg_def_q_op) {
return;
}
qd_arg_def_q_op.emitWarning()
<< " quantizer's output has another quantizer (" << q_op.getLoc()
<< ") as consumer - intentional?";
});
}
// Merges consecutive QuantizeCast ops. For example, the following case:
// %1 = tf.QuantizeCastOp(%0) : f32 -> qtype1
// %2 = tf.QuantizeCastOp(%1) : qtype1 -> qtype2
// %3 = tf.QuantizedOp1(%1)
// %4 = tf.QuantizedOp2(%2)
// will be tranformed to:
// %1 = tf.QuantizeCastOp(%0) : f32 -> qtype1
// %2 = tf.QuantizeCastOp(%0) : f32 -> qtype2
// %3 = tf.QuantizedOp1(%1)
// %4 = tf.QuantizedOp2(%2)
// Converting from f32 -> qtype1 -> qtype2 will add unexpected quantization
// lost for %2. This pattern avoids that by converting from f32 -> qtype2
// directly.
class MergeConsecutiveQuantizeCast
: public mlir::OpRewritePattern<QuantizeCastOp> {
public:
explicit MergeConsecutiveQuantizeCast(MLIRContext* context)
: OpRewritePattern<QuantizeCastOp>(context) {}
private:
LogicalResult matchAndRewrite(QuantizeCastOp q_op,
PatternRewriter& rewriter) const override {
auto preceding_qcast = q_op.getArg().getDefiningOp<QuantizeCastOp>();
if (!preceding_qcast) return failure();
auto new_qcast = QuantizeCastOp::create(
rewriter, q_op.getLoc(), q_op.getType(), preceding_qcast.getArg());
new_qcast->setAttr(kVolatileOpAttrName, rewriter.getUnitAttr());
q_op->replaceAllUsesWith(new_qcast);
return success();
}
};
bool PrepareQuantizePass::ContainsQuantizeOps(func::FuncOp func) {
for (const auto& op : func.getOps()) {
if (llvm::isa<DequantizeCastOp>(op)) return true;
}
return false;
}
using PrepareQuantStats =
quant::ConvertStatsToQDQs<QuantizeCastOp, DequantizeCastOp>;
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/prepare_quantize.inc"
void PrepareQuantizePass::runOnOperation() {
func::FuncOp func = getOperation();
MLIRContext* ctx = func.getContext();
quant_specs_.post_training_quantization = enable_post_training_quantize_;
if (quant_specs_.post_training_quantization) {
RemoveRedundantStats(func);
} else {
// Set the quantization parameters for the quantizable input nodes. If this
// failed, return the function immediately. This is only required for
// quantization aware training model conversion.
if (SetInputNodesQuantizationParams(func)) {
return;
}
}
bool is_signed = quant_specs_.IsSignedInferenceType();
int bit_width = quant_specs_.GetQuantizationTypeWidth();
// When this is true, the quantizer will try its best to extract the
// quantization parameters from the op quantization property and constant
// content. This is also set to true when the `quantize_allowlist` and
// `quantize_signed` test flags are enabled.
bool eager_quantize = ContainsQuantizeOps(func);
// Infer the tensor range for the activation ops and weight constants unless
// it is disabled explicitly.
bool infer_tensor_range =
(quant_specs_.post_training_quantization || eager_quantize) &&
!quant_specs_.disable_infer_tensor_range;
// During the legalization, unsigned quantized type is used, so we have to
// convert all of them to signed.
RewritePatternSet patterns(ctx);
populateWithGenerated(patterns);
patterns.add<ConvertUnsignedToSigned<QuantizeCastOp>>(ctx);
// Convert quant stats to int8 quantization parameters.
// Currently, only activation stats are imported, so narrow_range = false.
patterns.add<PrepareQuantStats>(bit_width, false, true,
/*legacy_float_scale=*/false, ctx);
if (failed(applyPatternsGreedily(func, std::move(patterns)))) {
signalPassFailure();
}
SanityCheckAndAdjustment(func);
// Finally, the quantization parameters can be propagated to the rest of the
// values (tensors).
ApplyQuantizationParamsPropagation(
func, is_signed, /*bit_width=*/8, !enable_per_channel_quantization_,
GetTFOpQuantSpec, GetTfQuantScaleSpec, infer_tensor_range,
quant_specs_.legacy_float_scale, /*is_qdq_conversion=*/false);
RewritePatternSet patterns2(ctx);
patterns2.add<MergeConsecutiveQuantizeCast>(ctx);
if (failed(applyPatternsGreedily(func, std::move(patterns2)))) {
signalPassFailure();
}
}
} // namespace
// Creates an instance of the TensorFlow dialect PrepareQuantize pass.
std::unique_ptr<OperationPass<func::FuncOp>> CreatePrepareQuantizePass(
const QuantizationSpecs& quant_specs, QuantMethod quantization_method) {
return std::make_unique<PrepareQuantizePass>(quant_specs,
quantization_method);
}
static PassRegistration<PrepareQuantizePass> pass;
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/passes/prepare_quantize.td
+28
View File
@@ -0,0 +1,28 @@
/* 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 "mlir/IR/OpBase.td"
include "mlir/IR/PatternBase.td"
include "mlir/Dialect/Func/IR/FuncOps.td"
include "mlir/Dialect/Arith/IR/ArithOps.td"
include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.td"
include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"
// Converts tf.Const to arith.constant for statically shaped, non-opaque constants.
// Needed for QuantizationDriver to recognize constants.
def ConvertTfConstToArithConst : Pat<
(TF_ConstOp:$res DenseElementsAttr:$value),
(Arith_ConstantOp $value),
[(AnyStaticShapeTensor $res)], [], (addBenefit 10)>;
tensorflow/compiler/mlir/quantization/tensorflow/passes/prepare_quantize_drq.cc
+312
View File
@@ -0,0 +1,312 @@
/* 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.
==============================================================================*/
// Copied and modified from
// //third_party/tensorflow/compiler/mlir/lite/transforms/prepare_quantize_dynamic_range.cc
// This transformation pass applies quantization propagation on TF dialect.
#include <memory>
#include <utility>
#include "llvm/Support/Casting.h"
#include "llvm/Support/CommandLine.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/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/BuiltinTypeInterfaces.h" // from @llvm-project
#include "mlir/IR/DialectRegistry.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Matchers.h" // from @llvm-project
#include "mlir/IR/Operation.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/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassRegistry.h" // from @llvm-project
#include "mlir/Rewrite/FrozenRewritePatternSet.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.h" // IWYU pragma: keep
#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"
#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_config.h"
#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/ops/tf_op_quant_spec.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
//===----------------------------------------------------------------------===//
// The prepare-quantize-drq Pass.
//
namespace mlir {
namespace quant {
namespace {
using QuantizationUnit = std::pair<Operation*, int>;
using QuantizationUnits = llvm::SetVector<QuantizationUnit>;
using ::tensorflow::quantization::OpSet;
// Applies prepare quantization on the model in TF dialect for dynamic range
// quantization case.
class PrepareQuantizeDRQPass
: public PassWrapper<PrepareQuantizeDRQPass, OperationPass<ModuleOp>> {
void getDependentDialects(DialectRegistry& registry) const override {
registry.insert<TF::TensorFlowDialect, ::mlir::quant::QuantDialect,
::mlir::quant::ir::TFQuantDialect>();
}
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PrepareQuantizeDRQPass)
// Constructor used by the PassRegistration and enforce int8 quantization.
// This is only used by test.
explicit PrepareQuantizeDRQPass() : op_set_(OpSet::UNIFORM_QUANTIZED) {
quant_specs_.inference_type = tensorflow::DT_QINT8;
}
// Constructor used by manually creating the pass.
explicit PrepareQuantizeDRQPass(const QuantizationSpecs& quant_specs,
OpSet op_set)
: quant_specs_(quant_specs), op_set_(op_set) {
enable_per_channel_quantization_ = !quant_specs_.disable_per_channel;
}
PrepareQuantizeDRQPass(const PrepareQuantizeDRQPass& other) {
quant_specs_ = other.quant_specs_;
op_set_ = other.op_set_;
enable_per_channel_quantization_ = !quant_specs_.disable_per_channel;
}
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-prepare-quantize-drq";
}
StringRef getDescription() const final {
// This is a brief description of the pass.
return "Prepare TF dialect for dynamic range quantization";
}
// The function might contain stats ops which are redundant for processing
// dynamic range quantization. And stats ops may cause conflict while
// processing the function for dynamic range quantization. Therefore, this
// method preprocess the function to remove all stats ops.
void removeAllStatsOp(func::FuncOp func);
void runOnOperation() override;
private:
QuantizationSpecs quant_specs_;
OpSet op_set_;
Option<bool> enable_per_channel_quantization_{
*this, "enable-per-channel-quantization", llvm::cl::init(false),
llvm::cl::desc("Whether enable per-channel quantized weights.")};
};
// If the weight is applicable to dynamic range quantization, insert Quantize
// and Dequantize ops with per-tensor scale.
class PrepareDRQQuantizableOp : public OpRewritePattern<arith::ConstantOp> {
public:
explicit PrepareDRQQuantizableOp(MLIRContext* context,
const quant::QuantizationSpecs& quant_specs,
OpSet op_set,
bool enable_per_channel_quantization)
: OpRewritePattern<arith::ConstantOp>(context),
quant_specs_(quant_specs),
op_set_(op_set),
enable_per_channel_quantization_(enable_per_channel_quantization) {}
LogicalResult matchAndRewrite(arith::ConstantOp op,
PatternRewriter& rewriter) const override {
QuantizationUnits quantizable_ops;
// 1. Collect quantizable ops.
if (!(getQuantizableOps(op, quantizable_ops))) {
return failure();
}
// 2. Quantize collected ops. It is immediately quantized by inserting Q-DQ
// pair for int8.
if (!(quantizeOps(rewriter, op, quantizable_ops))) {
return failure();
}
return success();
}
private:
// Mark users that are applicable for dynamic range quantization where the
// criteria for determining quantizable ops differs by the inference type.
bool getQuantizableOps(arith::ConstantOp op,
QuantizationUnits& quantizable_ops) const {
// Non-float tensors do not need quantization.
auto type = mlir::dyn_cast<ShapedType>(op.getType());
if (!type || !type.getElementType().isF32()) return false;
Value value = op.getResult();
// Check whether dynamic range quantization can be applied.
for (auto& use : value.getUses()) {
Operation* user = use.getOwner();
int operand_num = use.getOperandNumber();
std::unique_ptr<OpQuantSpec> spec = GetTFOpQuantSpec(user);
if (quant_specs_.inference_type == tensorflow::DT_QINT8 &&
spec->quantizable_operands.contains(operand_num)) {
quantizable_ops.insert({user, operand_num});
}
}
return !quantizable_ops.empty();
}
// Apply per-tensor quantization for int8 dynamic range quantization.
bool quantizeOpAsInt8(PatternRewriter& rewriter, arith::ConstantOp op,
QuantizationUnit quant_op) const {
auto [quantized_op, weight_idx] = quant_op;
const bool is_narrow_range = true;
const bool is_legacy_float = quant_specs_.legacy_float_scale;
const bool is_signed = quant_specs_.IsSignedInferenceType();
const int bit_width = quant_specs_.GetQuantizationTypeWidth();
std::unique_ptr<OpQuantSpec> spec = GetTFOpQuantSpec(quantized_op);
const int quant_dim = spec->coeff_op_quant_dim[weight_idx];
const bool is_per_channel_quantization =
enable_per_channel_quantization_ && quant_dim != -1;
QuantizedType quant_type;
DenseFPElementsAttr attr;
if (!matchPattern(op->getResult(0), m_Constant(&attr))) return false;
if (attr.size() < quant_specs_.minimum_elements_for_weights) {
op->emitRemark("Quantization is skipped for ")
<< quantized_op->getName().getStringRef().str() << " because it has "
<< mlir::dyn_cast<DenseFPElementsAttr>(attr).size()
<< " elements which is fewer than the threshold("
<< quant_specs_.minimum_elements_for_weights << " elements).";
return false;
}
if (is_per_channel_quantization) {
quant_type = mlir::dyn_cast<quant::QuantizedType>(
quant::GetUniformQuantizedPerAxisTypeForWeight(
attr, quant_dim,
/*symmetric=*/true, bit_width, is_signed, is_narrow_range,
is_legacy_float));
} else {
quant_type = mlir::dyn_cast<quant::QuantizedType>(
quant::GetUniformQuantizedTypeForWeight(
attr, is_narrow_range && is_signed, bit_width, is_signed,
is_narrow_range, is_legacy_float));
}
return insertQDQ(rewriter, op, quant_type, quant_op);
}
// Insert Quantize and Dequantize ops.
bool insertQDQ(PatternRewriter& rewriter, arith::ConstantOp op,
QuantizedType quant_type, QuantizationUnit quant_op) const {
if (!quant_type) return false;
Operation* quantize_op = quant_op.first;
int quantize_operand_num = quant_op.second;
Type expressed_type = op.getResult().getType();
Type cast_type = quant_type.castFromExpressedType(expressed_type);
// Insert DQ-op if it does not exist yet. Otherwise, just rewire without
// creating a new DQ-op.
for (auto connected_op : op->getUsers()) {
auto q_op =
llvm::dyn_cast_or_null<mlir::quant::ir::QuantizeCastOp>(connected_op);
if (q_op && q_op.getType() == cast_type) {
auto dq_op = llvm::cast<mlir::quant::ir::DequantizeCastOp>(
q_op.getResult().use_begin()->getOwner());
quantize_op->setOperand(quantize_operand_num, dq_op);
return false;
}
}
rewriter.setInsertionPointAfter(op);
auto q = mlir::quant::ir::QuantizeCastOp::create(rewriter, op->getLoc(),
cast_type, op.getResult());
auto dq = mlir::quant::ir::DequantizeCastOp::create(rewriter, op->getLoc(),
expressed_type, q);
quantize_op->setOperand(quantize_operand_num, dq.getResult());
return true;
}
// For each filtered user, apply quantization.
bool quantizeOps(PatternRewriter& rewriter, arith::ConstantOp op,
QuantizationUnits& quantizable_ops) const {
bool quantized = false;
for (auto& quant_op : quantizable_ops) {
if (quant_specs_.inference_type == tensorflow::DT_QINT8) {
quantized |= quantizeOpAsInt8(rewriter, op, quant_op);
}
}
return quantized;
}
protected:
QuantizationSpecs quant_specs_;
OpSet op_set_;
bool enable_per_channel_quantization_;
};
// Remove all the stats ops which are redundant for dynamic range quantizaiton.
void PrepareQuantizeDRQPass::removeAllStatsOp(func::FuncOp func) {
func.walk([&](mlir::quant::ir::StatisticsOp stats_op) {
stats_op.replaceAllUsesWith(stats_op.getArg());
stats_op.erase();
});
}
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/prepare_quantize.inc"
void PrepareQuantizeDRQPass::runOnOperation() {
MLIRContext* ctx = &getContext();
RewritePatternSet patterns(ctx);
ModuleOp module_op = getOperation();
populateWithGenerated(patterns);
patterns.add<PrepareDRQQuantizableOp>(ctx, quant_specs_, op_set_,
enable_per_channel_quantization_);
FrozenRewritePatternSet frozen_patterns(std::move(patterns));
for (auto func : module_op.getOps<func::FuncOp>()) {
removeAllStatsOp(func);
if (failed(applyPatternsGreedily(func, frozen_patterns))) {
func.emitError() << "quant-prepare-quantize-drq failed.";
signalPassFailure();
}
}
}
} // namespace
// Creates an instance of the TensorFlow dialect PrepareQuantizeDRQ
// pass.
std::unique_ptr<OperationPass<ModuleOp>> CreatePrepareQuantizeDRQPass(
const QuantizationSpecs& quant_specs, const OpSet op_set) {
return std::make_unique<PrepareQuantizeDRQPass>(quant_specs, op_set);
}
static PassRegistration<PrepareQuantizeDRQPass> pass;
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/passes/preprocess_op.cc
+276
View File
@@ -0,0 +1,276 @@
/* 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 transformation pass applies quantization propagation on TF dialect.
#include <cstdint>
#include <memory>
#include <utility>
#include "absl/container/flat_hash_set.h"
#include "llvm/Support/CommandLine.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/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/DialectRegistry.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/Matchers.h" // from @llvm-project
#include "mlir/IR/Operation.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/IR/SymbolTable.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/IR/ValueRange.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassRegistry.h" // from @llvm-project
#include "mlir/Rewrite/FrozenRewritePatternSet.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"
#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/ops/tf_op_quant_spec.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
//===----------------------------------------------------------------------===//
// The preprocess-op Pass.
//
namespace mlir {
namespace quant {
namespace {
using QuantMethod =
::tensorflow::quantization::QuantizationMethod::PresetMethod;
using QuantizationUnit = std::pair<Operation*, int>;
using QuantizationUnits = llvm::SetVector<QuantizationUnit>;
using ::tensorflow::quantization::OpSet;
// Preprocesses ops to allow multi-axis quantization, prior to quantization
// passes. Currently, per-channel quantization only supports 1D results.
class PreprocessOpPass
: public PassWrapper<PreprocessOpPass, OperationPass<ModuleOp>> {
void getDependentDialects(DialectRegistry& registry) const override {
registry.insert<TF::TensorFlowDialect, QuantDialect,
mlir::quant::ir::TFQuantDialect>();
}
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PreprocessOpPass)
explicit PreprocessOpPass() = default;
// Constructor used by manually creating the pass.
explicit PreprocessOpPass(OpSet op_set, const QuantMethod quantization_method,
bool enable_per_channel_quantization) {
op_set_ = op_set;
quantization_method_ = quantization_method;
enable_per_channel_quantization_ = enable_per_channel_quantization;
}
PreprocessOpPass(const PreprocessOpPass& other) {
op_set_ = other.op_set_;
quantization_method_ = other.quantization_method_;
enable_per_channel_quantization_ = other.enable_per_channel_quantization_;
}
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-preprocess-op";
}
StringRef getDescription() const final {
// This is a brief description of the pass.
return "Preprocess TF op prior to quantization";
}
void runOnOperation() override;
private:
Option<OpSet> op_set_{
*this, "target-opset", llvm::cl::init(OpSet::UNIFORM_QUANTIZED),
llvm::cl::desc("Choose target opset."),
llvm::cl::values(
clEnumValN(OpSet::TF, "TF",
"Uses TF ops that mimic quantization behavior"),
clEnumValN(OpSet::XLA, "XLA", "Uses TF XLA ops"),
clEnumValN(OpSet::UNIFORM_QUANTIZED, "UNIFORM_QUANTIZED",
"Uses TF Uniform Quantized ops"))};
Option<QuantMethod> quantization_method_{
*this, "quantization-method",
llvm::cl::init(tensorflow::quantization::QuantizationMethod::
METHOD_STATIC_RANGE_INT8),
llvm::cl::desc("Choose quantization method."),
llvm::cl::values(
clEnumValN(tensorflow::quantization::QuantizationMethod::
METHOD_STATIC_RANGE_INT8,
"ptq", "Post-training static-range quantization"),
clEnumValN(tensorflow::quantization::QuantizationMethod::
METHOD_DYNAMIC_RANGE_INT8,
"drq", "Post-training dynamic-range quantizaiton"),
clEnumValN(tensorflow::quantization::QuantizationMethod::
METHOD_STATIC_RANGE_WEIGHT_ONLY_INT8,
"weight_only", "Post-training weight-only quantizaiton"))};
Option<bool> enable_per_channel_quantization_{
*this, "enable-per-channel-quantization", llvm::cl::init(false),
llvm::cl::desc("Whether enable per-channel quantized weights.")};
};
// Apply constant transformations for the op_set.
class PreprocessConstantOp : public OpRewritePattern<TF::PartitionedCallOp> {
public:
explicit PreprocessConstantOp(MLIRContext* context, OpSet op_set,
QuantMethod quantization_method,
bool enable_per_channel_quantization)
: OpRewritePattern<TF::PartitionedCallOp>(context),
op_set_(op_set),
quantization_method_(quantization_method),
enable_per_channel_quantization_(enable_per_channel_quantization) {}
LogicalResult addReshapeOpToDepthwiseWeight(TF::PartitionedCallOp op,
PatternRewriter& rewriter,
StringRef function_name) const {
std::unique_ptr<OpQuantSpec> spec = GetTFOpQuantSpec(op);
const absl::flat_hash_set<int> operands = spec->quantizable_operands;
if (operands.size() != 1) return failure();
int weight_operand_idx = *operands.begin();
Operation* weight_op = op.getOperand(weight_operand_idx).getDefiningOp();
DenseFPElementsAttr attr;
if (!matchPattern(weight_op->getResult(0), m_Constant(&attr))) {
return failure();
}
// Get new shape.
llvm::ArrayRef<int64_t> cur_shape = attr.getType().getShape();
int cur_rank = cur_shape.size();
if (cur_rank != 4 || cur_shape[2] == 1) return failure();
TensorType new_shape = RankedTensorType::get(
{cur_shape[0], cur_shape[1], 1, cur_shape[2] * cur_shape[3]},
attr.getElementType());
// Inserts a reshape op.
auto shape_spec_type =
RankedTensorType::get({cur_rank}, rewriter.getIntegerType(64));
auto new_shape_const_attr =
DenseElementsAttr::get(shape_spec_type, new_shape.getShape());
rewriter.setInsertionPointAfter(weight_op);
auto new_shape_const = arith::ConstantOp::create(
rewriter, weight_op->getLoc(), shape_spec_type, new_shape_const_attr);
auto reshape_op =
TF::ReshapeOp::create(rewriter, weight_op->getLoc(), new_shape,
weight_op->getResult(0), new_shape_const);
op->setOperand(weight_operand_idx, reshape_op);
// Create a new function with preprocessed types.
ModuleOp module = op->getParentOfType<ModuleOp>();
SymbolTable symbol_table(module);
func::FuncOp float_func =
dyn_cast<func::FuncOp>(symbol_table.lookup(function_name));
OperandRange func_args = op.getArgs();
func::FuncOp new_float_func = float_func.clone();
SmallVector<Value> new_float_func_args{func_args.begin(), func_args.end()};
new_float_func_args[weight_operand_idx] = reshape_op;
new_float_func.getArgument(weight_operand_idx).setType(new_shape);
new_float_func.setType(FunctionType::get(
getContext(), TypeRange{ValueRange{new_float_func_args}},
new_float_func.getResultTypes()));
symbol_table.insert(new_float_func);
op->setAttr("f", SymbolRefAttr::get(rewriter.getContext(),
new_float_func.getName()));
return success();
}
LogicalResult matchAndRewrite(TF::PartitionedCallOp op,
PatternRewriter& rewriter) const override {
const auto f_attr = mlir::dyn_cast<FlatSymbolRefAttr>(op.getFAttr());
// Non-quantizable op
if (!op->hasAttr(kQuantTraitAttrName)) return failure();
StringRef function_name = f_attr.getValue();
// TODO(b/228928859): Improve the getter function to match attributes rather
// than function name.
if (!function_name.starts_with("composite_")) {
return failure();
}
if (function_name.contains("depthwise_conv2d")) {
// Uniform Quantized op requires weights of tf.DepthwiseConv2dNative to
// be transformed from [H,W,C,M] to [H,W,1,CxM] where
// H=height,W=width,C=channel,M=multiplier. Therefore, a reshape op is
// inserted between the constant op and the function op so that the
// constant is safely transformed for the multi-use cases as well. Note
// that bias doesn't need transformation as its shape is already in [CxM].
if (op_set_ == OpSet::UNIFORM_QUANTIZED ||
(op_set_ == OpSet::XLA && enable_per_channel_quantization_ &&
quantization_method_ ==
tensorflow::quantization::QuantizationMethod::
METHOD_STATIC_RANGE_WEIGHT_ONLY_INT8)) {
return addReshapeOpToDepthwiseWeight(op, rewriter, function_name);
}
}
return failure();
}
private:
const OpSet op_set_;
const QuantMethod quantization_method_;
const bool enable_per_channel_quantization_;
};
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/preprocess_op.inc"
void PreprocessOpPass::runOnOperation() {
MLIRContext* ctx = &getContext();
RewritePatternSet patterns(ctx);
ModuleOp module_op = getOperation();
populateWithGenerated(patterns);
patterns.add<PreprocessConstantOp>(ctx, op_set_, quantization_method_,
enable_per_channel_quantization_);
FrozenRewritePatternSet frozen_patterns(std::move(patterns));
for (auto func : module_op.getOps<func::FuncOp>()) {
if (failed(applyPatternsGreedily(func, frozen_patterns))) {
func.emitError() << "quant-preprocess-op failed.";
signalPassFailure();
}
}
}
} // namespace
// Creates an instance of the TensorFlow dialect PreprocessOp
// pass.
std::unique_ptr<OperationPass<ModuleOp>> CreatePreprocessOpPass(
const OpSet op_set, QuantMethod quantization_method,
const bool enable_per_channel_quantization) {
return std::make_unique<PreprocessOpPass>(op_set, quantization_method,
enable_per_channel_quantization);
}
static PassRegistration<PreprocessOpPass> pass;
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/passes/preprocess_op.td
+28
View File
@@ -0,0 +1,28 @@
/* 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 "mlir/IR/OpBase.td"
include "mlir/IR/PatternBase.td"
include "mlir/Dialect/Func/IR/FuncOps.td"
include "mlir/Dialect/Arith/IR/ArithOps.td"
include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.td"
include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"
// Converts tf.Const to arith.constant for statically shaped, non-opaque constants.
// Needed for QuantizationDriver to recognize constants.
def ConvertTfConstToArithConst : Pat<
(TF_ConstOp:$res DenseElementsAttr:$value),
(Arith_ConstantOp $value),
[(AnyStaticShapeTensor $res)], [], (addBenefit 10)>;
tensorflow/compiler/mlir/quantization/tensorflow/passes/propagate_quantize_type.cc
+172
View File
@@ -0,0 +1,172 @@
/* 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.
==============================================================================*/
#include <algorithm>
#include <memory>
#include <utility>
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/BuiltinTypeInterfaces.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassRegistry.h" // from @llvm-project
#include "mlir/Rewrite/FrozenRewritePatternSet.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "mlir/Transforms/DialectConversion.h" // from @llvm-project
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/ops/tf_op_quant_spec.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
namespace mlir {
namespace quant {
namespace {
constexpr StringRef kDequantizeFunctionName = "composite_dequantize";
class PropagateQuantizeType
: public PassWrapper<PropagateQuantizeType, OperationPass<ModuleOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PropagateQuantizeType)
// Constructor used by the PassRegistration. This will remove the adaptor ops.
explicit PropagateQuantizeType() = default;
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-propagate-quantize-type";
}
StringRef getDescription() const final {
// This is a brief description of the pass.
return "Propagate quantized type through allowed ops.";
}
void runOnOperation() override;
};
// Propagate dequantize op if the next op supports the data type.
// Given the below graph,
// op_before_dequantize -> dequantize_op -> user_op -> rest_op
// the transformation is applied to result the following graph:
// op_before_dequantize -> user_op -> new_dequantize_op -> rest_op
class PropagateDequantizeOpIfAllowed
: public OpRewritePattern<TF::PartitionedCallOp> {
public:
explicit PropagateDequantizeOpIfAllowed(MLIRContext* context)
: OpRewritePattern<TF::PartitionedCallOp>(context) {}
// Create a new dequantize op that is propagated.
void createNewDequantizeOp(PatternRewriter& rewriter,
TF::PartitionedCallOp original_dequantize_op,
Operation* user_op, int user_idx,
Type new_user_op_type) const {
auto op_before_dequantize = original_dequantize_op.getOperand(0);
// Create a new dequantize op that is propagated.
rewriter.setInsertionPointAfter(user_op);
TF::PartitionedCallOp new_dequantize_op =
cast<TF::PartitionedCallOp>(rewriter.clone(*original_dequantize_op));
// Skip the original dequant op and connect the op before dequantize to the
// user op.
user_op->setOperand(user_idx, op_before_dequantize);
// Wire input/output nodes.
new_dequantize_op->setOperand(0, user_op->getResult(0));
new_dequantize_op->getResult(0).setType(user_op->getResult(0).getType());
user_op->getResult(0).replaceAllUsesExcept(new_dequantize_op->getResult(0),
new_dequantize_op);
user_op->getResult(0).setType(new_user_op_type);
}
LogicalResult matchAndRewrite(TF::PartitionedCallOp op,
PatternRewriter& rewriter) const override {
const auto f_attr = mlir::dyn_cast<FlatSymbolRefAttr>(op.getFAttr());
StringRef function_name = f_attr.getValue();
if (!function_name.starts_with(kDequantizeFunctionName)) return failure();
llvm::SmallVector<Operation*> users(op->getUsers().begin(),
op->getUsers().end());
bool changed = false;
for (auto& use : op->getUses()) {
Operation* user_op = use.getOwner();
int user_idx = use.getOperandNumber();
if (!IsOpWithInt8TypeOperand(user_op)) continue;
// If the next op is terminator, function type needs to be changed so
// handle this case separately when propagating for function op is
// added.
if (std::any_of(user_op->getResult(0).getUsers().begin(),
user_op->getResult(0).getUsers().end(), [](Operation* y) {
return y->hasTrait<OpTrait::IsTerminator>();
}))
continue;
if (IsOpWithDataMovementTrait(user_op)) {
auto op_before_dequantize = op.getOperand(0);
// New user op type needs to be set since user_op can output integer
// type for the data movement case.
auto original_result_type = user_op->getResult(0).getType();
auto new_user_op_type = CloneTypeWithNewElementType(
original_result_type,
mlir::cast<ShapedType>(op_before_dequantize.getType())
.getElementType());
createNewDequantizeOp(rewriter, op, user_op, user_idx,
new_user_op_type);
} else {
createNewDequantizeOp(rewriter, op, user_op, user_idx,
user_op->getResult(0).getType());
}
changed = true;
}
return changed ? success() : failure();
}
};
void PropagateQuantizeType::runOnOperation() {
RewritePatternSet patterns(&getContext());
auto module_op = getOperation();
MLIRContext* ctx = &getContext();
patterns.add<PropagateDequantizeOpIfAllowed>(ctx);
FrozenRewritePatternSet frozen_patterns(std::move(patterns));
// Propagation can happen recursively with multiple functions so keep this
// module level.
for (auto func : module_op.getOps<func::FuncOp>()) {
if (failed(applyPatternsGreedily(func, frozen_patterns))) {
func.emitError() << "quant-propagate-quantize-type failed.";
signalPassFailure();
}
}
}
} // namespace
// Creates an instance of the TensorFlow dialect PropagateQuantizeType pass.
std::unique_ptr<OperationPass<ModuleOp>> CreatePropagateQuantizeTypePass() {
return std::make_unique<PropagateQuantizeType>();
}
static PassRegistration<PropagateQuantizeType> pass;
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/passes/quantize.cc
+576
View File
@@ -0,0 +1,576 @@
/* 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 <string>
#include <utility>
#include "absl/container/flat_hash_set.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/CommandLine.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/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/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/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/LogicalResult.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"
#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_config.h"
#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/ops/tf_op_quant_spec.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/core/framework/types.pb.h"
namespace mlir {
namespace quant {
//===----------------------------------------------------------------------===//
// The actual Quantize Pass.
//===----------------------------------------------------------------------===//
namespace {
using ::mlir::quant::ir::DequantizeCastOp;
using ::mlir::quant::ir::QuantizeCastOp;
using ::mlir::quant::ir::StorageCastOp;
using ::tensorflow::quantization::OpSet;
enum QuantizationTrait { kFullQuantization, kDynamicRangeQuantization };
// Base struct for quantization.
template <QuantizationTrait quantization_trait, typename ConcreteT,
typename RootOpT = DequantizeCastOp>
struct TFQuantizationBase
: public QuantizationPattern<ConcreteT, QuantizeCastOp, DequantizeCastOp,
/*VerifierT=*/void, RootOpT> {
explicit TFQuantizationBase(MLIRContext* ctx,
const QuantPassSpec& quant_params)
: QuantizationPattern<ConcreteT, QuantizeCastOp, DequantizeCastOp,
/*VerifierT=*/void, RootOpT>(ctx, quant_params) {}
// Custom op quantization is not supported.
static bool IsQuantizableCustomOp(Operation* op,
const CustomMap& custom_op_map) {
return false;
}
// All the quantized ops are supported if the quantization method is dynamic
// range quantization.
static bool AllowDynamicRangeQuantizedOperand(
Operation* quantized_op, const CustomMap& custom_op_map) {
auto call_op = cast<TF::PartitionedCallOp>(quantized_op);
StringRef function_name =
llvm::cast<FlatSymbolRefAttr>(call_op.getFAttr()).getValue();
// The below can be generalized as there are more read-only ops added such
// as slice.
const bool is_gather = function_name.contains("gather");
return quantization_trait != kFullQuantization || is_gather;
}
// All the quantized ops are supported if the quantization method is dynamic
// range quantization.
static bool AllowDynamicRangeQuantizedResult(Operation* quantized_op,
const CustomMap& custom_op_map) {
auto call_op = cast<TF::PartitionedCallOp>(quantized_op);
StringRef function_name =
llvm::cast<FlatSymbolRefAttr>(call_op.getFAttr()).getValue();
// The below can be generalized as there are more read-only ops added such
// as slice.
bool is_gather = false;
if (function_name.contains("gather")) is_gather = true;
return quantization_trait != kFullQuantization ||
(quantization_trait == kFullQuantization && is_gather);
}
// If weight_only_quantization is true, the legacy weight-only quantization is
// applied. The legacy weight-only graph has dequantization logic at the
// front.
static bool IsWeightOnlyOp(Operation* quantized_op,
absl::flat_hash_set<std::string>& ops_blocklist,
bool weight_only_quantization,
const CustomMap& custom_op_map) {
return weight_only_quantization;
}
};
// Full integer quantization rewrite pattern using DQ as the root op.
struct TFFullQuantization
: public TFQuantizationBase<kFullQuantization, TFFullQuantization> {
explicit TFFullQuantization(MLIRContext* ctx,
const QuantPassSpec& quant_params)
: TFQuantizationBase<kFullQuantization, TFFullQuantization>(
ctx, quant_params) {}
};
// Full integer quantization rewrite pattern using Q as the root op. This is for
// the quantizable ops without floating-point operands.
struct TFFullQuantizationReverse
: public TFQuantizationBase<kFullQuantization, TFFullQuantizationReverse,
QuantizeCastOp> {
explicit TFFullQuantizationReverse(MLIRContext* ctx,
const QuantPassSpec& quant_params)
: TFQuantizationBase<kFullQuantization, TFFullQuantizationReverse,
QuantizeCastOp>(ctx, quant_params) {}
};
// Dynamic range quantization rewrite pattern using DQ as the root op.
struct TFDynamicRangeQuantization
: public TFQuantizationBase<kDynamicRangeQuantization,
TFDynamicRangeQuantization> {
explicit TFDynamicRangeQuantization(MLIRContext* ctx,
const quant::QuantPassSpec& quant_params)
: TFQuantizationBase<kDynamicRangeQuantization,
TFDynamicRangeQuantization>(ctx, quant_params) {}
};
// Removes quantize-dequantize pairs that are not used in the quantization.
// The benefit of this pattern is set to lower value than other patterns, so
// that the other patterns can work on quantize/dequantize ops first.
class RemoveUnusedQdqPattern : public OpRewritePattern<DequantizeCastOp> {
public:
explicit RemoveUnusedQdqPattern(MLIRContext* context)
: OpRewritePattern<DequantizeCastOp>(context) {}
LogicalResult matchAndRewrite(DequantizeCastOp dq_op,
PatternRewriter& rewriter) const override {
auto q_op = dq_op.getArg().getDefiningOp<QuantizeCastOp>();
if (!q_op) return failure();
dq_op.replaceAllUsesWith(q_op.getArg());
return success();
}
};
class QuantizeSameScaleOpsPattern : public OpRewritePattern<DequantizeCastOp> {
public:
explicit QuantizeSameScaleOpsPattern(
MLIRContext* context, OpQuantScaleSpecGetter op_quant_scale_spec_getter,
OpSet target_opset)
// Set the score to a large number so it is always preferred, after
// quantization patterns.
: OpRewritePattern<DequantizeCastOp>(context, /*benefit=*/200),
op_quant_scale_spec_getter_(op_quant_scale_spec_getter),
target_opset_(target_opset) {}
LogicalResult matchAndRewrite(DequantizeCastOp op,
PatternRewriter& rewriter) const override {
SmallVector<Operation*, 4> quantizing_ops;
auto users = op.getResult().getUsers();
quantizing_ops.append(users.begin(), users.end());
bool changed = false;
// Rewrite the floating-point ops to the quantized version, by fusing
// preceding dequantize ops and succeding quantize ops.
for (Operation* quantizing_op : quantizing_ops) {
// If it is requantize op, we shouldn't rewrite this op.
if (llvm::isa<QuantizeCastOp, DequantizeCastOp>(quantizing_op)) {
return failure();
}
// If the op is terminator, not quantizable or any ops from the mlir quant
// ops dialect, we shouldn't rewrite.
if (quantizing_op->hasTrait<OpTrait::IsTerminator>()) {
return failure();
}
if (!op_quant_scale_spec_getter_(quantizing_op)
->has_same_scale_requirement) {
continue;
}
if (target_opset_ == OpSet::XLA &&
!IsConnectedWithCompsiteFunction(quantizing_op)) {
continue;
}
// Same scale op is not supported for Uniform Quantized ops.
if (target_opset_ == OpSet::UNIFORM_QUANTIZED) {
continue;
}
// Collect all the quantized inputs and "clone" the matched op by these
// inputs.
SmallVector<Value, 4> inputs;
inputs.reserve(quantizing_op->getNumOperands());
for (const auto& operand : quantizing_op->getOperands()) {
Type operand_type = operand.getType();
if (isa<NoneType>(operand_type)) {
inputs.push_back(operand);
continue;
}
Type elem_type = llvm::cast<TensorType>(operand_type).getElementType();
if (auto dq_op =
dyn_cast_or_null<DequantizeCastOp>(operand.getDefiningOp())) {
auto dq_arg_type = llvm::cast<TensorType>(dq_op.getArg().getType());
auto qtype = llvm::cast<QuantizedType>(dq_arg_type.getElementType());
auto scast_op = StorageCastOp::create(
rewriter, dq_op->getLoc(),
dq_arg_type.clone(qtype.getStorageType()), dq_op.getArg());
inputs.push_back(scast_op.getResult());
} else if (!elem_type.isF32()) {
// If the operand is an integer tensor, then it doesn't require the
// DQ op in the pattern.
inputs.push_back(operand);
} else {
return failure();
}
}
// 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();
if (isa<NoneType>(result_type)) {
outputs_replaced.insert({result, enumerated_result.index()});
output_types.push_back(result_type);
continue;
}
auto result_tensor_type = llvm::cast<TensorType>(result_type);
// If the user is the Quantize op, it must be the only user.
if (result.hasOneUse() &&
llvm::isa<QuantizeCastOp>(*result.user_begin())) {
auto user = llvm::cast<QuantizeCastOp>(*result.user_begin());
outputs_replaced.insert(
{user.getResult(), enumerated_result.index()});
auto qtype = llvm::cast<QuantizedType>(
llvm::cast<TensorType>(user.getType()).getElementType());
output_types.push_back(
result_tensor_type.clone(qtype.getStorageType()));
} else if (!result_tensor_type.getElementType().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 {
// TODO(b/224691264): separate matching and rewriting clearly.
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();
}
Operation* quantized_op = rewriter.create(new_state);
if (quantizing_op->getNumRegions() != 0) {
for (const auto& indexed_regions :
llvm::enumerate(quantizing_op->getRegions())) {
IRMapping mapping;
indexed_regions.value().cloneInto(
&quantized_op->getRegion(indexed_regions.index()), mapping);
}
}
for (const auto& output_index_pair : outputs_replaced) {
Value output = output_index_pair.getFirst();
int output_index = output_index_pair.getSecond();
auto scast_op =
StorageCastOp::create(rewriter, output.getLoc(), output.getType(),
quantized_op->getResult(output_index));
output.replaceAllUsesWith(scast_op);
}
changed = true;
}
return success(changed);
}
private:
// Checks whether the operation is connected with a composite function.
// If not, the same-scale op will not be quantized. This decision is based
// on the current assumption that the performance gain of the same-scale
// op itself could not beat the overhead of the quantize and dequantize
// routines need to be added around that op. When the assumption changes,
// this policy might change as well.
bool IsConnectedWithCompsiteFunction(Operation* same_scale_op) const {
for (const auto& operand : same_scale_op->getOperands()) {
auto dq_op = dyn_cast_or_null<DequantizeCastOp>(operand.getDefiningOp());
if (!dq_op) continue;
Operation* preceding_op = dq_op.getArg().getDefiningOp();
if (!preceding_op) continue;
// Check whether the preceding op is a quantized composite function.
if (llvm::isa<TF::PartitionedCallOp>(preceding_op)) {
auto call_op = llvm::cast<TF::PartitionedCallOp>(preceding_op);
if (!IsCompositeFunction(call_op)) continue;
return true;
}
// Check if the preceding op is a quantized same-scale op.
if (llvm::isa<StorageCastOp>(preceding_op)) {
auto sc_op = llvm::cast<StorageCastOp>(preceding_op);
auto sc_arg_type = llvm::dyn_cast<TensorType>(sc_op.getArg().getType());
if (sc_arg_type.getElementType().isInteger(8)) {
return true;
}
}
}
for (const auto& result : same_scale_op->getResults()) {
// If the user is the Quantize op, it must be the only user.
if (!result.hasOneUse() ||
!llvm::isa<QuantizeCastOp>(*result.user_begin())) {
continue;
}
auto q_op = llvm::cast<QuantizeCastOp>(*result.user_begin());
for (auto following_op : q_op->getUsers()) {
// Check whether the preceding op is a quantized composite function.
if (llvm::isa<TF::PartitionedCallOp>(following_op)) {
auto call_op = llvm::cast<TF::PartitionedCallOp>(following_op);
if (!IsCompositeFunction(call_op)) continue;
return true;
}
// Check if the preceding op is a quantized same-scale op.
if (llvm::isa<StorageCastOp>(following_op)) {
auto sc_op = llvm::cast<StorageCastOp>(following_op);
auto sc_arg_type =
llvm::dyn_cast<TensorType>(sc_op.getResult().getType());
if (sc_arg_type.getElementType().isInteger(8)) {
return true;
}
}
}
}
return false;
}
// Checks if op calls a composite function and all the inputs are quantized.
bool IsCompositeFunction(TF::PartitionedCallOp call_op) const {
if (!call_op->hasAttr(kQuantTraitAttrName)) {
return false;
}
const auto f_attr = llvm::dyn_cast<FlatSymbolRefAttr>(call_op.getFAttr());
if (!f_attr || !f_attr.getValue().starts_with("composite_")) {
return false;
}
bool has_quantized_types = false;
for (Value input : call_op.getArgs()) {
if (auto type = llvm::dyn_cast<TensorType>(input.getType())) {
if (isa<FloatType>(type.getElementType())) {
return false;
}
if (isa<QuantizedType>(type.getElementType())) {
has_quantized_types = true;
}
}
}
for (Value output : call_op.getOutput()) {
if (auto type = llvm::dyn_cast<TensorType>(output.getType())) {
if (isa<FloatType>(type.getElementType())) {
return false;
}
if (isa<QuantizedType>(type.getElementType())) {
has_quantized_types = true;
}
}
}
return has_quantized_types;
}
OpQuantScaleSpecGetter op_quant_scale_spec_getter_;
OpSet target_opset_;
};
// The AvgPool op is a same-scale op but it doesn't have int8 kernel, so
// we cast its input to float and its output to int8 as a workaround.
// TODO(b/229183248): Remove this workaround after int8 kernels have been
// added to TF and XLA.
struct QuantizeAvgPoolOpPattern : public OpRewritePattern<StorageCastOp> {
explicit QuantizeAvgPoolOpPattern(MLIRContext* context)
: OpRewritePattern<StorageCastOp>(context, /*benefit=*/100) {}
LogicalResult matchAndRewrite(StorageCastOp sc_op,
PatternRewriter& rewriter) const override {
auto avg_pool_op = sc_op.getArg().getDefiningOp<TF::AvgPoolOp>();
if (!avg_pool_op) return failure();
auto preceding_sc_op =
dyn_cast_or_null<StorageCastOp>(avg_pool_op.getValue().getDefiningOp());
if (!preceding_sc_op) return failure();
// Check if the same-scale requirement is met.
auto dq_arg_type =
llvm::cast<TensorType>(preceding_sc_op.getArg().getType());
auto qtype = llvm::cast<QuantizedType>(dq_arg_type.getElementType());
auto q_result_type = llvm::cast<TensorType>(sc_op.getType());
auto out_qtype = llvm::cast<QuantizedType>(q_result_type.getElementType());
if (qtype != out_qtype) {
avg_pool_op.emitError(
"The preceding StorageCastOp and the following "
"StorageCastOp must have the same quantized type");
return failure();
}
// Cast to float type before the AvgPool op.
OpBuilder::InsertionGuard g(rewriter);
rewriter.setInsertionPointAfter(preceding_sc_op);
auto fcast_op = TF::CastOp::create(rewriter, preceding_sc_op->getLoc(),
dq_arg_type.clone(rewriter.getF32Type()),
preceding_sc_op.getResult());
// Create a new AvgPool op with float type.
TF::AvgPoolOp float_avg_pool_op = TF::AvgPoolOp::create(
rewriter, avg_pool_op->getLoc(),
avg_pool_op.getType().clone(rewriter.getF32Type()),
/*operands=*/fcast_op.getResult(),
/*attributes=*/avg_pool_op->getAttrs());
// Cast back to the storage type after AvgPool op.
auto round_val = TF::RoundOp::create(rewriter, sc_op.getLoc(),
float_avg_pool_op.getOutput());
auto icast_op = TF::CastOp::create(
rewriter, sc_op.getLoc(), q_result_type.clone(qtype.getStorageType()),
round_val);
avg_pool_op.getResult().replaceAllUsesWith(icast_op.getResult());
return success();
}
};
// Applies quantization on the model in TF dialect.
class QuantizePass
: public PassWrapper<QuantizePass, OperationPass<func::FuncOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(QuantizePass)
// Constructor used by the PassRegistration and only used by test.
explicit QuantizePass() {
quant_specs_.inference_type = tensorflow::DT_QINT8;
}
// Constructor used by manually creating the pass.
explicit QuantizePass(const QuantizationSpecs& quant_specs,
OpSet target_opset)
: quant_specs_(quant_specs) {
weight_quantization_ = quant_specs.weight_quantization;
target_opset_ = target_opset;
}
QuantizePass(const QuantizePass& other) : quant_specs_(other.quant_specs_) {
weight_quantization_ = other.weight_quantization_;
target_opset_ = other.target_opset_;
}
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-quantize";
}
StringRef getDescription() const final {
// This is a brief description of the pass.
return "Apply quantization on models in TensorFlow dialect";
}
// Determine if the unused Q-DQ pairs need to be removed. For weight-only
// quantizable ops, Q-DQ ops need to be preserved.
bool shouldKeepUnusedQdqPattern();
void runOnOperation() override;
private:
QuantizationSpecs quant_specs_;
Option<bool> weight_quantization_{
*this, "weight-quantization", llvm::cl::init(false),
llvm::cl::desc("Whether to enable weight quantization.")};
Option<OpSet> target_opset_{
*this, "target-opset", llvm::cl::init(OpSet::TF),
llvm::cl::desc("Choose target opset."),
llvm::cl::values(
clEnumValN(OpSet::TF, "TF",
"Uses TF ops that mimic quantization behavior"),
clEnumValN(OpSet::XLA, "XLA", "Uses TF XLA ops"),
clEnumValN(OpSet::UNIFORM_QUANTIZED, "UNIFORM_QUANTIZED",
"Uses TF Uniform Quantized ops"))};
};
bool QuantizePass::shouldKeepUnusedQdqPattern() {
return target_opset_ == OpSet::XLA &&
(quant_specs_.weight_only_quantization ||
quant_specs_.weight_quantization);
}
void QuantizePass::runOnOperation() {
RewritePatternSet patterns(&getContext());
auto func = getOperation();
auto* ctx = func.getContext();
quant_specs_.weight_quantization = weight_quantization_;
const QuantPassSpec quant_params = {
{quant_specs_.verify_numeric, /*error_tolerance=*/5.0f,
quant_specs_.whole_model_verify, /*enable_log_if_failed=*/false},
quant_specs_};
if (quant_specs_.weight_quantization) {
patterns.add<TFDynamicRangeQuantization>(ctx, quant_params);
} else {
patterns.add<TFFullQuantization, TFFullQuantizationReverse>(ctx,
quant_params);
patterns.add<QuantizeSameScaleOpsPattern>(ctx, GetTfQuantScaleSpec,
target_opset_);
patterns.add<QuantizeAvgPoolOpPattern>(ctx);
}
if (failed(applyPatternsGreedily(func, std::move(patterns)))) {
func.emitWarning("Failed to converge pattern at QuantizePass.");
}
if (!shouldKeepUnusedQdqPattern()) {
RewritePatternSet patterns_2(&getContext());
patterns_2.add<RemoveUnusedQdqPattern>(ctx);
if (failed(applyPatternsGreedily(func, std::move(patterns_2)))) {
signalPassFailure();
}
}
}
} // namespace
// Creates an instance of the TensorFlow dialect Quantize pass.
std::unique_ptr<OperationPass<func::FuncOp>> CreateQuantizePass() {
QuantizationSpecs quant_specs;
return std::make_unique<QuantizePass>(quant_specs, OpSet::TF);
}
std::unique_ptr<OperationPass<func::FuncOp>> CreateQuantizePass(
QuantizationSpecs quant_specs, OpSet target_opset) {
return std::make_unique<QuantizePass>(quant_specs, target_opset);
}
static PassRegistration<QuantizePass> pass;
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/passes/quantize_composite_functions.cc
+1373
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/mlir/quantization/tensorflow/passes/quantize_composite_functions.td
+28
View File
@@ -0,0 +1,28 @@
/* 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 "mlir/IR/OpBase.td"
include "mlir/IR/PatternBase.td"
include "mlir/Dialect/Func/IR/FuncOps.td"
include "mlir/Dialect/Arith/IR/ArithOps.td"
include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.td"
include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"
// Converts reamaining arith.constant ops from quantization passes back to
// tf.Const ops.
def ConvertArithConstToTfConst : Pat<
(Arith_ConstantOp:$res DenseElementsAttr:$value),
(TF_ConstOp $value),
[(AnyStaticShapeTensor $res)], [], (addBenefit 20)>;
tensorflow/compiler/mlir/quantization/tensorflow/passes/quantize_weights.cc
+278
View File
@@ -0,0 +1,278 @@
/* 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.
==============================================================================*/
#include <memory>
#include <optional>
#include <utility>
#include "llvm/Support/Casting.h"
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/Quant/IR/Quant.h" // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/BuiltinTypeInterfaces.h" // from @llvm-project
#include "mlir/IR/DialectRegistry.h" // from @llvm-project
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "mlir/IR/OperationSupport.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
#include "mlir/IR/Verifier.h" // from @llvm-project
#include "mlir/IR/Visitors.h" // from @llvm-project
#include "mlir/Interfaces/CallInterfaces.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "mlir/Pass/PassRegistry.h" // from @llvm-project
#include "mlir/Rewrite/FrozenRewritePatternSet.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Support/LogicalResult.h" // from @llvm-project
#include "mlir/Support/TypeID.h" // from @llvm-project
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/tensorflow/ops/tf_op_quant_spec.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/ops/tf_quantize_op.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/passes.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
namespace mlir {
namespace quant {
namespace {
class QuantizeWeightsPass
: public mlir::PassWrapper<QuantizeWeightsPass, OperationPass<ModuleOp>> {
public:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(QuantizeWeightsPass)
explicit QuantizeWeightsPass() : test_mode_(true) { initializeForTest(); }
explicit QuantizeWeightsPass(
const tensorflow::quantization::QuantizationOptions& quant_options)
: test_mode_(false), quant_options_(quant_options) {}
QuantizeWeightsPass(const QuantizeWeightsPass& other) {
test_mode_ = other.test_mode_;
quant_options_ = other.quant_options_;
initializeForTest();
}
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-quantize-weights";
}
StringRef getDescription() const final {
// This is a brief description of the pass.
return "Quantize weights used by quantizable ops.";
}
void getDependentDialects(DialectRegistry& registry) const override {
registry.insert<TF::TensorFlowDialect, quant::QuantDialect>();
}
private:
void runOnOperation() override;
bool test_mode_;
tensorflow::quantization::QuantizationOptions quant_options_;
// Initialize for tests.
void initializeForTest() {
if (!test_mode_) return;
tensorflow::quantization::QuantizationComponentSpec quant_spec;
quant_spec.set_quantization_component(
tensorflow::quantization::QuantizationComponentSpec::COMPONENT_WEIGHT);
quant_spec.set_tensor_type(
tensorflow::quantization::QuantizationComponentSpec::TENSORTYPE_INT_8);
auto mutable_quant_method = quant_options_.mutable_quantization_method();
*mutable_quant_method->add_quantization_component_specs() = quant_spec;
}
};
// If a constant is connected to a quantizable op, quantize the constant to have
// the provided data type.
class QuantizeConstWeights : public OpRewritePattern<TF::ConstOp> {
public:
explicit QuantizeConstWeights(
MLIRContext* context,
const tensorflow::quantization::QuantizationOptions& quantization_options)
: OpRewritePattern<TF::ConstOp>(context),
quant_options_(quantization_options) {}
LogicalResult matchAndRewrite(TF::ConstOp op,
PatternRewriter& rewriter) const override {
auto weight_component_spec = GetWeightComponentSpec(quant_options_);
if (!weight_component_spec) return failure();
// 1. Check if the constant is quantizable.
if (failed((isQuantizableWeight(op)))) {
return failure();
}
// 2. Quantize the constant to the provided data type.
// After quantization, the graph will be transformed
// from:
// const -> some op -> quantizable_op
// to:
// q_const -> dequant_op -> some op -> quantizable_op
//
// A dequant_op will propagate to further quantize the next ops in another
// pass.
//
// Note that a constant can be used by multiple ops. For example, if a graph
// looks like below:
// const -> while -> quant_op
// -> not_quant_op
//
// the transformation will be:
// q_const -> dequant_op -> while -> quant_op
// -> not_quant_op
// And the dequant_op op will propagate towards quant_op only.
if (failed(quantizeOps(rewriter, op, weight_component_spec.value()))) {
return failure();
}
return success();
}
private:
// Check if op's user or op's user after an identity op is connected to a
// terminator.
bool checkIfAnyUserIsConnectedToTermiantor(BlockArgument op) const {
for (const auto& user : op.getUsers()) {
if (user->template hasTrait<OpTrait::IsTerminator>()) return true;
if (auto next_user = dyn_cast_or_null<TF::IdentityOp>(user)) {
return (*(next_user->getResult(0).getUsers().begin()))
->template hasTrait<OpTrait::IsTerminator>();
}
}
return false;
}
// Check if the constant op is connected to a quantizable op at some point.
bool hasUsageFromQuantizableOp(TF::ConstOp op) const {
llvm::SmallVector<mlir::Value> uses_at_current_level{op};
while (!uses_at_current_level.empty()) {
llvm::SmallVector<mlir::Value> next_values_to_visit;
for (auto cur_op : uses_at_current_level) {
for (auto& cur_op_use : cur_op.getUses()) {
Operation* next_op = cur_op_use.getOwner();
int next_op_operand_num = cur_op_use.getOperandNumber();
if (auto call_op = llvm::dyn_cast<mlir::CallOpInterface>(next_op)) {
mlir::func::FuncOp func =
llvm::dyn_cast<mlir::func::FuncOp>(call_op.resolveCallable());
if (!func) continue;
next_values_to_visit.push_back(
func.getArgument(next_op_operand_num));
} else if (auto while_op =
llvm::dyn_cast_or_null<TF::WhileOp>(next_op)) {
func::FuncOp func = while_op.body_function();
auto func_argument = func.getArgument(next_op_operand_num);
// Check if the op is returned without mutation. Returning values
// from a while op follow return or identity -> return pattern.
if (checkIfAnyUserIsConnectedToTermiantor(func_argument))
next_values_to_visit.push_back(
func.getArgument(next_op_operand_num));
} else if (IsOpWithQuantizableTrait(next_op)) {
// Check this before IsOpWithDataMovementTrait since some data
// movement ops are also quantizable ops.
return true;
} else if (IsOpWithDataMovementTrait(next_op)) {
next_values_to_visit.insert(next_values_to_visit.end(),
next_op->getResults().begin(),
next_op->getResults().end());
}
}
}
uses_at_current_level.swap(next_values_to_visit);
}
return false;
}
// List of conditions to check if a const op is quantizable.
LogicalResult isQuantizableWeight(TF::ConstOp op) const {
// Non-float tensors do not need quantization.
if (!IsValueWithQuantizablePrecision(op)) return failure();
// Check if quantizable ops are connected. Do this before num_elements check
// to avoid checking unnecessary constants which causes unintended remarks.
// This check also prevents quantizing unintended consts like scale.
if (!hasUsageFromQuantizableOp(op)) return failure();
// Check if the weight size is big enough.
int num_elements_threshold = quant_options_.min_num_elements_for_weights();
int num_elements = cast<ShapedType>(op.getType()).getNumElements();
if (num_elements < num_elements_threshold) {
op->emitRemark("Quantization is skipped because the op has ")
<< num_elements << " elements which is fewer than the threshold("
<< num_elements_threshold << " elements).";
return failure();
}
return success();
}
// Apply quantization with the provided spec.
LogicalResult quantizeOps(PatternRewriter& rewriter, TF::ConstOp op,
tensorflow::quantization::QuantizationComponentSpec&
weight_component_spec) const {
if (weight_component_spec.tensor_type() ==
tensorflow::quantization::QuantizationComponentSpec::TENSORTYPE_INT_8) {
// TODO - b/296535985: [Converter Component][TF-Quantizer] Factor out
// quant/dequant in QuantizeWeightsPass
auto dequantized_val =
ApplyUniformQuantization(rewriter, op, weight_component_spec);
if (!dequantized_val.has_value()) return failure();
op.getOutput().replaceAllUsesWith(dequantized_val.value().getResult(0));
return success();
}
op->emitRemark("Not supported quantization data type.");
return failure();
}
protected:
tensorflow::quantization::QuantizationOptions quant_options_;
};
static PassRegistration<QuantizeWeightsPass> pass;
void QuantizeWeightsPass::runOnOperation() {
MLIRContext* ctx = &getContext();
auto module_op = getOperation();
RewritePatternSet patterns(ctx);
patterns.add<QuantizeConstWeights>(ctx, quant_options_);
FrozenRewritePatternSet frozen_patterns(std::move(patterns));
// Apply transformation on each function. For recursive call case, another
// function can be modified at the same time so avoid running functions in
// parallel.
for (auto func : module_op.getOps<func::FuncOp>()) {
if (failed(applyPatternsGreedily(func, frozen_patterns))) {
func.emitError() << "quant-quantize-weights failed.";
signalPassFailure();
}
}
}
} // namespace
std::unique_ptr<OperationPass<ModuleOp>> CreateQuantizeWeightsPass(
const tensorflow::quantization::QuantizationOptions& quant_options) {
return std::make_unique<QuantizeWeightsPass>(quant_options);
}
} // namespace quant
} // namespace mlir
tensorflow/compiler/mlir/quantization/tensorflow/passes/quantized_function_library.mlir
+511
View File
@@ -0,0 +1,511 @@
// Copyright 2022 The TensorFlow Runtime Authors
//
// 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.
// Quantization as a function library.
//
// Internal functions should be marked as private. They will be inlined and
// deleted in `InsertQuantizedFunctionsPass`.
//
// Function template can generate functions with different parameters. Ex:
// ```
// parameters[
// {"key1": "value11", "key2": "value21"},
// {"key1": "value12", "key2": "value22"},
// ]
// func.func func_name_${key1}_fn (...) {
// ...${key2}...
// }
// ```
// The above template with generate two functions by substituting `key1` and
// `key2` with given values.
module {
// Rescales to the output scale and zero point.
func.func private @internal_rescale_fn(%accumulation : tensor<*xi32>,
%input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>,
%filter_scale : tensor<*xf32>, %filter_zp : tensor<*xi32>,
%out_scale : tensor<*xf32>, %out_zp : tensor<*xi32>) -> tensor<*xf32> {
%scale_prod = "tf.Mul"(%input_scale, %filter_scale) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
%rescale_factor = "tf.Div"(%scale_prod, %out_scale) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
%float_out_zp = "tf.Cast"(%out_zp) {Truncate = false} : (tensor<*xi32>) -> tensor<*xf32>
%cast = "tf.Cast"(%accumulation) {Truncate = false} : (tensor<*xi32>) -> tensor<*xf32>
%mul = "tf.Mul"(%cast, %rescale_factor) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
%add = "tf.AddV2"(%mul, %float_out_zp) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
func.return %add : tensor<*xf32>
}
func.func private @internal_dequantize_i8_fn(%input : tensor<*xi8>, %scale : tensor<*xf32>, %zp : tensor<*xi32>) -> tensor<*xf32> {
%input_i32 = "tf.Cast"(%input) : (tensor<*xi8>) -> tensor<*xi32>
%output = "tf.Sub"(%input_i32, %zp) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
%cast = "tf.Cast"(%output) : (tensor<*xi32>) -> tensor<*xf32>
%mul = "tf.Mul"(%cast, %scale) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
func.return %mul : tensor<*xf32>
}
// Requantizes and clips to the range of quantized type if there is no specific activation.
func.func private @internal_requantize_no_activation_fn(%accumulation : tensor<*xi32>,
%input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>,
%filter_scale : tensor<*xf32>, %filter_zp : tensor<*xi32>,
%out_scale : tensor<*xf32>, %out_zp : tensor<*xi32>) -> tensor<*xi8> {
%rescale = "tf.PartitionedCall"(%accumulation, %input_scale, %input_zp, %filter_scale, %filter_zp,
%out_scale, %out_zp) {
config = "", config_proto = "", executor_type = "", f=@internal_rescale_fn
} : (tensor<*xi32>, tensor<*xf32>, tensor<*xi32>, tensor<*xf32>, tensor<*xi32>,
tensor<*xf32>, tensor<*xi32>) -> tensor<*xf32>
%i8_min = "tf.Const"() {value = dense<-128.0> : tensor<f32>} : () -> tensor<f32>
%i8_max = "tf.Const"() {value = dense<127.0> : tensor<f32>} : () -> tensor<f32>
%clamp_max = "tf.Maximum"(%rescale, %i8_min) : (tensor<*xf32>, tensor<f32>) -> tensor<*xf32>
%clamp_min = "tf.Minimum"(%clamp_max, %i8_max) : (tensor<*xf32>, tensor<f32>) -> tensor<*xf32>
%round = "tf.Round"(%clamp_min) : (tensor<*xf32>) -> tensor<*xf32>
%cast = "tf.Cast"(%round) {Truncate = false} : (tensor<*xf32>) -> tensor<*xi8>
func.return %cast : tensor<*xi8>
}
// Requantizes and applies quantized Relu by clipping.
func.func private @internal_requantize_and_relu_fn(%accumulation : tensor<*xi32>,
%input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>,
%filter_scale : tensor<*xf32>, %filter_zp : tensor<*xi32>,
%out_scale : tensor<*xf32>, %out_zp : tensor<*xi32>) -> tensor<*xi8> {
%rescale = "tf.PartitionedCall"(%accumulation, %input_scale, %input_zp, %filter_scale, %filter_zp,
%out_scale, %out_zp) {
config = "", config_proto = "", executor_type = "", f=@internal_rescale_fn
} : (tensor<*xi32>, tensor<*xf32>, tensor<*xi32>, tensor<*xf32>, tensor<*xi32>,
tensor<*xf32>, tensor<*xi32>) -> tensor<*xf32>
%i8_min = "tf.Const"() {value = dense<-128.0> : tensor<f32>} : () -> tensor<f32>
%i8_max = "tf.Const"() {value = dense<127.0> : tensor<f32>} : () -> tensor<f32>
%float_out_zp = "tf.Cast"(%out_zp) {Truncate = false} : (tensor<*xi32>) -> tensor<*xf32>
%clip_min = "tf.Maximum"(%i8_min, %float_out_zp) : (tensor<f32>, tensor<*xf32>) -> tensor<*xf32>
%clamp_max = "tf.Maximum"(%rescale, %clip_min) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
%clamp_min = "tf.Minimum"(%clamp_max, %i8_max) : (tensor<*xf32>, tensor<f32>) -> tensor<*xf32>
%round = "tf.Round"(%clamp_min) : (tensor<*xf32>) -> tensor<*xf32>
%cast = "tf.Cast"(%round) {Truncate = false} : (tensor<*xf32>) -> tensor<*xi8>
func.return %cast : tensor<*xi8>
}
// Requantizes and applies quantized Relu6 by clipping.
func.func private @internal_requantize_and_relu6_fn(%accumulation : tensor<*xi32>,
%input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>,
%filter_scale : tensor<*xf32>, %filter_zp : tensor<*xi32>,
%out_scale : tensor<*xf32>, %out_zp : tensor<*xi32>) -> tensor<*xi8> {
%rescale = "tf.PartitionedCall"(%accumulation, %input_scale, %input_zp, %filter_scale, %filter_zp,
%out_scale, %out_zp) {
config = "", config_proto = "", executor_type = "", f=@internal_rescale_fn
} : (tensor<*xi32>, tensor<*xf32>, tensor<*xi32>, tensor<*xf32>, tensor<*xi32>,
tensor<*xf32>, tensor<*xi32>) -> tensor<*xf32>
%i8_min = "tf.Const"() {value = dense<-128.0> : tensor<f32>} : () -> tensor<f32>
%i8_max = "tf.Const"() {value = dense<127.0> : tensor<f32>} : () -> tensor<f32>
%act_max = "tf.Const"() {value = dense<6.0> : tensor<f32>} : () -> tensor<f32>
%i8_act_max_0 = "tf.PartitionedCall"(%act_max, %out_scale, %out_zp) {
config = "", config_proto = "", executor_type = "", f=@quantize_i8
} : (tensor<f32>, tensor<*xf32>, tensor<*xi32>) -> tensor<*xi8>
%i8_act_max_1 = "tf.Cast"(%i8_act_max_0) {Truncate = false} : (tensor<*xi8>) -> tensor<*xf32>
%float_out_zp = "tf.Cast"(%out_zp) {Truncate = false} : (tensor<*xi32>) -> tensor<*xf32>
%clip_min = "tf.Maximum"(%i8_min, %float_out_zp) : (tensor<f32>, tensor<*xf32>) -> tensor<*xf32>
%clip_max = "tf.Minimum"(%i8_max, %i8_act_max_1) : (tensor<f32>, tensor<*xf32>) -> tensor<*xf32>
%clamp_max = "tf.Maximum"(%rescale, %clip_min) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
%clamp_min = "tf.Minimum"(%clamp_max, %clip_max) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
%round = "tf.Round"(%clamp_min) : (tensor<*xf32>) -> tensor<*xf32>
%cast = "tf.Cast"(%round) {Truncate = false} : (tensor<*xf32>) -> tensor<*xi8>
func.return %cast : tensor<*xi8>
}
// Dequantizes and clips to the range of quantized type if there is no specific activation.
func.func private @internal_dequantize_no_activation_fn(%accumulation : tensor<*xi32>,
%input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>,
%filter_scale : tensor<*xf32>, %filter_zp : tensor<*xi32>,
%out_scale : tensor<*xf32>, %out_zp : tensor<*xi32>) -> tensor<*xf32> {
%accumulation_scale = "tf.Mul"(%input_scale, %filter_scale) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
%cast = "tf.Cast"(%accumulation) {Truncate = false} : (tensor<*xi32>) -> tensor<*xf32>
%dequantize = "tf.Mul"(%cast, %accumulation_scale) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
%i8_min = "tf.Const"() {value = dense<-128> : tensor<i8>} : () -> tensor<i8>
%i8_max = "tf.Const"() {value = dense<127> : tensor<i8>} : () -> tensor<i8>
%clip_min = "tf.PartitionedCall"(%i8_min, %out_scale, %out_zp) {
config = "", config_proto = "", executor_type = "", f=@internal_dequantize_i8_fn
} : (tensor<i8>, tensor<*xf32>, tensor<*xi32>) -> tensor<*xf32>
%clip_max = "tf.PartitionedCall"(%i8_max, %out_scale, %out_zp) {
config = "", config_proto = "", executor_type = "", f=@internal_dequantize_i8_fn
} : (tensor<i8>, tensor<*xf32>, tensor<*xi32>) -> tensor<*xf32>
%clamp_max = "tf.Maximum"(%dequantize, %clip_min) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
%clamp_min = "tf.Minimum"(%clamp_max, %clip_max) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
func.return %clamp_min : tensor<*xf32>
}
// Dequantizes and applies quantized Relu by clipping.
func.func private @internal_dequantize_and_relu_fn(%accumulation : tensor<*xi32>,
%input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>,
%filter_scale : tensor<*xf32>, %filter_zp : tensor<*xi32>,
%out_scale : tensor<*xf32>, %out_zp : tensor<*xi32>) -> tensor<*xf32> {
%accumulation_scale = "tf.Mul"(%input_scale, %filter_scale) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
%cast = "tf.Cast"(%accumulation) {Truncate = false} : (tensor<*xi32>) -> tensor<*xf32>
%dequantize = "tf.Mul"(%cast, %accumulation_scale) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
%i8_min = "tf.Const"() {value = dense<-128> : tensor<i8>} : () -> tensor<i8>
%i8_max = "tf.Const"() {value = dense<127> : tensor<i8>} : () -> tensor<i8>
%clip_min_0 = "tf.PartitionedCall"(%i8_min, %out_scale, %out_zp) {
config = "", config_proto = "", executor_type = "", f=@internal_dequantize_i8_fn
} : (tensor<i8>, tensor<*xf32>, tensor<*xi32>) -> tensor<*xf32>
%clip_max = "tf.PartitionedCall"(%i8_max, %out_scale, %out_zp) {
config = "", config_proto = "", executor_type = "", f=@internal_dequantize_i8_fn
} : (tensor<i8>, tensor<*xf32>, tensor<*xi32>) -> tensor<*xf32>
%clip_min = "tf.Relu"(%clip_min_0) : (tensor<*xf32>) -> tensor<*xf32>
%clamp_max = "tf.Maximum"(%dequantize, %clip_min) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
%clamp_min = "tf.Minimum"(%clamp_max, %clip_max) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
func.return %clamp_min : tensor<*xf32>
}
// Dequantizes and applies quantized Relu6 by clipping.
func.func private @internal_dequantize_and_relu6_fn(%accumulation : tensor<*xi32>,
%input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>,
%filter_scale : tensor<*xf32>, %filter_zp : tensor<*xi32>,
%out_scale : tensor<*xf32>, %out_zp : tensor<*xi32>) -> tensor<*xf32> {
%accumulation_scale = "tf.Mul"(%input_scale, %filter_scale) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
%cast = "tf.Cast"(%accumulation) {Truncate = false} : (tensor<*xi32>) -> tensor<*xf32>
%dequantize = "tf.Mul"(%cast, %accumulation_scale) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
%i8_min = "tf.Const"() {value = dense<-128> : tensor<i8>} : () -> tensor<i8>
%i8_max = "tf.Const"() {value = dense<127> : tensor<i8>} : () -> tensor<i8>
%relu6_upper = "tf.Const"() {value = dense<6.0>: tensor<f32>} : () -> tensor<f32>
%clip_min_0 = "tf.PartitionedCall"(%i8_min, %out_scale, %out_zp) {
config = "", config_proto = "", executor_type = "", f=@internal_dequantize_i8_fn
} : (tensor<i8>, tensor<*xf32>, tensor<*xi32>) -> tensor<*xf32>
%clip_max_0 = "tf.PartitionedCall"(%i8_max, %out_scale, %out_zp) {
config = "", config_proto = "", executor_type = "", f=@internal_dequantize_i8_fn
} : (tensor<i8>, tensor<*xf32>, tensor<*xi32>) -> tensor<*xf32>
%clip_min = "tf.Relu"(%clip_min_0) : (tensor<*xf32>) -> tensor<*xf32>
%clip_max = "tf.Minimum"(%clip_max_0, %relu6_upper) : (tensor<*xf32>, tensor<f32>) -> tensor<*xf32>
%clamp_max = "tf.Maximum"(%dequantize, %clip_min) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
%clamp_min = "tf.Minimum"(%clamp_max, %clip_max) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
func.return %clamp_min : tensor<*xf32>
}
// Conv2D with int32 accumulation.
func.func private @internal_conv2d_fn(
%input : tensor<*xi8>, %filter : tensor<*xi8>,
%input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>,
%filter_scale : tensor<*xf32>, %filter_zp : tensor<*xi32>) -> tensor<*xi32> {
%0 = "tf.Cast"(%input) {Truncate = false} : (tensor<*xi8>) -> tensor<*xi32>
%1 = "tf.Sub"(%0, %input_zp) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
// Use identity op to avoid the filter being constant-folded.
%identity = "tf.Identity"(%filter) : (tensor<*xi8>) -> tensor<*xi8>
%2 = "tf.Cast"(%identity) {Truncate = false} : (tensor<*xi8>) -> tensor<*xi32>
%3 = "tf.Sub"(%2, %filter_zp) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
%5 = "tf.Conv2D"(%1, %3) {
padding = "VALID", strides = [1, 1, 1, 1],
attr_map = "strides:0,use_cudnn_on_gpu:1,padding:2,explicit_paddings:3,dilations:4"
} : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
func.return %5 : tensor<*xi32>
}
// DepthwiseConv2D with (simulated) int32 accumulation.
func.func private @internal_depthwise_conv2d_fn(
%input : tensor<*xi8>, %filter : tensor<*xi8>,
%input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>,
%filter_scale : tensor<*xf32>, %filter_zp : tensor<*xi32>) -> tensor<*xi32> {
%0 = "tf.Cast"(%input) {Truncate = false} : (tensor<*xi8>) -> tensor<*xi32>
%1 = "tf.Sub"(%0, %input_zp) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
// Use identity op to avoid the filter being constant-folded.
%identity = "tf.Identity"(%filter) : (tensor<*xi8>) -> tensor<*xi8>
%2 = "tf.Cast"(%identity) {Truncate = false} : (tensor<*xi8>) -> tensor<*xi32>
%3 = "tf.Sub"(%2, %filter_zp) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
%cast_1_f32 = "tf.Cast"(%1) {Truncate = false} : (tensor<*xi32>) -> tensor<*xf32>
%cast_3_f32 = "tf.Cast"(%3) {Truncate = false} : (tensor<*xi32>) -> tensor<*xf32>
%5 = "tf.DepthwiseConv2dNative"(%cast_1_f32, %cast_3_f32) {
padding = "VALID", strides = [1, 1, 1, 1],
attr_map = "strides:0,padding:1,explicit_paddings:2,dilations:3"
} : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
%6 = "tf.Cast"(%5) : (tensor<*xf32>) -> tensor<*xi32>
func.return %6 : tensor<*xi32>
}
// Matmul with int32 accumulation.
func.func private @internal_matmul_fn(
%input : tensor<*xi8>, %weight : tensor<*xi8>,
%input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>,
%weight_scale : tensor<*xf32>, %weight_zp : tensor<*xi32>) -> tensor<*xi32> {
%0 = "tf.Cast"(%input) {Truncate = false} : (tensor<*xi8>) -> tensor<*xi32>
%1 = "tf.Sub"(%0, %input_zp) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
// Use identity op to avoid the weight being constant-folded.
%identity = "tf.Identity"(%weight) : (tensor<*xi8>) -> tensor<*xi8>
%2 = "tf.Cast"(%identity) {Truncate = false} : (tensor<*xi8>) -> tensor<*xi32>
%3 = "tf.Sub"(%2, %weight_zp) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
%5 = "tf.MatMul"(%1, %3) {
attr_map = "transpose_a:0,transpose_b:1"
} : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
func.return %5 : tensor<*xi32>
}
// Conv3D with int32 accumulation.
func.func private @internal_conv3d_fn(
%input : tensor<*xi8>, %filter : tensor<*xi8>,
%input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>,
%filter_scale : tensor<*xf32>, %filter_zp : tensor<*xi32>) -> tensor<*xi32> {
%0 = "tf.Cast"(%input) {Truncate = false} : (tensor<*xi8>) -> tensor<*xi32>
%1 = "tf.Sub"(%0, %input_zp) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
// Use identity op to avoid the filter being constant-folded.
%identity = "tf.Identity"(%filter) : (tensor<*xi8>) -> tensor<*xi8>
%2 = "tf.Cast"(%identity) {Truncate = false} : (tensor<*xi8>) -> tensor<*xi32>
%3 = "tf.Sub"(%2, %filter_zp) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
%cast_1_f32 = "tf.Cast"(%1) {Truncate = false} : (tensor<*xi32>) -> tensor<*xf32>
%cast_3_f32 = "tf.Cast"(%3) {Truncate = false} : (tensor<*xi32>) -> tensor<*xf32>
%5 = "tf.Conv3D"(%cast_1_f32, %cast_3_f32) {
padding = "VALID", strides = [1, 1, 1, 1, 1],
attr_map = "strides:0,padding:1,dilations:2"
} : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
%6 = "tf.Cast"(%5) : (tensor<*xf32>) -> tensor<*xi32>
func.return %6 : tensor<*xi32>
}
// BatchMatMul with int32 accumulation.
func.func private @internal_batch_matmul_fn(
%input : tensor<*xi8>, %weight : tensor<*xi8>,
%input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>,
%weight_scale : tensor<*xf32>, %weight_zp : tensor<*xi32>) -> tensor<*xi32> {
%0 = "tf.Cast"(%input) {Truncate = false} : (tensor<*xi8>) -> tensor<*xi32>
%1 = "tf.Sub"(%0, %input_zp) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
// Use identity op to avoid the weight being constant-folded.
%identity = "tf.Identity"(%weight) : (tensor<*xi8>) -> tensor<*xi8>
%2 = "tf.Cast"(%identity) {Truncate = false} : (tensor<*xi8>) -> tensor<*xi32>
%3 = "tf.Sub"(%2, %weight_zp) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
%5 = "tf.BatchMatMulV2"(%1, %3) {
attr_map = "adj_x:0,adj_y:1"
} : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
func.return %5 : tensor<*xi32>
}
// Einsum with int32 accumulation.
func.func private @internal_einsum_fn(
%input : tensor<*xi8>, %weight : tensor<*xi8>,
%input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>,
%weight_scale : tensor<*xf32>, %weight_zp : tensor<*xi32>) -> tensor<*xi32> {
// For Einsum, the input could also be constant, since we do the argument swapping.
%identity_input = "tf.Identity"(%input) : (tensor<*xi8>) -> tensor<*xi8>
%0 = "tf.Cast"(%identity_input) {Truncate = false} : (tensor<*xi8>) -> tensor<*xi32>
%1 = "tf.Sub"(%0, %input_zp) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
// Use identity op to avoid the weight being constant-folded.
%identity = "tf.Identity"(%weight) : (tensor<*xi8>) -> tensor<*xi8>
%2 = "tf.Cast"(%identity) {Truncate = false} : (tensor<*xi8>) -> tensor<*xi32>
%3 = "tf.Sub"(%2, %weight_zp) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
%4 = "tf.Einsum"(%1, %3) {
// Equation placeholder to prevent error during op creation.
equation = "",
attr_map = "equation:0"
} : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
func.return %4 : tensor<*xi32>
}
for main_op in ["Conv2D", "DepthwiseConv2D", "MatMul", "Conv3D", "BatchMatMul", "Einsum"] {
parameters[
{"quantized_ops": ["${main_op}", "BiasAdd"], "act_func": "internal_requantize_no_activation_fn", "output_type": "i8"},
{"quantized_ops": ["${main_op}", "BiasAdd", "Relu"], "act_func": "internal_requantize_and_relu_fn", "output_type": "i8"},
{"quantized_ops": ["${main_op}", "BiasAdd", "Relu6"], "act_func": "internal_requantize_and_relu6_fn", "output_type": "i8"},
{"quantized_ops": ["${main_op}", "BiasAdd"], "act_func": "internal_dequantize_no_activation_fn", "output_type": "f32"},
{"quantized_ops": ["${main_op}", "BiasAdd", "Relu"], "act_func": "internal_dequantize_and_relu_fn", "output_type": "f32"},
{"quantized_ops": ["${main_op}", "BiasAdd", "Relu6"], "act_func": "internal_dequantize_and_relu6_fn", "output_type": "f32"},
]
func.func @GenerateQuantizedFunctionName(${quantized_ops}, "${output_type}")(%input : tensor<*xi8>,
%filter : tensor<*xi8>, %bias : tensor<*xi32>,
%input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>,
%filter_scale : tensor<*xf32>, %filter_zp : tensor<*xi32>,
%bias_scale : tensor<*xf32>, %bias_zp : tensor<*xi32>,
%out_scale : tensor<*xf32>, %out_zp : tensor<*xi32>) -> tensor<*x${output_type}>
attributes {tf_quant.quantized_ops = ${quantized_ops}} {
%0 = "tf.PartitionedCall"(%input, %filter, %input_scale, %input_zp,
%filter_scale, %filter_zp) {
config = "", config_proto = "", executor_type = "", f=@GenerateImplFunctionName(${main_op})
} : (tensor<*xi8>, tensor<*xi8>, tensor<*xf32>, tensor<*xi32>,
tensor<*xf32>, tensor<*xi32>) -> tensor<*xi32>
%1 = "tf.AddV2"(%0, %bias) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
%2 = "tf.PartitionedCall"(%1, %input_scale, %input_zp, %filter_scale, %filter_zp,
%out_scale, %out_zp) {
config = "", config_proto = "", executor_type = "", f=@${act_func}
} : (tensor<*xi32>, tensor<*xf32>, tensor<*xi32>, tensor<*xf32>, tensor<*xi32>,
tensor<*xf32>, tensor<*xi32>) -> tensor<*x${output_type}>
func.return %2 : tensor<*x${output_type}>
}
parameters[
{"quantized_ops": ["${main_op}"], "act_func": "internal_requantize_no_activation_fn", "output_type": "i8"},
{"quantized_ops": ["${main_op}", "Relu"], "act_func": "internal_requantize_and_relu_fn", "output_type": "i8"},
{"quantized_ops": ["${main_op}", "Relu6"], "act_func": "internal_requantize_and_relu6_fn", "output_type": "i8"},
{"quantized_ops": ["${main_op}"], "act_func": "internal_dequantize_no_activation_fn", "output_type": "f32"},
{"quantized_ops": ["${main_op}", "Relu"], "act_func": "internal_dequantize_and_relu_fn", "output_type": "f32"},
{"quantized_ops": ["${main_op}", "Relu6"], "act_func": "internal_dequantize_and_relu6_fn", "output_type": "f32"},
]
func.func @GenerateQuantizedFunctionName(${quantized_ops}, "${output_type}")(
%input : tensor<*xi8>, %filter : tensor<*xi8>,
%input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>,
%filter_scale : tensor<*xf32>, %filter_zp : tensor<*xi32>,
%out_scale : tensor<*xf32>, %out_zp : tensor<*xi32>) -> tensor<*x${output_type}>
attributes {tf_quant.quantized_ops = ${quantized_ops}} {
%0 = "tf.PartitionedCall"(%input, %filter, %input_scale, %input_zp,
%filter_scale, %filter_zp) {
config = "", config_proto = "", executor_type = "", f=@GenerateImplFunctionName(${main_op})
} : (tensor<*xi8>, tensor<*xi8>, tensor<*xf32>, tensor<*xi32>,
tensor<*xf32>, tensor<*xi32>) -> tensor<*xi32>
%1 = "tf.PartitionedCall"(%0, %input_scale, %input_zp, %filter_scale, %filter_zp,
%out_scale, %out_zp) {
config = "", config_proto = "", executor_type = "", f=@${act_func}
} : (tensor<*xi32>, tensor<*xf32>, tensor<*xi32>, tensor<*xf32>, tensor<*xi32>,
tensor<*xf32>, tensor<*xi32>) -> tensor<*x${output_type}>
func.return %1 : tensor<*x${output_type}>
}
} // end for
// TODO(b/278493977): Create generic implementation of lifting any fused op
// with any reshaping op
for main_op in ["MatMul"] {
parameters[
{"quantized_ops": ["${main_op}", "Reshape", "BiasAdd"], "act_func": "internal_requantize_no_activation_fn", "output_type": "i8"},
{"quantized_ops": ["${main_op}", "Reshape", "BiasAdd"], "act_func": "internal_dequantize_no_activation_fn", "output_type": "f32"},
]
func.func @GenerateQuantizedFunctionName(${quantized_ops}, "${output_type}")(%input : tensor<*xi8>,
%filter : tensor<*xi8>, %bias : tensor<*xi32>, %shape : tensor<*xi32>,
%input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>,
%filter_scale : tensor<*xf32>, %filter_zp : tensor<*xi32>,
%bias_scale : tensor<*xf32>, %bias_zp : tensor<*xi32>,
%out_scale : tensor<*xf32>, %out_zp : tensor<*xi32>) -> tensor<*x${output_type}>
attributes {tf_quant.quantized_ops = ${quantized_ops}} {
%0 = "tf.PartitionedCall"(%input, %filter, %input_scale, %input_zp,
%filter_scale, %filter_zp) {
config = "", config_proto = "", executor_type = "", f=@GenerateImplFunctionName(${main_op})
} : (tensor<*xi8>, tensor<*xi8>, tensor<*xf32>, tensor<*xi32>,
tensor<*xf32>, tensor<*xi32>) -> tensor<*xi32>
%1 = "tf.Reshape"(%0, %shape) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
%2 = "tf.AddV2"(%1, %bias) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
%3 = "tf.PartitionedCall"(%2, %input_scale, %input_zp, %filter_scale, %filter_zp,
%out_scale, %out_zp) {
config = "", config_proto = "", executor_type = "", f=@${act_func}
} : (tensor<*xi32>, tensor<*xf32>, tensor<*xi32>, tensor<*xf32>, tensor<*xi32>,
tensor<*xf32>, tensor<*xi32>) -> tensor<*x${output_type}>
func.return %3 : tensor<*x${output_type}>
}
} // end for
func.func @quantize_i8(%input : tensor<*xf32>, %scale : tensor<*xf32>, %zp : tensor<*xi32>) -> tensor<*xi8> {
%float_zp = "tf.Cast"(%zp) {Truncate = false} : (tensor<*xi32>) -> tensor<*xf32>
%div = "tf.Div"(%input, %scale) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
%add = "tf.AddV2"(%div, %float_zp) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
%i8_min = "tf.Const"() {value = dense<-128.0> : tensor<f32>} : () -> tensor<f32>
%i8_max = "tf.Const"() {value = dense<127.0> : tensor<f32>} : () -> tensor<f32>
%clamp_max = "tf.Maximum"(%add, %i8_min) : (tensor<*xf32>, tensor<f32>) -> tensor<*xf32>
%clamp_min = "tf.Minimum"(%clamp_max, %i8_max) : (tensor<*xf32>, tensor<f32>) -> tensor<*xf32>
%round = "tf.Round"(%clamp_min) : (tensor<*xf32>) -> tensor<*xf32>
%i8 = "tf.Cast"(%round) : (tensor<*xf32>) -> tensor<*xi8>
func.return %i8 : tensor<*xi8>
}
func.func @dequantize_i8(%input : tensor<*xi8>, %scale : tensor<*xf32>, %zp : tensor<*xi32>) -> tensor<*xf32> {
// Use identity op to avoid the weight being constant-folded.
%identity = "tf.Identity"(%input) : (tensor<*xi8>) -> tensor<*xi8>
%input_i32 = "tf.Cast"(%identity) : (tensor<*xi8>) -> tensor<*xi32>
%output = "tf.Sub"(%input_i32, %zp) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
%cast = "tf.Cast"(%output) : (tensor<*xi32>) -> tensor<*xf32>
%mul = "tf.Mul"(%cast, %scale) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
func.return %mul : tensor<*xf32>
}
//===----------------------------------------------------------------------===//
// Weight-only functions.
//===----------------------------------------------------------------------===//
// Requantize to the output quantization parameters.
func.func private @internal_requantize_fn(%input : tensor<*xi8>,
%input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>,
%out_scale : tensor<*xf32>, %out_zp : tensor<*xi32>) -> tensor<*xi8> {
%rescale_factor = "tf.Div"(%input_scale, %out_scale) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
%float_in_zp = "tf.Cast"(%input_zp) {Truncate = false} : (tensor<*xi32>) -> tensor<*xf32>
%float_out_zp = "tf.Cast"(%out_zp) {Truncate = false} : (tensor<*xi32>) -> tensor<*xf32>
%cast = "tf.Cast"(%input) {Truncate = false} : (tensor<*xi8>) -> tensor<*xf32>
%sub = "tf.Sub"(%cast, %float_in_zp) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
%mul = "tf.Mul"(%sub, %rescale_factor) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
%add = "tf.AddV2"(%mul, %float_out_zp) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
%i8_min = "tf.Const"() {value = dense<-128.0> : tensor<f32>} : () -> tensor<f32>
%i8_max = "tf.Const"() {value = dense<127.0> : tensor<f32>} : () -> tensor<f32>
%clamp_max = "tf.Maximum"(%add, %i8_min) : (tensor<*xf32>, tensor<f32>) -> tensor<*xf32>
%clamp_min = "tf.Minimum"(%clamp_max, %i8_max) : (tensor<*xf32>, tensor<f32>) -> tensor<*xf32>
%round = "tf.Round"(%clamp_min) : (tensor<*xf32>) -> tensor<*xf32>
%cast2 = "tf.Cast"(%round) {Truncate = false} : (tensor<*xf32>) -> tensor<*xi8>
func.return %cast2 : tensor<*xi8>
}
// Note that input i64 type is also supported by this.
// As the output is quantized type, output scale/zp is required for the arguments.
func.func @quantized_gather_fn(
%weight : tensor<*xi8>, %input : tensor<*xi32>, %axis : tensor<i32>,
%weight_scale : tensor<*xf32>, %weight_zp : tensor<*xi32>,
%out_scale : tensor<*xf32>, %out_zp : tensor<*xi32>) -> tensor<*xi8>
attributes {tf_quant.quantized_ops = ["Gather"]} {
%out = "tf.GatherV2"(%weight, %input, %axis) {
batch_dims = 0 : i64, attr_map = "batch_dims:0"} : (tensor<*xi8>, tensor<*xi32>, tensor<i32>) -> tensor<*xi8>
// Requantize as the output quantization params can be different than the input for Gather ops.
// Ex: Input can be per-axis quantized while output can be per-tensor.
%requantize = "tf.PartitionedCall"(%out, %weight_scale, %weight_zp, %out_scale, %out_zp) {
config = "", config_proto = "", executor_type = "", f=@internal_requantize_fn
} : (tensor<*xi8>, tensor<*xf32>, tensor<*xi32>, tensor<*xf32>, tensor<*xi32>) -> tensor<*xi8>
func.return %requantize : tensor<*xi8>
}
func.func @quantized_gather_hybrid_fn(
%weight : tensor<*xi8>, %input : tensor<*xi32>, %axis : tensor<i32>,
%weight_scale : tensor<*xf32>, %weight_zp : tensor<*xi32>) -> tensor<*xf32>
attributes {tf_quant.quantized_ops = ["Gather"]} {
%accum_out = "tf.GatherV2"(%weight, %input, %axis) {
batch_dims = 0 : i64, attr_map = "batch_dims:0"} : (tensor<*xi8>, tensor<*xi32>, tensor<i32>) -> tensor<*xi8>
%out = "tf.PartitionedCall"(%accum_out, %weight_scale, %weight_zp) {
config = "", config_proto = "", executor_type = "", f=@internal_dequantize_i8_fn
} : (tensor<*xi8>, tensor<*xf32>, tensor<*xi32>) -> tensor<*xf32>
func.return %out : tensor<*xf32>
}
}
tensorflow/compiler/mlir/quantization/tensorflow/passes/quantized_function_library_tf_drq.mlir
+238
View File
@@ -0,0 +1,238 @@
// Copyright 2022 The TensorFlow Runtime Authors
//
// 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.
// Quantization as a function library with TF Ops for Dynamic PTQ
//
// Internal functions should be marked as private. They will be inlined and
// deleted in `InsertQuantizedFunctionsPass`.
//
// Function template can generate functions with different parameters. Ex:
// ```
// parameters[
// {"key1": "value11", "key2": "value21"},
// {"key1": "value12", "key2": "value22"},
// ]
// func.func func_name_${key1}_fn (...) {
// ...${key2}...
// }
// ```
// The above template with generate two functions by substituting `key1` and
// `key2` with given values.
module {
// Note: following functions won't handle per-channel quantization for now.
func.func private @internal_quantize_i8(%input : tensor<*xf32>, %scale : tensor<*xf32>, %zp : tensor<*xi32>) -> tensor<*xi8> {
// Uses tf.floor(x + 0.5) instead of tf.round(x) since tf.round generates
// a very expensive pattern.
%round_cst = "tf.Const"() {value = dense<0.5> : tensor<f32>} : () -> tensor<f32>
%float_zp = "tf.Cast"(%zp) {Truncate = false} : (tensor<*xi32>) -> tensor<*xf32>
%zp_plus_round_cst = "tf.AddV2"(%float_zp, %round_cst) : (tensor<*xf32>, tensor<f32>) -> tensor<*xf32>
%div = "tf.Div"(%input, %scale) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
%add = "tf.AddV2"(%div, %zp_plus_round_cst) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
%round = "tf.Floor"(%add) : (tensor<*xf32>) -> tensor<*xf32>
%i8_min = "tf.Const"() {value = dense<-128.0> : tensor<f32>} : () -> tensor<f32>
%i8_max = "tf.Const"() {value = dense<127.0> : tensor<f32>} : () -> tensor<f32>
%clip = "tf.ClipByValue"(%round, %i8_min, %i8_max) : (tensor<*xf32>, tensor<f32>, tensor<f32>) -> tensor<*xf32>
%i8 = "tf.Cast"(%clip) : (tensor<*xf32>) -> tensor<*xi8>
func.return %i8 : tensor<*xi8>
}
func.func private @internal_dequantize_i32(%input : tensor<*xi32>,
%input_scale : tensor<*xf32>,
%weight_scale : tensor<*xf32>) -> tensor<*xf32> {
%scale_prod = "tf.Mul"(%input_scale, %weight_scale) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
%cast = "tf.Cast"(%input) : (tensor<*xi32>) -> tensor<*xf32>
%mul = "tf.Mul"(%cast, %scale_prod) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
func.return %mul : tensor<*xf32>
}
// TODO(b/263199401): Support quantization options for activation quantization for DRQ
// Note: following function supports per-tensor, asymmetric, non_narrow_range.
func.func private @internal_calculate_quant_params(%input : tensor<*xf32>) -> (tensor<1xf32>, tensor<1xi32>) {
%zero = "tf.Const"() {value = dense<0.0> : tensor<1xf32>} : () -> tensor<1xf32>
%shape = "tf.Const"() {value = dense<[-1]> : tensor<1xi32>} : () -> tensor<1xi32>
%dim = "tf.Const"() { value = dense<0> : tensor<1xi64> } : () -> tensor<1xi64>
// Check and include zero in the range so that zero value can be correctly
// represented.
%input_1d = "tf.Reshape"(%input, %shape) : (tensor<*xf32>, tensor<1xi32>) -> tensor<?xf32>
%r_max_without_zero = "tf.Max"(%input_1d, %dim) { keep_dims = true }: (tensor<?xf32>, tensor<1xi64>) -> tensor<1xf32>
%r_max = "tf.Maximum"(%zero, %r_max_without_zero) : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
%r_min_without_zero = "tf.Min"(%input_1d, %dim) { keep_dims = true }: (tensor<?xf32>, tensor<1xi64>) -> tensor<1xf32>
%r_min = "tf.Minimum"(%zero, %r_min_without_zero) : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
%r_max_f64 = "tf.Cast"(%r_max) : (tensor<1xf32>) -> tensor<1xf64>
%r_min_f64 = "tf.Cast"(%r_min) : (tensor<1xf32>) -> tensor<1xf64>
%i8_min = "tf.Const"() {value = dense<-128.0> : tensor<f32>} : () -> tensor<f32>
%i8_max = "tf.Const"() {value = dense<127.0> : tensor<f32>} : () -> tensor<f32>
%i8_min_f64 = "tf.Cast"(%i8_min) : (tensor<f32>) -> tensor<f64>
%i8_max_f64 = "tf.Cast"(%i8_max) : (tensor<f32>) -> tensor<f64>
%range_nume = "tf.Sub"(%r_max_f64, %r_min_f64) : (tensor<1xf64>, tensor<1xf64>) -> tensor<1xf64>
%range_deno = "tf.Sub"(%i8_max_f64, %i8_min_f64) : (tensor<f64>, tensor<f64>) -> tensor<f64>
%scale_f64 = "tf.Div"(%range_nume, %range_deno) : (tensor<1xf64>, tensor<f64>) -> tensor<1xf64>
%scale = "tf.Cast"(%scale_f64) : (tensor<1xf64>) -> tensor<1xf32>
// Add comparison with minimum if needed
%intermediate_val = "tf.Div"(%r_max_f64, %scale_f64) : (tensor<1xf64>, tensor<1xf64>) -> tensor<1xf64>
%zp_from_max = "tf.Sub"(%i8_max_f64, %intermediate_val) : (tensor<f64>, tensor<1xf64>) -> tensor<1xf64>
%zp_fp32 = "tf.Cast"(%zp_from_max) : (tensor<1xf64>) -> tensor<1xf32>
%zp = "tf.Cast"(%zp_fp32) : (tensor<1xf32>) -> tensor<1xi32>
func.return %scale, %zp : tensor<1xf32>, tensor<1xi32>
}
// Matmul with int32 accumulation
func.func private @internal_matmul_fn(
%input : tensor<*xi8>, %filter : tensor<*xi8>,
%input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>,
%weight_scale : tensor<*xf32>, %weight_zp : tensor<*xi32>) -> tensor<*xi32> {
%0 = "tf.Cast"(%input) {Truncate = false} : (tensor<*xi8>) -> tensor<*xi32>
%1 = "tf.Sub"(%0, %input_zp) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
// Use identity op to avoid the weight being constant-folded.
%identity = "tf.Identity"(%filter) : (tensor<*xi8>) -> tensor<*xi8>
%2 = "tf.Cast"(%identity) {Truncate = false} : (tensor<*xi8>) -> tensor<*xi32>
%3 = "tf.Sub"(%2, %weight_zp) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
%5 = "tf.MatMul"(%1, %3) {
attr_map = "transpose_a:0,transpose_b:1"
} : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
func.return %5 : tensor<*xi32>
}
// Conv2D with int32 accumulation
func.func private @internal_conv2d_fn(
%input : tensor<*xi8>, %filter : tensor<*xi8>,
%input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>,
%filter_scale : tensor<*xf32>, %filter_zp : tensor<*xi32>) -> tensor<*xi32> {
%0 = "tf.Cast"(%input) {Truncate = false} : (tensor<*xi8>) -> tensor<*xi32>
%1 = "tf.Sub"(%0, %input_zp) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
// Use identity op to avoid the weight being constant-folded.
%identity = "tf.Identity"(%filter) : (tensor<*xi8>) -> tensor<*xi8>
%2 = "tf.Cast"(%identity) {Truncate = false} : (tensor<*xi8>) -> tensor<*xi32>
%3 = "tf.Sub"(%2, %filter_zp) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
%5 = "tf.Conv2D"(%1, %3) {
padding = "VALID", strides = [1, 1, 1, 1],
attr_map = "strides:0,use_cudnn_on_gpu:1,padding:2,explicit_paddings:3,dilations:4"
} : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
func.return %5 : tensor<*xi32>
}
// DepthwiseConv2D with float computation
func.func private @internal_depthwise_conv2d_fn(
%input : tensor<*xi8>, %filter : tensor<*xi8>,
%input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>,
%filter_scale : tensor<*xf32>, %filter_zp : tensor<*xi32>) -> tensor<*xi32> {
%0 = "tf.Cast"(%input) {Truncate = false} : (tensor<*xi8>) -> tensor<*xi32>
%1 = "tf.Sub"(%0, %input_zp) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
// Use identity op to avoid the weight being constant-folded.
%identity = "tf.Identity"(%filter) : (tensor<*xi8>) -> tensor<*xi8>
%2 = "tf.Cast"(%identity) {Truncate = false} : (tensor<*xi8>) -> tensor<*xi32>
%3 = "tf.Sub"(%2, %filter_zp) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
%cast_1_f32 = "tf.Cast"(%1) {Truncate = false} : (tensor<*xi32>) -> tensor<*xf32>
%cast_3_f32 = "tf.Cast"(%3) {Truncate = false} : (tensor<*xi32>) -> tensor<*xf32>
%5 = "tf.DepthwiseConv2dNative"(%cast_1_f32, %cast_3_f32) {
padding = "VALID", strides = [1, 1, 1, 1],
attr_map = "strides:0,padding:1,explicit_paddings:2,dilations:3"
} : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
%6 = "tf.Cast"(%5) : (tensor<*xf32>) -> tensor<*xi32>
func.return %6 : tensor<*xi32>
}
parameters[
{"quantized_ops": ["MatMul"], "internal_func_name": "internal_matmul_fn"},
{"quantized_ops": ["Conv2D"], "internal_func_name": "internal_conv2d_fn"},
{"quantized_ops": ["DepthwiseConv2D"], "internal_func_name": "internal_depthwise_conv2d_fn"}
]
func.func @GenerateQuantizedFunctionName(${quantized_ops})(
%input : tensor<*xf32>, %weight : tensor<*xi8>,
%weight_scale : tensor<*xf32>, %weight_zp : tensor<*xi32>) -> tensor<*xf32>
attributes {tf_quant.quantized_ops = ${quantized_ops}} {
%input_scale, %input_zp = "tf.PartitionedCall"(%input) {
config = "", config_proto = "", executor_type = "", f=@internal_calculate_quant_params
} : (tensor<*xf32>) -> (tensor<*xf32>, tensor<*xi32>)
%quantized_input = "tf.PartitionedCall"(%input, %input_scale, %input_zp) {
config = "", config_proto = "", executor_type = "", f=@internal_quantize_i8
} : (tensor<*xf32>, tensor<*xf32>, tensor<*xi32>) -> tensor<*xi8>
%accum_out = "tf.PartitionedCall"(%quantized_input, %weight, %input_scale, %input_zp,
%weight_scale, %weight_zp) {
config = "", config_proto = "", executor_type = "", f=@${internal_func_name}
} : (tensor<*xi8>, tensor<*xi8>, tensor<*xf32>, tensor<*xi32>,
tensor<*xf32>, tensor<*xi32>) -> tensor<*xi32>
%out = "tf.PartitionedCall"(%accum_out, %input_scale, %weight_scale) {
config = "", config_proto = "", executor_type = "", f=@internal_dequantize_i32
} : (tensor<*xi32>, tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
func.return %out : tensor<*xf32>
}
// For weight-only
func.func @dequantize_i8(%input : tensor<*xi8>, %scale : tensor<*xf32>, %zp : tensor<*xi32>) -> tensor<*xf32> {
// Use identity op to avoid the weight being constant-folded.
%identity = "tf.Identity"(%input) : (tensor<*xi8>) -> tensor<*xi8>
%input_i32 = "tf.Cast"(%identity) : (tensor<*xi8>) -> tensor<*xi32>
%output = "tf.Sub"(%input_i32, %zp) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
%cast = "tf.Cast"(%output) : (tensor<*xi32>) -> tensor<*xf32>
%mul = "tf.Mul"(%cast, %scale) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
func.return %mul : tensor<*xf32>
}
//===----------------------------------------------------------------------===//
// Weight-only functions.
//===----------------------------------------------------------------------===//
func.func private @internal_dequantize_f32(
%input : tensor<*xf32>, %weight_scale : tensor<*xf32>) -> tensor<*xf32> {
%mul = "tf.Mul"(%input, %weight_scale) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
func.return %mul : tensor<*xf32>
}
// Note that input i64 type is also supported by this.
parameters[
{"quantized_ops": ["Gather"]}
]
func.func @GenerateQuantizedFunctionName(${quantized_ops})(
%weight : tensor<*xi8>, %input : tensor<*xi32>, %axis : tensor<i32>,
%weight_scale : tensor<*xf32>, %weight_zp : tensor<*xi32>) -> tensor<*xf32>
attributes {tf_quant.quantized_ops = ${quantized_ops}}
{
%accum_out = "tf.GatherV2"(%weight, %input, %axis) {
batch_dims = 0 : i64, attr_map = "batch_dims:0"} : (tensor<*xi8>, tensor<*xi32>, tensor<i32>) -> tensor<*xi8>
%accum_out_new = "tf.Cast"(%accum_out) : (tensor<*xi8>) -> tensor<*xf32>
%out = "tf.PartitionedCall"(%accum_out_new, %weight_scale) {
config = "", config_proto = "", executor_type = "", f=@internal_dequantize_f32
} : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
func.return %out : tensor<*xf32>
}
}
tensorflow/compiler/mlir/quantization/tensorflow/passes/quantized_function_library_uniform_quantized.mlir
+313
View File
@@ -0,0 +1,313 @@
// Copyright 2022 The TensorFlow Runtime Authors
//
// 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.
// Quantization as a function library with Uniform Quantized Ops for Static
// PTQ
//
// Internal functions should be marked as private. They will be inlined and
// deleted in `InsertQuantizedFunctionsPass`.
//
// Function template can generate functions with different parameters. Ex:
// ```
// parameters[
// {"key1": "value11", "key2": "value21"},
// {"key1": "value12", "key2": "value22"},
// ]
// func.func func_name_${key1}_fn (...) {
// ...${key2}...
// }
// ```
// The above template with generate two functions by substituting `key1` and
// `key2` with given values.
module {
for main_op in ["Conv2D", "DepthwiseConv2D", "MatMul"] {
parameters[
{"quantized_ops": ["${main_op}", "BiasAdd"], "act_func": "internal_requantize_no_activation_fn", "output_type": "!tf_type.qint8"},
{"quantized_ops": ["${main_op}", "BiasAdd", "Relu"], "act_func": "internal_requantize_and_relu_fn", "output_type": "!tf_type.qint8"},
{"quantized_ops": ["${main_op}", "BiasAdd", "Relu6"], "act_func": "internal_requantize_and_relu6_fn", "output_type": "!tf_type.qint8"},
]
func.func @GenerateQuantizedFunctionName(${quantized_ops})(%input : tensor<*x!tf_type.qint8>,
%filter : tensor<*x!tf_type.qint8>, %bias : tensor<*x!tf_type.qint32>,
%input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>,
%filter_scale : tensor<*xf32>, %filter_zp : tensor<*xi32>,
%bias_scale : tensor<*xf32>, %bias_zp : tensor<*xi32>,
%out_scale : tensor<*xf32>, %out_zp : tensor<*xi32>) -> tensor<*x${output_type}>
attributes {tf_quant.quantized_ops = ${quantized_ops}} {
// Given the convolution takes 2 qint8 inputs and output a qint32.
// The accumulation scale is (input_scale * filter_scale).
// The accumulation zero point is 0.
%accum_scale = "tf.Mul"(%input_scale, %filter_scale) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
%accum_zp_f32 = "tf.ZerosLike"(%accum_scale) : (tensor<*xf32>) -> tensor<*xf32>
%accum_zp = "tf.Cast"(%accum_zp_f32) {Truncate = false} : (tensor<*xf32>) -> tensor<*xi32>
%main_out = "tf.PartitionedCall"(%input, %filter, %input_scale, %input_zp,
%filter_scale, %filter_zp, %accum_scale, %accum_zp) {
config = "", config_proto = "", executor_type = "", f=@GenerateImplFunctionName(${main_op})
} : (tensor<*x!tf_type.qint8>, tensor<*x!tf_type.qint8>, tensor<*xf32>, tensor<*xi32>, tensor<*xf32>, tensor<*xi32>, tensor<*xf32>, tensor<*xi32>) -> tensor<*x!tf_type.qint32>
%add = "tf.UniformQuantizedAdd"(%main_out, %bias, %accum_scale, %accum_zp, %bias_scale, %bias_zp, %accum_scale, %accum_zp) {
lhs_quantization_axis = -1,
lhs_quantization_min_val = -2147483648,
lhs_quantization_max_val = 2147483647,
rhs_quantization_axis = -1,
rhs_quantization_min_val = -2147483648,
rhs_quantization_max_val = 2147483647,
output_quantization_axis = -1,
output_quantization_min_val = -2147483648,
output_quantization_max_val = 2147483647,
T = "tfdtype$DT_QINT32",
attr_map = ""
} : (tensor<*x!tf_type.qint32>, tensor<*x!tf_type.qint32>, tensor<*xf32>, tensor<*xi32>, tensor<*xf32>, tensor<*xi32>, tensor<*xf32>, tensor<*xi32>) -> tensor<*x!tf_type.qint32>
%filter_shape = "tf.Shape" (%filter_scale) : (tensor<*xf32>) -> tensor<*xi32>
%out_scale_filled = "tf.Fill" (%filter_shape, %out_scale) : (tensor<*xi32>, tensor<*xf32>) -> tensor<*xf32>
%out_zp_filled = "tf.Fill" (%filter_shape, %out_zp) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
%act = "tf.PartitionedCall"(%add, %accum_scale, %accum_zp, %out_scale_filled, %out_zp_filled, %out_scale, %out_zp) {
config = "", config_proto = "", executor_type = "", f=@${act_func}
} : (tensor<*x!tf_type.qint32>, tensor<*xf32>, tensor<*xi32>, tensor<*xf32>, tensor<*xi32>, tensor<*xf32>, tensor<*xi32>) -> tensor<*x${output_type}>
func.return %act : tensor<*x${output_type}>
}
parameters[
{"quantized_ops": ["${main_op}"], "act_func": "internal_requantize_no_activation_fn", "output_type": "!tf_type.qint8"},
{"quantized_ops": ["${main_op}", "Relu"], "act_func": "internal_requantize_and_relu_fn", "output_type": "!tf_type.qint8"},
{"quantized_ops": ["${main_op}", "Relu6"], "act_func": "internal_requantize_and_relu6_fn", "output_type": "!tf_type.qint8"},
]
func.func @GenerateQuantizedFunctionName(${quantized_ops})(%input : tensor<*x!tf_type.qint8>, %filter : tensor<*x!tf_type.qint8>,
%input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>,
%filter_scale : tensor<*xf32>, %filter_zp : tensor<*xi32>,
%out_scale : tensor<*xf32>, %out_zp : tensor<*xi32>) -> tensor<*x${output_type}>
attributes {tf_quant.quantized_ops = ${quantized_ops}} {
// Given the convolution takes 2 qint8 inputs and output a qint32.
// The accumulation scale is (input_scale * filter_scale).
// The accumulation zero point is 0.
%accum_scale = "tf.Mul"(%input_scale, %filter_scale) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
%accum_zp_f32 = "tf.ZerosLike"(%accum_scale) : (tensor<*xf32>) -> tensor<*xf32>
%accum_zp = "tf.Cast"(%accum_zp_f32) {Truncate = false} : (tensor<*xf32>) -> tensor<*xi32>
%main_out = "tf.PartitionedCall"(%input, %filter, %input_scale, %input_zp,
%filter_scale, %filter_zp, %accum_scale, %accum_zp) {
config = "", config_proto = "", executor_type = "", f=@GenerateImplFunctionName(${main_op})
} : (tensor<*x!tf_type.qint8>, tensor<*x!tf_type.qint8>, tensor<*xf32>, tensor<*xi32>, tensor<*xf32>, tensor<*xi32>, tensor<*xf32>, tensor<*xi32>) -> tensor<*x!tf_type.qint32>
%filter_shape = "tf.Shape" (%filter_scale) : (tensor<*xf32>) -> tensor<*xi32>
%out_scale_filled = "tf.Fill" (%filter_shape, %out_scale) : (tensor<*xi32>, tensor<*xf32>) -> tensor<*xf32>
%out_zp_filled = "tf.Fill" (%filter_shape, %out_zp) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
%act = "tf.PartitionedCall"(%main_out, %accum_scale, %accum_zp, %out_scale_filled, %out_zp_filled, %out_scale, %out_zp) {
config = "", config_proto = "", executor_type = "", f=@${act_func}
} : (tensor<*x!tf_type.qint32>, tensor<*xf32>, tensor<*xi32>, tensor<*xf32>, tensor<*xi32>, tensor<*xf32>, tensor<*xi32>) -> tensor<*x${output_type}>
func.return %act : tensor<*x${output_type}>
}
} // end for
// Conv2d Convolution.
func.func private @internal_conv2d_fn(
%input : tensor<*x!tf_type.qint8>, %filter : tensor<*x!tf_type.qint8>,
%input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>,
%filter_scale : tensor<*xf32>, %filter_zp : tensor<*xi32>, %out_scale : tensor<*xf32>, %out_zp : tensor<*xi32>) -> tensor<*x!tf_type.qint32> {
%conv_out = "tf.UniformQuantizedConvolution"(%input, %filter,
%input_scale, %input_zp, %filter_scale, %filter_zp, %out_scale, %out_zp) {
Tin = "tfdtype$DT_QINT8",
Tout = "tfdtype$DT_QINT32",
window_strides = [1, 1],
padding = "SAME",
explicit_padding = [],
lhs_dilation = [],
rhs_dilation = [],
batch_group_count = 1,
feature_group_count = 1,
dimension_numbers = "",
lhs_quantization_axis = -1,
lhs_quantization_min_val = -128,
lhs_quantization_max_val = 127,
rhs_quantization_axis = -1,
rhs_quantization_min_val = -128,
rhs_quantization_max_val = 127,
output_quantization_axis = -1,
output_quantization_min_val = -128,
output_quantization_max_val = 127,
attr_map = ""
} : (tensor<*x!tf_type.qint8>, tensor<*x!tf_type.qint8>, tensor<*xf32>, tensor<*xi32>, tensor<*xf32>, tensor<*xi32>, tensor<*xf32>, tensor<*xi32>) -> tensor<*x!tf_type.qint32>
func.return %conv_out : tensor<*x!tf_type.qint32>
}
// Depthwise convolution. feature_group_count is set to 3rd dim of input shape.
func.func private @internal_depthwise_conv2d_fn(
%input : tensor<*x!tf_type.qint8>, %filter : tensor<*x!tf_type.qint8>,
%input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>,
%filter_scale : tensor<*xf32>, %filter_zp : tensor<*xi32>, %out_scale : tensor<*xf32>, %out_zp : tensor<*xi32>) -> tensor<*x!tf_type.qint32> {
%conv_out = "tf.UniformQuantizedConvolution"(%input, %filter,
%input_scale, %input_zp, %filter_scale, %filter_zp, %out_scale, %out_zp) {
Tin = "tfdtype$DT_QINT8",
Tout = "tfdtype$DT_QINT32",
window_strides = [1, 1],
padding = "SAME",
explicit_padding = [],
lhs_dilation = [],
rhs_dilation = [],
batch_group_count = 1,
feature_group_count = 1,
dimension_numbers = "",
lhs_quantization_axis = -1,
lhs_quantization_min_val = -128,
lhs_quantization_max_val = 127,
rhs_quantization_axis = -1,
rhs_quantization_min_val = -128,
rhs_quantization_max_val = 127,
output_quantization_axis = -1,
output_quantization_min_val = -2147483648,
output_quantization_max_val = 2147483647,
attr_map = ""
} : (tensor<*x!tf_type.qint8>, tensor<*x!tf_type.qint8>, tensor<*xf32>, tensor<*xi32>, tensor<*xf32>, tensor<*xi32>, tensor<*xf32>, tensor<*xi32>) -> tensor<*x!tf_type.qint32>
func.return %conv_out : tensor<*x!tf_type.qint32>
}
// MatMul.
func.func private @internal_matmul_fn(%input : tensor<*x!tf_type.qint8>, %filter : tensor<*x!tf_type.qint8>,
%input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>,
%filter_scale : tensor<*xf32>, %filter_zp : tensor<*xi32>, %out_scale : tensor<*xf32>, %out_zp : tensor<*xi32>) -> tensor<*x!tf_type.qint32> {
%dot_out = "tf.UniformQuantizedDot"(%input, %filter,
%input_scale, %input_zp, %filter_scale, %filter_zp, %out_scale, %out_zp) {
Tin = "tfdtype$DT_QINT8",
Tout = "tfdtype$DT_QINT32",
lhs_quantization_axis = -1,
lhs_quantization_min_val = -128,
lhs_quantization_max_val = 127,
rhs_quantization_axis = -1,
rhs_quantization_min_val = -128,
rhs_quantization_max_val = 127,
output_quantization_axis = -1,
output_quantization_min_val = -2147483648,
output_quantization_max_val = 2147483647,
attr_map = ""
} : (tensor<*x!tf_type.qint8>, tensor<*x!tf_type.qint8>, tensor<*xf32>, tensor<*xi32>, tensor<*xf32>, tensor<*xi32>, tensor<*xf32>, tensor<*xi32>) -> tensor<*x!tf_type.qint32>
func.return %dot_out : tensor<*x!tf_type.qint32>
}
// Quantize initial input at the start of the graph. Output is qint8.
func.func @quantize_i8(%input : tensor<*xf32>, %input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>) -> tensor<*x!tf_type.qint8> {
%quantize = "tf.UniformQuantize"(%input, %input_scale, %input_zp) {
Tin = "tfdtype$DT_FLOAT",
Tout = "tfdtype$DT_QINT8",
quantization_axis = -1,
quantization_min_val = -128,
quantization_max_val = 127,
attr_map = ""
} : (tensor<*xf32>, tensor<*xf32>, tensor<*xi32>) -> tensor<*x!tf_type.qint8>
func.return %quantize : tensor<*x!tf_type.qint8>
}
// Requantize a qint32 tensor to qint8 tensor for the next input.
func.func private @internal_requantize_qi8_fn(%input : tensor<*x!tf_type.qint32>, %input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>, %out_scale: tensor<*xf32>, %out_zp: tensor<*xi32>) -> tensor<*x!tf_type.qint8> {
%requantize = "tf.UniformRequantize"(%input, %input_scale, %input_zp, %out_scale, %out_zp) {
Tin = "tfdtype$DT_QINT32",
Tout = "tfdtype$DT_QINT8",
input_quantization_axis = -1,
input_quantization_min_val = -2147483648,
input_quantization_max_val = 2147483647,
output_quantization_axis = -1,
output_quantization_min_val = -128,
output_quantization_max_val = 127,
attr_map = ""
} : (tensor<*x!tf_type.qint32>, tensor<*xf32>, tensor<*xi32>, tensor<*xf32>, tensor<*xi32>) -> tensor<*x!tf_type.qint8>
func.return %requantize : tensor<*x!tf_type.qint8>
}
// Quantize initial input at the start of the graph. Output is qint32.
func.func @quantize_i32(%input : tensor<*xf32>, %input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>) -> tensor<*x!tf_type.qint32> {
%quantize = "tf.UniformQuantize"(%input, %input_scale, %input_zp) {
Tin = "tfdtype$DT_FLOAT",
Tout = "tfdtype$DT_QINT32",
quantization_axis = -1,
quantization_min_val = -2147483648,
quantization_max_val = 2147483647,
attr_map = ""
} : (tensor<*xf32>, tensor<*xf32>, tensor<*xi32>) -> tensor<*x!tf_type.qint32>
func.return %quantize : tensor<*x!tf_type.qint32>
}
// Dequantize final graph output back to f32. Input is qint8.
func.func @dequantize_i8(%input : tensor<*x!tf_type.qint8>, %input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>) -> tensor<*xf32> {
%dequantize = "tf.UniformDequantize"(%input, %input_scale, %input_zp) {
Tin = "tfdtype$DT_QINT8",
Tout = "tfdtype$DT_FLOAT",
quantization_axis = -1,
quantization_min_val = -128,
quantization_max_val = 127,
attr_map = ""
} : (tensor<*x!tf_type.qint8>, tensor<*xf32>, tensor<*xi32>) -> tensor<*xf32>
func.return %dequantize : tensor<*xf32>
}
// Requantizes and applies quantized Relu by clipping.
func.func private @internal_requantize_no_activation_fn(%input : tensor<*x!tf_type.qint32>, %input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>,
%out_scale : tensor<*xf32>, %out_zp : tensor<*xi32>, %out_scale_single : tensor<*xf32>, %out_zp_single : tensor<*xi32>) -> tensor<*x!tf_type.qint8> {
%q_out = "tf.PartitionedCall"(%input, %input_scale, %input_zp, %out_scale_single, %out_zp_single) {
config = "", config_proto = "", executor_type = "", f=@internal_requantize_qi8_fn
} : (tensor<*x!tf_type.qint32>, tensor<*xf32>, tensor<*xi32>, tensor<*xf32>, tensor<*xi32>) -> tensor<*x!tf_type.qint8>
func.return %q_out : tensor<*x!tf_type.qint8>
}
// Requantizes and applies quantized Relu6 by clipping.
func.func private @internal_requantize_and_relu_fn(%input : tensor<*x!tf_type.qint32>, %input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>,
%out_scale : tensor<*xf32>, %out_zp : tensor<*xi32>, %out_scale_single : tensor<*xf32>, %out_zp_single : tensor<*xi32>) -> tensor<*x!tf_type.qint8> {
%filter_shape = "tf.Shape" (%input_scale) : (tensor<*xf32>) -> tensor<*xi32>
%i32_min = "tf.Const"() {value = dense<-2147483648> : tensor<i32>} : () -> tensor<i32>
%i32_max = "tf.Const"() {value = dense<2147483647> : tensor<i32>} : () -> tensor<i32>
%i32_min_filled = "tf.Fill" (%filter_shape, %i32_min) : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
%i32_max_filled = "tf.Fill" (%filter_shape, %i32_max) : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
%clip_min = "tf.Maximum"(%i32_min_filled, %input_zp) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
%qclip_min = "tf.Cast"(%clip_min) {Truncate = false} : (tensor<*xi32>) -> tensor<*x!tf_type.qint32>
%qclip_max = "tf.Cast"(%i32_max_filled) {Truncate = false} : (tensor<*xi32>) -> tensor<*x!tf_type.qint32>
%relu = "tf.UniformQuantizedClipByValue"(%input, %qclip_min, %qclip_max, %input_scale, %input_zp) {
T = "tfdtype$DT_QINT32",
quantization_axis = -1,
quantization_min_val = -2147483648,
quantization_max_val = 2147483647,
attr_map = ""
} : (tensor<*x!tf_type.qint32>, tensor<*x!tf_type.qint32>, tensor<*x!tf_type.qint32>, tensor<*xf32>, tensor<*xi32>) -> tensor<*x!tf_type.qint32>
%requantize = "tf.PartitionedCall"(%relu, %input_scale, %input_zp, %out_scale_single, %out_zp_single) {
config = "", config_proto = "", executor_type = "", f=@internal_requantize_qi8_fn
} : (tensor<*x!tf_type.qint32>, tensor<*xf32>, tensor<*xi32>, tensor<*xf32>, tensor<*xi32>) -> tensor<*x!tf_type.qint8>
func.return %requantize : tensor<*x!tf_type.qint8>
}
// Apply requantization and relu6.
func.func private @internal_requantize_and_relu6_fn(%input : tensor<*x!tf_type.qint32>, %input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>,
%out_scale : tensor<*xf32>, %out_zp : tensor<*xi32>, %out_scale_single : tensor<*xf32>, %out_zp_single : tensor<*xi32>) -> tensor<*x!tf_type.qint8> {
%input_shape = "tf.Shape" (%input_scale) : (tensor<*xf32>) -> tensor<*xi32>
%i32_min = "tf.Const"() {value = dense<-2147483648> : tensor<i32>} : () -> tensor<i32>
%i32_max = "tf.Const"() {value = dense<2147483647> : tensor<i32>} : () -> tensor<i32>
%act_max = "tf.Const"() {value = dense<6.0> : tensor<f32>} : () -> tensor<f32>
%i32_act_max_q32 = "tf.PartitionedCall"(%act_max, %input_scale, %input_zp) {
config = "", config_proto = "", executor_type = "", f=@quantize_i32
} : (tensor<f32>, tensor<*xf32>, tensor<*xi32>) -> tensor<*x!tf_type.qint32>
%i32_act_max_f32 = "tf.Cast"(%i32_act_max_q32) {Truncate = false} : (tensor<*x!tf_type.qint32>) -> tensor<i32>
%i32_min_filled = "tf.Fill" (%input_shape, %i32_min) : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
%i32_max_filled = "tf.Fill" (%input_shape, %i32_max) : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
%i32_act_max_f32_filled = "tf.Fill" (%input_shape, %i32_act_max_f32) : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
%clip_min = "tf.Maximum"(%i32_min_filled, %input_zp) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
%clip_max = "tf.Minimum"(%i32_max_filled, %i32_act_max_f32_filled) : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
%qclip_min = "tf.Cast"(%clip_min) {Truncate = false} : (tensor<*xi32>) -> tensor<*x!tf_type.qint32>
%qclip_max = "tf.Cast"(%clip_max) {Truncate = false} : (tensor<*xi32>) -> tensor<*x!tf_type.qint32>
%relu = "tf.UniformQuantizedClipByValue"(%input, %qclip_min, %qclip_max, %input_scale, %input_zp) {
T = "tfdtype$DT_QINT32",
quantization_axis = -1,
quantization_min_val = -2147483648,
quantization_max_val = 2147483647,
attr_map = ""
} : (tensor<*x!tf_type.qint32>, tensor<*x!tf_type.qint32>, tensor<*x!tf_type.qint32>, tensor<*xf32>, tensor<*xi32>) -> tensor<*x!tf_type.qint32>
%requantize = "tf.PartitionedCall"(%relu, %input_scale, %input_zp, %out_scale_single, %out_zp_single) {
config = "", config_proto = "", executor_type = "", f=@internal_requantize_qi8_fn
} : (tensor<*x!tf_type.qint32>, tensor<*xf32>, tensor<*xi32>, tensor<*xf32>, tensor<*xi32>) -> tensor<*x!tf_type.qint8>
func.return %requantize : tensor<*x!tf_type.qint8>
}
}

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