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/stablehlo/cc/BUILD
+465
View File
@@ -0,0 +1,465 @@
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 = [
# For TFLite Converter integration.
"//tensorflow/compiler/mlir/lite:__subpackages__",
"//tensorflow/compiler/mlir/quantization:__subpackages__",
"//tensorflow/compiler/mlir/tensorflow_to_stablehlo:__subpackages__",
],
licenses = ["notice"],
)
cc_library(
name = "types",
hdrs = ["types.h"],
compatible_with = get_compatible_with_portable(),
deps = [],
)
cc_library(
name = "component",
hdrs = ["component.h"],
compatible_with = get_compatible_with_portable(),
deps = [
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
"@com_google_absl//absl/status:statusor",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
)
cc_library(
name = "config",
srcs = ["config.cc"],
hdrs = ["config.h"],
compatible_with = get_compatible_with_portable(),
deps = [
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/strings:string_view",
],
)
tf_cc_test(
name = "config_test",
srcs = ["config_test.cc"],
deps = [
":config",
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
"@com_google_googletest//:gtest_main",
],
)
cc_library(
name = "io",
srcs = ["io.cc"],
hdrs = ["io.h"],
compatible_with = get_compatible_with_portable(),
deps = [
"@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",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:statusor",
],
)
tf_cc_test(
name = "io_test",
srcs = ["io_test.cc"],
deps = [
":io",
"@com_google_absl//absl/functional:any_invocable",
"@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_absl//absl/strings:string_view",
"@com_google_googletest//:gtest_main",
"@xla//xla/tsl/platform:env",
],
)
cc_library(
name = "graph_def",
srcs = [],
hdrs = ["graph_def.h"],
compatible_with = get_compatible_with_portable(),
deps = [
"//tensorflow/core:protos_all_cc",
],
)
tf_cc_test(
name = "graph_def_test",
srcs = ["graph_def_test.cc"],
deps = [
":graph_def",
"//tensorflow/core:protos_all_cc",
"@com_google_googletest//:gtest_main",
"@tsl//tsl/platform:protobuf",
],
)
cc_library(
name = "debugger",
srcs = ["debugger.cc"],
hdrs = ["debugger.h"],
compatible_with = get_compatible_with_portable(),
deps = [
"//tensorflow/compiler/mlir/quantization/tensorflow:tf_quant_ops",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
)
cc_library(
name = "permutation",
hdrs = ["permutation.h"],
compatible_with = get_compatible_with_portable(),
deps = [
"@llvm-project//llvm:Support",
"@llvm-project//mlir:Support",
],
)
tf_cc_test(
name = "permutation_test",
srcs = ["permutation_test.cc"],
deps = [
":permutation",
"@com_google_googletest//:gtest_main",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:Support",
],
)
cc_library(
name = "saved_model_export",
srcs = ["saved_model_export.cc"],
hdrs = ["saved_model_export.h"],
compatible_with = get_compatible_with_portable(),
deps = [
":io",
":pass_pipeline",
":saved_model_import",
":types",
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
"//tensorflow/compiler/mlir/quantization/tensorflow:exported_model_proto_cc",
"//tensorflow/compiler/mlir/quantization/tensorflow:passes",
"//tensorflow/compiler/mlir/quantization/tensorflow/cc:convert_asset_args",
"//tensorflow/compiler/mlir/quantization/tensorflow/cc:run_passes",
"//tensorflow/compiler/mlir/quantization/tensorflow/python:unfreeze_constants",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:mlir_roundtrip_flags",
"//tensorflow/compiler/mlir/tensorflow/transforms:verify_no_outside_compilation_markers_pass",
"//tensorflow/compiler/mlir/tf2xla/api/v2:tf_executor_to_graph",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/protobuf:for_core_protos_cc",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/base:core_headers",
"@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:string_view",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:statusor",
],
)
tf_cc_test(
name = "saved_model_export_test",
srcs = ["saved_model_export_test.cc"],
deps = [
":saved_model_export",
"//tensorflow/compiler/mlir/quantization/common:test_base",
"//tensorflow/compiler/mlir/quantization/tensorflow:exported_model_proto_cc",
"//tensorflow/core:all_kernels", # buildcleaner: keep Required to export to GraphDef
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:ops", # buildcleaner: keep Required to export to GraphDef
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:status_matchers",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:string_view",
"@com_google_googletest//:gtest_main",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
"@tsl//tsl/platform:protobuf",
"@xla//xla/tsl/platform:statusor",
],
)
cc_library(
name = "saved_model_import",
srcs = ["saved_model_import.cc"],
hdrs = ["saved_model_import.h"],
compatible_with = get_compatible_with_portable(),
deps = [
":types",
"//tensorflow/cc/saved_model:loader",
"//tensorflow/cc/saved_model:reader",
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
"//tensorflow/compiler/mlir/quantization/tensorflow:quantize_preprocess",
"//tensorflow/compiler/mlir/tensorflow:mlir_import_options",
"//tensorflow/compiler/mlir/tensorflow:translate_lib",
"//tensorflow/core/protobuf:for_core_protos_cc",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:string_view",
"@com_google_absl//absl/types:span",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:statusor",
],
)
tf_cc_test(
name = "saved_model_import_test",
srcs = ["saved_model_import_test.cc"],
deps = [
":saved_model_import",
":types",
"//tensorflow/compiler/mlir/quantization/common:test_base",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:IR",
],
)
cc_library(
name = "pass_pipeline",
srcs = ["pass_pipeline.cc"],
hdrs = ["pass_pipeline.h"],
compatible_with = get_compatible_with_portable(),
deps = [
"//tensorflow/compiler/mlir/quantization/stablehlo:bridge_passes",
"//tensorflow/compiler/mlir/quantization/stablehlo:passes",
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
"//tensorflow/compiler/mlir/quantization/tensorflow:passes",
"//tensorflow/compiler/mlir/quantization/tensorflow:quantization_options_proto_cc",
"//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:ReconcileUnrealizedCasts",
"@llvm-project//mlir:Transforms",
"@stablehlo//:stablehlo_passes",
"@xla//xla/mlir_hlo:mhlo_passes",
],
)
cc_library(
name = "pre_calibration",
srcs = ["pre_calibration.cc"],
hdrs = ["pre_calibration.h"],
compatible_with = get_compatible_with_portable(),
visibility = [
"//tensorflow/compiler/mlir/quantization/stablehlo:__subpackages__",
"//tensorflow/compiler/mlir/quantization/tensorflow:__subpackages__",
],
deps = [
":component",
":pass_pipeline",
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
"//tensorflow/compiler/mlir/quantization/tensorflow:passes",
"//tensorflow/compiler/mlir/quantization/tensorflow:quantization_options_proto_cc",
"//tensorflow/compiler/mlir/quantization/tensorflow/cc:run_passes",
"@com_google_absl//absl/base:nullability",
"@com_google_absl//absl/log:die_if_null",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:string_view",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@xla//xla/tsl/platform:errors",
],
)
tf_cc_test(
name = "pre_calibration_test",
srcs = ["pre_calibration_test.cc"],
deps = [
":config",
":pre_calibration",
"//tensorflow/compiler/mlir/quantization/common:test_base",
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
"//tensorflow/compiler/mlir/quantization/tensorflow:quantization_options_proto_cc",
"//tensorflow/compiler/mlir/quantization/tensorflow:tf_quant_ops",
"//tensorflow/compiler/mlir/tensorflow",
"@com_google_absl//absl/status:status_matchers",
"@com_google_absl//absl/status:statusor",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Support",
],
)
cc_library(
name = "report",
srcs = ["report.cc"],
hdrs = ["report.h"],
compatible_with = get_compatible_with_portable(),
deps = [
":io",
"//tensorflow/compiler/mlir/quantization/common:lift_as_function_call",
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
"//tensorflow/compiler/mlir/tensorflow",
"@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:Support",
"@tsl//tsl/platform:protobuf",
],
)
tf_cc_test(
name = "report_test",
srcs = ["report_test.cc"],
deps = [
":io",
":report",
"//tensorflow/compiler/mlir/quantization/common:test_base",
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
"@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_absl//absl/strings:string_view",
"@com_google_googletest//:gtest_main",
"@llvm-project//mlir:IR",
"@tsl//tsl/platform:protobuf",
],
)
cc_library(
name = "context",
srcs = [],
hdrs = ["context.h"],
compatible_with = get_compatible_with_portable(),
visibility = ["//visibility:public"],
deps = [
"@llvm-project//mlir:FuncExtensions",
"@llvm-project//mlir:IR",
],
)
cc_library(
name = "post_calibration",
srcs = ["post_calibration.cc"],
hdrs = ["post_calibration.h"],
compatible_with = get_compatible_with_portable(),
deps = [
":component",
":config",
":pass_pipeline",
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
"//tensorflow/compiler/mlir/quantization/stablehlo/instrumentations:save_report",
"//tensorflow/compiler/mlir/quantization/tensorflow:passes",
"//tensorflow/compiler/mlir/quantization/tensorflow/cc:run_passes",
"@com_google_absl//absl/base:nullability",
"@com_google_absl//absl/log:die_if_null",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:string_view",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@xla//xla/mlir_hlo:mhlo_passes",
"@xla//xla/tsl/platform:errors",
],
)
cc_library(
name = "static_range_ptq",
srcs = ["static_range_ptq.cc"],
hdrs = ["static_range_ptq.h"],
compatible_with = get_compatible_with_portable(),
deps = [
":component",
":context",
":post_calibration",
":pre_calibration",
":saved_model_export",
":saved_model_import",
":types",
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
"//tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration:component",
"//tensorflow/compiler/mlir/quantization/tensorflow:exported_model_proto_cc",
"//tensorflow/compiler/mlir/quantization/tensorflow:quantization_options_proto_cc",
"//tensorflow/compiler/mlir/quantization/tensorflow/python:py_function_lib",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/protobuf:for_core_protos_cc",
"@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/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:string_view",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@xla//xla/tsl/platform:statusor",
],
)
cc_library(
name = "weight_only_ptq",
srcs = ["weight_only_ptq.cc"],
hdrs = ["weight_only_ptq.h"],
compatible_with = get_compatible_with_portable(),
deps = [
":component",
":config",
":context",
":pass_pipeline",
":saved_model_export",
":saved_model_import",
":types",
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
"//tensorflow/compiler/mlir/quantization/stablehlo/instrumentations:save_report",
"//tensorflow/compiler/mlir/quantization/tensorflow:exported_model_proto_cc",
"//tensorflow/compiler/mlir/quantization/tensorflow:passes",
"//tensorflow/compiler/mlir/quantization/tensorflow/cc:run_passes",
"//tensorflow/compiler/mlir/quantization/tensorflow/python:py_function_lib",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core/protobuf:for_core_protos_cc",
"@com_google_absl//absl/base:nullability",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/log:die_if_null",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:string_view",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@xla//xla/mlir_hlo:mhlo_passes",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:statusor",
],
)
tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/BUILD
+133
View File
@@ -0,0 +1,133 @@
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/stablehlo:__subpackages__",
"//tensorflow/compiler/mlir/quantization/tensorflow:__subpackages__",
],
licenses = ["notice"],
)
cc_library(
name = "min_max_value",
srcs = [],
hdrs = ["min_max_value.h"],
compatible_with = get_compatible_with_portable(),
deps = [],
)
cc_library(
name = "statistics",
srcs = ["statistics.cc"],
hdrs = ["statistics.h"],
compatible_with = get_compatible_with_portable(),
deps = [
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
"//tensorflow/compiler/mlir/quantization/stablehlo/cc:io",
"//tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration:min_max_value",
"//tensorflow/compiler/mlir/quantization/tensorflow:tf_quant_ops",
"//tensorflow/compiler/mlir/quantization/tensorflow/calibrator:calibration_statistics_proto_cc",
"//tensorflow/compiler/mlir/quantization/tensorflow/python:py_function_lib",
"@com_google_absl//absl/container:flat_hash_map",
"@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:Support",
"@tsl//tsl/platform:path",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:statusor",
],
)
cc_library(
name = "component",
srcs = ["component.cc"],
hdrs = ["component.h"],
compatible_with = get_compatible_with_portable(),
deps = [
":representative_dataset",
":statistics",
"//tensorflow/compiler/mlir/quantization/stablehlo:passes",
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
"//tensorflow/compiler/mlir/quantization/stablehlo/cc:component",
"//tensorflow/compiler/mlir/quantization/stablehlo/cc:debugger",
"//tensorflow/compiler/mlir/quantization/stablehlo/cc:io",
"//tensorflow/compiler/mlir/quantization/stablehlo/cc:saved_model_export",
"//tensorflow/compiler/mlir/quantization/stablehlo/cc:types",
"//tensorflow/compiler/mlir/quantization/tensorflow:exported_model_proto_cc",
"//tensorflow/compiler/mlir/quantization/tensorflow:quantization_options_proto_cc",
"//tensorflow/compiler/mlir/quantization/tensorflow/calibrator:calibration_statistics_proto_cc",
"//tensorflow/compiler/mlir/quantization/tensorflow/cc:run_passes",
"//tensorflow/compiler/mlir/quantization/tensorflow/python:py_function_lib",
"//tensorflow/core/protobuf:for_core_protos_cc",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/base:nullability",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:die_if_null",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:string_view",
"@com_google_absl//absl/types:span",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:statusor",
],
)
cc_library(
name = "representative_dataset",
srcs = ["representative_dataset.cc"],
hdrs = ["representative_dataset.h"],
compatible_with = get_compatible_with_portable(),
deps = [
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
"//tensorflow/compiler/mlir/quantization/tensorflow:quantization_options_proto_cc",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:span",
],
)
tf_cc_test(
name = "representative_dataset_test",
srcs = ["representative_dataset_test.cc"],
deps = [
":representative_dataset",
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
"//tensorflow/compiler/mlir/quantization/tensorflow:quantization_options_proto_cc",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:status_matchers",
"@com_google_absl//absl/status:statusor",
"@com_google_googletest//:gtest_main",
],
)
cc_library(
name = "calibration_parameters",
srcs = [],
hdrs = ["calibration_parameters.h"],
compatible_with = get_compatible_with_portable(),
deps = ["//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc"],
)
tf_cc_test(
name = "calibration_parameters_test",
srcs = ["calibration_parameters_test.cc"],
deps = [
":calibration_parameters",
"@com_google_googletest//:gtest_main",
],
)
tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/calibration_parameters.h
+77
View File
@@ -0,0 +1,77 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_CALIBRATION_CALIBRATION_PARAMETERS_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_CALIBRATION_CALIBRATION_PARAMETERS_H_
#include <algorithm>
#include <cmath>
#include <cstdint>
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
namespace stablehlo::quantization {
// Calculates the bin width from the range and expected number of bins. The
// bin width is formalized to the form of 2^n. As a consequence, the actual
// number of bins might be smaller than the given `num_bins`.
inline float CalculateBinWidth(const float min_value, const float max_value,
const int32_t num_bins) {
const float raw_bin_width = (max_value - min_value) / num_bins;
return std::pow(2, std::ceil(std::log2(raw_bin_width)));
}
// Calculates the lower bound of the histogram. The lower bound is in form of
// `N * bin_width`.
inline float CalculateLowerBound(const float min_value, const float bin_width) {
return std::floor(min_value / bin_width) * bin_width;
}
// Calculates the bin index of the current value.
inline int32_t CalculateBinIndex(const float value, const float lower_bound,
const float bin_width) {
return std::floor((value - lower_bound) / bin_width);
}
// Same as `CalculateBinIndex` but clamps to avoid out-of-bound.
inline int32_t CalculateBinIndexSafe(const float value, const float lower_bound,
const float bin_width,
const int32_t num_bins) {
const int32_t bin_index = CalculateBinIndex(value, lower_bound, bin_width);
return std::clamp(bin_index, 0, num_bins - 1);
}
// Checks if the given method is a histogram-based calibration method.
inline bool IsHistogramCalibration(
const CalibrationOptions::CalibrationMethod method) {
return method ==
CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_PERCENTILE ||
method ==
CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_MSE_BRUTEFORCE ||
method == CalibrationOptions::
CALIBRATION_METHOD_HISTOGRAM_MSE_MAX_FREQUENCY ||
method ==
CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_MSE_SYMMETRIC;
}
// Gets the number of bins for the given calibration method.
inline int32_t GetNumBins(const CalibrationOptions& calib_opts) {
return IsHistogramCalibration(calib_opts.calibration_method())
? calib_opts.calibration_parameters().num_bins()
: 0;
}
} // namespace stablehlo::quantization
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_CALIBRATION_CALIBRATION_PARAMETERS_H_
tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/calibration_parameters_test.cc
+101
View File
@@ -0,0 +1,101 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/calibration_parameters.h"
#include <cmath>
#include <cstdint>
#include <gtest/gtest.h>
namespace stablehlo::quantization {
namespace {
// Calculates the number of bins from the range and bin width.
inline int32_t CalculateActualNumBins(const float min_value,
const float max_value,
const float bin_width) {
const float lower_bound = CalculateLowerBound(min_value, bin_width);
return std::ceil((max_value - lower_bound) / bin_width);
}
TEST(CalibrationParametersTest, CalculateBinWidthSmallerThanOne) {
float bin_width = CalculateBinWidth(/*min_value=*/0.0, /*max_value=*/25.0,
/*num_bins=*/256);
EXPECT_FLOAT_EQ(bin_width, 0.125);
int32_t actual_num_bins =
CalculateActualNumBins(/*min_value=*/0.0, /*max_value=*/25.0, bin_width);
EXPECT_EQ(actual_num_bins, 200);
// Calculate the bin width with the actual num bins.
float raw_bin_width = 25.0 / actual_num_bins;
EXPECT_FLOAT_EQ(bin_width, raw_bin_width);
}
TEST(CalibrationParametersTest, CalculateBinWidthLargerThanOne) {
float bin_width = CalculateBinWidth(/*min_value=*/0.0, /*max_value=*/360.0,
/*num_bins=*/256);
EXPECT_FLOAT_EQ(bin_width, 2.0);
int32_t actual_num_bins =
CalculateActualNumBins(/*min_value=*/0.0, /*max_value=*/360.0, bin_width);
EXPECT_EQ(actual_num_bins, 180);
// Calculate the bin width with the actual num bins.
float raw_bin_width = 360.0 / actual_num_bins;
EXPECT_FLOAT_EQ(bin_width, raw_bin_width);
}
TEST(CalibrationParametersTest, CalculateBinWidthDivisible) {
float bin_width = CalculateBinWidth(/*min_value=*/0.0, /*max_value=*/256.0,
/*num_bins=*/256);
EXPECT_FLOAT_EQ(bin_width, 1.0);
int32_t actual_num_bins =
CalculateActualNumBins(/*min_value=*/0.0, /*max_value=*/256.0, bin_width);
EXPECT_EQ(actual_num_bins, 256);
// Calculate the bin width with the actual num bins.
float raw_bin_width = 256.0 / actual_num_bins;
EXPECT_FLOAT_EQ(bin_width, raw_bin_width);
}
TEST(CalibrationParametersTest, CalculateNumBinsDivisible) {
int32_t num_bins = CalculateActualNumBins(
/*min_value=*/0.0, /*max_value=*/4.0, /*bin_width=*/2.0);
// Expect 2 bins: [0, 2), [2, 4].
EXPECT_EQ(num_bins, 2);
}
TEST(CalibrationParametersTest, CalculateNumBinsNotDivisible) {
int32_t num_bins = CalculateActualNumBins(
/*min_value=*/0.0, /*max_value=*/5.0, /*bin_width=*/2.0);
// Expect 3 bins: [0, 2), [2, 4), [4, 6].
EXPECT_EQ(num_bins, 3);
}
TEST(CalibrationParametersTest, CalculateBinIndex) {
int32_t bin_index = CalculateBinIndexSafe(/*value=*/3.0, /*lower_bound=*/0.0,
/*bin_width=*/2.0, /*num_bins=*/2);
EXPECT_EQ(bin_index, 1);
}
TEST(CalibrationParametersTest, CalculateBinIndexMaxValue) {
int32_t bin_index = CalculateBinIndexSafe(/*value=*/4.0, /*lower_bound=*/0.0,
/*bin_width=*/2.0, /*num_bins=*/2);
EXPECT_EQ(bin_index, 1);
}
} // namespace
} // namespace stablehlo::quantization
tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/component.cc
+215
View File
@@ -0,0 +1,215 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/component.h"
#include <optional>
#include <string>
#include <unordered_set>
#include <utility>
#include <vector>
#include "absl/algorithm/container.h"
#include "absl/base/nullability.h"
#include "absl/container/flat_hash_map.h"
#include "absl/log/die_if_null.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#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/Pass/PassManager.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/representative_dataset.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/statistics.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/debugger.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/io.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/saved_model_export.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/types.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.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/cc/run_passes.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/exported_model.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
namespace mlir::quant::stablehlo {
namespace {
using ::stablehlo::quantization::AddCalibrationStatistics;
using ::stablehlo::quantization::CreateRepresentativeDatasetFileMap;
using ::stablehlo::quantization::DisableDebugging;
using ::stablehlo::quantization::IsCalibrationRequired;
using ::stablehlo::quantization::QuantizationConfig;
using ::stablehlo::quantization::ReadStatistics;
using ::stablehlo::quantization::RepresentativeDatasetConfig;
using ::stablehlo::quantization::io::CreateTmpDir;
using ::stablehlo::quantization::io::GetLocalTmpFileName;
using ::stablehlo::quantization::io::ListDirectory;
using ::tensorflow::AssetFileDef;
using ::tensorflow::SignatureDef;
using ::tensorflow::calibrator::CalibrationStatistics;
using ::tensorflow::quantization::ExportedModel;
using ::tensorflow::quantization::PyFunctionLibrary;
using ::tensorflow::quantization::RunPasses;
using CalibrationStatisticsFlatMap =
absl::flat_hash_map<std::string, CalibrationStatistics>;
} // namespace
absl::Status RunCalibrationPasses(
mlir::ModuleOp module_op, MLIRContext& ctx,
absl::string_view calibration_data_dir,
const bool force_regenerate_calibration_data) {
// Disable DumpTensor ops when running calibration.
DisableDebugging(module_op);
std::vector<std::string> skipping_aggregator_ops;
if (!force_regenerate_calibration_data) {
TF_ASSIGN_OR_RETURN(const CalibrationStatisticsFlatMap statistics_map,
ReadStatistics(calibration_data_dir));
absl::c_for_each(statistics_map, [&](const auto& iter) {
return skipping_aggregator_ops.push_back(iter.first);
});
}
return RunPasses(
/*name=*/
CalibrationComponent::kName,
/*add_passes_func=*/
[calibration_data_dir, &skipping_aggregator_ops](PassManager& pm) {
pm.addPass(CreateInsertCalibrationStatisticsSaverPass(
calibration_data_dir, skipping_aggregator_ops));
},
ctx, module_op);
}
CalibrationComponent::CalibrationComponent(
MLIRContext* absl_nonnull ctx,
const PyFunctionLibrary* absl_nonnull py_function_lib,
const absl::string_view src_saved_model_path,
absl::flat_hash_map<FunctionName, FunctionAlias> function_aliases,
std::unordered_set<std::string> tags,
absl::flat_hash_map<std::string, SignatureDef> signature_def_map,
std::vector<std::string> signature_keys)
: ctx_(ABSL_DIE_IF_NULL(ctx)), // Crash OK
py_function_lib_(ABSL_DIE_IF_NULL(py_function_lib)), // Crash OK
src_saved_model_path_(src_saved_model_path),
function_aliases_(std::move(function_aliases)),
tags_(std::move(tags)),
signature_def_map_(std::move(signature_def_map)),
signature_keys_(std::move(signature_keys)) {}
absl::Status CalibrationComponent::ExportToSavedModel(
ModuleOp module_op, absl::string_view calibration_data_dir,
const bool force_regenerate_calibration_data,
const absl::string_view dst_saved_model_path) {
TF_ASSIGN_OR_RETURN(const std::string checkpoint_dir, GetLocalTmpFileName());
// Clone ModuleOp and function aliases so changes in this pipeline won't
// be reflected in the original values.
mlir::OwningOpRef<mlir::ModuleOp> cloned_module_ref(module_op.clone());
TF_RETURN_IF_ERROR(RunCalibrationPasses(*cloned_module_ref, *ctx_,
calibration_data_dir,
force_regenerate_calibration_data));
const bool is_calibration_required =
IsCalibrationRequired(*cloned_module_ref);
if (!is_calibration_required) return absl::OkStatus();
// `duplicate_shape_determining_constants = false` because the
// resulting graph of this step is not expected to be loaded on TPU.
const quant::stablehlo::ExportOptions export_opts = {
/*duplicate_shape_determining_constants=*/false,
/*unfreeze_constants=*/false, checkpoint_dir,
/*debug_name=*/
absl::StrCat(kName, quant::stablehlo::kExportStepSuffix)};
TF_ASSIGN_OR_RETURN(const SmallVector<AssetFileDef> asset_file_defs,
RunExportPasses(export_opts, *ctx_, *cloned_module_ref));
TF_ASSIGN_OR_RETURN(ExportedModel exported_model,
quant::stablehlo::ConvertMlirModuleToExportedModel(
*cloned_module_ref, checkpoint_dir, function_aliases_,
{asset_file_defs.begin(), asset_file_defs.end()}));
py_function_lib_->SaveExportedModel(dst_saved_model_path, exported_model,
src_saved_model_path_, tags_,
signature_def_map_);
return absl::OkStatus();
}
absl::StatusOr<ModuleOp> CalibrationComponent::Run(
ModuleOp module_op, const QuantizationConfig& config) {
// Export the calibration model to SavedModel.
TF_ASSIGN_OR_RETURN(const std::string calibration_saved_model_dir,
CreateTmpDir());
std::string calibration_data_dir =
config.calibration_options().calibration_data_dir();
if (calibration_data_dir.empty()) {
TF_ASSIGN_OR_RETURN(calibration_data_dir, CreateTmpDir());
}
TF_RETURN_IF_ERROR(ExportToSavedModel(
module_op, calibration_data_dir,
config.calibration_options().force_regenerate_calibration_data(),
calibration_saved_model_dir));
TF_ASSIGN_OR_RETURN(std::vector<std::string> calibration_saved_model_files,
ListDirectory(calibration_saved_model_dir));
if (!calibration_saved_model_files.empty()) {
// Translate `RepresentativeDatasetConfig`s to signature key ->
// `RepresentativeDatasetFile` mapping.
const auto dataset_configs =
config.calibration_options().representative_datasets();
const std::vector<RepresentativeDatasetConfig> dataset_config_vector(
dataset_configs.begin(), dataset_configs.end());
TF_ASSIGN_OR_RETURN(
const auto representative_dataset_file_map,
CreateRepresentativeDatasetFileMap(dataset_config_vector));
// Run calibration on the exported model.
if (py_function_lib_->RunCalibration(
calibration_saved_model_dir, signature_keys_, tags_,
/*force_graph_mode_calibration=*/true,
representative_dataset_file_map) == std::nullopt) {
return absl::InternalError(
"CalibrationComponent error: Failed to run calibration.");
}
}
if (absl::Status status = AddCalibrationStatistics(
module_op, calibration_data_dir, config.calibration_options(),
*py_function_lib_);
!status.ok()) {
LOG(WARNING) << "Some CustomAggregator ops do not have min or max "
"values. Parts of the graph are not quantized. "
<< status;
}
return module_op;
}
} // namespace mlir::quant::stablehlo
tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/component.h
+122
View File
@@ -0,0 +1,122 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_CALIBRATION_COMPONENT_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_CALIBRATION_COMPONENT_H_
#include <string>
#include <unordered_set>
#include <vector>
#include "absl/base/nullability.h"
#include "absl/container/flat_hash_map.h"
#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/IR/MLIRContext.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/component.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/types.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/exported_model.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
namespace mlir::quant::stablehlo {
// Performs post-calibration graph transformation as part of post-training
// static-range quantization.
//
// The resulting `ModuleOp` contains quantized StableHLO ops serialized in
// `TF::XlaCallModuleOp`s. They are quantized using the statistics collected
// after the calibration step, corresponding to each `TF::CustomAggregatorOp`s
// in the input module op.
//
// TODO: b/320607042 - Add tests for this component on the python layer.
class CalibrationComponent : public Component {
public:
// Name of the post-training quantization post-calibration step. Used for
// debugging purposes.
static constexpr absl::string_view kName = "quant_ptq_calibration";
// `CalibrationComponent` ctor with necessary information required to run
// calibration on a `ModuleOp`. Meta information like `function_aliases`,
// `tags`, `signature_def_map`, and `signature_keys` are required to properly
// save and load the module_op to and from SavedModel.
// `representative_dataset_file_map` contains information about the
// calibration dataset.
CalibrationComponent(
MLIRContext* absl_nonnull ctx,
const tensorflow::quantization::PyFunctionLibrary* absl_nonnull
py_function_lib,
absl::string_view src_saved_model_path,
absl::flat_hash_map<FunctionName, FunctionAlias> function_aliases,
std::unordered_set<std::string> tags,
absl::flat_hash_map<std::string, tensorflow::SignatureDef>
signature_def_map,
std::vector<std::string> signature_keys);
// Runs calibration on `module_op` and returns a calibrated ModuleOp with
// calibrated statistics embedded.
absl::StatusOr<ModuleOp> Run(
ModuleOp module_op,
const ::stablehlo::quantization::QuantizationConfig& config) override;
private:
// Exports `module_op` to SavedModel at `dst_saved_model_path`. This is used
// to export the pre-calibrated `module_op` to SavedModel so that the
// calibration process can use it to load and run the graph with the
// representative dataset. Returns a failure status if the export fails.
absl::Status ExportToSavedModel(ModuleOp module_op,
absl::string_view calibration_data_dir,
bool force_regenerate_calibration_data,
absl::string_view dst_saved_model_path);
// Imports the SavedModel at `calibrated_saved_model_path` to `ModuleOp` after
// running calibration.
absl::StatusOr<ModuleOp> ImportCalibratedSavedModel(
absl::string_view calibrated_saved_model_path);
MLIRContext* absl_nonnull ctx_;
// Contains function implementations from the python layer. Should be injected
// from the python level using pybind11.
const tensorflow::quantization::PyFunctionLibrary* absl_nonnull
py_function_lib_;
// Path to the pre-calibrated SavedModel.
std::string src_saved_model_path_;
// Function alias mapping for pre-calibrated SavedModel. Used to preserve
// aliased functions.
absl::flat_hash_map<FunctionName, FunctionAlias> function_aliases_;
// Tags to identify the MetaGraphDef to load from a SavedModel.
const std::unordered_set<std::string> tags_;
const absl::flat_hash_map<std::string, tensorflow::SignatureDef>
signature_def_map_;
// Signature keys to identify the functions to load & quantize.
const std::vector<std::string> signature_keys_;
};
// Runs passes to prepare the calibration model.
absl::Status RunCalibrationPasses(mlir::ModuleOp module_op, MLIRContext& ctx,
absl::string_view calibration_data_dir,
bool force_regenerate_calibration_data);
} // namespace mlir::quant::stablehlo
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_CALIBRATION_COMPONENT_H_
tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/min_max_value.h
+28
View File
@@ -0,0 +1,28 @@
/* 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_STABLEHLO_CC_CALIBRATION_MIN_MAX_VALUE_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_CALIBRATION_MIN_MAX_VALUE_H_
#include <utility>
namespace stablehlo::quantization {
// Represents the (min, max) value pair, representing the range of values after
// calibrating for quantization.
using MinMaxValue = std::pair<float, float>;
} // namespace stablehlo::quantization
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_CALIBRATION_MIN_MAX_VALUE_H_
tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/representative_dataset.cc
+60
View File
@@ -0,0 +1,60 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/representative_dataset.h"
#include <string>
#include <utility>
#include "absl/container/flat_hash_map.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/types/span.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
namespace stablehlo::quantization {
using ::tensorflow::quantization::RepresentativeDatasetFile;
absl::StatusOr<absl::flat_hash_map<std::string, RepresentativeDatasetFile>>
CreateRepresentativeDatasetFileMap(absl::Span<const RepresentativeDatasetConfig>
representative_dataset_configs) {
absl::flat_hash_map<std::string, RepresentativeDatasetFile>
repr_dataset_file_map{};
for (const RepresentativeDatasetConfig& dataset_config :
representative_dataset_configs) {
RepresentativeDatasetFile repr_dataset_file;
repr_dataset_file.set_tfrecord_file_path(dataset_config.tf_record().path());
// If the signature_key has not been explicitly specified, use the default
// value of "serving_default".
const std::string signature_key = dataset_config.has_signature_key()
? dataset_config.signature_key()
: "serving_default";
if (repr_dataset_file_map.contains(signature_key)) {
return absl::InvalidArgumentError(
absl::StrCat("RepresentativeDatasetConfig should not contain "
"duplicate signature key: ",
signature_key));
}
repr_dataset_file_map[signature_key] = std::move(repr_dataset_file);
}
return repr_dataset_file_map;
}
} // namespace stablehlo::quantization
tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/representative_dataset.h
+41
View File
@@ -0,0 +1,41 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_CALIBRATION_REPRESENTATIVE_DATASET_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_CALIBRATION_REPRESENTATIVE_DATASET_H_
#include <string>
#include "absl/container/flat_hash_map.h"
#include "absl/status/statusor.h"
#include "absl/types/span.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
namespace stablehlo::quantization {
// Translates a set of `RepresentativeDatsetConfig` to signature key ->
// `RepresentativeDatasetFile` mapping. This is useful when using
// `RepresentativeDatasetConfig`s at places that accept the legacy
// `RepresentativeDatasetFile` mapping.
// Returns a non-OK status when there is a duplicate signature key among
// `representative_dataset_configs`.
absl::StatusOr<absl::flat_hash_map<
std::string, tensorflow::quantization::RepresentativeDatasetFile>>
CreateRepresentativeDatasetFileMap(absl::Span<const RepresentativeDatasetConfig>
representative_dataset_configs);
} // namespace stablehlo::quantization
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_CALIBRATION_REPRESENTATIVE_DATASET_H_
tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/representative_dataset_test.cc
+110
View File
@@ -0,0 +1,110 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/representative_dataset.h"
#include <string>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/container/flat_hash_map.h"
#include "absl/status/status.h"
#include "absl/status/status_matchers.h"
#include "absl/status/statusor.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
namespace mlir::quant::stablehlo {
namespace {
using ::stablehlo::quantization::RepresentativeDatasetConfig;
using ::tensorflow::quantization::RepresentativeDatasetFile;
using ::testing::Contains;
using ::testing::HasSubstr;
using ::testing::Key;
using ::testing::SizeIs;
using ::testing::StrEq;
TEST(CreateRepresentativeDatasetFileMapTest,
ConfigWithoutExplicitSignatureKeyMappedToServingDefault) {
std::vector<RepresentativeDatasetConfig> representative_dataset_configs;
RepresentativeDatasetConfig config{};
*(config.mutable_tf_record()->mutable_path()) = "test_path";
representative_dataset_configs.push_back(config);
const absl::StatusOr<
absl::flat_hash_map<std::string, RepresentativeDatasetFile>>
representative_dataset_file_map =
CreateRepresentativeDatasetFileMap(representative_dataset_configs);
ASSERT_THAT(representative_dataset_file_map, absl_testing::IsOk());
ASSERT_THAT(*representative_dataset_file_map, SizeIs(1));
EXPECT_THAT(*representative_dataset_file_map,
Contains(Key("serving_default")));
EXPECT_THAT(representative_dataset_file_map->at("serving_default")
.tfrecord_file_path(),
StrEq("test_path"));
}
TEST(CreateRepresentativeDatasetFileMapTest, ConfigWithExplicitSignatureKey) {
std::vector<RepresentativeDatasetConfig> representative_dataset_configs;
RepresentativeDatasetConfig config{};
config.set_signature_key("test_signature_key");
*(config.mutable_tf_record()->mutable_path()) = "test_path";
representative_dataset_configs.push_back(config);
const absl::StatusOr<
absl::flat_hash_map<std::string, RepresentativeDatasetFile>>
representative_dataset_file_map =
CreateRepresentativeDatasetFileMap(representative_dataset_configs);
ASSERT_THAT(representative_dataset_file_map, absl_testing::IsOk());
ASSERT_THAT(*representative_dataset_file_map, SizeIs(1));
EXPECT_THAT(*representative_dataset_file_map,
Contains(Key(StrEq("test_signature_key"))));
EXPECT_THAT(representative_dataset_file_map->at("test_signature_key")
.tfrecord_file_path(),
StrEq("test_path"));
}
TEST(CreateRepresentativeDatasetFileMapTest,
ConfigWithDuplicateSignatureKeyReturnsInvalidArgumentError) {
std::vector<RepresentativeDatasetConfig> representative_dataset_configs;
RepresentativeDatasetConfig config_1{};
config_1.set_signature_key("serving_default");
*(config_1.mutable_tf_record()->mutable_path()) = "test_path_1";
representative_dataset_configs.push_back(config_1);
// Signature key is implicitly "serving_default".
RepresentativeDatasetConfig config_2{};
*(config_2.mutable_tf_record()->mutable_path()) = "test_path_2";
representative_dataset_configs.push_back(config_2);
const absl::StatusOr<
absl::flat_hash_map<std::string, RepresentativeDatasetFile>>
representative_dataset_file_map =
CreateRepresentativeDatasetFileMap(representative_dataset_configs);
EXPECT_THAT(representative_dataset_file_map,
absl_testing::StatusIs(
absl::StatusCode::kInvalidArgument,
HasSubstr("duplicate signature key: serving_default")));
}
} // namespace
} // namespace mlir::quant::stablehlo
tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/statistics.cc
+115
View File
@@ -0,0 +1,115 @@
/* 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/stablehlo/cc/calibration/statistics.h"
#include <optional>
#include <string>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_format.h"
#include "absl/strings/string_view.h"
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/min_max_value.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/io.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/passes/tf_quant_ops.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h"
#include "xla/tsl/platform/statusor.h"
#include "tsl/platform/path.h"
namespace stablehlo::quantization {
namespace {
using ::stablehlo::quantization::CalibrationOptions;
using ::tensorflow::calibrator::CalibrationStatistics;
using ::tensorflow::calibrator::CalibrationStatisticsMap;
using ::tensorflow::quantization::PyFunctionLibrary;
using CalibrationStatisticsFlatMap =
absl::flat_hash_map<std::string, CalibrationStatistics>;
} // namespace
// Reads the calibration statistics from the given directory.
absl::StatusOr<CalibrationStatisticsFlatMap> ReadStatistics(
absl::string_view calibration_data_dir) {
TF_ASSIGN_OR_RETURN(std::vector<std::string> statistics_files,
io::ListDirectory(calibration_data_dir));
CalibrationStatisticsFlatMap statistics_map;
for (const std::string& statistics_file : statistics_files) {
TF_ASSIGN_OR_RETURN(
const auto single_map,
io::ReadBinaryProto<CalibrationStatisticsMap>(
tsl::io::JoinPath(calibration_data_dir, statistics_file)));
statistics_map.insert(single_map.statistics().begin(),
single_map.statistics().end());
}
return statistics_map;
}
absl::Status AddCalibrationStatistics(
mlir::ModuleOp module_op, absl::string_view calibration_data_dir,
const CalibrationOptions& calibration_options,
const PyFunctionLibrary& py_function_library) {
TF_ASSIGN_OR_RETURN(const CalibrationStatisticsFlatMap statistics_map,
ReadStatistics(calibration_data_dir));
absl::Status status = absl::OkStatus();
module_op.walk([&py_function_library, &calibration_options, &status,
&statistics_map](mlir::TF::CustomAggregatorOp aggregator_op) {
mlir::StringRef id = aggregator_op.getId();
auto iter = statistics_map.find(id);
if (iter == statistics_map.end()) {
status = absl::InternalError(
absl::StrFormat("Calibrated data does not exist. Cannot find "
"statistics. value for id: %s",
id));
return;
}
const std::optional<MinMaxValue> min_max_values =
py_function_library.GetCalibrationMinMaxValue(iter->second,
calibration_options);
if (min_max_values == std::nullopt) {
status = absl::InternalError(
"Cannot find min/max values for calibration statistics.");
return;
}
const auto [min_value, max_value] = *min_max_values;
mlir::OpBuilder builder(aggregator_op);
aggregator_op->setAttr("min", builder.getF32FloatAttr(min_value));
aggregator_op->setAttr("max", builder.getF32FloatAttr(max_value));
});
return status;
}
bool IsCalibrationRequired(mlir::ModuleOp module_op) {
bool calibration_required = false;
module_op.walk(
[&calibration_required](
mlir::TF::CalibrationStatisticsSaverOp statistics_saver_op) {
calibration_required = true;
});
return calibration_required;
}
} // namespace stablehlo::quantization
tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/statistics.h
+50
View File
@@ -0,0 +1,50 @@
/* 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_STABLEHLO_CC_CALIBRATION_STATISTICS_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_CALIBRATION_STATISTICS_H_
#include <string>
#include "absl/container/flat_hash_map.h"
#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 "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/python/py_function_lib.h"
namespace stablehlo::quantization {
// Reads the calibration statistics from the given directory.
absl::StatusOr<absl::flat_hash_map<
std::string, tensorflow::calibrator::CalibrationStatistics>>
ReadStatistics(absl::string_view calibration_data_dir);
// Adds calibrated min / max values to CustomAggregator nodes in `graph_def`.
// The min and max values will be added to the "min" and "max" attributes,
// respectively. `calibration_options` provides the strategy to retrieve min and
// max values.
absl::Status AddCalibrationStatistics(
mlir::ModuleOp module_op, absl::string_view calibration_data_dir,
const stablehlo::quantization::CalibrationOptions& calibration_options,
const tensorflow::quantization::PyFunctionLibrary& py_function_library);
// Checks if the model required calibration.
bool IsCalibrationRequired(mlir::ModuleOp module_op);
} // namespace stablehlo::quantization
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_CALIBRATION_STATISTICS_H_
tensorflow/compiler/mlir/quantization/stablehlo/cc/component.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_STABLEHLO_CC_COMPONENT_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_COMPONENT_H_
#include "absl/status/statusor.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
namespace mlir::quant::stablehlo {
// Component is a public abstraction for StableHLO Quantizer that represents the
// most basic unit of action applied to the StableHLO graph. Derived classes
// should override the `Run` method to implement the action.
class Component {
public:
virtual ~Component() = default;
// Runs the action to the StableHLO graph, passed by the `module_op`. `config`
// should provide information necessary to configure the action's behavior.
virtual absl::StatusOr<ModuleOp> Run(
ModuleOp module_op,
const ::stablehlo::quantization::QuantizationConfig& config) = 0;
};
} // namespace mlir::quant::stablehlo
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_COMPONENT_H_
tensorflow/compiler/mlir/quantization/stablehlo/cc/config.cc
+234
View File
@@ -0,0 +1,234 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/config.h"
#include <utility>
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
namespace stablehlo::quantization {
namespace {
// Populate `CalibrationOptions` with default fields.
void PopulateDefaultCalibrationOptions(QuantizationConfig& quant_config) {
if (!quant_config.has_calibration_options() ||
quant_config.calibration_options().calibration_method() ==
CalibrationOptions::CALIBRATION_METHOD_UNSPECIFIED) {
quant_config.mutable_calibration_options()->set_calibration_method(
CalibrationOptions::CALIBRATION_METHOD_MIN_MAX);
}
switch (quant_config.calibration_options().calibration_method()) {
case CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_PERCENTILE:
case CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_MSE_BRUTEFORCE:
case CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_MSE_MAX_FREQUENCY:
case CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_MSE_SYMMETRIC:
if (quant_config.calibration_options()
.calibration_parameters()
.num_bins() == 0) {
quant_config.mutable_calibration_options()
->mutable_calibration_parameters()
->set_num_bins(512);
}
if (quant_config.calibration_options().calibration_method() ==
CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_PERCENTILE) {
if (quant_config.calibration_options()
.calibration_parameters()
.min_percentile() == 0) {
quant_config.mutable_calibration_options()
->mutable_calibration_parameters()
->set_min_percentile(0.001);
}
if (quant_config.calibration_options()
.calibration_parameters()
.max_percentile() == 0) {
quant_config.mutable_calibration_options()
->mutable_calibration_parameters()
->set_max_percentile(99.999);
}
}
break;
default:
break;
}
}
// Returns a default `QuantizationSpec` for performing static-range PTQ on all
// ops.
//
// In textproto, the spec corresponds to:
//
// {
// {matcher {function_name {regex: ".*"}}
// {method {static_range_ptq {}}}
// }
QuantizationSpec GetDefaultStaticRangePtqSpec(StaticRangePtqPreset preset) {
QuantizationSpec spec{};
// Default for all ops.
spec.mutable_matcher()->mutable_function_name()->set_regex(
preset.enable_full_int_quantization() ? ".*"
: "^.*(dot_general|gather).*");
spec.mutable_method()->mutable_static_range_ptq();
return spec;
}
QuantizationSpec GetDefaultWeightOnlyPtqSpec() {
QuantizationSpec spec{};
spec.mutable_matcher()->mutable_function_name()->set_regex(
"^.*(conv|dot_general).*");
WeightOnlyPtq& weight_only_ptq_spec =
*spec.mutable_method()->mutable_weight_only_ptq();
if (auto [iter, inserted] =
weight_only_ptq_spec.mutable_input_quantized_types()->try_emplace(1);
inserted) {
iter->second.mutable_dimension_specs();
}
return spec;
}
// Returns a `QuantizationSpec` for performing static-range PTQ on the
// convolution quantizable unit family. Enables per-channel quantization for
// weights, on the channel dimension.
//
// In textproto, the spec corresponds to:
//
// {
// {matcher {function_name {regex: "composite_conv.*"}}}
// {method {static_range_ptq
// {input_quantized_types {
// key: 1,
// value {dimension_specs {dimension: 3}}}}
// }}
// }
QuantizationSpec GetPtqSpecForConvolution(Method::MethodCase method_case) {
QuantizationSpec spec{};
if (method_case != Method::kStaticRangePtq) {
return spec;
}
// Matches all convolution quantizable unit family.
spec.mutable_matcher()->mutable_function_name()->set_regex(
"composite_conv.*");
// Enable per-channel quantization for convolution weights.
QuantizedType conv_weight_quantized_type{};
// Assumes NHWC format, specifying the channel dimension (3) as the
// quantized axis.
conv_weight_quantized_type.mutable_dimension_specs()->set_dimension(3);
// The index of weight operands passed to lifted functions for convolution
// is 1.
StaticRangePtq& static_range_ptq_spec =
*spec.mutable_method()->mutable_static_range_ptq();
static_range_ptq_spec.mutable_input_quantized_types()->try_emplace(
1, std::move(conv_weight_quantized_type));
return spec;
};
void ExpandStaticRangePtqPreset(const StaticRangePtqPreset& preset,
QuantizationConfig& config) {
// Populate with preset's representative dataset configs if the user didn't
// explicitly specify other representative dataset configs to the top-level
// `CalibrationOptions`.
if (config.calibration_options().representative_datasets().empty()) {
auto preset_datasets = preset.representative_datasets();
config.mutable_calibration_options()
->mutable_representative_datasets()
->Add(preset_datasets.begin(), preset_datasets.end());
}
// Create a new `QuantizationSpecs` to replace the existing one. The
// expansion from `StaticRangePtqPreset` gets populated first and then
// user-provided explicit `QuantizationSpec`s will be appended.
QuantizationSpecs new_specs{};
*new_specs.add_specs() =
GetDefaultStaticRangePtqSpec(/*preset=*/config.static_range_ptq_preset());
*new_specs.add_specs() =
GetPtqSpecForConvolution(Method::MethodCase::kStaticRangePtq);
// Append user-provided specs to override existing specs.
const QuantizationSpecs& previous_specs = config.specs();
new_specs.mutable_specs()->Add(previous_specs.specs().begin(),
previous_specs.specs().end());
config.clear_static_range_ptq_preset();
config.mutable_specs()->Swap(&new_specs);
}
void ExpandWeightOnlyPtqPreset(QuantizationConfig& config) {
// Create a new `QuantizationSpecs` to replace the existing one. The
// expansion from `WeightOnlyPtqPreset` gets populated first and then
// user-provided explicit `QuantizationSpec`s will be appended.
QuantizationSpecs new_specs{};
*new_specs.add_specs() = GetDefaultWeightOnlyPtqSpec();
// Append user-provided specs to override existing specs.
const QuantizationSpecs& previous_specs = config.specs();
new_specs.mutable_specs()->Add(previous_specs.specs().begin(),
previous_specs.specs().end());
config.clear_weight_only_ptq_preset();
config.mutable_specs()->Swap(&new_specs);
}
} // namespace
QuantizationConfig ExpandPresets(const QuantizationConfig& config) {
QuantizationConfig new_config = config;
// Update the `new_config` with each preset's expansions.
switch (config.preset_case()) {
case QuantizationConfig::kStaticRangePtqPreset:
ExpandStaticRangePtqPreset(config.static_range_ptq_preset(), new_config);
break;
case QuantizationConfig::kWeightOnlyPtqPreset:
ExpandWeightOnlyPtqPreset(new_config);
break;
default:
// Preset has not been specified. The expansion is a no-op.
break;
}
return new_config;
}
bool HasQuantizationMethod(const QuantizationSpecs& specs,
Method::MethodCase method_case) {
for (const auto& spec : specs.specs()) {
if (spec.method().method_case() == method_case) {
return true;
}
}
return false;
}
QuantizationConfig PopulateDefaults(
const QuantizationConfig& user_provided_config) {
QuantizationConfig config = user_provided_config;
PopulateDefaultCalibrationOptions(config);
PipelineConfig& pipeline_config = *config.mutable_pipeline_config();
if (!pipeline_config.has_unpack_quantized_types()) {
pipeline_config.set_unpack_quantized_types(true);
}
return config;
}
} // namespace stablehlo::quantization
tensorflow/compiler/mlir/quantization/stablehlo/cc/config.h
+65
View File
@@ -0,0 +1,65 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_CONFIG_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_CONFIG_H_
#include <optional>
#include "absl/base/attributes.h"
#include "absl/strings/string_view.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
namespace stablehlo::quantization {
// Returns a copy of `user_provided_config` with default values populated where
// the user did not explicitly specify.
QuantizationConfig PopulateDefaults(
const QuantizationConfig& user_provided_config);
// Returns a copy of `QuantizationConfig` where presets are expanded and
// transformed into other fields in `QuantizationConfig`.
//
// The expansion rules are as follows:
// * StaticRangePtqPreset
// - The preset's `representative_datasets` field will be transferred to
// `QuantizationConfig.calibration_options.representative_datasets`, unless
// the user explicitly provided representative dataset configs to
// `calibration_options`. In that case, the explicit configs take precedence
// and the preset's configs are ignored.
// - For `QuantizationSpecs`, the expanded `QuantizationSpec`s will be
// populated first and user-provided `QuantizationSpec`s, if any, will be
// appended. This expresses the fact that user-provided specs take precedence.
// * Preset unspecified
// - No-op.
QuantizationConfig ExpandPresets(const QuantizationConfig& config);
// Returns whether a given QuantizationSpecs has the given quantization method.
bool HasQuantizationMethod(const QuantizationSpecs& specs,
Method::MethodCase method_case);
// Convenience function for converting the optional `report_file_path` field to
// `std::optional<absl::string_view>`, where `std::nullopt` represents that the
// field is not explicitly set. The returned value is a reference type
// (`absl::string_view`) so its lifetime is bound to the input `config`.
inline std::optional<absl::string_view> GetReportFilePath(
const QuantizationConfig& config ABSL_ATTRIBUTE_LIFETIME_BOUND) {
return config.has_report_file_path()
? std::make_optional<absl::string_view>(config.report_file_path())
: std::nullopt;
}
} // namespace stablehlo::quantization
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_CONFIG_H_
tensorflow/compiler/mlir/quantization/stablehlo/cc/config_test.cc
+298
View File
@@ -0,0 +1,298 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/config.h"
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
namespace stablehlo::quantization {
namespace {
using ::testing::Eq;
using ::testing::SizeIs;
using ::testing::StrEq;
using ::testing::Truly;
TEST(PopulateDefaultsTest, PopulateDefaultsForEmptyConfig) {
QuantizationConfig config{};
const QuantizationConfig new_config = PopulateDefaults(config);
EXPECT_TRUE(new_config.pipeline_config().unpack_quantized_types());
}
TEST(PopulateDefaultsTest, PopulateDefaultsForConfigWithUnpackQuantizedTypes) {
QuantizationConfig config{};
config.mutable_pipeline_config()->set_unpack_quantized_types(false);
// Test that if the user explicitly provided `unpack_quantized_types`, it is
// not overridden.
const QuantizationConfig new_config = PopulateDefaults(config);
EXPECT_FALSE(new_config.pipeline_config().unpack_quantized_types());
}
TEST(PopulateDefaultsTest, DefaultCalibrationOptionsPopulated) {
QuantizationConfig config{};
const QuantizationConfig new_config = PopulateDefaults(config);
EXPECT_THAT(new_config.calibration_options().calibration_method(),
Eq(CalibrationOptions::CALIBRATION_METHOD_MIN_MAX));
}
TEST(PopulateDefaultsTest,
DefaultCalibrationOptionsPopulatedForUnspecifiedMethod) {
QuantizationConfig config{};
CalibrationOptions& calibration_options =
*config.mutable_calibration_options();
calibration_options.set_calibration_method(
CalibrationOptions::CALIBRATION_METHOD_UNSPECIFIED);
const QuantizationConfig new_config = PopulateDefaults(config);
EXPECT_THAT(new_config.calibration_options().calibration_method(),
Eq(CalibrationOptions::CALIBRATION_METHOD_MIN_MAX));
}
TEST(PopulateDefaultsTest, ExplicitCalibrationOptionsNotOverridden) {
QuantizationConfig config{};
CalibrationOptions& calibration_options =
*config.mutable_calibration_options();
calibration_options.set_calibration_method(
CalibrationOptions::CALIBRATION_METHOD_AVERAGE_MIN_MAX);
calibration_options.mutable_calibration_parameters()->set_num_bins(512);
// Test that if the user explicitly provided `calibration_options`, it is not
// overridden.
const QuantizationConfig new_config = PopulateDefaults(config);
EXPECT_THAT(new_config.calibration_options().calibration_method(),
Eq(CalibrationOptions::CALIBRATION_METHOD_AVERAGE_MIN_MAX));
EXPECT_THAT(
new_config.calibration_options().calibration_parameters().num_bins(),
Eq(512));
}
TEST(PopulateDefaultsTest, DefaultNumbersPopulatedForPartOfCalibrationOptions) {
QuantizationConfig config{};
CalibrationOptions& calibration_options =
*config.mutable_calibration_options();
calibration_options.set_calibration_method(
CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_PERCENTILE);
calibration_options.mutable_calibration_parameters()->set_num_bins(512);
// Test that if the user explicitly provided part of the
// `calibration_options`, it is not overridden, rest of the data are default.
const QuantizationConfig new_config = PopulateDefaults(config);
EXPECT_THAT(new_config.calibration_options().calibration_method(),
Eq(CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_PERCENTILE));
EXPECT_THAT(
new_config.calibration_options().calibration_parameters().num_bins(),
Eq(512));
EXPECT_THAT(new_config.calibration_options()
.calibration_parameters()
.min_percentile(),
Eq(0.001f));
EXPECT_THAT(new_config.calibration_options()
.calibration_parameters()
.max_percentile(),
Eq(99.999f));
}
TEST(PopulateDefaultsTest,
DefaultNumbersPopulatedForCalibrationOptionsOfHistogramMseBruteforce) {
QuantizationConfig config{};
CalibrationOptions& calibration_options =
*config.mutable_calibration_options();
calibration_options.set_calibration_method(
CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_MSE_BRUTEFORCE);
const QuantizationConfig new_config = PopulateDefaults(config);
EXPECT_THAT(
new_config.calibration_options().calibration_method(),
Eq(CalibrationOptions::CALIBRATION_METHOD_HISTOGRAM_MSE_BRUTEFORCE));
EXPECT_THAT(
new_config.calibration_options().calibration_parameters().num_bins(),
Eq(512));
EXPECT_THAT(new_config.calibration_options()
.calibration_parameters()
.min_percentile(),
Eq(0.0f));
EXPECT_THAT(new_config.calibration_options()
.calibration_parameters()
.max_percentile(),
Eq(0.0f));
}
TEST(ExpandPresetsTest, ExpandUnspecifiedPreset) {
QuantizationConfig config{};
const QuantizationConfig new_config = ExpandPresets(config);
// Test that nothing has been changed.
EXPECT_FALSE(new_config.has_specs());
EXPECT_FALSE(new_config.has_calibration_options());
EXPECT_FALSE(new_config.has_pipeline_config());
}
TEST(ExpandPresetsTest, ExpandStaticRangePtqEnableFullIntquantization) {
QuantizationConfig config{};
RepresentativeDatasetConfig& preset_dataset_config =
*config.mutable_static_range_ptq_preset()->add_representative_datasets();
config.mutable_static_range_ptq_preset()->set_enable_full_int_quantization(
true);
preset_dataset_config.mutable_tf_record()->set_path("/test/path");
const QuantizationConfig new_config = ExpandPresets(config);
ASSERT_THAT(new_config.specs().specs(), SizeIs(2));
const QuantizationSpec& default_spec = new_config.specs().specs(0);
EXPECT_THAT(default_spec.matcher().function_name().regex(), StrEq(".*"));
EXPECT_TRUE(default_spec.method().has_static_range_ptq());
// Test that the expansion for convolution ops is done.
const QuantizationSpec& conv_spec = new_config.specs().specs(1);
EXPECT_THAT(conv_spec.matcher().function_name().regex(),
StrEq("composite_conv.*"));
ASSERT_TRUE(conv_spec.method().has_static_range_ptq());
const StaticRangePtq& srq_spec = conv_spec.method().static_range_ptq();
ASSERT_THAT(srq_spec.input_quantized_types(), SizeIs(1));
ASSERT_TRUE(srq_spec.input_quantized_types().contains(1));
ASSERT_TRUE(srq_spec.input_quantized_types().at(1).has_dimension_specs());
const QuantizedDimension& dimension_specs =
srq_spec.input_quantized_types().at(1).dimension_specs();
ASSERT_TRUE(dimension_specs.has_dimension());
EXPECT_THAT(dimension_specs.dimension(), Eq(3));
// Test that representative dataset config has been transferred to the
// `CalibrationOptions`.
ASSERT_THAT(new_config.calibration_options().representative_datasets(),
SizeIs(1));
EXPECT_THAT(new_config.calibration_options()
.representative_datasets(0)
.tf_record()
.path(),
StrEq("/test/path"));
}
TEST(ExpandPresetsTest, ExpandStaticRangePtqPresetDefault) {
QuantizationConfig config{};
RepresentativeDatasetConfig& preset_dataset_config =
*config.mutable_static_range_ptq_preset()->add_representative_datasets();
preset_dataset_config.mutable_tf_record()->set_path("/test/path");
const QuantizationConfig new_config = ExpandPresets(config);
ASSERT_THAT(new_config.specs().specs(), SizeIs(2));
const QuantizationSpec& spec = new_config.specs().specs(0);
EXPECT_THAT(spec.matcher().function_name().regex(),
StrEq("^.*(dot_general|gather).*"));
EXPECT_TRUE(spec.method().has_static_range_ptq());
}
TEST(ExpandPresetsTest,
ExpandStaticRangePtqPresetWithTopLevelRepresentativeDataset) {
// Test the scenario where both
// `config.calibration_options.representative_datasets` and
// `config.static_range_ptq_preset.representative_datasets` are both
// specified. In this case, the one set to the `calibration_options` takes
// precedence.
QuantizationConfig config{};
RepresentativeDatasetConfig& top_level_dataset_config =
*config.mutable_calibration_options()->add_representative_datasets();
top_level_dataset_config.mutable_tf_record()->set_path("/test/path/1");
RepresentativeDatasetConfig& preset_dataset_config =
*config.mutable_static_range_ptq_preset()->add_representative_datasets();
preset_dataset_config.mutable_tf_record()->set_path("/test/path/2");
const QuantizationConfig new_config = ExpandPresets(config);
// Test that representative dataset config has not been transferred to the
// `CalibrationOptions`. Top-level config takes precedence.
ASSERT_THAT(new_config.calibration_options().representative_datasets(),
SizeIs(1));
EXPECT_THAT(new_config.calibration_options()
.representative_datasets(0)
.tf_record()
.path(),
StrEq("/test/path/1"));
}
TEST(ExpandPresetsTest, ExpandStaticRangePtqPresetThenAppendExplicitSpecs) {
QuantizationConfig config{};
config.mutable_static_range_ptq_preset()->set_enable_full_int_quantization(
true);
QuantizationSpec& user_provided_spec = *config.mutable_specs()->add_specs();
user_provided_spec.mutable_matcher()->mutable_function_name()->set_regex(
"composite_dot_general_fn_1");
user_provided_spec.mutable_method()->mutable_no_quantization();
// Test that the expanded `QuantizationSpec`s are populated first and then
// user-provided specs are appended.
//
// It should look like:
//
// specs {matcher {function_name {regex: ".*"}} method {static_range_ptq {}}}
// specs {
// matcher {function_name {regex: "composite_conv.*"}}
// method {static_range_ptq {...}}}
// }
// specs {
// matcher {function_name {regex: "composite_dot_general_fn_1"}}
// method {no_quantization {}}
// }
const QuantizationConfig new_config = ExpandPresets(config);
ASSERT_THAT(new_config.specs().specs(), SizeIs(3));
const QuantizationSpec& first_spec = new_config.specs().specs(0);
EXPECT_THAT(first_spec.matcher().function_name().regex(), StrEq(".*"));
EXPECT_TRUE(first_spec.method().has_static_range_ptq());
const QuantizationSpec& second_spec = new_config.specs().specs(1);
EXPECT_THAT(second_spec.matcher().function_name().regex(),
StrEq("composite_conv.*"));
EXPECT_TRUE(second_spec.method().has_static_range_ptq());
// This corresponds to `user_provided_spec`.
const QuantizationSpec& third_spec = new_config.specs().specs(2);
EXPECT_THAT(third_spec.matcher().function_name().regex(),
StrEq("composite_dot_general_fn_1"));
EXPECT_TRUE(third_spec.method().has_no_quantization());
}
TEST(ExpandPresetsTest, ExpandWeightOnlyPtqPresetDefault) {
QuantizationConfig config{};
*config.mutable_weight_only_ptq_preset() = WeightOnlyPtqPreset();
const QuantizationConfig new_config = ExpandPresets(config);
ASSERT_THAT(new_config.specs().specs(), SizeIs(1));
const QuantizationSpec& spec = new_config.specs().specs(0);
EXPECT_THAT(spec.matcher().function_name().regex(),
StrEq("^.*(conv|dot_general).*"));
EXPECT_TRUE(spec.method().has_weight_only_ptq());
const WeightOnlyPtq& weight_only_ptq_spec = spec.method().weight_only_ptq();
EXPECT_THAT(weight_only_ptq_spec.input_quantized_types(),
UnorderedElementsAre(Pair(
1, Truly([](const auto& quantized_type) {
return quantized_type.has_dimension_specs() &&
!quantized_type.dimension_specs().has_dimension();
}))));
}
} // namespace
} // namespace stablehlo::quantization
tensorflow/compiler/mlir/quantization/stablehlo/cc/context.h
+36
View File
@@ -0,0 +1,36 @@
/* 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_STABLEHLO_CC_CONTEXT_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_CONTEXT_H_
#include <memory>
#include "mlir/Dialect/Func/Extensions/AllExtensions.h" // from @llvm-project
#include "mlir/IR/DialectRegistry.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
namespace mlir::quant::stablehlo {
// Creates an MLIRContext with the extensions required for quantization are
// registered.
inline std::unique_ptr<MLIRContext> CreateMlirContextForQuantization() {
DialectRegistry registry{};
func::registerAllExtensions(registry);
return std::make_unique<MLIRContext>(registry);
}
} // namespace mlir::quant::stablehlo
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_CONTEXT_H_
tensorflow/compiler/mlir/quantization/stablehlo/cc/debugger.cc
+34
View File
@@ -0,0 +1,34 @@
/* 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/stablehlo/cc/debugger.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/tf_quant_ops.h"
namespace stablehlo::quantization {
void DisableDebugging(mlir::ModuleOp module_op) {
module_op.walk(
[](mlir::TF::DumpTensorOp dump_op) { dump_op.setEnabled(false); });
}
void ChangeToQuantizedFilename(mlir::ModuleOp module_op) {
module_op.walk([](mlir::TF::DumpTensorOp dump_op) {
dump_op.setFileName("quantized_tensor_data.pb");
});
}
} // namespace stablehlo::quantization
tensorflow/compiler/mlir/quantization/stablehlo/cc/debugger.h
+31
View File
@@ -0,0 +1,31 @@
/* 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_STABLEHLO_CC_DEBUGGER_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_DEBUGGER_H_
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
namespace stablehlo::quantization {
// Disables debugging on `DumpTensor` ops.
void DisableDebugging(mlir::ModuleOp module_op);
// Changes the filename from `unquantized_tensor_data.pb` to
// `quantized_tensor_data.pb`.
void ChangeToQuantizedFilename(mlir::ModuleOp module_op);
} // namespace stablehlo::quantization
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_DEBUGGER_H_
tensorflow/compiler/mlir/quantization/stablehlo/cc/graph_def.h
+46
View File
@@ -0,0 +1,46 @@
/* 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_STABLEHLO_CC_GRAPH_DEF_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_GRAPH_DEF_H_
#include <type_traits>
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
namespace stablehlo::quantization {
// Mutates all `NodeDef`s in `graph_def` by applying `func`. It modifies the
// top-level `NodeDef`s as well as all `NodeDef`s in the function library.
// `func` should accept a `NodeDef` reference.
template <typename FuncT, typename = std::enable_if_t<std::is_invocable_r_v<
void, FuncT, tensorflow::NodeDef&>>>
void MutateNodeDefs(tensorflow::GraphDef& graph_def, FuncT&& func) {
for (tensorflow::NodeDef& node_def : *graph_def.mutable_node()) {
func(node_def);
}
for (tensorflow::FunctionDef& function_def :
*graph_def.mutable_library()->mutable_function()) {
for (tensorflow::NodeDef& node_def : *function_def.mutable_node_def()) {
func(node_def);
}
}
}
} // namespace stablehlo::quantization
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_GRAPH_DEF_H_
tensorflow/compiler/mlir/quantization/stablehlo/cc/graph_def_test.cc
+63
View File
@@ -0,0 +1,63 @@
/* 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/stablehlo/cc/graph_def.h"
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tsl/platform/protobuf.h" // IWYU pragma: keep
namespace stablehlo::quantization {
namespace {
using ::tensorflow::GraphDef;
using ::tensorflow::NodeDef;
using ::testing::SizeIs;
using ::testing::StrEq;
using ::tsl::protobuf::TextFormat;
TEST(GraphDefTest, MutateNodeDefsMutatesTopLevelNodeDefs) {
GraphDef graph_def;
ASSERT_TRUE(TextFormat::ParseFromString(R"pb(
node { name: "foo" }
)pb",
&graph_def));
MutateNodeDefs(graph_def,
[](NodeDef& node_def) { node_def.set_name("bar"); });
ASSERT_THAT(graph_def.node(), SizeIs(1));
EXPECT_THAT(graph_def.node()[0].name(), StrEq("bar"));
}
TEST(GraphDefTest, MutateNodeDefsMutatesFunctionNodeDefs) {
GraphDef graph_def;
ASSERT_TRUE(TextFormat::ParseFromString(
R"pb(
library { function { node_def { name: "foo" } } }
)pb",
&graph_def));
MutateNodeDefs(graph_def,
[](NodeDef& node_def) { node_def.set_name("bar"); });
ASSERT_THAT(graph_def.library().function(), SizeIs(1));
ASSERT_THAT(graph_def.library().function()[0].node_def(), SizeIs(1));
EXPECT_THAT(graph_def.library().function()[0].node_def()[0].name(),
StrEq("bar"));
}
} // namespace
} // namespace stablehlo::quantization
tensorflow/compiler/mlir/quantization/stablehlo/cc/io.cc
+87
View File
@@ -0,0 +1,87 @@
/* 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/stablehlo/cc/io.h"
#include <string>
#include <vector>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_format.h"
#include "absl/strings/string_view.h"
#include "xla/tsl/platform/env.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/statusor.h"
namespace stablehlo::quantization::io {
absl::StatusOr<std::string> GetLocalTmpFileName(tsl::Env* const env) {
std::string tmp_fname{};
if (!env->LocalTempFilename(&tmp_fname)) {
return absl::InternalError("Failed to create tmp file name.");
}
return tmp_fname;
}
absl::StatusOr<std::string> GetLocalTmpFileName() {
return GetLocalTmpFileName(tsl::Env::Default());
}
absl::StatusOr<std::string> CreateTmpDir(tsl::Env* const env) {
TF_ASSIGN_OR_RETURN(std::string tmp_dir, GetLocalTmpFileName(env));
if (!env->RecursivelyCreateDir(tmp_dir).ok()) {
return absl::InternalError(
absl::StrFormat("Failed to create tmp dir: '%s'", tmp_dir));
}
return tmp_dir;
}
absl::StatusOr<std::string> CreateTmpDir() {
// The overloaded function uses the default env.
return CreateTmpDir(tsl::Env::Default());
}
absl::Status WriteStringToFile(const absl::string_view file_path,
const absl::string_view data) {
auto* env = tsl::Env::Default();
return WriteStringToFile(env, std::string(file_path), data);
}
absl::StatusOr<std::string> ReadFileToString(
const absl::string_view file_path) {
auto* env = tsl::Env::Default();
std::string data{};
absl::Status read_status =
ReadFileToString(env, std::string(file_path), &data);
if (read_status.ok()) {
return data;
} else {
return read_status;
}
}
absl::StatusOr<std::vector<std::string>> ListDirectory(
absl::string_view directory) {
std::vector<std::string> children;
TF_RETURN_IF_ERROR(
tsl::Env::Default()->GetChildren(std::string(directory), &children));
return children;
}
} // namespace stablehlo::quantization::io
tensorflow/compiler/mlir/quantization/stablehlo/cc/io.h
+73
View File
@@ -0,0 +1,73 @@
/* 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_STABLEHLO_CC_IO_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_IO_H_
#include <string>
#include <vector>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "xla/tsl/platform/env.h"
#include "xla/tsl/platform/errors.h"
namespace stablehlo::quantization::io {
// Generates a unique local tmp file name. This function only generates the name
// (path) and doesn't actually creates the file.
absl::StatusOr<std::string> GetLocalTmpFileName(tsl::Env* env);
// Generates a unique local tmp file name. This function only generates the name
// (path) and doesn't actually creates the file. The default environment
// `tsl::Env::Default` is used to generate the name.
absl::StatusOr<std::string> GetLocalTmpFileName();
// Creates a temporary directory on an environment defined by the implementation
// of `tsl::Env` and returns its path. Returns an InternalError status if
// failed.
absl::StatusOr<std::string> CreateTmpDir(tsl::Env* env);
// Creates a temporary directory and returns its path. Returns an InternalError
// status if failed. The file system used will be the default environment
// returned by `tsl::Env::Default`.
absl::StatusOr<std::string> CreateTmpDir();
// Convenience function for writing string `data` to file without the need to
// pass `tsl::Env` instance. Internally it uses the default `tsl::Env::Default`.
absl::Status WriteStringToFile(absl::string_view file_path,
absl::string_view data);
// Convenience function for reading string data from file at `file_path` without
// the need to pass `tsl::Env` instance. Internally it uses the default
// `tsl::Env::Default`. Returns an OK status with string data containing file
// contents. Returns non-ok status upon error, e.g. file doesn't exist.
absl::StatusOr<std::string> ReadFileToString(absl::string_view file_path);
// Lists all files and directories under the given directory.
absl::StatusOr<std::vector<std::string>> ListDirectory(
absl::string_view directory);
template <class MessageT>
absl::StatusOr<MessageT> ReadBinaryProto(const std::string& binary_file_path) {
MessageT message;
TF_RETURN_IF_ERROR(
tsl::ReadBinaryProto(tsl::Env::Default(), binary_file_path, &message));
return message;
}
} // namespace stablehlo::quantization::io
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_IO_H_
tensorflow/compiler/mlir/quantization/stablehlo/cc/io_test.cc
+206
View File
@@ -0,0 +1,206 @@
/* 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/stablehlo/cc/io.h"
#include <cstdint>
#include <fstream>
#include <string>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/functional/any_invocable.h"
#include "absl/status/status.h"
#include "absl/status/status_matchers.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "xla/tsl/platform/env.h"
#include "xla/tsl/platform/file_system.h"
namespace stablehlo::quantization::io {
namespace {
using ::testing::Eq;
using ::testing::HasSubstr;
using ::testing::IsEmpty;
using ::testing::Not;
using ::testing::SizeIs;
using ::testing::UnorderedElementsAre;
// A test-only derived class of `tsl::Env` which is broken. Used to cause
// failure for the `CreateTmpDir` function. Each of the overridden member
// functions implements a dummy functionality just to be able to create an
// instance of this class.
class TestEnvBrokenFileSystem : public tsl::Env {
public:
TestEnvBrokenFileSystem() = default;
bool MatchPath(absl::string_view path, absl::string_view pattern) override {
return false;
}
void SleepForMicroseconds(int64_t micros) override {}
std::string GetRunfilesDir() override { return std::string("dummy_path"); }
int64_t GetCurrentThreadId() override { return 0; }
tsl::Thread* StartThread(const tsl::ThreadOptions& thread_options,
const std::string& name,
absl::AnyInvocable<void()> fn) override {
return nullptr;
}
bool GetCurrentThreadName(std::string* name) override { return false; }
void SchedClosure(absl::AnyInvocable<void()> closure) override {}
void SchedClosureAfter(int64_t micros,
absl::AnyInvocable<void()> closure) override {}
absl::Status LoadDynamicLibrary(const char* library_filename,
void** handle) override {
return absl::OkStatus();
}
absl::Status GetSymbolFromLibrary(void* handle, const char* symbol_name,
void** symbol) override {
return absl::OkStatus();
}
std::string FormatLibraryFileName(const std::string& name,
const std::string& version) override {
return std::string("dummy_path");
}
// This is the part that would break the `CreateTmpDir` function because it
// fails to provide a valid file system.
absl::Status GetFileSystemForFile(absl::string_view fname,
tsl::FileSystem** result) override {
return absl::InternalError("Broken file system");
}
private:
void GetLocalTempDirectories(std::vector<std::string>* list) override {
list->push_back("/tmp");
}
};
// Represents an environment with broken file system and no available local tmp
// directories.
class TestEnvBrokenFileSystemAndNoLocalTempDirs
: public TestEnvBrokenFileSystem {
private:
// This is the part that essentially breaks the `GetLocalTmpFileName` function
// because it doesn't provide any available temp dirs.
void GetLocalTempDirectories(std::vector<std::string>* list) override {}
};
TEST(IoTest, GetLocalTmpFileNameGivesValidFileName) {
absl::StatusOr<std::string> tmp_file_name = GetLocalTmpFileName();
ASSERT_THAT(tmp_file_name, absl_testing::IsOk());
EXPECT_THAT(*tmp_file_name, Not(IsEmpty()));
}
TEST(IoTest, GetLocalTmpFileNameWhenNoTempDirsReturnsInternalError) {
TestEnvBrokenFileSystemAndNoLocalTempDirs broken_env;
absl::StatusOr<std::string> tmp_file_name = GetLocalTmpFileName(&broken_env);
EXPECT_THAT(tmp_file_name, absl_testing::StatusIs(
absl::StatusCode::kInternal,
HasSubstr("Failed to create tmp file name")));
}
TEST(IoTest, CreateTmpDirReturnsValidTmpPath) {
absl::StatusOr<std::string> tmp_dir = CreateTmpDir();
ASSERT_THAT(tmp_dir, absl_testing::IsOk());
auto* const env = tsl::Env::Default();
EXPECT_THAT(env->FileExists(*tmp_dir), absl_testing::IsOk());
}
TEST(IoTest, CreateTmpDirWhenInvalidPathReturnsInternalError) {
TestEnvBrokenFileSystem test_env{};
absl::StatusOr<std::string> tmp_dir = CreateTmpDir(&test_env);
EXPECT_THAT(tmp_dir,
absl_testing::StatusIs(absl::StatusCode::kInternal,
HasSubstr("Failed to create tmp dir")));
}
TEST(IoTest, WriteStringToFile) {
const std::string dst_file_path =
absl::StrCat(testing::TempDir(), "/tmp_file");
const absl::Status write_status =
WriteStringToFile(dst_file_path, "test_string");
ASSERT_THAT(write_status, absl_testing::IsOk());
auto* const env = tsl::Env::Default();
ASSERT_THAT(env->FileExists(dst_file_path), absl_testing::IsOk());
std::string data{};
ASSERT_THAT(tsl::ReadFileToString(env, dst_file_path, &data),
absl_testing::IsOk());
EXPECT_THAT(data, Eq("test_string"));
}
TEST(IoTest, ReadFileToString) {
// Prepare a temp file and write some string to it.
const std::string src_file_path =
absl::StrCat(testing::TempDir(), "/tmp_file");
{
std::ofstream ofs(src_file_path);
ofs << "test_string";
}
// Test that the contents match.
const absl::StatusOr<std::string> read_status =
ReadFileToString(src_file_path);
ASSERT_THAT(read_status, absl_testing::IsOk());
EXPECT_THAT(*read_status, Eq("test_string"));
}
TEST(IoTest, ListChildrenInDirectory) {
absl::StatusOr<std::string> tmp_dir = CreateTmpDir();
ASSERT_THAT(tmp_dir, absl_testing::IsOk());
auto* const env = tsl::Env::Default();
EXPECT_THAT(env->FileExists(*tmp_dir), absl_testing::IsOk());
ASSERT_THAT(
WriteStringToFile(absl::StrCat(*tmp_dir, "/tmp_file1"), "test_string"),
absl_testing::IsOk());
ASSERT_THAT(
WriteStringToFile(absl::StrCat(*tmp_dir, "/tmp_file2"), "test_string"),
absl_testing::IsOk());
ASSERT_THAT(env->RecursivelyCreateDir(absl::StrCat(*tmp_dir, "/subdir")),
absl_testing::IsOk());
absl::StatusOr<std::vector<std::string>> children = ListDirectory(*tmp_dir);
EXPECT_THAT(children, absl_testing::IsOk());
EXPECT_THAT(children.value(), SizeIs(3));
EXPECT_THAT(children.value(),
UnorderedElementsAre("subdir", "tmp_file1", "tmp_file2"));
}
} // namespace
} // namespace stablehlo::quantization::io
tensorflow/compiler/mlir/quantization/stablehlo/cc/pass_pipeline.cc
+176
View File
@@ -0,0 +1,176 @@
/* 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/stablehlo/cc/pass_pipeline.h"
#include "mlir/Conversion/ReconcileUnrealizedCasts/ReconcileUnrealizedCasts.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "mlir/Pass/PassRegistry.h" // from @llvm-project
#include "mlir/Transforms/Passes.h" // from @llvm-project
#include "stablehlo/transforms/Passes.h" // from @stablehlo
#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/passes.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "xla/mlir_hlo/mhlo/transforms/passes.h"
namespace mlir::quant::stablehlo {
using ::stablehlo::quantization::CalibrationOptions;
using ::stablehlo::quantization::DebuggerConfig;
using ::stablehlo::quantization::PipelineConfig;
using ::stablehlo::quantization::QuantizationSpecs;
void AddPreCalibrationPasses(OpPassManager& pm,
const CalibrationOptions& calibration_options,
const QuantizationSpecs& quantization_specs,
const DebuggerConfig& debugger_config) {
// Convert NCHW tensors to NHWC at along with extra optimizations as
// downstream passes perform better optimizations when dealing with NHWC
// formatted tensors.
AddProcessNchwTensorPasses(pm);
pm.addPass(CreateLiftQuantizableSpotsAsFunctionsPass(quantization_specs));
if (debugger_config.debugger_type() !=
DebuggerConfig::DEBUGGER_TYPE_UNSPECIFIED) {
pm.addPass(CreateAddDumpTensorOpPass(debugger_config.debugger_type(),
debugger_config.log_dir_path()));
}
pm.addNestedPass<func::FuncOp>(
CreateInsertCustomAggregationOpsPass(calibration_options));
}
void AddPostCalibrationPasses(OpPassManager& pm,
const PipelineConfig& pipeline_config,
const QuantizationSpecs& specs) {
QuantizeCompositeFunctionsPassOptions options;
// TODO: b/331120943 - Temporarily set below to true, signaling per-channel
// quantization will be applied for all where applicable. This will be
// replaced by individual `Method` in `QuantizationSpecs`.
options.enable_per_channel_quantized_weight_ = true;
// For debugging purposes.
options.mlir_dump_file_name_ = "quantize_composite_functions";
options.merge_fusion_with_dequantize_ =
pipeline_config.merge_fusion_with_dequantize();
AddShapeLegalizationPasses(pm);
pm.addNestedPass<func::FuncOp>(
CreateConvertCustomAggregationOpToQuantStatsPass());
pm.addPass(createQuantizeCompositeFunctionsPass(options));
// Add an inliner pass to inline quantized StableHLO functions.
pm.addPass(createInlinerPass());
if (pipeline_config.unpack_quantized_types()) {
AddStablehloQuantToIntPasses(pm);
}
}
void AddWeightOnlyQuantizationPasses(
OpPassManager& pm, const QuantizationSpecs& quantization_specs,
const PipelineConfig& pipeline_config,
const DebuggerConfig& debugger_config) {
// For models with NCHW convolution format. This pass is required because
// downstream pipeline handles NHWC convolution better for most cases.
pm.addNestedPass<func::FuncOp>(createNchwConvolutionToNhwcPass());
// Folds `stablehlo.constant`->`stablehlo.transpose` patterns, which is often
// generated as by-products after optimizing dimension numbers (e.g.
// NCHW->NHWC convolution conversion).
pm.addNestedPass<func::FuncOp>(createFoldConstantTransposePass());
pm.addPass(CreateLiftQuantizableSpotsAsFunctionsPass(quantization_specs));
if (debugger_config.debugger_type() !=
DebuggerConfig::DEBUGGER_TYPE_UNSPECIFIED) {
pm.addPass(CreateAddDumpTensorOpPass(debugger_config.debugger_type(),
debugger_config.log_dir_path()));
}
AddShapeLegalizationPasses(pm);
QuantizeCompositeFunctionsPassOptions options;
// For debugging purposes.
options.mlir_dump_file_name_ = "quantize_composite_functions";
pm.addPass(createQuantizeCompositeFunctionsPass(options));
// Add an inliner pass to inline quantized StableHLO functions.
pm.addPass(createInlinerPass());
if (pipeline_config.unpack_quantized_types()) {
AddStablehloQuantToIntPasses(pm);
}
}
void AddXlaCallModuleOpDeserializationPasses(OpPassManager& pm) {
pm.addPass(TF::CreateXlaCallModuleDeserializationPass());
pm.addPass(createRestoreFunctionNamePass());
pm.addPass(createUnwrapXlaCallModuleOpPass());
pm.addPass(createSymbolDCEPass());
}
void AddShapeLegalizationPasses(OpPassManager& pm) {
// TODO: We may need to make a parent pass here that does
// shape->StableHLO+cstr because the stablehlo pass requires that the ops made
// by cstr are legal.
pm.addNestedPass<func::FuncOp>(
createConvertShapeToStablehloWithConstraintsPass());
pm.addPass(createReconcileUnrealizedCastsPass());
pm.addNestedPass<func::FuncOp>(mlir::createCanonicalizerPass());
}
void AddStablehloQuantToIntPasses(OpPassManager& pm) {
pm.addNestedPass<func::FuncOp>(
mlir::stablehlo::createStablehloLegalizeQuantToMathPass());
// StableHLO -> MHLO legalization.
pm.addPass(mhlo::createStablehloLegalizeToHloPass());
pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
// Integer graph optimization relies on chlo broadcast ops for easier handling
// of dynamic shapes. Therefore we lower chlo ops after optimization.
pm.addNestedPass<func::FuncOp>(CreateOptimizeIntGraphPass());
pm.addNestedPass<func::FuncOp>(mhlo::createChloLegalizeToHloPass());
pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
pm.addPass(createSymbolDCEPass());
// MHLO -> StableHLO legalization.
pm.addPass(mhlo::createHloLegalizeToStablehloPass());
}
// NOMUTANTS -- Add tests for individual passes with migration below.
void AddCallModuleSerializationPasses(OpPassManager& pm) {
AddShapeLegalizationPasses(pm);
pm.addPass(createReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPass());
// ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPass may create
// duplicate constants. Add canonicalizer to deduplicate.
pm.addNestedPass<func::FuncOp>(mlir::createCanonicalizerPass());
pm.addPass(TF::CreateXlaCallModuleSerializationPass());
}
void AddProcessNchwTensorPasses(OpPassManager& pm) {
// For models with NCHW convolution format. This pass is required because
// downstream pipeline handles NHWC convolution better for most cases.
pm.addNestedPass<func::FuncOp>(createNchwConvolutionToNhwcPass());
// Recursively push down the `stablehlo.transpose` ops for activations
// generated by the `NchwConvolutionToNhwc` pass.
pm.addNestedPass<func::FuncOp>(createDeferActivationTransposePass());
// Folds `stablehlo.constant`->`stablehlo.transpose` patterns, which is often
// generated as by-products after optimizing dimension numbers (e.g.
// NCHW->NHWC convolution conversion).
pm.addNestedPass<func::FuncOp>(createFoldConstantTransposePass());
}
void RegisterPassPipelines() {
static PassPipelineRegistration<> nchw_tensor_format_processing_pipeline(
/*arg=*/"stablehlo-process-nchw-tensor",
/*description=*/"Optimizes tensors with NCHW format.",
AddProcessNchwTensorPasses);
}
} // namespace mlir::quant::stablehlo
tensorflow/compiler/mlir/quantization/stablehlo/cc/pass_pipeline.h
+75
View File
@@ -0,0 +1,75 @@
/* 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_STABLEHLO_CC_PASS_PIPELINE_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_PASS_PIPELINE_H_
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
namespace mlir::quant::stablehlo {
// Adds passes for static-range quantization pre-calibration. Inserts ops
// required to collect tensor statistics.
void AddPreCalibrationPasses(
OpPassManager& pm,
const ::stablehlo::quantization::CalibrationOptions& calibration_options,
const ::stablehlo::quantization::QuantizationSpecs& specs,
const ::stablehlo::quantization::DebuggerConfig& debugger_config);
// Adds passes for static-range quantization post-calibration. Utilizes tensor
// statistics collected from the calibration step and performs quantization.
void AddPostCalibrationPasses(
OpPassManager& pm,
const ::stablehlo::quantization::PipelineConfig& pipeline_config,
const ::stablehlo::quantization::QuantizationSpecs& specs);
// Adds passes for weight-only quantization.
void AddWeightOnlyQuantizationPasses(
OpPassManager& pm,
const ::stablehlo::quantization::QuantizationSpecs& quantization_specs,
const ::stablehlo::quantization::PipelineConfig& pipeline_config,
const ::stablehlo::quantization::DebuggerConfig& debugger_config);
// Deserializes StableHLO functions serialized and embedded in XlaCallModuleOps.
void AddXlaCallModuleOpDeserializationPasses(OpPassManager& pm);
// Legalizes shape/tensor/arith dialect ops to StableHLO for handling dynamic
// shapes, by going through a round-trip to MHLO.
void AddShapeLegalizationPasses(OpPassManager& pm);
// Serializes the StableHLO module into a tf.XlaCallModuleOp for compatibility
// with passes that expect TF format. This also allows the StableHLO ops to be
// exported as a TF SavedModel.
void AddCallModuleSerializationPasses(OpPassManager& pm);
// Passes for unpacking quantized ops to int valued StableHLO ops. This is
// useful when uniform quantized types are suboptimal for the hardware. It goes
// through a StableHLO <-> MHLO roundtrip to utilize the MHLOQuantToInt pass.
void AddStablehloQuantToIntPasses(OpPassManager& pm);
// Processes tensors with NCHW format (== (batch, channel, height, weight)) by
// converting them to NHWC formats along with extra optimizations such as
// constant folding the transpose->convolution pattern. This is useful when
// downstream pipeline (e.g. XLA) is more optimized when accepting NHWC formats.
void AddProcessNchwTensorPasses(OpPassManager& pm);
// Registers quantization pass pipelines. This is only required when running
// MLIR opt binaries and not required when adding passes programmatically.
void RegisterPassPipelines();
} // namespace mlir::quant::stablehlo
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_PASS_PIPELINE_H_
tensorflow/compiler/mlir/quantization/stablehlo/cc/permutation.h
+44
View File
@@ -0,0 +1,44 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_PERMUTATION_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_PERMUTATION_H_
#include <cstdint>
#include <type_traits>
#include "llvm/ADT/ArrayRef.h" // IWYU pragma: keep; required to include the definition of ArrayRef
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h" // IWYU pragma: keep; required to include the definition of SmallVector
#include "mlir/Support/LLVM.h" // from @llvm-project
namespace mlir::quant {
// Permutes `values` with `permutation`. Returns the permuted values. Sizes of
// `values` and `permutation` must be equal, and the elements of `permutation`
// should be less than `values.size()`.
template <typename T,
typename = std::enable_if_t<std::is_default_constructible_v<T>, void>>
SmallVector<T> Permute(const ArrayRef<T> values,
const ArrayRef<int64_t> permutation) {
SmallVector<T> permuted_values(/*Size=*/values.size(), /*Value=*/T{});
for (auto [i, permutation_idx] : llvm::enumerate(permutation)) {
permuted_values[i] = std::move(values[permutation_idx]);
}
return permuted_values;
}
} // namespace mlir::quant
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_PERMUTATION_H_
tensorflow/compiler/mlir/quantization/stablehlo/cc/permutation_test.cc
+64
View File
@@ -0,0 +1,64 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/permutation.h"
#include <cstdint>
#include <string>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "mlir/Support/LLVM.h" // from @llvm-project
namespace mlir::quant {
namespace {
using testing::ElementsAre;
using testing::IsEmpty;
TEST(PermutationTest, PermuteEmptyArray) {
const SmallVector<int> permutation_result =
Permute<int>(SmallVector<int>{}, SmallVector<int64_t>{});
EXPECT_THAT(permutation_result, IsEmpty());
}
TEST(PermutationTest, PermuteOneElement) {
const SmallVector<int> single_element_array = {8};
const SmallVector<int64_t> permutation = {0};
const SmallVector<int> permutation_result =
Permute<int>(single_element_array, permutation);
EXPECT_THAT(permutation_result, ElementsAre(8));
}
TEST(PermutationTest, PermuteFourElements) {
const SmallVector<int> arr = {0, 3, 1, 2};
// Permutation inverse of {0, 3, 1, 2}.
const SmallVector<int64_t> permutation = {0, 2, 3, 1};
const SmallVector<int> permutation_result = Permute<int>(arr, permutation);
EXPECT_THAT(permutation_result, ElementsAre(0, 1, 2, 3));
}
TEST(PermutationTest, PermuteFourStringElements) {
const SmallVector<std::string> arr = {"a", "b", "c", "d"};
const SmallVector<int64_t> permutation = {0, 2, 3, 1};
const SmallVector<std::string> permutation_result =
Permute<std::string>(arr, permutation);
EXPECT_THAT(permutation_result, ElementsAre("a", "c", "d", "b"));
}
} // namespace
} // namespace mlir::quant
tensorflow/compiler/mlir/quantization/stablehlo/cc/post_calibration.cc
+67
View File
@@ -0,0 +1,67 @@
/* 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/stablehlo/cc/post_calibration.h"
#include <memory>
#include "absl/base/nullability.h"
#include "absl/log/die_if_null.h"
#include "absl/status/statusor.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/config.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/pass_pipeline.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/instrumentations/save_report.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/cc/run_passes.h"
#include "xla/mlir_hlo/mhlo/transforms/passes.h"
#include "xla/tsl/platform/errors.h"
namespace mlir::quant::stablehlo {
using ::stablehlo::quantization::GetReportFilePath;
using ::stablehlo::quantization::PipelineConfig;
using ::stablehlo::quantization::QuantizationConfig;
using ::stablehlo::quantization::QuantizationSpecs;
using ::tensorflow::quantization::RunPasses;
PostCalibrationComponent::PostCalibrationComponent(
MLIRContext* absl_nonnull ctx)
: ctx_(ABSL_DIE_IF_NULL(ctx)) {} // Crash OK
absl::StatusOr<ModuleOp> PostCalibrationComponent::Run(
ModuleOp module_op, const QuantizationConfig& config) {
TF_RETURN_IF_ERROR(RunPasses(
kName, /*add_passes_func=*/
[&config](PassManager& pm) {
// Add instrumentation to save quantization report after quantization.
pm.addInstrumentation(
std::make_unique<SaveQuantizationReportInstrumentation>(
GetReportFilePath(config)));
quant::stablehlo::AddPostCalibrationPasses(pm, config.pipeline_config(),
config.specs());
},
*ctx_, module_op));
return module_op;
}
void PostCalibrationComponent::AddPasses(
OpPassManager& pm, const QuantizationSpecs& specs,
const PipelineConfig& pipeline_config) const {
quant::stablehlo::AddPostCalibrationPasses(pm, pipeline_config, specs);
}
} // namespace mlir::quant::stablehlo
tensorflow/compiler/mlir/quantization/stablehlo/cc/post_calibration.h
+59
View File
@@ -0,0 +1,59 @@
/* 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_STABLEHLO_CC_POST_CALIBRATION_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_POST_CALIBRATION_H_
#include "absl/base/nullability.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/component.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
namespace mlir::quant::stablehlo {
// Performs post-calibration graph transformation as part of post-training
// static-range quantization.
//
// The resulting `ModuleOp` contains quantized StableHLO ops serialized in
// `TF::XlaCallModuleOp`s. They are quantized using the statistics collected
// after the calibration step, corresponding to each `TF::CustomAggregatorOp`s
// in the input module op.
class PostCalibrationComponent : public Component {
public:
// Name of the post-training quantization post-calibration step. Used for
// debugging purposes.
static constexpr absl::string_view kName = "quant_ptq_post_calibration";
explicit PostCalibrationComponent(MLIRContext* absl_nonnull ctx);
absl::StatusOr<ModuleOp> Run(
ModuleOp module_op,
const ::stablehlo::quantization::QuantizationConfig& config) override;
void AddPasses(
OpPassManager& pm,
const ::stablehlo::quantization::QuantizationSpecs& specs,
const ::stablehlo::quantization::PipelineConfig& pipeline_config) const;
private:
MLIRContext* absl_nonnull ctx_;
};
} // namespace mlir::quant::stablehlo
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_POST_CALIBRATION_H_
tensorflow/compiler/mlir/quantization/stablehlo/cc/pre_calibration.cc
+49
View File
@@ -0,0 +1,49 @@
/* 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/stablehlo/cc/pre_calibration.h"
#include "absl/base/nullability.h"
#include "absl/log/die_if_null.h"
#include "absl/status/statusor.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/Pass/PassManager.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/pass_pipeline.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/cc/run_passes.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
#include "xla/tsl/platform/errors.h"
namespace mlir::quant::stablehlo {
using ::stablehlo::quantization::QuantizationConfig;
using ::tensorflow::quantization::RunPasses;
PreCalibrationComponent::PreCalibrationComponent(MLIRContext* absl_nonnull ctx)
: ctx_(ABSL_DIE_IF_NULL(ctx)) {} // Crash OK
absl::StatusOr<ModuleOp> PreCalibrationComponent::Run(
ModuleOp module_op, const QuantizationConfig& config) {
TF_RETURN_IF_ERROR(RunPasses(
kName, /*add_passes_func=*/
[&config](PassManager& pm) {
quant::stablehlo::AddPreCalibrationPasses(
pm, config.calibration_options(), config.specs(),
config.debugger_config());
},
*ctx_, module_op));
return module_op;
}
} // namespace mlir::quant::stablehlo
tensorflow/compiler/mlir/quantization/stablehlo/cc/pre_calibration.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_STABLEHLO_CC_PRE_CALIBRATION_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_PRE_CALIBRATION_H_
#include "absl/base/nullability.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/component.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
namespace mlir::quant::stablehlo {
// Performs pre-calibration graph transformation as part of post-training
// static-range quantization.
// The resulting `ModuleOp` contains `TF::CustomAggregatorOp`s for collecting
// quantization statistics, along with `TF::XlaCallModuleOp`s that correspond to
// lifted quantizable functions.
class PreCalibrationComponent : public Component {
public:
// Name of the post-training quantization pre-calibration step. Used for
// debugging purposes.
static constexpr absl::string_view kName = "quant_ptq_pre_calibration";
explicit PreCalibrationComponent(MLIRContext* absl_nonnull ctx);
absl::StatusOr<ModuleOp> Run(
ModuleOp,
const ::stablehlo::quantization::QuantizationConfig& config) override;
private:
MLIRContext* absl_nonnull ctx_;
};
} // namespace mlir::quant::stablehlo
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_PRE_CALIBRATION_H_
tensorflow/compiler/mlir/quantization/stablehlo/cc/pre_calibration_test.cc
+140
View File
@@ -0,0 +1,140 @@
/* 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/stablehlo/cc/pre_calibration.h"
#include <type_traits>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/status_matchers.h"
#include "absl/status/statusor.h"
#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/OwningOpRef.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/common/test_base.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/config.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.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_ops.h"
namespace mlir::quant::stablehlo {
namespace {
using ::stablehlo::quantization::ExpandPresets;
using ::stablehlo::quantization::PopulateDefaults;
using ::stablehlo::quantization::QuantizationConfig;
using ::testing::Contains;
using ::testing::SizeIs;
using ::testing::StartsWith;
using ::testing::StrEq;
// Matches an operation whose `getSymName` equals `name`.
MATCHER_P(HasSymName, name, "") {
auto non_const_arg = const_cast<std::remove_const_t<decltype(arg)>>(arg);
*result_listener << "where the name is " << non_const_arg.getSymName().str();
return non_const_arg.getSymName() == name;
}
// Matches an operation that has a StringAttr whose name is `name` and value
// matches `value_matcher`.
MATCHER_P2(HasStringAttr, name, value_matcher,
absl::StrCat(negation ? "doesn't have" : "has",
"string attribute: ", name, ", with desirable value")) {
auto non_const_arg = const_cast<std::remove_const_t<decltype(arg)>>(arg);
return non_const_arg->template hasAttrOfType<StringAttr>(name) &&
ExplainMatchResult(
value_matcher,
non_const_arg->template getAttrOfType<StringAttr>(name).str(),
result_listener);
}
// Matches an operation that has a FlatSymbolRefAttr whose name is `name` and
// value matches `value_matcher`.
MATCHER_P2(HasSymNameAttr, name, value_matcher,
absl::StrCat(negation ? "doesn't have" : "has",
"string attribute: ", name, ", with desirable value")) {
auto non_const_arg = const_cast<std::remove_const_t<decltype(arg)>>(arg);
return non_const_arg->template hasAttrOfType<FlatSymbolRefAttr>(name) &&
ExplainMatchResult(
value_matcher,
non_const_arg->template getAttrOfType<FlatSymbolRefAttr>(name)
.getValue()
.str(),
result_listener);
}
using PreCalibrationComponentTest = ::mlir::quant::QuantizationTestBase;
TEST_F(PreCalibrationComponentTest,
HasCustomAggregatorOpAndQuantizableFuncForSimpleDotGeneral) {
PreCalibrationComponent component(ctx_.get());
OwningOpRef<ModuleOp> module_op = ParseModuleOpString(R"mlir(
module attributes {} {
func.func @main(%arg0: tensor<1x4xf32>) -> tensor<1x3xf32> attributes {} {
%0 = stablehlo.constant dense<1.0> : tensor<4x3xf32>
%1 = stablehlo.dot_general %arg0, %0, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x4xf32>, tensor<4x3xf32>) -> tensor<1x3xf32>
return %1 : tensor<1x3xf32>
}
}
)mlir");
ASSERT_TRUE(module_op);
QuantizationConfig quantization_config{};
quantization_config.mutable_static_range_ptq_preset();
quantization_config = ExpandPresets(PopulateDefaults(quantization_config));
absl::StatusOr<ModuleOp> pre_calibration_result =
component.Run(*module_op, quantization_config);
EXPECT_THAT(pre_calibration_result, absl_testing::IsOk());
SmallVector<func::FuncOp> func_ops;
for (auto func_op : pre_calibration_result->getOps<func::FuncOp>()) {
func_ops.push_back(func_op);
}
ASSERT_THAT(func_ops, SizeIs(2));
EXPECT_THAT(func_ops, Contains(HasSymName("main")));
EXPECT_THAT(func_ops, Contains(HasSymName("composite_dot_general_fn_1")));
// Tests that there is a XlaCallModuleOp that calls the composite quantizable
// function.
SmallVector<TF::XlaCallModuleOp> xla_call_module_ops;
for (auto xla_call_module_op : func_ops[0].getOps<TF::XlaCallModuleOp>()) {
xla_call_module_ops.push_back(xla_call_module_op);
}
ASSERT_THAT(xla_call_module_ops, SizeIs(1));
auto xla_call_module_op = xla_call_module_ops[0];
EXPECT_THAT(xla_call_module_op,
HasStringAttr("_tfl_quant_trait", StrEq("fully_quantizable")));
EXPECT_THAT(xla_call_module_op,
HasSymNameAttr("_entry_function",
StartsWith("composite_dot_general_fn")));
EXPECT_THAT(xla_call_module_op,
HasStringAttr("_original_entry_function",
StartsWith("composite_dot_general_fn")));
// Tests that there are CustomAggregatorOps inserted.
SmallVector<TF::CustomAggregatorOp> custom_aggregator_ops;
for (auto custom_aggregator_op :
func_ops[0].getOps<TF::CustomAggregatorOp>()) {
custom_aggregator_ops.push_back(custom_aggregator_op);
}
EXPECT_THAT(custom_aggregator_ops, SizeIs(2));
}
} // namespace
} // namespace mlir::quant::stablehlo
tensorflow/compiler/mlir/quantization/stablehlo/cc/report.cc
+174
View File
@@ -0,0 +1,174 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/report.h"
#include <optional>
#include <string>
#include <utility>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "llvm/Support/raw_ostream.h"
#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/Visitors.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/common/lift_as_function_call.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/io.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tsl/platform/protobuf.h" // IWYU pragma: keep
namespace mlir::quant::stablehlo {
namespace {
using ::stablehlo::quantization::Method;
using ::stablehlo::quantization::QuantizationResult;
using ::stablehlo::quantization::QuantizationResults;
using ::stablehlo::quantization::io::WriteStringToFile;
using ::tsl::protobuf::TextFormat;
// Given a `quantized_func_name` that starts with `kQuantizedFuncPrefix`,
// converts `kQuantizedFuncPrefix` to `kCompositeFuncPrefix`.
std::string GetCompositeFunctionName(const StringRef quantized_func_name) {
return Twine(kCompositeFuncPrefix)
.concat(quantized_func_name.rsplit(kQuantizedFuncPrefix).second)
.str();
}
// Retrieves `QuantizationResult` from `call_op`. If the callee's name starts
// with `kQuantizedFuncPrefix` then a `QuantizationResult` will be returned with
// its `name` field set to the callee's name reverted back to the lifted
// function's name. Also, `call_op` must have the `kQuantizationMethodAttr`
// attribute, which is deserialized as `Method` and set in the returned
// `QuantizationResult`. Otherwise, it returns `std::nullopt`.
std::optional<QuantizationResult> GetQuantizationResult(func::CallOp call_op) {
const StringRef callee_name = call_op.getCalleeAttr().getValue();
if (!callee_name.starts_with(kQuantizedFuncPrefix)) {
return std::nullopt; // `call_op` is not a quantized function call.
}
absl::StatusOr<Method> method = GetQuantizationMethod(call_op);
if (!method.ok()) {
call_op->emitError() << "Failed to get quantization method: "
<< method.status().ToString();
return std::nullopt;
}
QuantizationResult result{};
result.mutable_quantizable_unit()->set_name(
GetCompositeFunctionName(callee_name));
*result.mutable_method() = std::move(*method);
return result;
}
// Retrieves `QuantizationResult` from `xla_call_module_op`. If
// `xla_call_module_op` is a quantizable unit, then a `QuantizationResult` will
// be returned with its `name` field set to the callee's name. The `method`
// field will be set to `NoQuantization` because remaining `xla_call_module_op`s
// means they are not quantized. Returns `std::nullopt` if `xla_call_module_op`
// is not a quantizable unit.
std::optional<QuantizationResult> GetQuantizationResult(
TF::XlaCallModuleOp xla_call_module_op) {
const StringAttr callee_name_attr =
mlir::dyn_cast_or_null<StringAttr>(xla_call_module_op->getDiscardableAttr(
kOriginalStablehloEntryFunctionAttrName));
// `TF::XlaCallModuleOp` without the `_original_entry_function` means it is
// not a quantizable unit.
if (callee_name_attr == nullptr) return std::nullopt;
if (callee_name_attr.getValue().starts_with(kCompositeFuncPrefix)) {
QuantizationResult result{};
result.mutable_quantizable_unit()->set_name(
callee_name_attr.getValue().str());
result.mutable_method()->mutable_no_quantization();
return result;
} else {
return std::nullopt;
}
}
// Populates quantized ops from `module_op` to `results`. After going through
// the quantization passes, quantized ops are represented as `func::CallOp` with
// a callee's prefix of `quantized_`.
void PopulateQuantizedResults(ModuleOp module_op,
QuantizationResults& results) {
module_op.walk([&results](func::CallOp call_op) {
std::optional<QuantizationResult> result = GetQuantizationResult(call_op);
if (result == std::nullopt) return WalkResult::skip();
*results.add_results() = std::move(*result);
return WalkResult::advance();
});
}
// Populates non-quantized ops from `module_op` to `results`. After going
// through the quantization passes, non-quantized quantizable units remain as
// `TF::XlaCallModuleOp` with a callee's prefix of `composite_`.
void PopulateNonQuantizedResults(ModuleOp module_op,
QuantizationResults& results) {
module_op.walk([&results](TF::XlaCallModuleOp xla_call_module_op) {
std::optional<QuantizationResult> result =
GetQuantizationResult(xla_call_module_op);
if (result == std::nullopt) return WalkResult::skip();
*results.add_results() = std::move(*result);
return WalkResult::advance();
});
}
} // namespace
QuantizationReport::QuantizationReport(ModuleOp module_op)
: quantization_results_(CollectResultsFromModuleOp(module_op)) {}
QuantizationResults QuantizationReport::CollectResultsFromModuleOp(
ModuleOp module_op) const {
QuantizationResults results{};
PopulateQuantizedResults(module_op, results);
PopulateNonQuantizedResults(module_op, results);
return results;
}
void QuantizationReport::AddQuantizationResult(QuantizationResult&& result) {
*quantization_results_.add_results() = std::move(result);
}
std::string QuantizationReport::ToString() const {
std::string results_str{};
TextFormat::PrintToString(quantization_results_, &results_str);
return absl::StrCat("===== Quantization Report =====\n\n", results_str,
"\n===== Quantization Report End =====\n\n");
}
void QuantizationReport::Print() const {
llvm::outs() << ToString();
llvm::outs().flush(); // Show the report immediately.
}
absl::Status QuantizationReport::Save(const StringRef file_path) const {
std::string results_str{};
TextFormat::PrintToString(GetQuantizationResults(), &results_str);
return WriteStringToFile(file_path, results_str);
}
} // namespace mlir::quant::stablehlo
tensorflow/compiler/mlir/quantization/stablehlo/cc/report.h
+71
View File
@@ -0,0 +1,71 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_REPORT_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_REPORT_H_
#include <string>
#include "absl/status/status.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
namespace mlir::quant::stablehlo {
// A class that manages information about `QuantizableUnit`s post-quantization,
// internally in the form of `QuantizationUnits`. It is used to collect
// quantization summary from a quantized `ModuleOp` and emit it in a human- and
// machine-readable format.
class QuantizationReport {
public:
QuantizationReport() = default;
// Initializes `QuantizationReport` by collecting `QuantizationResults` from
// `module_op`.
explicit QuantizationReport(ModuleOp module_op);
// Adds a `QuantizationResult` to the report.
void AddQuantizationResult(
::stablehlo::quantization::QuantizationResult&& result);
// Returns `QuantizationResults` that are registered in this report.
const ::stablehlo::quantization::QuantizationResults& GetQuantizationResults()
const {
return quantization_results_;
}
// Returns a human-readable string representation of this report.
std::string ToString() const;
// Prints a human-readable report to stdout.
void Print() const;
// Saves the report to `file_path`. The textproto representation of
// `QuantizationResults` will be written to the file. Returns non-ok status
// when the file write fails.
absl::Status Save(StringRef file_path) const;
private:
::stablehlo::quantization::QuantizationResults CollectResultsFromModuleOp(
ModuleOp module_op) const;
// Quantization results that are registered in this report. A quantization
// result may be added manually by calling `AddQuantizationResult`.
::stablehlo::quantization::QuantizationResults quantization_results_;
};
} // namespace mlir::quant::stablehlo
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_REPORT_H_
tensorflow/compiler/mlir/quantization/stablehlo/cc/report_test.cc
+325
View File
@@ -0,0 +1,325 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/report.h"
#include <string>
#include <utility>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "absl/status/status_matchers.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/common/test_base.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/io.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tsl/platform/protobuf.h" // IWYU pragma: keep
namespace mlir::quant::stablehlo {
namespace {
using ::stablehlo::quantization::Method;
using ::stablehlo::quantization::QuantizableUnit;
using ::stablehlo::quantization::QuantizationResult;
using ::stablehlo::quantization::QuantizationResults;
using ::stablehlo::quantization::io::ReadFileToString;
using ::testing::HasSubstr;
using ::testing::IsEmpty;
using ::testing::SizeIs;
using ::testing::StrEq;
using ::testing::TempDir;
using ::tsl::protobuf::TextFormat;
using QuantizationReportTest = ::mlir::quant::QuantizationTestBase;
TEST_F(QuantizationReportTest, GetQuantizationResultsReturnsEmptyResults) {
QuantizationReport report{};
const QuantizationResults& results = report.GetQuantizationResults();
ASSERT_THAT(results.results(), IsEmpty());
}
TEST_F(QuantizationReportTest, AddQuantizationResult) {
// Construct a `QuantizationResult` to add, representing a unit named
// `quantized_my_function` that is not quantized.
QuantizationResult result{};
QuantizableUnit& quantizable_unit = *result.mutable_quantizable_unit();
quantizable_unit.set_name("quantized_my_function");
Method& method = *result.mutable_method();
method.mutable_no_quantization();
QuantizationReport report{};
report.AddQuantizationResult(std::move(result));
const QuantizationResults& results = report.GetQuantizationResults();
ASSERT_THAT(results.results(), SizeIs(1));
const QuantizationResult& first_result = results.results(0);
EXPECT_THAT(first_result.quantizable_unit().name(),
StrEq("quantized_my_function"));
EXPECT_TRUE(first_result.method().has_no_quantization());
}
TEST_F(QuantizationReportTest, InitializeWithModuleOp) {
constexpr absl::string_view kQuantizedDotGeneral = R"mlir(
func.func @main(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> {
%0 = stablehlo.constant() {value = dense<127> : tensor<2x3xi8>} : () -> tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {1.000000e+0,2.000000e+0,3.000000e+0}>>
%1 = stablehlo.uniform_quantize %arg0 : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<i8:f32, 4.000000e+0>>
%2 = call @quantized_dot_general_fn(%1, %0) {_quantization_method = "static_range_ptq { }"} : (tensor<1x2x!quant.uniform<i8:f32, 4.000000e+0>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {1.000000e+0,2.000000e+0,3.000000e+0}>>) -> tensor<1x3x!quant.uniform<i8:f32, 5.000000e+0>>
%3 = stablehlo.uniform_dequantize %2 : (tensor<1x3x!quant.uniform<i8:f32, 5.000000e+0>>) -> tensor<1x3xf32>
return %3 : tensor<1x3xf32>
}
func.func private @quantized_dot_general_fn(%arg0: tensor<1x2x!quant.uniform<i8:f32, 4.000000e+0>>, %arg1: tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {1.000000e+0,2.000000e+0,3.000000e+0}>>) -> tensor<1x3x!quant.uniform<i8:f32, 5.000000e+0>> {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2x!quant.uniform<i8:f32, 4.000000e+0>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {1.000000e+0,2.000000e+0,3.000000e+0}>>) -> tensor<1x3x!quant.uniform<i32:f32:1, {6.000000e+0,7.000000e+0,8.000000e+0}>>
%1 = stablehlo.uniform_quantize %0 : (tensor<1x3x!quant.uniform<i32:f32:1, {6.000000e+0,7.000000e+0,8.000000e+0}>>) -> tensor<1x3x!quant.uniform<i8:f32, 5.000000e+0>>
return %1 : tensor<1x3x!quant.uniform<i8:f32, 5.000000e+0>>
}
)mlir";
const OwningOpRef<ModuleOp> module_op =
ParseModuleOpString(kQuantizedDotGeneral);
ASSERT_TRUE(module_op);
const QuantizationReport report(*module_op);
const QuantizationResults& results = report.GetQuantizationResults();
ASSERT_THAT(results.results(), SizeIs(1));
// Test that the quantized `QuantizableUnit` corresponding to
// `composite_dot_general_fn` is captured.
const QuantizationResult& result = results.results(0);
EXPECT_THAT(result.quantizable_unit().name(),
StrEq("composite_dot_general_fn"));
EXPECT_TRUE(result.method().has_static_range_ptq());
}
TEST_F(QuantizationReportTest,
InitializeWithModuleOpWithoutQuantizationMethodAttribute) {
// A quantized dot_general op but the `CallOp` is missing the
// `_quantization_method` attribute.
constexpr absl::string_view
kQuantizedDotGeneralMissingQuantizationMethodAttr = R"mlir(
func.func @main(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> {
%0 = stablehlo.constant() {value = dense<127> : tensor<2x3xi8>} : () -> tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {1.000000e+0,2.000000e+0,3.000000e+0}>>
%1 = stablehlo.uniform_quantize %arg0 : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<i8:f32, 4.000000e+0>>
%2 = call @quantized_dot_general_fn(%1, %0) : (tensor<1x2x!quant.uniform<i8:f32, 4.000000e+0>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {1.000000e+0,2.000000e+0,3.000000e+0}>>) -> tensor<1x3x!quant.uniform<i8:f32, 5.000000e+0>>
%3 = stablehlo.uniform_dequantize %2 : (tensor<1x3x!quant.uniform<i8:f32, 5.000000e+0>>) -> tensor<1x3xf32>
return %3 : tensor<1x3xf32>
}
func.func private @quantized_dot_general_fn(%arg0: tensor<1x2x!quant.uniform<i8:f32, 4.000000e+0>>, %arg1: tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {1.000000e+0,2.000000e+0,3.000000e+0}>>) -> tensor<1x3x!quant.uniform<i8:f32, 5.000000e+0>> {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2x!quant.uniform<i8:f32, 4.000000e+0>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {1.000000e+0,2.000000e+0,3.000000e+0}>>) -> tensor<1x3x!quant.uniform<i32:f32:1, {6.000000e+0,7.000000e+0,8.000000e+0}>>
%1 = stablehlo.uniform_quantize %0 : (tensor<1x3x!quant.uniform<i32:f32:1, {6.000000e+0,7.000000e+0,8.000000e+0}>>) -> tensor<1x3x!quant.uniform<i8:f32, 5.000000e+0>>
return %1 : tensor<1x3x!quant.uniform<i8:f32, 5.000000e+0>>
}
)mlir";
const OwningOpRef<ModuleOp> module_op =
ParseModuleOpString(kQuantizedDotGeneralMissingQuantizationMethodAttr);
ASSERT_TRUE(module_op);
const QuantizationReport report(*module_op);
const QuantizationResults& results = report.GetQuantizationResults();
// The quantized call op without the _quantization_method attribute is not
// captured as a `QuantizationResult`.
ASSERT_THAT(results.results(), IsEmpty());
}
TEST_F(QuantizationReportTest, InitializeWithModuleOpWithInvalidCalleeName) {
// A quantized dot_general op but the callee function has an invalid name. It
// is expected to start with `quantized_`.
constexpr absl::string_view kQuantizedDotGeneralWithInvalidCalleeName =
R"mlir(
func.func @main(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> {
%0 = stablehlo.constant() {value = dense<127> : tensor<2x3xi8>} : () -> tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {1.000000e+0,2.000000e+0,3.000000e+0}>>
%1 = stablehlo.uniform_quantize %arg0 : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<i8:f32, 4.000000e+0>>
%2 = call @invalid_quantized_dot_general_fn(%1, %0) {_quantization_method = "static_range_ptq { }"} : (tensor<1x2x!quant.uniform<i8:f32, 4.000000e+0>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {1.000000e+0,2.000000e+0,3.000000e+0}>>) -> tensor<1x3x!quant.uniform<i8:f32, 5.000000e+0>>
%3 = stablehlo.uniform_dequantize %2 : (tensor<1x3x!quant.uniform<i8:f32, 5.000000e+0>>) -> tensor<1x3xf32>
return %3 : tensor<1x3xf32>
}
func.func private @invalid_quantized_dot_general_fn(%arg0: tensor<1x2x!quant.uniform<i8:f32, 4.000000e+0>>, %arg1: tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {1.000000e+0,2.000000e+0,3.000000e+0}>>) -> tensor<1x3x!quant.uniform<i8:f32, 5.000000e+0>> {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2x!quant.uniform<i8:f32, 4.000000e+0>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {1.000000e+0,2.000000e+0,3.000000e+0}>>) -> tensor<1x3x!quant.uniform<i32:f32:1, {6.000000e+0,7.000000e+0,8.000000e+0}>>
%1 = stablehlo.uniform_quantize %0 : (tensor<1x3x!quant.uniform<i32:f32:1, {6.000000e+0,7.000000e+0,8.000000e+0}>>) -> tensor<1x3x!quant.uniform<i8:f32, 5.000000e+0>>
return %1 : tensor<1x3x!quant.uniform<i8:f32, 5.000000e+0>>
}
)mlir";
const OwningOpRef<ModuleOp> module_op =
ParseModuleOpString(kQuantizedDotGeneralWithInvalidCalleeName);
ASSERT_TRUE(module_op);
const QuantizationReport report(*module_op);
const QuantizationResults& results = report.GetQuantizationResults();
// The quantized call op whose callee doesn't start with `quantized_` is not
// captured as a `QuantizationResult`.
ASSERT_THAT(results.results(), IsEmpty());
}
TEST_F(QuantizationReportTest, InitializeWithModuleOpWithNonQuantizedOp) {
constexpr absl::string_view kNonQuantizedDotGeneral = R"mlir(
func.func @main(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> {
%0 = stablehlo.constant dense<3.000000e+0> : tensor<2x3xf32>
%1 = "tf.XlaCallModule"(%arg0, %0) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn, _stablehlo_version = "1.0.0", _original_entry_function = "composite_dot_general_fn", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
return %1 : tensor<1x3xf32>
}
func.func private @composite_dot_general_fn(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<1x3xf32> {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
}
)mlir";
const OwningOpRef<ModuleOp> module_op =
ParseModuleOpString(kNonQuantizedDotGeneral);
ASSERT_TRUE(module_op);
const QuantizationReport report(*module_op);
const QuantizationResults& results = report.GetQuantizationResults();
ASSERT_THAT(results.results(), SizeIs(1));
// Test that the unquantized `QuantizableUnit` corresponding to
// `composite_dot_general_fn` is captured. The `Method` contains
// `NoQuantization`.
const QuantizationResult& result = results.results(0);
EXPECT_THAT(result.quantizable_unit().name(),
StrEq("composite_dot_general_fn"));
EXPECT_TRUE(result.method().has_no_quantization());
}
TEST_F(QuantizationReportTest,
InitializeWithModuleOpWithQuantizedAndNonQuantizedOps) {
constexpr absl::string_view kQuantizedDotGeneralAndNonQuantizedDotGeneral =
R"mlir(
func.func @main(%arg0: tensor<1x2xf32>, %arg1: tensor<1x2xf32>) -> tensor<1x3xf32> {
// Non-quantized dot_general.
%0 = stablehlo.constant dense<3.000000e+0> : tensor<2x3xf32>
%1 = "tf.XlaCallModule"(%arg0, %0) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn_1, _stablehlo_verison = "1.0.0", _original_entry_function = "composite_dot_general_fn_1", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
// Quantized dot_general.
%2 = stablehlo.constant() {value = dense<127> : tensor<2x3xi8>} : () -> tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {1.000000e+0,2.000000e+0,3.000000e+0}>>
%3 = stablehlo.uniform_quantize %arg1 : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<i8:f32, 4.000000e+0>>
%4 = call @quantized_dot_general_fn_2(%3, %2) {_quantization_method = "static_range_ptq { }"} : (tensor<1x2x!quant.uniform<i8:f32, 4.000000e+0>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {1.000000e+0,2.000000e+0,3.000000e+0}>>) -> tensor<1x3x!quant.uniform<i8:f32, 5.000000e+0>>
%5 = stablehlo.uniform_dequantize %4 : (tensor<1x3x!quant.uniform<i8:f32, 5.000000e+0>>) -> tensor<1x3xf32>
// Add is there to prevent from dot_generals from being DCEed.
%6 = stablehlo.add %1, %5 : tensor<1x3xf32>
return %6 : tensor<1x3xf32>
}
// Callee of non-quantized op.
func.func private @composite_dot_general_fn_1(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<1x3xf32> {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
return %0 : tensor<1x3xf32>
}
// Callee of quantized op.
func.func private @quantized_dot_general_fn_2(%arg0: tensor<1x2x!quant.uniform<i8:f32, 4.000000e+0>>, %arg1: tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {1.000000e+0,2.000000e+0,3.000000e+0}>>) -> tensor<1x3x!quant.uniform<i8:f32, 5.000000e+0>> {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2x!quant.uniform<i8:f32, 4.000000e+0>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {1.000000e+0,2.000000e+0,3.000000e+0}>>) -> tensor<1x3x!quant.uniform<i32:f32:1, {6.000000e+0,7.000000e+0,8.000000e+0}>>
%1 = stablehlo.uniform_quantize %0 : (tensor<1x3x!quant.uniform<i32:f32:1, {6.000000e+0,7.000000e+0,8.000000e+0}>>) -> tensor<1x3x!quant.uniform<i8:f32, 5.000000e+0>>
return %1 : tensor<1x3x!quant.uniform<i8:f32, 5.000000e+0>>
}
)mlir";
const OwningOpRef<ModuleOp> module_op =
ParseModuleOpString(kQuantizedDotGeneralAndNonQuantizedDotGeneral);
ASSERT_TRUE(module_op);
const QuantizationReport report(*module_op);
const QuantizationResults& results = report.GetQuantizationResults();
ASSERT_THAT(results.results(), SizeIs(2));
// Test that the quantized op is captured in `results`.
const QuantizationResult& quantized_result = results.results(0);
EXPECT_THAT(quantized_result.quantizable_unit().name(),
StrEq("composite_dot_general_fn_2"));
EXPECT_TRUE(quantized_result.method().has_static_range_ptq());
// Test that the non-quantized op is captured in `results`.
const QuantizationResult& non_quantized_result = results.results(1);
EXPECT_THAT(non_quantized_result.quantizable_unit().name(),
StrEq("composite_dot_general_fn_1"));
EXPECT_TRUE(non_quantized_result.method().has_no_quantization());
}
TEST_F(QuantizationReportTest, ToString) {
QuantizationResult result{};
QuantizableUnit& quantizable_unit = *result.mutable_quantizable_unit();
quantizable_unit.set_name("quantized_my_function");
Method& method = *result.mutable_method();
method.mutable_no_quantization();
QuantizationReport report{};
report.AddQuantizationResult(std::move(result));
// Check that the report string is equivalent to the textproto representation
// of the `QuantizationResults`.
std::string result_str{};
TextFormat::PrintToString(report.GetQuantizationResults(), &result_str);
EXPECT_THAT(report.ToString(), HasSubstr("Quantization Report"));
EXPECT_THAT(report.ToString(), HasSubstr(result_str));
EXPECT_THAT(report.ToString(), HasSubstr("Quantization Report End"));
}
TEST_F(QuantizationReportTest, Save) {
constexpr absl::string_view kQuantizedDotGeneral = R"mlir(
func.func @main(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> {
%0 = stablehlo.constant() {value = dense<127> : tensor<2x3xi8>} : () -> tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {1.000000e+0,2.000000e+0,3.000000e+0}>>
%1 = stablehlo.uniform_quantize %arg0 : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<i8:f32, 4.000000e+0>>
%2 = call @quantized_dot_general_fn(%1, %0) {_quantization_method = "static_range_ptq { }"} : (tensor<1x2x!quant.uniform<i8:f32, 4.000000e+0>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {1.000000e+0,2.000000e+0,3.000000e+0}>>) -> tensor<1x3x!quant.uniform<i8:f32, 5.000000e+0>>
%3 = stablehlo.uniform_dequantize %2 : (tensor<1x3x!quant.uniform<i8:f32, 5.000000e+0>>) -> tensor<1x3xf32>
return %3 : tensor<1x3xf32>
}
func.func private @quantized_dot_general_fn(%arg0: tensor<1x2x!quant.uniform<i8:f32, 4.000000e+0>>, %arg1: tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {1.000000e+0,2.000000e+0,3.000000e+0}>>) -> tensor<1x3x!quant.uniform<i8:f32, 5.000000e+0>> {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2x!quant.uniform<i8:f32, 4.000000e+0>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {1.000000e+0,2.000000e+0,3.000000e+0}>>) -> tensor<1x3x!quant.uniform<i32:f32:1, {6.000000e+0,7.000000e+0,8.000000e+0}>>
%1 = stablehlo.uniform_quantize %0 : (tensor<1x3x!quant.uniform<i32:f32:1, {6.000000e+0,7.000000e+0,8.000000e+0}>>) -> tensor<1x3x!quant.uniform<i8:f32, 5.000000e+0>>
return %1 : tensor<1x3x!quant.uniform<i8:f32, 5.000000e+0>>
}
)mlir";
const OwningOpRef<ModuleOp> module_op =
ParseModuleOpString(kQuantizedDotGeneral);
ASSERT_TRUE(module_op);
const QuantizationReport report(*module_op);
const std::string dst_file_path =
absl::StrCat(TempDir(), "/quantization_report.txtpb");
const absl::Status save_status = report.Save(dst_file_path);
ASSERT_THAT(save_status, absl_testing::IsOk());
const absl::StatusOr<std::string> file_data = ReadFileToString(dst_file_path);
ASSERT_THAT(file_data, absl_testing::IsOk());
// Test that the file data can be parsed as `QuantizationResults`.
QuantizationResults results{};
ASSERT_TRUE(TextFormat::ParseFromString(*file_data, &results));
// Check that `results` reflects the information of the quantized units
// properly.
ASSERT_THAT(results.results(), SizeIs(1));
EXPECT_THAT(results.results(0).quantizable_unit().name(),
StrEq("composite_dot_general_fn"));
EXPECT_TRUE(results.results(0).method().has_static_range_ptq());
}
} // namespace
} // namespace mlir::quant::stablehlo
tensorflow/compiler/mlir/quantization/stablehlo/cc/saved_model_export.cc
+290
View File
@@ -0,0 +1,290 @@
/* 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/stablehlo/cc/saved_model_export.h"
#include <memory>
#include <optional>
#include <string>
#include <unordered_set>
#include <utility>
#include <vector>
#include "absl/algorithm/container.h"
#include "absl/base/attributes.h"
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/match.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_join.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/Pass/PassManager.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/stablehlo/cc/io.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/pass_pipeline.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/types.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/cc/convert_asset_args.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/cc/run_passes.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/exported_model.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/constants.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/passes.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/python/unfreeze_constants.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/translate/mlir_roundtrip_flags.h"
#include "tensorflow/compiler/mlir/tf2xla/api/v2/tf_executor_to_graph.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include "tensorflow/core/protobuf/saver.pb.h"
namespace mlir::quant::stablehlo {
namespace {
using ::mlir::tf_saved_model::kTfSavedModelIndexPathAttr;
using ::mlir::tf_saved_model::kTfSavedModelInitializerInitType;
using ::mlir::tf_saved_model::kTfSavedModelInitializerRestoreType;
using ::stablehlo::quantization::QuantizationConfig;
using ::stablehlo::quantization::io::GetLocalTmpFileName;
using ::tensorflow::AssetFileDef;
using ::tensorflow::FunctionDefLibrary;
using ::tensorflow::FunctionLibraryDefinition;
using ::tensorflow::Graph;
using ::tensorflow::GraphDef;
using ::tensorflow::Node;
using ::tensorflow::NodeDef;
using ::tensorflow::OpRegistry;
using ::tensorflow::SaverDef;
using ::tensorflow::quantization::ExportedModel;
using ::tensorflow::quantization::RunPasses;
using ::tensorflow::quantization::UnfreezeConstantsAndSaveVariables;
// Finds and returns the name of the node from a set of control output nodes.
// The name should contain the string `contains`. Returns an empty string if no
// node whose name contains `contains` is found. Assumes there is at most one
// such a node.
std::string GetNodeName(const std::vector<std::string>& control_ret_node_names,
const absl::string_view contains) {
for (const std::string& node_name : control_ret_node_names) {
if (absl::StrContains(node_name, contains)) {
VLOG(1) << "Node found: " << node_name << ", contains: " << contains;
return node_name;
}
}
VLOG(1) << "Could not find node whose name conatins: " << contains;
return "";
}
// Returns the file prefix tensor name. An empty string is returned if no such a
// tensor is found (when there are no variables to restore, it is expected that
// the file prefix tensor does not exist). The file prefix tensor is found among
// the "_Arg" nodes, as it is translated from the MLIR @main function's
// argument. It also must have the attribute `tf_saved_model.index_path =
// ["__tf_file_prefix"]`.
//
// See `MergeSaveFunctionOpsToMainPass` for details how the file prefix tensor
// ends up at the MLIR @main function's argument.
std::string FindFilePrefixTensorName(const GraphDef& graph_def) {
for (const NodeDef& node_def : graph_def.node()) {
if (node_def.op() == FunctionLibraryDefinition::kArgOp) {
// Matches the `tf_saved_model.index_path = ["__tf_file_prefix"]`.
const auto index_path_attr_itr =
node_def.attr().find(kTfSavedModelIndexPathAttr.str());
if (index_path_attr_itr != node_def.attr().end()) {
const auto& index_paths = index_path_attr_itr->second.list().s();
if (absl::c_find(index_paths, kTfFilePrefix.str()) !=
index_paths.end()) {
// ":0" appended to indicate that it is a tensor, not an Operation.
return absl::StrCat(node_def.name(), ":0");
}
}
}
}
return "";
}
} // namespace
absl::StatusOr<ExportedModel> CreateExportedModel(
const std::vector<std::string>& signature_keys,
const std::unordered_set<std::string>& tags,
const QuantizationConfig& quantization_config,
absl::string_view debug_name_prefix,
const absl::flat_hash_map<FunctionName, FunctionAlias>& function_aliases,
MLIRContext& ctx ABSL_ATTRIBUTE_LIFETIME_BOUND, ModuleOp module_op) {
TF_ASSIGN_OR_RETURN(const std::string checkpoint_dir, GetLocalTmpFileName());
const ExportOptions export_opts = {
/*duplicate_shape_determining_constants=*/true,
/*unfreeze_constants=*/false, checkpoint_dir,
/*debug_name=*/
absl::StrCat(debug_name_prefix, kExportStepSuffix)};
TF_ASSIGN_OR_RETURN(const SmallVector<AssetFileDef> asset_file_defs,
RunExportPasses(export_opts, ctx, module_op));
return ConvertMlirModuleToExportedModel(
module_op, checkpoint_dir, function_aliases,
{asset_file_defs.begin(), asset_file_defs.end()});
}
ExportedModel CreateExportedModelFromGraphDef(
GraphDef&& graph_def, const absl::string_view init_node_name,
const absl::string_view checkpoint_dir,
const std::optional<SaverDef> saver_def,
const absl::flat_hash_map<FunctionName, FunctionAlias>& function_aliases,
const std::vector<AssetFileDef>& asset_file_defs) {
ExportedModel exported_model{};
*exported_model.mutable_graph_def() = graph_def;
exported_model.set_init_node_name(init_node_name);
exported_model.set_checkpoint_dir(checkpoint_dir);
exported_model.mutable_function_aliases()->insert(function_aliases.begin(),
function_aliases.end());
for (const AssetFileDef& asset_file_def : asset_file_defs) {
*exported_model.mutable_asset_file_defs()->Add() = asset_file_def;
}
if (saver_def != std::nullopt) {
*exported_model.mutable_saver_def() = *std::move(saver_def);
}
return exported_model;
}
void AddExportPasses(mlir::PassManager& pm,
const bool duplicate_shape_determining_constants) {
quant::stablehlo::AddCallModuleSerializationPasses(pm);
if (duplicate_shape_determining_constants) {
pm.addNestedPass<mlir::func::FuncOp>(
mlir::quant::CreateDuplicateShapeDeterminingConstantsPass());
}
pm.addPass(mlir::quant::CreateInsertMainFunctionPass());
pm.addPass(mlir::quant::CreateLiftHashTableOpsAsArgsPass());
pm.addNestedPass<mlir::func::FuncOp>(
mlir::CreateFunctionalToExecutorDialectConversionPass());
pm.addPass(mlir::CreateBreakUpIslandsPass());
pm.addPass(mlir::quant::CreateMergeInitializerFunctionOpsToMainPass());
pm.addPass(mlir::quant::CreateMergeSaveFunctionOpsToMainPass());
pm.addNestedPass<mlir::func::FuncOp>(
mlir::quant::CreateMergeDuplicateResourceOpsPass());
// Used to clean up the "tf._noinliner" attribute that is previously used to
// prevent certain functions from being inlined (see
// `MarkFunctionsNoinlinePass`). InlinerPass must not come after this pass.
pm.addPass(mlir::TF::CreateStripNoinlineAttributePass());
}
absl::StatusOr<std::optional<SaverDef>> CreateSaverDef(
const std::vector<std::string>& control_ret_node_names,
const GraphDef& graph_def) {
const std::string filename_tensor_name = FindFilePrefixTensorName(graph_def);
const std::string restore_op_name =
GetNodeName(control_ret_node_names, kTfSavedModelInitializerRestoreType);
const std::string save_node_name =
GetNodeName(control_ret_node_names, kTfQuantSaveOpName);
const std::vector<absl::string_view> fields = {
filename_tensor_name, restore_op_name, save_node_name};
const auto is_empty_predicate = [](const absl::string_view s) {
return s.empty();
};
if (absl::c_all_of(fields, is_empty_predicate)) {
return std::nullopt;
} else if (absl::c_none_of(fields, is_empty_predicate)) {
SaverDef saver_def{};
saver_def.set_version(SaverDef::V2);
saver_def.set_filename_tensor_name(filename_tensor_name);
saver_def.set_restore_op_name(restore_op_name);
// :0 attached to indicate the first result tensor. This saves the model
// checkpoint when fetched.
saver_def.set_save_tensor_name(absl::StrCat(save_node_name, ":0"));
return saver_def;
} else {
return absl::InternalError(
absl::StrCat("Failed to create SaverDef. Fields should be either all "
"empty strings or all non-empty strings. Got fields: ",
absl::StrJoin(fields, ",")));
}
}
absl::StatusOr<ExportedModel> ConvertMlirModuleToExportedModel(
const mlir::ModuleOp module_op, const absl::string_view checkpoint_dir,
const absl::flat_hash_map<FunctionName, FunctionAlias>& function_aliases,
const std::vector<AssetFileDef>& asset_file_defs) {
const tensorflow::GraphExportConfig config{};
FunctionLibraryDefinition flib_def{OpRegistry::Global(),
FunctionDefLibrary()};
std::unique_ptr<Graph> graph;
absl::flat_hash_set<Node*> control_ret_nodes{};
TF_RETURN_IF_ERROR(tensorflow::tf2xla::v2::ConvertTfExecutorToGraph(
module_op, config, &graph, &flib_def, &control_ret_nodes));
GraphDef graph_def{};
graph->ToGraphDef(&graph_def);
std::vector<std::string> control_ret_node_names{};
for (Node* node : control_ret_nodes) {
control_ret_node_names.push_back(node->name());
}
const std::string init_node_name =
GetNodeName(control_ret_node_names, kTfSavedModelInitializerInitType);
TF_ASSIGN_OR_RETURN(const std::optional<SaverDef> saver_def,
CreateSaverDef(control_ret_node_names, graph_def));
return CreateExportedModelFromGraphDef(std::move(graph_def), init_node_name,
checkpoint_dir, std::move(saver_def),
function_aliases, asset_file_defs);
}
absl::StatusOr<SmallVector<AssetFileDef>> RunExportPasses(
const ExportOptions& export_opts, MLIRContext& ctx, ModuleOp module_op) {
if (export_opts.unfreeze_constants) {
TF_RETURN_IF_ERROR(UnfreezeConstantsAndSaveVariables(
export_opts.checkpoint_dir, ctx, module_op));
LOG(INFO) << "Unfrozen constants and saved variables to checkpoint file: "
<< export_opts.checkpoint_dir;
}
TF_RETURN_IF_ERROR(RunPasses(
/*name=*/
export_opts.debug_name,
/*add_passes_func=*/
[dup_constants = export_opts.duplicate_shape_determining_constants](
PassManager& pm) { AddExportPasses(pm, dup_constants); },
ctx, module_op));
FailureOr<SmallVector<AssetFileDef>> asset_file_defs =
quant::ConvertAssetArgs(module_op);
if (failed(asset_file_defs)) {
return absl::InternalError("Failed to convert asset args.");
}
return *asset_file_defs;
}
} // namespace mlir::quant::stablehlo
tensorflow/compiler/mlir/quantization/stablehlo/cc/saved_model_export.h
+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.
==============================================================================*/
// Functionalities for exporting MLIR ModuleOp to TensorFlow SavedModel.
#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_SAVED_MODEL_EXPORT_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_SAVED_MODEL_EXPORT_H_
#include <optional>
#include <string>
#include <unordered_set>
#include <vector>
#include "absl/base/attributes.h"
#include "absl/container/flat_hash_map.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 "mlir/Support/LLVM.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/types.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/exported_model.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include "tensorflow/core/protobuf/saver.pb.h"
namespace mlir::quant::stablehlo {
// Suffix string for the module export step. Used for debugging.
constexpr absl::string_view kExportStepSuffix = "_export";
// Options when running passes for exporting an MLIR ModuleOp.
struct ExportOptions {
// If set to `true`, it runs `DuplicateShapeDeterminingConstantsPass` before
// lowering to tf_executor dialect.
bool duplicate_shape_determining_constants = true;
// If set to `true`, unfreezes constants into variables and saves them to a
// checkpoint file. Setting this to `true` is an experimental feature that has
// no stability guarantees.
bool unfreeze_constants = false;
// Path to the directory where checkpoint files are saved.
std::string checkpoint_dir = "";
// Name used to identify the ModuleOp this is exporting. Only used for
// debugging and does not modify the behavior of the export.
std::string debug_name = "stablehlo_quant";
};
// Creates `ExportedModel` from `module_op`. `module_op` goes through post
// process passes before an `ExportModel` is created.
// TODO: b/329206105 - Add unit tests after decomposing post processing passes.
absl::StatusOr<tensorflow::quantization::ExportedModel> CreateExportedModel(
const std::vector<std::string>& signature_keys,
const std::unordered_set<std::string>& tags,
const ::stablehlo::quantization::QuantizationConfig& quantization_config,
absl::string_view debug_name_prefix,
const absl::flat_hash_map<FunctionName, FunctionAlias>& function_aliases,
MLIRContext& ctx ABSL_ATTRIBUTE_LIFETIME_BOUND, ModuleOp module_op);
// Factory function for `ExportedModel`.
[[nodiscard]] tensorflow::quantization::ExportedModel
CreateExportedModelFromGraphDef(
tensorflow::GraphDef&& graph_def, absl::string_view init_node_name,
absl::string_view checkpoint_dir,
std::optional<tensorflow::SaverDef> saver_def,
const absl::flat_hash_map<std::string, std::string>& function_aliases,
const std::vector<tensorflow::AssetFileDef>& asset_file_defs);
// Creates a new `SaverDef` instance, which contains information regarding
// checkpoint saving and restoring. This function returns a `SaverDef` instance
// with four fields populated: `version`, `filename_tensor_name`,
// `restore_op_name` and `save_tensor_name`. For valid quantized `graph_def` and
// `control_ret_node_names`, it should be able to retrieve the last three fields
// if there is at lest one variable in the graph.
//
// Returns a `std::nullopt` if there are no variables in the graph and no saving
// & restoring are required. Returns an `InternalError` status for when the
// required fields are only partially provided.
absl::StatusOr<std::optional<tensorflow::SaverDef>> CreateSaverDef(
const std::vector<std::string>& control_ret_node_names,
const tensorflow::GraphDef& graph_def);
// Adds passes for transforming the MLIR module op so that it can be exported
// back to GraphDef. Roughly, this consists of:
// 1) Inserting the @main function, which will become the main Graph.
// 2) Duplicating shape-determining constants.
// 3) Converting TF dialect -> tf_executor dialect.
// 4) Adding initializer function's ops into @main function for correct
// resource initialization when loading the exported model.
//
// Duplicating shape-determining constants is required to place constants that
// affect the shape of a tensor to be placed in the TPU graph instead of in the
// CPU graph, when the graph gets converted for TPU inference. This allows these
// constants to be known at XLA compilation time.
void AddExportPasses(mlir::PassManager& pm,
bool duplicate_shape_determining_constants);
// Converts MLIR ModuleOp to `ExportedModel`. Returns `InternalError` status
// when the conversion fails.
//
// * `checkpoint_dir` is the directory where checkpoints where variable values
// are stored. This value will be fed to the "file_prefix" tensor to restore the
// variables.
// * `function_aliases` maps the actual function name to the function alias.
// This associates the quantized functions to the original functions' aliases.
// If there were no function aliases in the input model, this should be empty.
// * `asset_file_defs` include information about the assets, if any, that are
// used directly to initialize resources (like hash tables). If no assets are
// used in the model, this should be empty.
absl::StatusOr<tensorflow::quantization::ExportedModel>
ConvertMlirModuleToExportedModel(
mlir::ModuleOp module_op, absl::string_view checkpoint_dir,
const absl::flat_hash_map<std::string, std::string>& function_aliases,
const std::vector<tensorflow::AssetFileDef>& asset_file_defs);
// Sets up and runs the passes for exporting `module_op`. The behavior of the
// exporting passes is controlled by `export_opts`. Returns `AssetFileDef`s that
// associate the input arguments of @main and the asset file names. Asset file
// names will be used to feed the corresponding tensors during initialization
// upon model loading.
// TODO: b/329206105 - Add unit tests after decomposing post processing passes.
absl::StatusOr<SmallVector<::tensorflow::AssetFileDef>> RunExportPasses(
const ExportOptions& export_opts, MLIRContext& ctx, ModuleOp module_op);
} // namespace mlir::quant::stablehlo
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_SAVED_MODEL_EXPORT_H_
tensorflow/compiler/mlir/quantization/stablehlo/cc/saved_model_export_test.cc
+438
View File
@@ -0,0 +1,438 @@
/* 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/stablehlo/cc/saved_model_export.h"
#include <optional>
#include <string>
#include <utility>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "absl/status/status_matchers.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "llvm/ADT/STLExtras.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/common/test_base.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/exported_model.pb.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include "tensorflow/core/protobuf/saver.pb.h"
#include "tsl/platform/protobuf.h" // IWYU pragma: keep
namespace mlir::quant::stablehlo {
namespace {
using ::tensorflow::AssetFileDef;
using ::tensorflow::GraphDef;
using ::tensorflow::NodeDef;
using ::tensorflow::SaverDef;
using ::tensorflow::quantization::ExportedModel;
using ::testing::HasSubstr;
using ::testing::IsEmpty;
using ::testing::SizeIs;
using ::testing::StrEq;
using ::tsl::protobuf::TextFormat;
TEST(CreateExportedModelTest, CreateExportedModelBasicFieldsSet) {
GraphDef graph_def{};
ASSERT_TRUE(
TextFormat::ParseFromString(R"pb(node { name: "foo" })pb", &graph_def));
const ExportedModel exported_model = CreateExportedModelFromGraphDef(
std::move(graph_def), "init_node_name", "checkpoint_dir",
/*saver_def=*/std::nullopt,
/*function_aliases=*/{}, /*asset_file_defs=*/{});
ASSERT_THAT(exported_model.graph_def().node(), SizeIs(1));
EXPECT_THAT(exported_model.graph_def().node()[0].name(), StrEq("foo"));
EXPECT_THAT(exported_model.init_node_name(), StrEq("init_node_name"));
EXPECT_THAT(exported_model.checkpoint_dir(), StrEq("checkpoint_dir"));
EXPECT_FALSE(exported_model.has_saver_def());
EXPECT_THAT(exported_model.function_aliases(), IsEmpty());
EXPECT_THAT(exported_model.asset_file_defs(), IsEmpty());
}
TEST(CreateExportedModelTest, CreateExportedModelWithAddedFunctionAliases) {
const ExportedModel exported_model = CreateExportedModelFromGraphDef(
GraphDef(), /*init_node_name=*/"", /*checkpoint_dir=*/"",
/*saver_def=*/std::nullopt,
/*function_aliases=*/{{"func1", "alias1"}, {"func2", "alias2"}},
/*asset_file_defs=*/{});
ASSERT_THAT(exported_model.function_aliases(), SizeIs(2));
EXPECT_TRUE(exported_model.function_aliases().contains("func1"));
EXPECT_THAT(exported_model.function_aliases().at("func1"), StrEq("alias1"));
EXPECT_TRUE(exported_model.function_aliases().contains("func2"));
EXPECT_THAT(exported_model.function_aliases().at("func2"), StrEq("alias2"));
}
TEST(CreateExportedModelTest, CreateExportedModelWithAddedAssetFileDefs) {
AssetFileDef asset1;
ASSERT_TRUE(
TextFormat::ParseFromString(R"pb(filename: "fname1")pb", &asset1));
AssetFileDef asset2;
ASSERT_TRUE(
TextFormat::ParseFromString(R"pb(filename: "fname2")pb", &asset2));
const ExportedModel exported_model = CreateExportedModelFromGraphDef(
GraphDef(), /*init_node_name=*/"", /*checkpoint_dir=*/"",
/*saver_def=*/std::nullopt, /*function_aliases=*/{},
/*asset_file_defs=*/{asset1, asset2});
ASSERT_THAT(exported_model.asset_file_defs(), SizeIs(2));
EXPECT_THAT(exported_model.asset_file_defs()[0].filename(), StrEq("fname1"));
EXPECT_THAT(exported_model.asset_file_defs()[1].filename(), StrEq("fname2"));
}
TEST(CreateExportedModelTest, CreateExportedModelWithAddedSaverDef) {
SaverDef saver_def;
ASSERT_TRUE(TextFormat::ParseFromString(
R"pb(filename_tensor_name: "my_file")pb", &saver_def));
const ExportedModel exported_model = CreateExportedModelFromGraphDef(
GraphDef(), /*init_node_name=*/"", /*checkpoint_dir=*/"", saver_def,
/*function_aliases=*/{}, /*asset_file_defs=*/{});
EXPECT_THAT(exported_model.saver_def().filename_tensor_name(), "my_file");
}
TEST(CreateSaverDefTest, CreateValidSaverDef) {
// Needs to have a _Arg node with an attribute "tf_saved_model.index_path" =
// ["__tf_file_prefix"].
GraphDef graph_def;
ASSERT_TRUE(TextFormat::ParseFromString(
R"pb(node {
name: "foo",
op: "_Arg",
attr {
key: "tf_saved_model.index_path",
value { list { s: "__tf_file_prefix" } }
}
})pb",
&graph_def));
// Restore op's name should start with "restore_op" and the save op's name
// should start with "tf_quant__save_op".
const std::vector<std::string> control_ret_node_names = {
"restore_op_0", "tf_quant__save_op_0"};
TF_ASSERT_OK_AND_ASSIGN(const std::optional<SaverDef> saver_def,
CreateSaverDef(control_ret_node_names, graph_def));
ASSERT_NE(saver_def, std::nullopt);
EXPECT_THAT(saver_def->version(), SaverDef::V2);
EXPECT_THAT(saver_def->restore_op_name(), "restore_op_0");
EXPECT_THAT(saver_def->filename_tensor_name(), "foo:0");
EXPECT_THAT(saver_def->save_tensor_name(), "tf_quant__save_op_0:0");
}
TEST(CreateSaverDefTest, ReturnsNulloptIfNoSaverDefRelatedNodesExist) {
TF_ASSERT_OK_AND_ASSIGN(
const std::optional<SaverDef> saver_def,
CreateSaverDef(/*control_ret_node_names=*/{}, GraphDef()));
EXPECT_EQ(saver_def, std::nullopt);
}
TEST(CreateSaverDefTest, ReturnsErrorStatusIfSaverDefNodesPartiallyExist) {
// An _Arg node missing the attribute "tf_saved_model.index_path" =
// ["__tf_file_prefix"].
GraphDef graph_def;
ASSERT_TRUE(TextFormat::ParseFromString(
R"pb(node { name: "foo", op: "_Arg" })pb", &graph_def));
// Restore op's name should start with "restore_op" and the save op's name
// should start with "tf_quant__save_op".
const std::vector<std::string> control_ret_node_names = {
"restore_op_0", "tf_quant__save_op_0"};
const absl::StatusOr<std::optional<SaverDef>> saver_def =
CreateSaverDef(control_ret_node_names, graph_def);
EXPECT_THAT(
saver_def,
absl_testing::StatusIs(
absl::StatusCode::kInternal,
HasSubstr(
"should be either all empty strings or all non-empty strings")));
}
// Testing ConvertMlirModuleToExportedModel requires parsing MLIR string to
// ModuleOp.
using ConvertMlirModuleToExportedModelTest =
::mlir::quant::QuantizationTestBase;
TEST_F(ConvertMlirModuleToExportedModelTest, SimpleGraphDefSet) {
// Define a module a no-op main function.
mlir::OwningOpRef<mlir::ModuleOp> module_op = ParseModuleOpString(R"mlir(
module attributes {tf_saved_model.semantics} {
func.func @main(%arg: tensor<1x2xf32> {tf_saved_model.index_path = ["input_tensor:0"]}) -> (tensor<1x2xf32> {tf_saved_model.index_path = ["output_tensor:0"]}) attributes {tf.entry_function = {inputs = "input_tensor:0", outputs = "output_tensor:0"}, tf_saved_model.exported_names = ["main"]} {
%0 = tf_executor.graph {
tf_executor.fetch %arg : tensor<1x2xf32>
}
return %0 : tensor<1x2xf32>
}
}
)mlir");
ASSERT_TRUE(module_op);
const absl::StatusOr<ExportedModel> exported_model =
ConvertMlirModuleToExportedModel(*module_op, /*checkpoint_dir=*/"",
/*function_aliases=*/{},
/*asset_file_defs=*/{});
ASSERT_THAT(exported_model, absl_testing::IsOk());
// There are 2 nodes in the graph, one for arg and another for retval.
ASSERT_THAT(exported_model->graph_def().node(), SizeIs(2));
// Match the `_Arg` node that corresponds to the argument of @main.
const auto arg_node_itr =
llvm::find_if(exported_model->graph_def().node(),
[](const NodeDef& node) { return node.op() == "_Arg"; });
ASSERT_NE(arg_node_itr, exported_model->graph_def().node().end());
EXPECT_THAT(arg_node_itr->name(), StrEq("input_tensor"));
ASSERT_TRUE(arg_node_itr->attr().contains("tf_saved_model.index_path"));
ASSERT_THAT(arg_node_itr->attr().at("tf_saved_model.index_path").list().s(),
SizeIs(1));
EXPECT_THAT(
arg_node_itr->attr().at("tf_saved_model.index_path").list().s()[0],
StrEq("input_tensor:0"));
// Match the `_Retval` node that corresponds to the return value of @main.
const auto retval_node_itr =
llvm::find_if(exported_model->graph_def().node(),
[](const NodeDef& node) { return node.op() == "_Retval"; });
ASSERT_NE(retval_node_itr, exported_model->graph_def().node().end());
EXPECT_THAT(retval_node_itr->name(), StrEq("output_tensor"));
ASSERT_TRUE(retval_node_itr->attr().contains("tf_saved_model.index_path"));
ASSERT_THAT(
retval_node_itr->attr().at("tf_saved_model.index_path").list().s(),
SizeIs(1));
EXPECT_THAT(
retval_node_itr->attr().at("tf_saved_model.index_path").list().s()[0],
StrEq("output_tensor:0"));
}
TEST_F(ConvertMlirModuleToExportedModelTest, CheckpointDirSet) {
// Define a module a no-op main function.
mlir::OwningOpRef<mlir::ModuleOp> module_op = ParseModuleOpString(R"mlir(
module attributes {tf_saved_model.semantics} {
func.func @main() -> () attributes {tf_saved_model.exported_names = ["main"]} {
tf_executor.graph {
tf_executor.fetch
}
return
}
}
)mlir");
ASSERT_TRUE(module_op);
const absl::StatusOr<ExportedModel> exported_model =
ConvertMlirModuleToExportedModel(*module_op, "my_checkpoint_dir",
/*function_aliases=*/{},
/*asset_file_defs=*/{});
ASSERT_THAT(exported_model, absl_testing::IsOk());
EXPECT_THAT(exported_model->checkpoint_dir(), StrEq("my_checkpoint_dir"));
}
TEST_F(ConvertMlirModuleToExportedModelTest, FunctionAliasesSet) {
// Define a module with 2 function calls, function_1 and function_2.
mlir::OwningOpRef<mlir::ModuleOp> module_op = ParseModuleOpString(R"mlir(
module attributes {tf_saved_model.semantics} {
func.func private @function_1() -> () attributes {tf._original_func_name = "__func_1"} {
tf_executor.graph {
%control_0 = tf_executor.island wraps "tf.NoOp"() : () -> ()
}
return
}
func.func private @function_2() -> () attributes {tf._original_func_name = "__func_2"} {
tf_executor.graph {
%control_0 = tf_executor.island wraps "tf.NoOp"() : () -> ()
}
return
}
func.func @main() -> () attributes {tf_saved_model.exported_names = ["main"]} {
tf_executor.graph {
%control_0 = tf_executor.island wraps "tf.PartitionedCall"() <{config = "", config_proto = "", executor_type = "", f = @function_1}> : () -> ()
%control_1 = tf_executor.island wraps "tf.PartitionedCall"() <{config = "", config_proto = "", executor_type = "", f = @function_2}> : () -> ()
tf_executor.fetch %control_0, %control_1 : !tf_executor.control, !tf_executor.control
}
return
}
}
)mlir");
ASSERT_TRUE(module_op);
const absl::StatusOr<ExportedModel> exported_model =
ConvertMlirModuleToExportedModel(
*module_op, /*checkpoint_dir=*/"",
/*function_aliases=*/
{{"alias_1", "function_1"}, {"alias_2", "function_2"}},
/*asset_file_defs=*/{});
ASSERT_THAT(exported_model, absl_testing::IsOk());
ASSERT_THAT(exported_model->function_aliases(), SizeIs(2));
EXPECT_THAT(exported_model->function_aliases().at("alias_1"),
StrEq("function_1"));
EXPECT_THAT(exported_model->function_aliases().at("alias_2"),
StrEq("function_2"));
}
TEST_F(ConvertMlirModuleToExportedModelTest, AssetFileDefSet) {
// Define a module a no-op main function.
mlir::OwningOpRef<mlir::ModuleOp> module_op = ParseModuleOpString(R"mlir(
module attributes {tf_saved_model.semantics} {
func.func @main() -> () attributes {tf_saved_model.exported_names = ["main"]} {
tf_executor.graph {
tf_executor.fetch
}
return
}
}
)mlir");
ASSERT_TRUE(module_op);
AssetFileDef asset_file_def{};
ASSERT_TRUE(
TextFormat::ParseFromString(R"pb(filename: "vocab_file.txt",
tensor_info { name: "arg_0:0" })pb",
&asset_file_def));
const std::vector<AssetFileDef> asset_file_defs = {asset_file_def};
const absl::StatusOr<ExportedModel> exported_model =
ConvertMlirModuleToExportedModel(*module_op, /*checkpoint_dir=*/"",
/*function_aliases=*/{},
/*asset_file_defs=*/asset_file_defs);
ASSERT_THAT(exported_model, absl_testing::IsOk());
ASSERT_THAT(exported_model->asset_file_defs(), SizeIs(1));
EXPECT_THAT(exported_model->asset_file_defs()[0].filename(),
StrEq("vocab_file.txt"));
EXPECT_THAT(exported_model->asset_file_defs()[0].tensor_info().name(),
StrEq("arg_0:0"));
}
TEST_F(ConvertMlirModuleToExportedModelTest,
InitNodeNameSetToLocOfControlOutput) {
// Define a module that initializes a tf.HashTableV2 whose control output node
// for the initialization is named "init_op_init_all_tables".
mlir::OwningOpRef<mlir::ModuleOp> module_op = ParseModuleOpString(R"mlir(
module attributes {tf_saved_model.semantics} {
"tf_saved_model.session_initializer"() <{initializers = []}> : () -> ()
"tf_saved_model.asset"() <{filename = "assets/vocab_file.txt", sym_name = "__tf_saved_model_asset0_vocab_file.txt"}> : () -> ()
func.func @main(%arg1: tensor<!tf_type.string> {tf_saved_model.index_path = ["arg_0:0"]}) -> (tensor<1x2xf32> {tf_saved_model.index_path = ["output:0"]}) attributes {tf.entry_function = {inputs = "arg_0:0", outputs = "output:0"}, tf_saved_model.exported_names = ["main"]} {
%0 = tf_executor.graph {
%o_0, %c_0 = tf_executor.island wraps "tf.Const"() <{value = dense<1.0> : tensor<1x2xf32>}> : () -> tensor<1x2xf32>
%o, %c = tf_executor.island wraps "tf.HashTableV2"() <{container = "", key_dtype = !tf_type.string, shared_name = "vocab_file.txt", use_node_name_sharing = false, value_dtype = i64}> {device = ""} : () -> tensor<!tf_type.resource>
%c_9 = tf_executor.island wraps "tf.InitializeTableFromTextFileV2"(%o, %arg1) <{delimiter = "\09", key_index = -2 : i64, value_index = -1 : i64, vocab_size = -1 : i64}> {_has_manual_control_dependencies = true, device = ""} : (tensor<!tf_type.resource>, tensor<!tf_type.string>) -> ()
// Location of this control output op becomes the name of the init_op.
%c_10 = tf_executor.island(%c_9) wraps "tf.NoOp"() : () -> () loc("init_op_init_all_tables")
tf_executor.fetch %o_0, %c_10 : tensor<1x2xf32>, !tf_executor.control
}
return %0 : tensor<1x2xf32>
}
}
)mlir");
ASSERT_TRUE(module_op);
const absl::StatusOr<ExportedModel> exported_model =
ConvertMlirModuleToExportedModel(*module_op, /*checkpoint_dir=*/"",
/*function_aliases=*/{},
/*asset_file_defs=*/{});
ASSERT_THAT(exported_model, absl_testing::IsOk());
EXPECT_THAT(exported_model->init_node_name(),
StrEq("init_op_init_all_tables"));
// Match the init node, which is a NoOp that has control dependency to
// HashTableV2 initialization. Fetching this node in TF Session will
// initialize the hash table.
const auto init_node_itr = llvm::find_if(
exported_model->graph_def().node(), [](const NodeDef& node) {
return node.name() == "init_op_init_all_tables";
});
ASSERT_NE(init_node_itr, exported_model->graph_def().node().end());
EXPECT_THAT(init_node_itr->op(), StrEq("NoOp"));
ASSERT_THAT(init_node_itr->input(), SizeIs(1));
// "^" means control input.
EXPECT_THAT(init_node_itr->input()[0],
StrEq("^tf.InitializeTableFromTextFileV2"));
}
TEST_F(ConvertMlirModuleToExportedModelTest, InitNodeNotSetIfLocNameMismatch) {
// Define a module that initializes a tf.HashTableV2 whose control output node
// for the initialization is named "init_ok". Since the output control node
// name does not begin with "init_op" the init node could not have been found
// after the conversion.
mlir::OwningOpRef<mlir::ModuleOp> module_op = ParseModuleOpString(R"mlir(
module attributes {tf_saved_model.semantics} {
"tf_saved_model.session_initializer"() <{initializers = []}> : () -> ()
"tf_saved_model.asset"() <{filename = "assets/vocab_file.txt", sym_name = "__tf_saved_model_asset0_vocab_file.txt"}> : () -> ()
func.func @main(%arg1: tensor<!tf_type.string> {tf_saved_model.index_path = ["arg_0:0"]}) -> (tensor<1x2xf32> {tf_saved_model.index_path = ["output:0"]}) attributes {tf.entry_function = {inputs = "arg_0:0", outputs = "output:0"}, tf_saved_model.exported_names = ["main"]} {
%0 = tf_executor.graph {
%output_0, %control_0 = tf_executor.island wraps "tf.Const"() <{value = dense<1.0> : tensor<1x2xf32>}> : () -> tensor<1x2xf32>
%output_1, %control_1 = tf_executor.island wraps "tf.HashTableV2"() <{container = "", key_dtype = !tf_type.string, shared_name = "vocab_file.txt", use_node_name_sharing = false, value_dtype = i64}> {device = ""} : () -> tensor<!tf_type.resource>
%control_2 = tf_executor.island wraps "tf.InitializeTableFromTextFileV2"(%output_1, %arg1) <{delimiter = "\09", key_index = -2 : i64, value_index = -1 : i64, vocab_size = -1 : i64}> {_has_manual_control_dependencies = true, device = ""} : (tensor<!tf_type.resource>, tensor<!tf_type.string>) -> ()
// Location of this control output op becomes the name of the init_op.
%control_3 = tf_executor.island(%control_2) wraps "tf.NoOp"() : () -> () loc("init_ok")
tf_executor.fetch %output_0, %control_3 : tensor<1x2xf32>, !tf_executor.control
}
return %0 : tensor<1x2xf32>
}
}
)mlir");
ASSERT_TRUE(module_op);
const absl::StatusOr<ExportedModel> exported_model =
ConvertMlirModuleToExportedModel(*module_op, /*checkpoint_dir=*/"",
/*function_aliases=*/{},
/*asset_file_defs=*/{});
ASSERT_THAT(exported_model, absl_testing::IsOk());
EXPECT_THAT(exported_model->init_node_name(), IsEmpty());
}
TEST_F(ConvertMlirModuleToExportedModelTest,
ConversionFailureWhenNoMainFunction) {
// Define a module a function whose name is not @main.
mlir::OwningOpRef<mlir::ModuleOp> module_op = ParseModuleOpString(R"mlir(
module attributes {tf_saved_model.semantics} {
func.func @not_main() -> () attributes {tf_saved_model.exported_names = ["not_main"]} {
tf_executor.graph {
tf_executor.fetch
}
return
}
}
)mlir");
ASSERT_TRUE(module_op);
const absl::StatusOr<ExportedModel> exported_model =
ConvertMlirModuleToExportedModel(*module_op, "my_checkpoint_dir",
/*function_aliases=*/{},
/*asset_file_defs=*/{});
EXPECT_THAT(exported_model,
absl_testing::StatusIs(
absl::StatusCode::kFailedPrecondition,
HasSubstr("entry function `main` must be present")));
}
} // namespace
} // namespace mlir::quant::stablehlo
tensorflow/compiler/mlir/quantization/stablehlo/cc/saved_model_import.cc
+153
View File
@@ -0,0 +1,153 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/saved_model_import.h"
#include <memory>
#include <string>
#include <unordered_set>
#include <utility>
#include <vector>
#include "absl/algorithm/container.h"
#include "absl/base/attributes.h"
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#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/OwningOpRef.h" // from @llvm-project
#include "tensorflow/cc/saved_model/loader.h"
#include "tensorflow/cc/saved_model/reader.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/types.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantize_preprocess.h"
#include "tensorflow/compiler/mlir/tensorflow/translate/mlir_import_options.h"
#include "tensorflow/compiler/mlir/tensorflow/translate/tf_mlir_translate.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
namespace mlir::quant::stablehlo {
using ::stablehlo::quantization::QuantizationConfig;
using ::tensorflow::MLIRImportOptions;
using ::tensorflow::SavedModelBundle;
using ::tensorflow::SavedModelSignatureDefsToMlirImport;
using ::tensorflow::quantization::PreprocessAndFreezeGraph;
absl::StatusOr<ImportedMlirModuleOp> SavedModelToMlirModuleOp(
const absl::string_view saved_model_path,
const std::unordered_set<std::string>& tags,
const std::vector<std::string>& signature_keys,
MLIRContext& ctx ABSL_ATTRIBUTE_LIFETIME_BOUND) {
MLIRImportOptions import_options;
import_options.upgrade_legacy = true;
import_options.lift_variables = false;
import_options.include_variables_in_initializers = true;
auto bundle = std::make_unique<SavedModelBundle>();
// Copy to eliminate the `const` qualifier so that `absl::MakeSpan` can be
// called on it.
std::vector<std::string> exported_names = signature_keys;
absl::StatusOr<OwningOpRef<ModuleOp>> module_op =
SavedModelSignatureDefsToMlirImport(saved_model_path, tags,
absl::MakeSpan(exported_names), &ctx,
import_options, &bundle);
if (!module_op.status().ok()) {
return absl::InternalError(absl::StrCat("Failed to import SavedModel: ",
module_op.status().ToString()));
}
return std::make_pair(std::move(*module_op), std::move(bundle));
}
absl::StatusOr<absl::flat_hash_map<FunctionName, FunctionAlias>>
GetFunctionAliases(absl::string_view saved_model_path,
const std::unordered_set<std::string>& tags) {
tensorflow::MetaGraphDef meta_graph;
TF_RETURN_IF_ERROR(tensorflow::ReadMetaGraphDefFromSavedModel(
saved_model_path, tags, &meta_graph));
absl::flat_hash_map<FunctionName, FunctionAlias> function_aliases(
meta_graph.meta_info_def().function_aliases().begin(),
meta_graph.meta_info_def().function_aliases().end());
return function_aliases;
}
void UpdateFunctionAliases(
absl::flat_hash_map<FunctionName, FunctionAlias>& function_aliases,
ModuleOp module_op) {
absl::flat_hash_set<FunctionName> existing_func_names;
module_op->walk([&](func::FuncOp func_op) {
FunctionName func_name = func_op.getSymName().str();
existing_func_names.insert(func_name);
// We may retrieve the original function's name from the attribute.
// Functions without this attribute are ignored.
auto original_func_name =
func_op->getAttrOfType<StringAttr>("tf._original_func_name");
if (original_func_name) {
if (auto alias_itr = function_aliases.find(original_func_name.str());
alias_itr != function_aliases.end()) {
const FunctionAlias alias = alias_itr->second;
function_aliases[func_name] = alias;
}
}
});
// Remove aliases to function that no-longer exists.
absl::erase_if(function_aliases, [&existing_func_names](const auto& item) {
return !existing_func_names.contains(item.first);
});
}
absl::StatusOr<OwningOpRef<ModuleOp>> ImportSavedModel(
const absl::string_view saved_model_path,
const std::vector<std::string>& signature_keys,
const std::unordered_set<std::string>& tags,
const QuantizationConfig& quantization_config,
const absl::string_view mlir_dump_file_prefix,
absl::flat_hash_map<FunctionName, FunctionAlias>& function_aliases,
MLIRContext& ctx ABSL_ATTRIBUTE_LIFETIME_BOUND) {
TF_ASSIGN_OR_RETURN(
ImportedMlirModuleOp imported_module,
SavedModelToMlirModuleOp(saved_model_path, tags, signature_keys, ctx));
auto [module_op, saved_model_bundle] = std::move(imported_module);
UpdateFunctionAliases(function_aliases, *module_op);
// Collect the names of the functions that have aliases so that they may not
// be inlined.
absl::flat_hash_set<std::string> aliased_function_names;
absl::c_for_each(function_aliases, [&](const auto& aliases) {
return aliased_function_names.insert(aliases.first);
});
TF_RETURN_IF_ERROR(PreprocessAndFreezeGraph(
mlir_dump_file_prefix, /*is_inliner_run=*/true,
/*noinline_functions=*/aliased_function_names, *module_op, &ctx,
saved_model_bundle == nullptr ? nullptr
: saved_model_bundle->GetSession(),
/*run_tf_to_stablehlo=*/true, /*deserialize_xla_call_module=*/false));
return std::move(module_op);
}
} // namespace mlir::quant::stablehlo
tensorflow/compiler/mlir/quantization/stablehlo/cc/saved_model_import.h
+90
View File
@@ -0,0 +1,90 @@
/* 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.
==============================================================================*/
// Functionalities for importing MLIR ModuleOp from TensorFlow SavedModel.
#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_SAVED_MODEL_IMPORT_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_SAVED_MODEL_IMPORT_H_
#include <memory>
#include <string>
#include <unordered_set>
#include <utility>
#include <vector>
#include "absl/base/attributes.h"
#include "absl/container/flat_hash_map.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#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 "tensorflow/cc/saved_model/loader.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/types.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
namespace mlir::quant::stablehlo {
// Represents a pair of `mlir::ModuleOp` and `tensorflow::SavedModelBundle`. The
// SavedModelBundle complements the imported ModuleOp by providing access to
// `tensorflow::Session` which may be useful when reading values from resources
// (e.g. `TF::VarHandleOp`s).
using ImportedMlirModuleOp =
std::pair<OwningOpRef<ModuleOp>,
std::unique_ptr<::tensorflow::SavedModelBundle>>;
// Loads a SavedModel at `saved_model_path` and converts it to `mlir::ModuleOp`.
//
// `tags` identify the `tensorflow::MetaGraphDef` to load from the SavedModel.
// Similarly, `signature_keys` identify the functions (`SignatureDef`s) to load
// within the `MetaGraphDef`. `ctx` is the `MLIRContext`, which should outlive
// the returned `ModuleOp`, thus marked with the lifetime bound attribute.
// TODO: b/329206105 - Add unit tests after decomposing preprocessing passes.
absl::StatusOr<ImportedMlirModuleOp> SavedModelToMlirModuleOp(
absl::string_view saved_model_path,
const std::unordered_set<std::string>& tags,
const std::vector<std::string>& signature_keys,
MLIRContext& ctx ABSL_ATTRIBUTE_LIFETIME_BOUND);
// Gets the function aliases from the SavedModel.
absl::StatusOr<absl::flat_hash_map<FunctionName, FunctionAlias>>
GetFunctionAliases(absl::string_view saved_model_path,
const std::unordered_set<std::string>& tags);
// Updates the function aliases. `module_op` may have different
// function names from the original model, so it re-associates the aliases
// with the new function names. Both the input `function_aliases` and the
// returned value are function name -> alias mappings. `function_aliases` is
// the function alias mapping of the original function. The original function's
// name is retrieved by looking at the "tf._original_func_name" string attribute
// attached to a `func::FuncOp`.
void UpdateFunctionAliases(
absl::flat_hash_map<FunctionName, FunctionAlias>& function_aliases,
ModuleOp module_op);
// Loads a SavedModel to `mlir::ModuleOp` and performs preprocesses including
// shape inference and graph freezing.
// TODO: b/329206105 - Add unit tests after decomposing preprocessing passes.
absl::StatusOr<OwningOpRef<ModuleOp>> ImportSavedModel(
absl::string_view saved_model_path,
const std::vector<std::string>& signature_keys,
const std::unordered_set<std::string>& tags,
const ::stablehlo::quantization::QuantizationConfig& quantization_config,
absl::string_view mlir_dump_file_prefix,
absl::flat_hash_map<FunctionName, FunctionAlias>& function_aliases,
MLIRContext& ctx ABSL_ATTRIBUTE_LIFETIME_BOUND);
} // namespace mlir::quant::stablehlo
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_SAVED_MODEL_IMPORT_H_
tensorflow/compiler/mlir/quantization/stablehlo/cc/saved_model_import_test.cc
+120
View File
@@ -0,0 +1,120 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/saved_model_import.h"
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/container/flat_hash_map.h"
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/common/test_base.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/types.h"
namespace mlir::quant::stablehlo {
namespace {
using ::testing::IsEmpty;
using ::testing::Pair;
using ::testing::UnorderedElementsAre;
using UpdateFunctionAliasesTest = ::mlir::quant::QuantizationTestBase;
TEST_F(UpdateFunctionAliasesTest, NoAliasesReturnsEmptyMap) {
// MLIR @main function corresponds to the TF function "main_original".
OwningOpRef<ModuleOp> module_op = ParseModuleOpString(R"mlir(
func.func private @main(%arg: tensor<1x2xf32>) -> (tensor<1x2xf32>) attributes {tf._original_func_name = "main_original"} {
return %arg : tensor<1x2xf32>
}
)mlir");
ASSERT_TRUE(module_op);
absl::flat_hash_map<FunctionName, FunctionAlias> function_aliases;
UpdateFunctionAliases(function_aliases, *module_op);
EXPECT_THAT(function_aliases, IsEmpty());
}
TEST_F(UpdateFunctionAliasesTest, AliasUpdatedByMlirFunctionName) {
// MLIR @main function corresponds to the TF function "main_original".
OwningOpRef<ModuleOp> module_op = ParseModuleOpString(R"mlir(
func.func private @main(%arg: tensor<1x2xf32>) -> (tensor<1x2xf32>) attributes {tf._original_func_name = "main_original"} {
return %arg : tensor<1x2xf32>
}
)mlir");
ASSERT_TRUE(module_op);
absl::flat_hash_map<FunctionName, FunctionAlias> function_aliases{
{"main_original", "main_alias"}};
UpdateFunctionAliases(function_aliases, *module_op);
EXPECT_THAT(function_aliases,
UnorderedElementsAre(Pair("main", "main_alias")));
}
TEST_F(UpdateFunctionAliasesTest, IgnoresUnmatchedFunctions) {
// MLIR @main function corresponds to the TF function "main_original".
OwningOpRef<ModuleOp> module_op = ParseModuleOpString(R"mlir(
func.func private @main(%arg: tensor<1x2xf32>) -> (tensor<1x2xf32>) attributes {tf._original_func_name = "main_original"} {
return %arg : tensor<1x2xf32>
}
)mlir");
ASSERT_TRUE(module_op);
// There is no alias corresponding to "main_original". The existing entry
// without a corresponding function is ignored.
absl::flat_hash_map<FunctionName, FunctionAlias> function_aliases{
{"not_main", "not_main_alias"}};
UpdateFunctionAliases(function_aliases, *module_op);
EXPECT_THAT(function_aliases, IsEmpty());
}
TEST_F(UpdateFunctionAliasesTest,
SkipsFunctionsWithNoOriginalFuncNameAttribute) {
// @main does not have the "tf._original_func_name" attribute.
OwningOpRef<ModuleOp> module_op = ParseModuleOpString(R"mlir(
func.func private @main(%arg: tensor<1x2xf32>) -> (tensor<1x2xf32>) {
return %arg : tensor<1x2xf32>
}
)mlir");
ASSERT_TRUE(module_op);
// The existing entry without a corresponding function is ignored.
absl::flat_hash_map<FunctionName, FunctionAlias> function_aliases{
{"main_original", "main_alias"}};
UpdateFunctionAliases(function_aliases, *module_op);
EXPECT_THAT(function_aliases, IsEmpty());
}
TEST_F(UpdateFunctionAliasesTest, FunctionNameNotChanged) {
// @main does not have the "tf._original_func_name" attribute.
OwningOpRef<ModuleOp> module_op = ParseModuleOpString(R"mlir(
func.func private @main_original(%arg: tensor<1x2xf32>) -> (tensor<1x2xf32>) {
return %arg : tensor<1x2xf32>
}
)mlir");
ASSERT_TRUE(module_op);
// The existing entry without a corresponding function is ignored.
absl::flat_hash_map<FunctionName, FunctionAlias> function_aliases{
{"main_original", "main_alias"}};
UpdateFunctionAliases(function_aliases, *module_op);
EXPECT_THAT(function_aliases,
UnorderedElementsAre(Pair("main_original", "main_alias")));
}
} // namespace
} // namespace mlir::quant::stablehlo
tensorflow/compiler/mlir/quantization/stablehlo/cc/static_range_ptq.cc
+137
View File
@@ -0,0 +1,137 @@
/* 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/stablehlo/cc/static_range_ptq.h"
#include <memory>
#include <string>
#include <unordered_set>
#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/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#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/Pass/PassManager.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/component.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/component.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/context.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/post_calibration.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/pre_calibration.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/saved_model_export.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/saved_model_import.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/types.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/exported_model.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include "tensorflow/core/protobuf/saver.pb.h"
namespace mlir::quant::stablehlo {
using ::stablehlo::quantization::QuantizationConfig;
using ::tensorflow::SignatureDef;
using ::tensorflow::quantization::ExportedModel;
using ::tensorflow::quantization::PyFunctionLibrary;
StaticRangePtqComponent::StaticRangePtqComponent(
MLIRContext* absl_nonnull ctx,
const PyFunctionLibrary* absl_nonnull py_function_library,
const absl::string_view src_saved_model_path,
std::vector<std::string> signature_keys,
std::unordered_set<std::string> tags,
absl::flat_hash_map<std::string, SignatureDef> signature_def_map,
absl::flat_hash_map<FunctionName, FunctionAlias> function_aliases)
: ctx_(ctx) {
// Initialize the three sub-components.
sub_components_[0] = std::make_unique<PreCalibrationComponent>(ctx_);
sub_components_[1] = std::make_unique<CalibrationComponent>(
ctx_, py_function_library, src_saved_model_path,
std::move(function_aliases), std::move(tags),
std::move(signature_def_map), std::move(signature_keys));
sub_components_[2] = std::make_unique<PostCalibrationComponent>(ctx_);
}
absl::StatusOr<ModuleOp> StaticRangePtqComponent::Run(
ModuleOp module_op, const QuantizationConfig& config) {
// Runs sub-components in sequence: PreCalibrationComponent ->
// CalibrationComponent -> PostCalibrationComponents.
for (std::unique_ptr<Component>& sub_component : sub_components_) {
TF_ASSIGN_OR_RETURN(module_op, sub_component->Run(module_op, config));
}
return module_op;
}
// TODO: b/317167427 - Enable debugger.
absl::Status QuantizeStaticRangePtq(
const absl::string_view src_saved_model_path,
const absl::string_view dst_saved_model_path,
const QuantizationConfig& quantization_config,
const std::vector<std::string>& signature_keys,
const absl::flat_hash_map<std::string, SignatureDef>& signature_def_map,
const PyFunctionLibrary& py_function_library) {
std::unordered_set<std::string> tags;
tags.insert(quantization_config.tf_saved_model().tags().begin(),
quantization_config.tf_saved_model().tags().end());
std::unique_ptr<MLIRContext> ctx = CreateMlirContextForQuantization();
absl::StatusOr<absl::flat_hash_map<FunctionName, FunctionAlias>>
function_aliases = GetFunctionAliases(src_saved_model_path, tags);
if (!function_aliases.ok()) {
return absl::InternalError(absl::StrCat(
"Failed to get function alias: ", function_aliases.status().message()));
}
TF_ASSIGN_OR_RETURN(
OwningOpRef<ModuleOp> module,
ImportSavedModel(src_saved_model_path, signature_keys, tags,
quantization_config, PreCalibrationComponent::kName,
*function_aliases, *ctx));
StaticRangePtqComponent static_range_ptq_component(
ctx.get(), &py_function_library, src_saved_model_path, signature_keys,
tags, signature_def_map, *function_aliases);
TF_ASSIGN_OR_RETURN(
*module, static_range_ptq_component.Run(*module, quantization_config));
TF_ASSIGN_OR_RETURN(
const ExportedModel post_calibrated_exported_model,
CreateExportedModel(signature_keys, tags, quantization_config,
PostCalibrationComponent::kName, *function_aliases,
*ctx, *module));
// Remove the `tpu` tag for exporting because the output quantized model is
// essentially a CPU model.
tags.erase("tpu");
py_function_library.SaveExportedModel(
dst_saved_model_path, post_calibrated_exported_model,
src_saved_model_path, tags, signature_def_map);
return absl::OkStatus();
}
} // namespace mlir::quant::stablehlo
tensorflow/compiler/mlir/quantization/stablehlo/cc/static_range_ptq.h
+103
View File
@@ -0,0 +1,103 @@
/* 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_STABLEHLO_CC_STATIC_RANGE_PTQ_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_STATIC_RANGE_PTQ_H_
#include <array>
#include <memory>
#include <string>
#include <unordered_set>
#include <vector>
#include "absl/base/nullability.h"
#include "absl/container/flat_hash_map.h"
#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/IR/MLIRContext.h" // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/component.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/types.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/exported_model.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
namespace mlir::quant::stablehlo {
// Component for static-range post-training quantization (PTQ).
// TODO: b/320607042 - Add tests in python level.
class StaticRangePtqComponent : public Component {
public:
// Name of this component. Used for debugging purposes.
static constexpr absl::string_view kName = "quant_static_range_ptq";
// Constructs `StaticRangePtqComponent` by creating three sub-components:
// `PreCalibrationComponent`, `CalibrationComponent`, and
// `PostCalibrationComponent`. These are stored in `sub_components_` in
// sequence. All arguments except `ctx` is used to initialize
// `CalibrationComponent`. For detailed explanation of each argument, see the
// comment of `CalibrationComponent`'s constructor.
StaticRangePtqComponent(
MLIRContext* absl_nonnull ctx,
const tensorflow::quantization::PyFunctionLibrary* absl_nonnull
py_function_library,
absl::string_view src_saved_model_path,
std::vector<std::string> signature_keys,
std::unordered_set<std::string> tags,
absl::flat_hash_map<std::string, tensorflow::SignatureDef>
signature_def_map,
absl::flat_hash_map<FunctionName, FunctionAlias> function_aliases);
// Runs the static-range post-training quantization (PTQ) on `module_op`.
absl::StatusOr<ModuleOp> Run(
ModuleOp module_op,
const ::stablehlo::quantization::QuantizationConfig& config) override;
private:
// A non-owning `MLIRContext`. This `MLIRContext` should exceed the lifetime
// of `StaticRangePtqComponent`.
MLIRContext* absl_nonnull ctx_;
// This component consists of three sub-components, `PreCalibrationComponent`,
// `CalibrationComponent`, and `PostCalibrationComponent`.
std::array<std::unique_ptr<Component>, 3> sub_components_;
};
// Runs static-range post-training quantization (PTQ) on a SavedModel at
// `src_saved_model_path` and saves the resulting model to
// `dst_saved_model_path`.
//
// `quantization_config` configures the quantization behavior for the
// static-range PTQ.
//
// `signature_keys` specify the signatures that correspond to functions to be
// quantized. `signature_def_map` connects the signature keys to
// `SignatureDef`s.
//
// Returns a non-OK status when the quantization is not successful.
// LINT.IfChange
absl::Status QuantizeStaticRangePtq(
absl::string_view src_saved_model_path,
absl::string_view dst_saved_model_path,
const ::stablehlo::quantization::QuantizationConfig& quantization_config,
const std::vector<std::string>& signature_keys,
const absl::flat_hash_map<std::string, tensorflow::SignatureDef>&
signature_def_map,
const tensorflow::quantization::PyFunctionLibrary& py_function_library);
// LINT.ThenChange(../python/pywrap_quantization.cc:static_range_ptq)
} // namespace mlir::quant::stablehlo
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_STATIC_RANGE_PTQ_H_
tensorflow/compiler/mlir/quantization/stablehlo/cc/types.h
+31
View File
@@ -0,0 +1,31 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_TYPES_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_TYPES_H_
#include <string>
namespace mlir::quant::stablehlo {
// Introduces aliases for `std::string` to distinguish btw. function name and
// its alias, to prevent confusion when used together in a container. For
// example, it is easy to confuse function name -> alias mapping with alias ->
// function name mapping when both are just represented as `std::string`.
using FunctionAlias = std::string;
using FunctionName = std::string;
} // namespace mlir::quant::stablehlo
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_TYPES_H_
tensorflow/compiler/mlir/quantization/stablehlo/cc/weight_only_ptq.cc
+125
View File
@@ -0,0 +1,125 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/weight_only_ptq.h"
#include <memory>
#include <string>
#include <unordered_set>
#include <vector>
#include "absl/base/nullability.h"
#include "absl/container/flat_hash_map.h"
#include "absl/log/die_if_null.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.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/stablehlo/cc/config.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/context.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/pass_pipeline.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/saved_model_export.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/saved_model_import.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/types.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/instrumentations/save_report.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/cc/run_passes.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/exported_model.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h"
#include "xla/mlir_hlo/mhlo/transforms/passes.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
namespace mlir::quant::stablehlo {
using ::mlir::quant::stablehlo::AddWeightOnlyQuantizationPasses;
using ::stablehlo::quantization::GetReportFilePath;
using ::stablehlo::quantization::QuantizationConfig;
using ::tensorflow::SignatureDef;
using ::tensorflow::quantization::ExportedModel;
using ::tensorflow::quantization::PyFunctionLibrary;
using ::tensorflow::quantization::RunPasses;
WeightOnlyPtqComponent::WeightOnlyPtqComponent(MLIRContext* absl_nonnull ctx)
: ctx_(ABSL_DIE_IF_NULL(ctx)) {} // Crash OK
absl::StatusOr<ModuleOp> WeightOnlyPtqComponent::Run(
ModuleOp module_op, const QuantizationConfig& config) {
TF_RETURN_IF_ERROR(RunPasses(
kName, /*add_passes_func=*/
[&config](PassManager& pm) {
// Add instrumentation to save quantization report after quantization.
pm.addInstrumentation(
std::make_unique<SaveQuantizationReportInstrumentation>(
GetReportFilePath(config)));
AddWeightOnlyQuantizationPasses(pm, config.specs(),
config.pipeline_config(),
config.debugger_config());
},
*ctx_, module_op));
return module_op;
}
absl::Status QuantizeWeightOnlyPtq(
const absl::string_view src_saved_model_path,
const absl::string_view dst_saved_model_path,
QuantizationConfig quantization_config,
const std::vector<std::string>& signature_keys,
const absl::flat_hash_map<std::string, SignatureDef>& signature_def_map,
const PyFunctionLibrary& py_function_library) {
std::unordered_set<std::string> tags;
tags.insert(quantization_config.tf_saved_model().tags().begin(),
quantization_config.tf_saved_model().tags().end());
std::unique_ptr<MLIRContext> ctx = CreateMlirContextForQuantization();
absl::StatusOr<absl::flat_hash_map<FunctionName, FunctionAlias>>
function_aliases = GetFunctionAliases(src_saved_model_path, tags);
if (!function_aliases.ok()) {
return absl::InternalError(absl::StrCat(
"Failed to get function alias: ", function_aliases.status().message()));
}
TF_ASSIGN_OR_RETURN(
auto module,
ImportSavedModel(src_saved_model_path, signature_keys, tags,
quantization_config, WeightOnlyPtqComponent::kName,
*function_aliases, *ctx));
WeightOnlyPtqComponent weight_only_ptq_component(ctx.get());
TF_ASSIGN_OR_RETURN(
*module, weight_only_ptq_component.Run(*module, quantization_config));
TF_ASSIGN_OR_RETURN(
const ExportedModel post_calibrated_exported_model,
CreateExportedModel(signature_keys, tags, quantization_config,
WeightOnlyPtqComponent::kName, *function_aliases,
*ctx, *module));
// Remove the `tpu` tag for exporting because the output quantized model is
// essentially a CPU model.
tags.erase("tpu");
py_function_library.SaveExportedModel(
dst_saved_model_path, post_calibrated_exported_model,
src_saved_model_path, tags, signature_def_map);
return absl::OkStatus();
}
} // namespace mlir::quant::stablehlo
tensorflow/compiler/mlir/quantization/stablehlo/cc/weight_only_ptq.h
+80
View File
@@ -0,0 +1,80 @@
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_WEIGHT_ONLY_PTQ_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_WEIGHT_ONLY_PTQ_H_
#include <string>
#include <vector>
#include "absl/base/nullability.h"
#include "absl/container/flat_hash_map.h"
#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/stablehlo/cc/component.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
namespace mlir::quant::stablehlo {
// Performs int8 weight-only quantization on dot_general ops.
//
// The resulting `ModuleOp` contains quantized StableHLO ops serialized in
// `TF::XlaCallModuleOp`s. They are quantized using the weight constants, not
// relying on calibration.
class WeightOnlyPtqComponent : public Component {
public:
// Used for debugging purposes.
static constexpr absl::string_view kName = "quant_ptq_weight_only";
explicit WeightOnlyPtqComponent(MLIRContext* absl_nonnull ctx);
absl::StatusOr<ModuleOp> Run(
ModuleOp module_op,
const ::stablehlo::quantization::QuantizationConfig& config) override;
private:
MLIRContext* absl_nonnull ctx_;
};
// Runs weight-only quantization on a SavedModel at
// `src_saved_model_path` and saves the resulting model to
// `dst_saved_model_path`.
//
// `quantization_config` configures the quantization behavior for the
// weight-only quantization.
//
// `signature_keys` specify the signatures that correspond to functions to be
// quantized. `signature_def_map` connects the signature keys to
// `SignatureDef`s.
//
// Returns a non-OK status when the quantization is not successful.
// LINT.IfChange
absl::Status QuantizeWeightOnlyPtq(
absl::string_view src_saved_model_path,
absl::string_view dst_saved_model_path,
::stablehlo::quantization::QuantizationConfig quantization_config,
const std::vector<std::string>& signature_keys,
const absl::flat_hash_map<std::string, tensorflow::SignatureDef>&
signature_def_map,
const tensorflow::quantization::PyFunctionLibrary& py_function_library);
// LINT.ThenChange(../python/pywrap_quantization.cc:weight_only_ptq)
} // namespace mlir::quant::stablehlo
#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_CC_WEIGHT_ONLY_PTQ_H_