chore: import upstream snapshot with attribution

tensorflow/compiler/mlir/quantization/stablehlo/BUILD

@@ -0,0 +1,735 @@
+load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library")
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@rules_python//python:proto.bzl", "py_proto_library")
+load("//tensorflow:tensorflow.bzl", "tf_cc_binary", "tf_cc_test")
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
+load("//tensorflow/compiler/mlir/quantization/stablehlo:internal_visibility_allowlist.bzl", "internal_visibility_allowlist")
+load("//tensorflow/core/platform:build_config.bzl", "tf_proto_library")
+
+package_group(
+    name = "internal_visibility_allowlist_package",
+    packages = [
+        "//learning/brain/mlir/quantization/stablehlo/python/integration_test/...",
+        "//tensorflow/compiler/mlir/lite/...",
+        "//tensorflow/compiler/mlir/quantization/...",
+        "//tensorflow/compiler/mlir/tf2xla/transforms/...",
+        "//tensorflow/lite/...",
+        "//third_party/cloud_tpu/inference_converter/...",  # TPU Inference Converter V1
+        "//waymo/ml/compiler/mlir/...",
+    ] + internal_visibility_allowlist(),
+)
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["@stablehlo//:license"],
+    default_visibility = [
+        ":internal_visibility_allowlist_package",
+        "//tensorflow:__pkg__",
+        "//third_party/py/orbax:__subpackages__",
+    ],
+    licenses = ["notice"],
+)
+
+gentbl_cc_library(
+    name = "stablehlo_passes_inc_gen",
+    compatible_with = get_compatible_with_portable(),
+    tbl_outs = {"passes/passes.h.inc": [
+        "-gen-pass-decls",
+    ]},
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "passes/passes.td",
+    deps = [
+        "@llvm-project//mlir:PassBaseTdFiles",
+    ],
+)
+
+cc_library(
+    name = "passes",
+    srcs = [
+        "passes/convert_func_to_bfloat16.cc",
+        "passes/convert_shape_constraint_to_assert.cc",
+        "passes/convert_xla_call_module_op_to_bfloat16.cc",
+        "passes/defer_activation_transpose.cc",
+        "passes/fold_constant_transpose.cc",
+        "passes/insert_calibration_statistics_saver.cc",
+        "passes/insert_weight_param.cc",
+        "passes/lift_quantizable_spots_as_functions.cc",
+        "passes/lift_quantizable_spots_as_functions_fusion.inc",
+        "passes/lift_quantizable_spots_as_functions_simple.inc",
+        "passes/merge_fusion_with_dequantize.cc",
+        "passes/nchw_convolution_to_nhwc.cc",
+        "passes/optimize_graph.cc",
+        "passes/post_quantize.cc",
+        "passes/prepare_quantize.cc",
+        "passes/quantize.cc",
+        "passes/quantize_composite_functions.cc",
+        "passes/quantize_weight.cc",
+        "passes/remove_sharding_custom_call.cc",
+        "passes/remove_sharding_custom_call.inc",
+        "passes/replace_stablehlo_ops_in_main_function_with_xla_call_module_ops.cc",
+        "passes/restore_function_name.cc",
+        "passes/unfuse_mhlo_batch_norm.cc",
+        "passes/unwrap_xla_call_module_op.cc",
+        "passes/xla_call_module_to_call.cc",
+    ],
+    hdrs = [
+        "passes/passes.h",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        ":bfloat16_type",
+        ":fill_quantization_options",
+        ":lift_quantizable_spots_as_functions_fusion_inc_gen",
+        ":lift_quantizable_spots_as_functions_simple_inc_gen",
+        ":optimize_graph_inc_gen",
+        ":quantization_config_proto_cc",
+        ":quantization_options_proto_cc",
+        ":quantization_patterns",
+        ":remove_sharding_custom_call_inc_gen",
+        ":stablehlo_passes_inc_gen",
+        ":stablehlo_type_utils",
+        "//tensorflow/compiler/mlir/quantization/common:attrs_and_constraints",
+        "//tensorflow/compiler/mlir/quantization/common:func",
+        "//tensorflow/compiler/mlir/quantization/common:lift_as_function_call",
+        "//tensorflow/compiler/mlir/quantization/common:uniform_quantized_types",
+        "//tensorflow/compiler/mlir/quantization/common/ir:QuantOps",
+        "//tensorflow/compiler/mlir/quantization/common/quantization_lib",
+        "//tensorflow/compiler/mlir/quantization/common/quantization_lib:quantization_config",
+        "//tensorflow/compiler/mlir/quantization/stablehlo/cc:permutation",
+        "//tensorflow/compiler/mlir/quantization/stablehlo/ops:stablehlo_op_quant_spec",
+        "//tensorflow/compiler/mlir/quantization/tensorflow:quantization_options_proto_cc",
+        "//tensorflow/compiler/mlir/quantization/tensorflow:tf_quant_ops",
+        "//tensorflow/compiler/mlir/quantization/tensorflow/cc:run_passes",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow:xla_call_module_attrs",
+        "//tensorflow/core:portable_gif_internal",
+        "//tensorflow/core/ir/types:Dialect",
+        "@com_google_absl//absl/base:nullability",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings:string_view",
+        "@eigen_archive//:eigen3",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:QuantOps",
+        "@llvm-project//mlir:Rewrite",
+        "@llvm-project//mlir:ShapeDialect",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TensorDialect",
+        "@llvm-project//mlir:TransformUtils",
+        "@shardy//shardy/dialect/sdy/ir:register",
+        "@stablehlo//:chlo_ops",
+        "@stablehlo//:stablehlo_ops",
+        "@stablehlo//:stablehlo_passes",
+        "@stablehlo//:stablehlo_portable_api",
+        "@stablehlo//:stablehlo_serialization",
+        "@stablehlo//:version",
+        "@tsl//tsl/platform:path",
+        "@tsl//tsl/platform:protobuf",
+        "@tsl//tsl/platform:regexp",
+        "@xla//xla/mlir_hlo",
+        "@xla//xla/mlir_hlo:mhlo_passes",
+    ],
+)
+
+cc_library(
+    name = "quantization_patterns",
+    srcs = ["passes/quantization_patterns.cc"],
+    hdrs = [
+        "passes/quantization_patterns.h",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        ":quantization_config_proto_cc",
+        "//tensorflow/compiler/mlir/quantization/common:attrs_and_constraints",
+        "//tensorflow/compiler/mlir/quantization/common:lift_as_function_call",
+        "//tensorflow/compiler/mlir/quantization/common:uniform_quantized_types",
+        "//tensorflow/compiler/mlir/quantization/common/ir:QuantOps",
+        "//tensorflow/compiler/mlir/quantization/common/quantization_lib",
+        "//tensorflow/compiler/mlir/quantization/stablehlo/ops:stablehlo_op_quant_spec",
+        "//tensorflow/compiler/mlir/quantization/tensorflow:quantization_options_proto_cc",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/platform:path",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:QuantOps",
+        "@llvm-project//mlir:Support",
+        "@stablehlo//:stablehlo_ops",
+    ],
+)
+
+td_library(
+    name = "quant_td_files",
+    srcs = [
+        "passes/lift_quantizable_spots_as_functions_fusion.td",
+        "passes/lift_quantizable_spots_as_functions_simple.td",
+        "passes/optimize_graph.td",
+        "passes/remove_sharding_custom_call.td",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        "//tensorflow/compiler/mlir/quantization/common:lift_as_function_call_td_files",
+        "//tensorflow/compiler/mlir/quantization/tensorflow:quant_td_files",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_ops_td_files",
+        "@llvm-project//mlir:ArithOpsTdFiles",
+        "@llvm-project//mlir:FuncTdFiles",
+        "@stablehlo//:stablehlo_ops_td_files",
+    ],
+)
+
+gentbl_cc_library(
+    name = "lift_quantizable_spots_as_functions_simple_inc_gen",
+    compatible_with = get_compatible_with_portable(),
+    tbl_outs = {"passes/lift_quantizable_spots_as_functions_simple.inc": ["-gen-rewriters"]},
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "passes/lift_quantizable_spots_as_functions_simple.td",
+    deps = [
+        ":quant_td_files",
+        "//tensorflow/compiler/mlir/quantization/common:quant_td_files",
+    ],
+)
+
+gentbl_cc_library(
+    name = "lift_quantizable_spots_as_functions_fusion_inc_gen",
+    compatible_with = get_compatible_with_portable(),
+    tbl_outs = {"passes/lift_quantizable_spots_as_functions_fusion.inc": ["-gen-rewriters"]},
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "passes/lift_quantizable_spots_as_functions_fusion.td",
+    deps = [
+        ":quant_td_files",
+        "//tensorflow/compiler/mlir/quantization/common:quant_td_files",
+    ],
+)
+
+gentbl_cc_library(
+    name = "optimize_graph_inc_gen",
+    compatible_with = get_compatible_with_portable(),
+    tbl_outs = {"passes/optimize_graph.inc": ["-gen-rewriters"]},
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "passes/optimize_graph.td",
+    deps = [
+        ":quant_td_files",
+        "//tensorflow/compiler/mlir/quantization/common:quant_td_files",
+    ],
+)
+
+gentbl_cc_library(
+    name = "remove_sharding_custom_call_inc_gen",
+    compatible_with = get_compatible_with_portable(),
+    tbl_outs = {"passes/remove_sharding_custom_call.inc": ["-gen-rewriters"]},
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "passes/remove_sharding_custom_call.td",
+    deps = [
+        ":quant_td_files",
+    ],
+)
+
+gentbl_cc_library(
+    name = "bridge_passes_inc_gen",
+    compatible_with = get_compatible_with_portable(),
+    tbl_outs = {"passes/bridge/passes.h.inc": [
+        "-gen-pass-decls",
+        "-name=Bridge",
+    ]},
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "passes/bridge/passes.td",
+    deps = [
+        "@llvm-project//mlir:PassBaseTdFiles",
+    ],
+)
+
+cc_library(
+    name = "bridge_passes",
+    srcs = [
+        "passes/bridge/convert_tf_quant_ops_to_mhlo.cc",
+        "passes/bridge/convert_tf_quant_types.cc",
+        "passes/bridge/optimize.cc",
+        "passes/bridge/passes.cc",
+        "passes/bridge/verify_quant_legalization.cc",
+    ],
+    hdrs = [
+        "passes/bridge/passes.h",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    visibility = ["//visibility:public"],
+    deps = [
+        ":bridge_passes_inc_gen",
+        ":optimize_inc_gen",
+        ":tf_type_utils",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow:mangling_util",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_attributes",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
+        "//tensorflow/compiler/mlir/tf2xla/transforms:legalize_utils",
+        "//tensorflow/compiler/mlir/tf2xla/transforms:xla_legalize_targets",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/framework:numeric_types",
+        "//tensorflow/core/util/quantization:uniform_quant_ops_attr_proto_cc",
+        "//tensorflow/core/util/quantization:uniform_quant_ops_params",
+        "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/strings",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:FuncExtensions",
+        "@llvm-project//mlir:FuncTransforms",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:QuantOps",
+        "@llvm-project//mlir:ShapeDialect",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+        "@llvm-project//mlir:Transforms",
+        "@stablehlo//:chlo_ops",
+        "@stablehlo//:stablehlo_passes",
+        "@xla//xla:xla_data_proto_cc",
+        "@xla//xla/hlo/translate/hlo_to_mhlo:attribute_importer",
+        "@xla//xla/mlir_hlo",
+        "@xla//xla/mlir_hlo:mhlo_passes",
+    ],
+)
+
+td_library(
+    name = "optimize_td_files",
+    srcs = [
+        "passes/bridge/optimize.td",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_ops_td_files",
+        "@llvm-project//mlir:ArithOpsTdFiles",
+        "@llvm-project//mlir:FuncTdFiles",
+        "@xla//xla/mlir_hlo:hlo_ops_td_files",
+    ],
+)
+
+gentbl_cc_library(
+    name = "optimize_inc_gen",
+    compatible_with = get_compatible_with_portable(),
+    tbl_outs = {"passes/bridge/optimize.inc": ["-gen-rewriters"]},
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "passes/bridge/optimize.td",
+    deps = [":optimize_td_files"],
+)
+
+tf_cc_test(
+    name = "convert_tf_quant_to_mhlo_int_test",
+    srcs = [
+        "passes/bridge/convert_tf_quant_to_mhlo_int_test.cc",
+    ],
+    tags = ["nomac"],  # TODO(b/297362678): re-enable mac test.
+    deps = [
+        ":bridge_passes",
+        "//tensorflow/compiler/mlir/quantization/tensorflow/cc:constant_fold",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow:convert_tensor",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
+        "//tensorflow/compiler/tf2xla:common",
+        "//tensorflow/core:framework",
+        "//tensorflow/core/kernels:math",
+        "//tensorflow/core/kernels:nn",
+        "//tensorflow/core/kernels/uniform_quant_ops:kernels",
+        "//tensorflow/core/ops",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/random",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Parser",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:QuantOps",
+        "@llvm-project//mlir:Support",
+        "@stablehlo//:chlo_ops",
+        "@xla//xla:error_spec",
+        "@xla//xla:literal",
+        "@xla//xla:literal_util",
+        "@xla//xla:shape_util",
+        "@xla//xla:xla_data_proto_cc",
+        "@xla//xla/mlir_hlo",
+        "@xla//xla/pjrt:maybe_owning_mlir_module",
+        "@xla//xla/pjrt:pjrt_client",
+        "@xla//xla/pjrt:pjrt_executable",
+        "@xla//xla/pjrt/plugin/xla_cpu:cpu_client_options",
+        "@xla//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client",
+        "@xla//xla/tests:literal_test_util",
+        "@xla//xla/tsl/platform:errors",
+        "@xla//xla/tsl/platform:statusor",
+    ],
+)
+
+tf_cc_test(
+    name = "legalize_tf_quant_test",
+    srcs = ["passes/bridge/legalize_tf_quant_test.cc"],
+    # TODO(b/299290907): Fix OSS test.
+    tags = ["no_oss"],
+    deps = [
+        "//tensorflow/compiler/jit",
+        "//tensorflow/compiler/mlir/tf2xla/api/v2:legalize_tf",
+        "//tensorflow/compiler/tf2xla:xla_helpers",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core/protobuf:for_core_protos_cc",
+        "//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
+        "//tensorflow/core/tpu/kernels:tpu_compile_op_support",
+        "@com_google_googletest//:gtest",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:Pass",
+        "@xla//xla:shape_util",
+        "@xla//xla/client:client_library",
+        "@xla//xla/stream_executor:platform",
+        "@xla//xla/stream_executor:platform_manager",
+    ],
+)
+
+tf_cc_test(
+    name = "convert_tf_quant_types_test",
+    srcs = ["passes/bridge/convert_tf_quant_types_test.cc"],
+    deps = [
+        ":bridge_passes",
+        "//tensorflow/compiler/mlir:register_common_dialects",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow:deserialize_mlir_module_utils",
+        "//tensorflow/core:test",
+        "//tensorflow/core/lib/monitoring:cell_reader",
+        "@com_google_absl//absl/strings",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@xla//xla/tsl/platform:statusor",
+    ],
+)
+
+cc_library(
+    name = "quantize_passes",
+    srcs = [
+        "quantize_passes.cc",
+    ],
+    hdrs = [
+        "quantize_passes.h",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    visibility = ["//visibility:public"],
+    deps = [
+        ":fill_quantization_options",
+        ":passes",
+        ":quantization_options_proto_cc",
+        "//tensorflow/compiler/mlir/tensorflow/transforms:tf_dialect_passes",
+        "//tensorflow/core/platform:path",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:Pass",
+    ],
+)
+
+gentbl_cc_library(
+    name = "stablehlo_test_passes_inc_gen",
+    compatible_with = get_compatible_with_portable(),
+    tbl_outs = {"passes/testing/passes.h.inc": [
+        "-gen-pass-decls",
+        "-name=Test",
+    ]},
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "passes/testing/passes.td",
+    deps = [
+        "@llvm-project//mlir:PassBaseTdFiles",
+    ],
+)
+
+cc_library(
+    name = "test_passes",
+    srcs = [
+        "passes/testing/test_lift_quantizable_spots_as_functions_with_quantization_specs.cc",
+        "passes/testing/test_post_calibration_component.cc",
+        "passes/testing/test_pre_calibration_component.cc",
+        "passes/testing/test_tf_to_stablehlo_pass.cc",
+    ],
+    hdrs = [
+        "passes/testing/passes.h",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        ":passes",
+        ":quantization_config_proto_cc",
+        ":stablehlo_test_passes_inc_gen",
+        "//tensorflow/compiler/mlir/quantization/common/ir:QuantOps",
+        "//tensorflow/compiler/mlir/quantization/stablehlo/cc:config",
+        "//tensorflow/compiler/mlir/quantization/stablehlo/cc:post_calibration",
+        "//tensorflow/compiler/mlir/quantization/stablehlo/cc:pre_calibration",
+        "//tensorflow/compiler/mlir/quantization/tensorflow:quantization_options_proto_cc",
+        "//tensorflow/compiler/mlir/quantization/tensorflow:quantize_preprocess",
+        "//tensorflow/compiler/mlir/quantization/tensorflow/cc:run_passes",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings:string_view",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:QuantOps",
+        "@llvm-project//mlir:SparseTensorDialect",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:UBDialect",
+        "@stablehlo//:chlo_ops",
+        "@stablehlo//:stablehlo_ops",
+        "@stablehlo//:vhlo_ops",
+        "@tsl//tsl/platform:protobuf",
+        "@xla//xla/mlir_hlo",
+    ],
+)
+
+cc_library(
+    name = "fill_quantization_options",
+    srcs = ["utils/fill_quantization_options.cc"],
+    hdrs = ["utils/fill_quantization_options.h"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        ":quantization_options_proto_cc",
+        "//tensorflow/compiler/mlir/tensorflow/transforms:tf_dialect_passes",
+        "//tensorflow/core/platform:path",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
+cc_library(
+    name = "math_utils",
+    srcs = ["utils/math_utils.cc"],
+    hdrs = ["utils/math_utils.h"],
+    compatible_with = get_compatible_with_portable(),
+    deps = ["@llvm-project//mlir:Support"],
+)
+
+tf_cc_test(
+    name = "math_utils_test",
+    srcs = ["utils/math_utils_test.cc"],
+    deps = [
+        ":math_utils",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
+cc_library(
+    name = "stablehlo_type_utils",
+    hdrs = ["utils/stablehlo_type_utils.h"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow:mangling_util",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/framework:numeric_types",
+        "//tensorflow/core/platform:status",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+        "@stablehlo//:stablehlo_ops",
+    ],
+)
+
+tf_cc_test(
+    name = "stablehlo_type_utils_test",
+    srcs = ["utils/stablehlo_type_utils_test.cc"],
+    deps = [
+        ":stablehlo_type_utils",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@stablehlo//:stablehlo_ops",
+    ],
+)
+
+cc_library(
+    name = "tf_type_utils",
+    srcs = ["utils/tf_type_utils.cc"],
+    hdrs = ["utils/tf_type_utils.h"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow:mangling_util",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/framework:numeric_types",
+        "@com_google_absl//absl/status",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+    ],
+)
+
+tf_cc_test(
+    name = "tf_type_utils_test",
+    srcs = [
+        "utils/tf_type_utils_test.cc",
+    ],
+    deps = [
+        ":bridge_passes",
+        ":tf_type_utils",
+        "//tensorflow/compiler/mlir:register_common_dialects",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow:mangling_util",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
+        "//tensorflow/compiler/mlir/tensorflow/transforms:tf_dialect_passes",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/ir/types:Dialect",
+        "//tensorflow/core/platform:path",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:QuantOps",
+        "@llvm-project//mlir:SparseTensorDialect",
+        "@llvm-project//mlir:Support",
+        "@xla//xla/mlir_hlo",
+        "@xla//xla/tsl/framework:numeric_types",
+    ],
+)
+
+cc_library(
+    name = "bfloat16_type",
+    srcs = ["utils/bfloat16_type.cc"],
+    hdrs = ["utils/bfloat16_type.h"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow:mangling_util",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/framework:numeric_types",
+        "//tensorflow/core/platform:status",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+    ],
+)
+
+tf_cc_test(
+    name = "bfloat16_type_test",
+    srcs = [
+        "utils/bfloat16_type_test.cc",
+    ],
+    deps = [
+        ":bfloat16_type",
+        "//tensorflow/compiler/mlir:register_common_dialects",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow:mangling_util",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
+        "//tensorflow/compiler/mlir/tensorflow/transforms:tf_dialect_passes",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/ir/types:Dialect",
+        "//tensorflow/core/platform:path",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:QuantOps",
+        "@llvm-project//mlir:SparseTensorDialect",
+        "@llvm-project//mlir:Support",
+        "@xla//xla/mlir_hlo",
+        "@xla//xla/tsl/framework:numeric_types",
+    ],
+)
+
+# OSS only: This target is header-only. Link `quantization_options_proto_cc_impl` only to
+# `libtensorflow_framework.so` via `lib_internal_impl`. Do NOT link
+# `quantization_options_proto_cc_impl` directly unless the target does not link
+# `libtensorflow_framework.so`.
+tf_proto_library(
+    name = "quantization_options_proto",
+    srcs = ["quantization_options.proto"],
+    make_default_target_header_only = True,
+    visibility = ["//visibility:public"],
+)
+
+# copybara:uncomment_begin(google-only)
+# py_proto_library(
+#     name = "quantization_options_py_pb2",
+#     visibility = [":internal_visibility_allowlist_package"],
+#     deps = [":quantization_options_proto"],
+# )
+# copybara:uncomment_end
+
+# OSS only: This target is header-only. Link `quantization_config_proto_cc_impl` only to
+# `libtensorflow_framework.so` via `lib_internal_impl`. Do NOT link
+# `quantization_config_proto_cc_impl` directly unless the target does not link
+# `libtensorflow_framework.so`.
+tf_proto_library(
+    name = "quantization_config_proto",
+    srcs = ["quantization_config.proto"],
+    make_default_target_header_only = True,
+    visibility = ["//visibility:public"],
+)
+
+# copybara:uncomment_begin(google-only)
+# py_proto_library(
+#     name = "quantization_config_py_pb2",
+#     visibility = [
+#         ":internal_visibility_allowlist_package",
+#     ],
+#     deps = [":quantization_config_proto"],
+# )
+# copybara:uncomment_end
+
+exports_files([
+    "run_lit.sh",
+])
+
+tf_cc_binary(
+    name = "stablehlo-quant-opt",
+    srcs = ["tools/stablehlo_quant_opt.cc"],
+    visibility = [":internal_visibility_allowlist_package"],
+    deps = [
+        ":bridge_passes",
+        ":passes",
+        ":test_passes",
+        "//tensorflow/compiler/mlir:init_mlir",
+        "//tensorflow/compiler/mlir/quantization/common/ir:QuantOps",
+        "//tensorflow/compiler/mlir/quantization/stablehlo/cc:pass_pipeline",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
+        "//tensorflow/compiler/mlir/tensorflow/transforms:tf_dialect_passes",
+        "//tensorflow/core/ir/types:Dialect",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:AllPassesAndDialects",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:FuncExtensions",
+        "@llvm-project//mlir:MlirOptLib",
+        "@llvm-project//mlir:QuantOps",
+        "@llvm-project//mlir:RegisterAllPasses",
+        "@llvm-project//mlir:SCFDialect",
+        "@llvm-project//mlir:ShapeDialect",
+        "@llvm-project//mlir:TensorDialect",
+        "@stablehlo//:stablehlo_ops",
+        "@stablehlo//:stablehlo_passes",
+        "@stablehlo//:vhlo_ops",
+        "@xla//xla/mlir_hlo:hlo_dialect_registration",
+        "@xla//xla/mlir_hlo:mhlo_passes",
+    ],
+)

tensorflow/compiler/mlir/quantization/stablehlo/cc/BUILD

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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_

tensorflow/compiler/mlir/quantization/stablehlo/instrumentations/BUILD

@@ -0,0 +1,48 @@
+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 = ["@stablehlo//:license"],
+    default_visibility = [
+        "//tensorflow/compiler/mlir/quantization/stablehlo:__subpackages__",
+    ],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "save_report",
+    srcs = ["save_report.cc"],
+    hdrs = ["save_report.h"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        "//tensorflow/compiler/mlir/quantization/stablehlo/cc:report",
+        "@com_google_absl//absl/base:nullability",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/strings:string_view",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
+tf_cc_test(
+    name = "save_report_test",
+    srcs = ["save_report_test.cc"],
+    deps = [
+        ":save_report",
+        "//tensorflow/compiler/mlir/quantization/common:test_base",
+        "//tensorflow/compiler/mlir/quantization/stablehlo:passes",
+        "//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
+        "//tensorflow/compiler/mlir/quantization/stablehlo/cc:io",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@tsl//tsl/platform:protobuf",
+    ],
+)

tensorflow/compiler/mlir/quantization/stablehlo/instrumentations/save_report.cc

@@ -0,0 +1,97 @@
+/* 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/instrumentations/save_report.h"
+
+#include <optional>
+#include <string>
+
+#include "absl/base/nullability.h"
+#include "absl/log/log.h"
+#include "absl/strings/string_view.h"
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/report.h"
+
+namespace mlir::quant::stablehlo {
+namespace {
+
+// Converts `std::optional<absl::string_view>` to `std::optional<std::string>`.
+// A `std::nullopt` is returned when `view` is `std::nullopt`.
+std::optional<std::string> OptionalStringViewToOptionalString(
+    std::optional<absl::string_view> view) {
+  if (view == std::nullopt) return std::nullopt;
+  return std::make_optional<std::string>(*view);
+}
+
+using quant::stablehlo::QuantizationReport;
+
+// Whether the pass is `QuantizeCompositeFunctionPass`.
+bool IsQuantizeCompositeFunctionPass(Pass* absl_nullable pass,
+                                     Operation* absl_nullable op) {
+  // It is known that `op` is `ModuleOp` when `pass` is
+  // `QuantizeCompositeFunctionPass`, but the check is still performed to be
+  // defensive.
+  return pass != nullptr &&
+         pass->getArgument() == "tf-stablehlo-quantize-composite-functions" &&
+         isa_and_nonnull<ModuleOp>(op);
+}
+
+// Report is saved only when:
+// * After running `QuantizeCompositeFunctionPass`.
+// * The pass is run on `ModuleOp`.
+// * `file_path` is not `nullopt`.
+bool ShouldSaveReport(Pass* absl_nullable pass, Operation* absl_nullable op,
+                      const std::optional<std::string>& file_path) {
+  return file_path != std::nullopt && IsQuantizeCompositeFunctionPass(pass, op);
+}
+
+void SaveReport(const QuantizationReport& report,
+                const absl::string_view file_path) {
+  if (const absl::Status save_status = report.Save(file_path);
+      save_status.ok()) {
+    LOG(INFO) << "Successfully saved quantization report to: " << file_path;
+  } else {
+    LOG(ERROR) << "Failed to save quantization report to: " << file_path
+               << " with status: " << save_status;
+  }
+}
+
+}  // namespace
+
+SaveQuantizationReportInstrumentation::SaveQuantizationReportInstrumentation(
+    std::optional<absl::string_view> file_path)
+    : file_path_(OptionalStringViewToOptionalString(file_path)) {}
+
+void SaveQuantizationReportInstrumentation::runAfterPass(Pass* pass,
+                                                         Operation* op) {
+  // Only run after `QuantizeCompositeFunctionPass`.
+  if (!IsQuantizeCompositeFunctionPass(pass, op)) return;
+
+  auto module_op = cast<ModuleOp>(op);
+  const QuantizationReport report(module_op);
+
+  // Print a human-readable report to stdout regardless of whether the report
+  // is saved to file.
+  report.Print();
+
+  // Exit early if the report should not be saved to file.
+  if (!ShouldSaveReport(pass, op, file_path_)) return;
+
+  SaveReport(report, *file_path_);
+}
+
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/instrumentations/save_report.h

@@ -0,0 +1,52 @@
+/* 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_INSTRUMENTATIONS_SAVE_REPORT_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_INSTRUMENTATIONS_SAVE_REPORT_H_
+
+#include <optional>
+#include <string>
+
+#include "absl/strings/string_view.h"
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Pass/PassInstrumentation.h"  // from @llvm-project
+
+namespace mlir::quant::stablehlo {
+
+// A `PassInstrumentation` that saves quantization report to file after
+// `QuantizeCompositeFunctionsPass` is run. It inspects the `ModuleOp` after
+// quantization and analyzes the quantizable units and quantization methods
+// used. The report file will be saved at the `file_path`. The report file
+// contains textproto of `QuantizationResults`. `file_path`'s base directories
+// should exist (this pass instrumentation will not `mkdir` them).
+//
+// See `QuantizationReport` for further details on the quantization report.
+class SaveQuantizationReportInstrumentation : public PassInstrumentation {
+ public:
+  // `file_path` is the path to save the report file. The report file is in
+  // textproto format so a `.txtpb` extension is preferred but it doesn't result
+  // in error if other extension is used. This instrumentation will not be run
+  // if `file_path` is a `nullopt`.
+  explicit SaveQuantizationReportInstrumentation(
+      std::optional<absl::string_view> file_path);
+
+  void runAfterPass(Pass* pass, Operation* op) override;
+
+ private:
+  std::optional<std::string> file_path_;  // Path to file to save the report.
+};
+
+}  // namespace mlir::quant::stablehlo
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_INSTRUMENTATIONS_SAVE_REPORT_H_

tensorflow/compiler/mlir/quantization/stablehlo/instrumentations/save_report_test.cc

@@ -0,0 +1,185 @@
+/* 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/instrumentations/save_report.h"
+
+#include <memory>
+#include <optional>
+#include <string>
+
+#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 "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.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/passes/passes.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::QuantizationResults;
+using ::stablehlo::quantization::io::ReadFileToString;
+using ::testing::SizeIs;
+using ::testing::StrEq;
+using ::tsl::protobuf::TextFormat;
+
+using SaveQuantizationReportInstrumentationTest = QuantizationTestBase;
+
+TEST_F(SaveQuantizationReportInstrumentationTest, SaveReport) {
+  constexpr absl::string_view kModuleWithCompositeDotGeneral = R"mlir(
+    func.func @main(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> {
+      %cst = "tf.Const"() {value = dense<3.00000000e-1> : tensor<2x3xf32>} : () -> tensor<2x3xf32>
+      %0 = "quantization.stats"(%arg0) {layerStats = dense<[6.00000000e-6, 9.00000000e-1]> : tensor<2xf32>} : (tensor<1x2xf32>) -> tensor<1x2xf32>
+      %1 = "tf.XlaCallModule"(%0, %cst) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn, _stablehlo_version = "1.0.0", _original_entry_function = "composite_dot_general_fn", _quantization_method = "static_range_ptq { }", _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>
+      %2 = "quantization.stats"(%1) {layerStats = dense<[5.00000000e-6, 7.00000000e-1]> : tensor<2xf32>} : (tensor<1x3xf32>) -> tensor<1x3xf32>
+      return %2 : tensor<1x3xf32>
+    }
+
+    func.func private @composite_dot_general_fn(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
+      %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(kModuleWithCompositeDotGeneral);
+  ASSERT_TRUE(module_op);
+
+  // Create a pass manager with `SaveQuantizationReportInstrumentation` and
+  // `QuantizeCompositeFunctionsPass`. Run the passes against `module_op`.
+  PassManager pm(ctx_.get());
+
+  QuantizeCompositeFunctionsPassOptions options;
+  pm.addPass(createQuantizeCompositeFunctionsPass(options));
+
+  const std::string report_file_path =
+      absl::StrCat(testing::TempDir(), "/save_report.txtpb");
+  pm.addInstrumentation(std::make_unique<SaveQuantizationReportInstrumentation>(
+      report_file_path));
+
+  const LogicalResult run_result = pm.run(*module_op);
+  ASSERT_TRUE(succeeded(run_result));
+
+  // Check that the report file contains `QuantizationResults` textproto,
+  // reflecting the quantization results, in this case the
+  // `composite_dot_general_fn` with quantized with `static_range_ptq` method.
+  const absl::StatusOr<std::string> file_data =
+      ReadFileToString(report_file_path);
+  ASSERT_THAT(file_data, absl_testing::IsOk());
+
+  /*
+  results {
+    quantizable_unit {
+      name: "composite_dot_general_fn"
+    }
+    method { static_range_ptq { } }
+  }
+  */
+  QuantizationResults results{};
+  ASSERT_TRUE(TextFormat::ParseFromString(*file_data, &results));
+  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());
+}
+
+TEST_F(SaveQuantizationReportInstrumentationTest,
+       ReportNotSavedWhenNoQuantizeCompositeFunctionsPass) {
+  constexpr absl::string_view kModuleWithCompositeDotGeneral = R"mlir(
+    func.func @main(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> {
+      %cst = "stablehlo.constant"() {value = dense<3.00000000e-1> : tensor<2x3xf32>} : () -> tensor<2x3xf32>
+      %0 = "quantization.stats"(%arg0) {layerStats = dense<[6.00000000e-6, 9.00000000e-1]> : tensor<2xf32>} : (tensor<1x2xf32>) -> tensor<1x2xf32>
+      %1 = "tf.XlaCallModule"(%0, %cst) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn, _stablehlo_version = "1.0.0", _original_entry_function = "composite_dot_general_fn", _quantization_method = "static_range_ptq { }", _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>
+      %2 = "quantization.stats"(%1) {layerStats = dense<[5.00000000e-6, 7.00000000e-1]> : tensor<2xf32>} : (tensor<1x3xf32>) -> tensor<1x3xf32>
+      return %2 : tensor<1x3xf32>
+    }
+
+    func.func private @composite_dot_general_fn(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
+      %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(kModuleWithCompositeDotGeneral);
+  ASSERT_TRUE(module_op);
+
+  // Create a pass manager with `SaveQuantizationReportInstrumentation` a pass
+  // that is not `QuantizeCompositeFunctionsPass`. Run the passes against
+  // `module_op`.
+  PassManager pm(ctx_.get());
+
+  pm.addPass(createPrepareQuantizePass());
+
+  const std::string report_file_path = absl::StrCat(
+      testing::TempDir(),
+      "/report_not_saved_no_quantize_composite_functions_pass.txtpb");
+  pm.addInstrumentation(std::make_unique<SaveQuantizationReportInstrumentation>(
+      report_file_path));
+
+  const LogicalResult run_result = pm.run(*module_op);
+  ASSERT_TRUE(succeeded(run_result));
+
+  // The report file is not created because `QuantizeCompositeFunctionsPass` was
+  // not run.
+  EXPECT_THAT(ReadFileToString(report_file_path),
+              absl_testing::StatusIs(absl::StatusCode::kNotFound));
+}
+
+TEST_F(SaveQuantizationReportInstrumentationTest,
+       ReportNotSavedWhenReportFilePathIsNullopt) {
+  constexpr absl::string_view kModuleWithCompositeDotGeneral = R"mlir(
+    func.func @main(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> {
+      %cst = "stablehlo.constant"() {value = dense<3.00000000e-1> : tensor<2x3xf32>} : () -> tensor<2x3xf32>
+      %0 = "quantization.stats"(%arg0) {layerStats = dense<[6.00000000e-6, 9.00000000e-1]> : tensor<2xf32>} : (tensor<1x2xf32>) -> tensor<1x2xf32>
+      %1 = "tf.XlaCallModule"(%0, %cst) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn, _stablehlo_version = "1.0.0", _original_entry_function = "composite_dot_general_fn", _quantization_method = "static_range_ptq { }", _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>
+      %2 = "quantization.stats"(%1) {layerStats = dense<[5.00000000e-6, 7.00000000e-1]> : tensor<2xf32>} : (tensor<1x3xf32>) -> tensor<1x3xf32>
+      return %2 : tensor<1x3xf32>
+    }
+
+    func.func private @composite_dot_general_fn(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
+      %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(kModuleWithCompositeDotGeneral);
+  ASSERT_TRUE(module_op);
+
+  PassManager pm(ctx_.get());
+
+  QuantizeCompositeFunctionsPassOptions options;
+  pm.addPass(createQuantizeCompositeFunctionsPass(options));
+  pm.addInstrumentation(std::make_unique<SaveQuantizationReportInstrumentation>(
+      /*file_path=*/std::nullopt));
+
+  // The report file is not created and `SaveQuantizationReportInstrumentation`
+  // is not run, but the passes still run without errors.
+  const LogicalResult run_result = pm.run(*module_op);
+  ASSERT_TRUE(succeeded(run_result));
+}
+
+}  // namespace
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/internal_visibility_allowlist.bzl

@@ -0,0 +1,5 @@
+"""Internal visibility rules."""
+
+def internal_visibility_allowlist():
+    """Returns a list of the packages that can depend on internal targets."""
+    return []

tensorflow/compiler/mlir/quantization/stablehlo/ops/BUILD

@@ -0,0 +1,57 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("//tensorflow:tensorflow.bzl", "tf_cc_test")
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [
+        "//tensorflow/compiler/mlir/quantization/tensorflow:internal_visibility_allowlist_package",
+    ],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "stablehlo_op_quant_spec",
+    srcs = [
+        "stablehlo_op_quant_spec.cc",
+    ],
+    hdrs = ["stablehlo_op_quant_spec.h"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        "//tensorflow/compiler/mlir/quantization/common:attrs_and_constraints",
+        "//tensorflow/compiler/mlir/quantization/common:lift_as_function_call",
+        "//tensorflow/compiler/mlir/quantization/common/ir:QuantOps",
+        "//tensorflow/compiler/mlir/quantization/common/quantization_lib",
+        "//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
+        "//tensorflow/compiler/mlir/quantization/tensorflow:quantization_options_proto_cc",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "@com_google_absl//absl/status:statusor",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@stablehlo//:stablehlo_ops",
+        "@tsl//tsl/platform:protobuf",
+    ],
+)
+
+tf_cc_test(
+    name = "stablehlo_op_quant_spec_test",
+    srcs = ["stablehlo_op_quant_spec_test.cc"],
+    deps = [
+        ":stablehlo_op_quant_spec",
+        "//tensorflow/compiler/mlir/quantization/common:attrs_and_constraints",
+        "//tensorflow/compiler/mlir/quantization/common:func",
+        "//tensorflow/compiler/mlir/quantization/common:test_base",
+        "//tensorflow/compiler/mlir/quantization/common/ir:QuantOps",
+        "//tensorflow/compiler/mlir/quantization/common/quantization_lib",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/core:test",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@stablehlo//:stablehlo_ops",
+    ],
+)

tensorflow/compiler/mlir/quantization/stablehlo/ops/stablehlo_op_quant_spec.cc

@@ -0,0 +1,184 @@
+/* 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/ops/stablehlo_op_quant_spec.h"
+
+#include <memory>
+
+#include "absl/status/statusor.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/Debug.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/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo
+#include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.h"
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"
+#include "tensorflow/compiler/mlir/quantization/common/lift_as_function_call.h"
+#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
+#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tsl/platform/protobuf.h"  // IWYU pragma: keep
+
+// To be used with LLVM_DEBUG.
+#define DEBUG_TYPE "stablehlo_opt_quant_spec"
+
+namespace mlir::quant::stablehlo {
+namespace {
+
+using ::mlir::stablehlo::DotGeneralOp;
+using ::stablehlo::quantization::Method;
+using ::stablehlo::quantization::StaticRangePtq;
+
+// Whether it represents a lifted function (i.e. `op` is the corresponding
+// `XlaCallModuleOp`) that is explicitly marked `NoQuantization`.
+bool IsDenylistedLiftedFunction(Operation* op) {
+  if (auto xla_call_module_op = dyn_cast_or_null<TF::XlaCallModuleOp>(op);
+      xla_call_module_op != nullptr) {
+    absl::StatusOr<Method> method = GetQuantizationMethod(xla_call_module_op);
+    if (method.ok() && method->has_no_quantization()) {
+      return true;
+    }
+  }
+  return false;
+}
+
+// Populates `spec.coeff_op_quant_dim` according to `xla_call_module_op`'s
+// `_quantization_method` attribute. If there is an input `QuantizedType` with
+// `dimension_specs` set, which represents the quantization dimension for the
+// input, then the corresponding operand index -> quantization dimension mapping
+// is set for `spec`.
+// TODO: b/323478683 - Duplicate tracking of config will be eliminated.
+// `OpQuantSpec` will be deprecated and `Method` will be used instead.
+void PopulateCoeffOpQuantDimIfPerChannelQuantized(
+    TF::XlaCallModuleOp xla_call_module_op, OpQuantSpec& spec) {
+  absl::StatusOr<Method> method = GetQuantizationMethod(xla_call_module_op);
+  if (method.ok() && method->has_static_range_ptq()) {
+    // TODO: b/331145946 - Use `Method` accessors.
+    const StaticRangePtq& static_range_ptq_spec = method->static_range_ptq();
+    // Look for quantized dimension specs for each quantized type and
+    // populate `coeff_op_quant_dim`.
+    for (const auto& [operand_idx, quantized_type] :
+         static_range_ptq_spec.input_quantized_types()) {
+      if (quantized_type.has_dimension_specs()) {
+        spec.coeff_op_quant_dim[operand_idx] =
+            quantized_type.dimension_specs().dimension();
+      }
+    }
+  }
+}
+
+}  // namespace
+
+std::unique_ptr<OpQuantSpec> GetStableHloOpQuantSpec(Operation* op) {
+  auto spec = std::make_unique<OpQuantSpec>();
+  if (auto call_op = dyn_cast_or_null<TF::XlaCallModuleOp>(op)) {
+    auto entry_function =
+        call_op->getAttrOfType<FlatSymbolRefAttr>("_entry_function");
+    StringRef function_name = entry_function.getValue();
+    if (!function_name.starts_with("composite_")) {
+      return spec;
+    }
+
+    if (function_name.contains("conv")) {
+      // Looks up `Method` to see if it should be per-channel quantized and
+      // populates the spec accordingly.
+      PopulateCoeffOpQuantDimIfPerChannelQuantized(call_op, *spec);
+
+      if (function_name.contains("with_bias")) {
+        spec->biases_params[2] = {{0, 1}, GetUniformQuantizedTypeForBias};
+      }
+    } else if (function_name.contains("dot_general")) {
+      const auto module_op = call_op->getParentOfType<ModuleOp>();
+
+      const SymbolTable symbol_table(module_op);
+      auto entry_func_op =
+          dyn_cast_or_null<func::FuncOp>(symbol_table.lookup(function_name));
+      auto dot_general_op = *entry_func_op.getOps<DotGeneralOp>().begin();
+      if (auto optional_dim = GetDotGeneralQuantizationDim(dot_general_op);
+          optional_dim) {
+        spec->coeff_op_quant_dim[1] = optional_dim.value();
+      } else {
+        spec->coeff_op_quant_dim[1] = -1;
+      }
+      if (function_name.contains("with_bias")) {
+        spec->biases_params[2] = {{0, 1}, GetUniformQuantizedTypeForBias};
+      }
+    }
+    for (const auto [operand_idx, per_channel_dim] : spec->coeff_op_quant_dim) {
+      spec->quantizable_operands.insert(operand_idx);
+    }
+  }
+  return spec;
+}
+
+std::unique_ptr<OpQuantScaleSpec> GetStableHloQuantConstraints(Operation* op) {
+  auto scale_spec = std::make_unique<OpQuantScaleSpec>();
+  if (llvm::isa<mlir::stablehlo::BroadcastInDimOp,
+                mlir::stablehlo::ConcatenateOp,
+                mlir::stablehlo::DynamicReshapeOp,
+                mlir::stablehlo::DynamicSliceOp, mlir::stablehlo::GatherOp,
+                mlir::stablehlo::PadOp, mlir::stablehlo::ReduceWindowOp,
+                mlir::stablehlo::ReshapeOp, mlir::stablehlo::SelectOp,
+                mlir::stablehlo::SliceOp, mlir::stablehlo::TransposeOp>(op)) {
+    scale_spec->has_same_scale_requirement = true;
+  }
+  if (llvm::isa<mlir::stablehlo::DynamicSliceOp, mlir::stablehlo::GatherOp,
+                mlir::stablehlo::PadOp, mlir::stablehlo::SliceOp>(op)) {
+    scale_spec->has_same_operand_and_result_type_requirement = true;
+  }
+  return scale_spec;
+}
+
+bool IsOpQuantizableStableHlo(Operation* op) {
+  if (isa<func::ConstantOp, mlir::stablehlo::ConstantOp>(op)) {
+    // Constant ops do not have QuantizableResult attribute but can be
+    // quantized.
+    return true;
+  } else if (op->hasTrait<OpTrait::IsTerminator>() ||
+             isa<mlir::quant::ir::QuantizeCastOp,
+                 mlir::quant::ir::DequantizeCastOp>(op)) {
+    // Terminators, qcast and decast are not quantizable.
+    return false;
+  }
+
+  // `op` is not quantizable when it is an `XlaCallModuleOp` representing lifted
+  // function whose `_quantization_method` attribute is marked `NoQuantization`.
+  // This means this quantizable unit has been explicitly denylisted by the
+  // user.
+  if (IsDenylistedLiftedFunction(op)) {
+    LLVM_DEBUG(llvm::errs() << "Denylisted quantizable unit: \n" << op << "\n");
+    return false;
+  }
+
+  if (GetStableHloQuantConstraints(op)->has_same_scale_requirement) {
+    return true;
+  }
+
+  const bool attr_enforced_quantizable =
+      op->hasAttrOfType<StringAttr>(kQuantTraitAttrName) &&
+      op->getAttrOfType<StringAttr>(kQuantTraitAttrName).getValue().str() ==
+          QuantTraitValues[QuantizationTrait::FullyQuantizable];
+  return attr_enforced_quantizable;
+}
+
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/ops/stablehlo_op_quant_spec.h

@@ -0,0 +1,41 @@
+/* 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_OPS_STABLEHLO_OP_QUANT_SPEC_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_OPS_STABLEHLO_OP_QUANT_SPEC_H_
+
+#include <memory>
+
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"
+#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
+
+namespace mlir::quant::stablehlo {
+
+// Returns StableHLO quantization specs for an op.
+std::unique_ptr<OpQuantSpec> GetStableHloOpQuantSpec(Operation* op);
+
+// Returns quantization constraints (ex: fixed output, same scale) given
+// a StableHLO op.
+std::unique_ptr<OpQuantScaleSpec> GetStableHloQuantConstraints(Operation* op);
+
+// Checks if an op is quantizable in StableHLO quantizer. Argument op is not
+// necessarily a StableHLO op.
+bool IsOpQuantizableStableHlo(Operation* op);
+
+}  // namespace mlir::quant::stablehlo
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_OPS_STABLEHLO_OP_QUANT_SPEC_H_

tensorflow/compiler/mlir/quantization/stablehlo/ops/stablehlo_op_quant_spec_test.cc

@@ -0,0 +1,332 @@
+/* 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/ops/stablehlo_op_quant_spec.h"
+
+#include <memory>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/strings/string_view.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo
+#include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.h"
+#include "tensorflow/compiler/mlir/quantization/common/func.h"
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"
+#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"
+#include "tensorflow/compiler/mlir/quantization/common/test_base.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace mlir::quant::stablehlo {
+namespace {
+
+using ::mlir::stablehlo::GatherOp;
+using ::testing::IsEmpty;
+using ::testing::IsTrue;
+using ::testing::NotNull;
+using ::testing::Pair;
+using ::testing::UnorderedElementsAre;
+
+using IsOpQuantizableStableHloTest = ::mlir::quant::QuantizationTestBase;
+
+// Quantizable ops: constants
+// Non-quantizable ops: normal StableHLO ops and terminators
+constexpr absl::string_view kModuleConstantAdd = R"mlir(
+  module {
+    func.func @constant_add() -> (tensor<3x2xf32>) {
+      %cst1 = stablehlo.constant dense<2.4> : tensor<3x2xf32>
+      %cst2 = stablehlo.constant dense<5.7> : tensor<3x2xf32>
+      %add = stablehlo.add %cst1, %cst2 : (tensor<3x2xf32>, tensor<3x2xf32>) -> tensor<3x2xf32>
+      func.return %add : tensor<3x2xf32>
+    }
+  }
+)mlir";
+
+// Quantizable ops: XlaCallModule op with "fully_quantizable" attribute and
+// same-scale StableHLO ops
+// Non-quantizable ops: quantize/dequantize ops
+constexpr absl::string_view kModuleCompositeSameScale = R"mlir(
+  module {
+    func.func @same_scale_after_composite() -> tensor<3x1xf32> {
+      %0 = "tf.XlaCallModule"() {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn_1, _stablehlo_version = "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<1x3xf32>
+      %1 = "quantization.qcast"(%0) {volatile} : (tensor<1x3xf32>) -> tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
+      %2 = "quantization.dcast"(%1) : (tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<1x3xf32>
+      %3 = stablehlo.reshape %2 : (tensor<1x3xf32>) -> tensor<3x1xf32>
+      %4 = "quantization.qcast"(%3) {volatile} : (tensor<3x1xf32>) -> tensor<3x1x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
+      %5 = "quantization.dcast"(%4) : (tensor<3x1x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<3x1xf32>
+      return %5 : tensor<3x1xf32>
+    }
+  }
+)mlir";
+
+// Non-quantizable ops: XlaCallModule op without "fully_quantizable" attribute
+constexpr absl::string_view kModuleCompositeNoAttr = R"mlir(
+  module {
+    func.func @composite_without_attr() -> tensor<1x3xf32> {
+      %0 = "tf.XlaCallModule"() {Sout = [#tf_type.shape<1x3>], _entry_function = @non_quantizable_composite, _stablehlo_version = "1.0.0", _original_entry_function = "non_quantizable_composite", _stablehlo_module_attrs = {}, device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : () -> tensor<1x3xf32>
+      return %0 : tensor<1x3xf32>
+    }
+  }
+)mlir";
+
+TEST_F(IsOpQuantizableStableHloTest, ConstantOpQuantizable) {
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleConstantAdd);
+  ASSERT_TRUE(module_op);
+
+  auto test_func = module_op->lookupSymbol<func::FuncOp>("constant_add");
+  ASSERT_THAT(test_func, NotNull());
+
+  auto constant_op =
+      FindOperationOfType<mlir::stablehlo::ConstantOp>(test_func);
+  EXPECT_TRUE(IsOpQuantizableStableHlo(constant_op));
+}
+
+TEST_F(IsOpQuantizableStableHloTest, TerminatorOpNotQuantizable) {
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleConstantAdd);
+  ASSERT_TRUE(module_op);
+
+  auto test_func = module_op->lookupSymbol<func::FuncOp>("constant_add");
+  ASSERT_THAT(test_func, NotNull());
+
+  auto return_op = FindOperationOfType<func::ReturnOp>(test_func);
+  EXPECT_FALSE(IsOpQuantizableStableHlo(return_op));
+}
+
+TEST_F(IsOpQuantizableStableHloTest, SameScaleOpQuantizable) {
+  OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kModuleCompositeSameScale);
+  ASSERT_TRUE(module_op);
+
+  auto test_func =
+      module_op->lookupSymbol<func::FuncOp>("same_scale_after_composite");
+  ASSERT_THAT(test_func, NotNull());
+
+  auto reshape_op = FindOperationOfType<mlir::stablehlo::ReshapeOp>(test_func);
+  EXPECT_TRUE(IsOpQuantizableStableHlo(reshape_op));
+}
+
+TEST_F(IsOpQuantizableStableHloTest, NonSameScaleOpNotQuantizable) {
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleConstantAdd);
+  ASSERT_TRUE(module_op);
+
+  auto test_func = module_op->lookupSymbol<func::FuncOp>("constant_add");
+  ASSERT_THAT(test_func, NotNull());
+
+  auto add_op = FindOperationOfType<mlir::stablehlo::AddOp>(test_func);
+  EXPECT_FALSE(IsOpQuantizableStableHlo(add_op));
+}
+
+TEST_F(IsOpQuantizableStableHloTest, ValidXlaCallModuleOpQuantizable) {
+  OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kModuleCompositeSameScale);
+  ASSERT_TRUE(module_op);
+
+  auto test_func =
+      module_op->lookupSymbol<func::FuncOp>("same_scale_after_composite");
+  ASSERT_THAT(test_func, NotNull());
+
+  auto xla_call_module_op = FindOperationOfType<TF::XlaCallModuleOp>(test_func);
+  EXPECT_TRUE(IsOpQuantizableStableHlo(xla_call_module_op));
+}
+
+TEST_F(IsOpQuantizableStableHloTest, InvalidXlaCallModuleOpNotQuantizable) {
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleCompositeNoAttr);
+  ASSERT_TRUE(module_op);
+
+  auto test_func =
+      module_op->lookupSymbol<func::FuncOp>("composite_without_attr");
+  ASSERT_THAT(test_func, NotNull());
+
+  auto xla_call_module_op = FindOperationOfType<TF::XlaCallModuleOp>(test_func);
+  EXPECT_FALSE(IsOpQuantizableStableHlo(xla_call_module_op));
+}
+
+TEST_F(IsOpQuantizableStableHloTest, QuantizeDequantizeOpNotQuantizable) {
+  OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kModuleCompositeSameScale);
+  ASSERT_TRUE(module_op);
+
+  auto test_func =
+      module_op->lookupSymbol<func::FuncOp>("same_scale_after_composite");
+  ASSERT_THAT(test_func, NotNull());
+
+  auto quantize_op =
+      FindOperationOfType<mlir::quant::ir::QuantizeCastOp>(test_func);
+  EXPECT_FALSE(IsOpQuantizableStableHlo(quantize_op));
+
+  auto dequantize_op =
+      FindOperationOfType<mlir::quant::ir::DequantizeCastOp>(test_func);
+  EXPECT_FALSE(IsOpQuantizableStableHlo(dequantize_op));
+}
+
+TEST_F(IsOpQuantizableStableHloTest,
+       XlaCallModuleOpQuantizableWhenNotDenylisted) {
+  // A `TF::XlaCallModuleOp` with `_quantization_method = ""`.
+  constexpr absl::string_view
+      kModuleXlaCallModuleOpWithDefaultQuantizationMethod = R"mlir(
+    func.func @xla_call_module_default_quantization_method(%arg0: tensor<1x1x3xf32>, %arg1: tensor<3x4xf32>) -> tensor<1x1x4xf32> {
+      %0 = "tf.XlaCallModule"(%arg0, %arg1) <{Sout = [#tf_type.shape<1x1x4>], dim_args_spec = [], disabled_checks = [], function_list = [], has_token_input_output = false, module = "", platforms = ["CPU"], version = 9 : i64}> {_entry_function = @composite_dot_general_fn_1, _stablehlo_version = "1.0.0", _original_entry_function = "composite_dot_general_fn_1", _quantization_method = "", _stablehlo_module_attrs = {jax.uses_shape_polymorphism = true}, _tfl_quant_trait = "fully_quantizable"} : (tensor<1x1x3xf32>, tensor<3x4xf32>) -> tensor<1x1x4xf32>
+      return %0 : tensor<1x1x4xf32>
+    }
+  )mlir";
+
+  OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kModuleXlaCallModuleOpWithDefaultQuantizationMethod);
+  ASSERT_TRUE(module_op);
+
+  auto test_func = module_op->lookupSymbol<func::FuncOp>(
+      "xla_call_module_default_quantization_method");
+  ASSERT_THAT(test_func, NotNull());
+
+  auto xla_call_module_op = FindOperationOfType<TF::XlaCallModuleOp>(test_func);
+  EXPECT_TRUE(IsOpQuantizableStableHlo(xla_call_module_op));
+}
+
+TEST_F(IsOpQuantizableStableHloTest, DenylistedXlaCallModuleOpNotQuantizable) {
+  // A `TF::XlaCallModuleOp` with `_quantization_method = "no_quantization {}"`,
+  // indicating it has been explicitly denylisted by the user.
+  constexpr absl::string_view kModuleDenylistedXlaCallModuleOp = R"mlir(
+    func.func @xla_call_module_denylisted(%arg0: tensor<1x1x3xf32>, %arg1: tensor<3x4xf32>) -> tensor<1x1x4xf32> {
+      %0 = "tf.XlaCallModule"(%arg0, %arg1) <{Sout = [#tf_type.shape<1x1x4>], dim_args_spec = [], disabled_checks = [], function_list = [], has_token_input_output = false, module = "", platforms = ["CPU"], version = 9 : i64}> {_entry_function = @composite_dot_general_fn_1, _stablehlo_version = "1.0.0", _original_entry_function = "composite_dot_general_fn_1", _quantization_method = "no_quantization {}", _stablehlo_module_attrs = {jax.uses_shape_polymorphism = true}, _tfl_quant_trait = "fully_quantizable"} : (tensor<1x1x3xf32>, tensor<3x4xf32>) -> tensor<1x1x4xf32>
+      return %0 : tensor<1x1x4xf32>
+    }
+  )mlir";
+
+  OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kModuleDenylistedXlaCallModuleOp);
+  ASSERT_TRUE(module_op);
+
+  auto test_func =
+      module_op->lookupSymbol<func::FuncOp>("xla_call_module_denylisted");
+  ASSERT_THAT(test_func, NotNull());
+
+  auto xla_call_module_op = FindOperationOfType<TF::XlaCallModuleOp>(test_func);
+  EXPECT_FALSE(IsOpQuantizableStableHlo(xla_call_module_op));
+}
+
+using GetStableHloOpQuantSpecTest = ::mlir::quant::QuantizationTestBase;
+
+TEST_F(GetStableHloOpQuantSpecTest,
+       EmptyCoeffOpQuantDimForPerTensorQuantizedConvolution) {
+  // A `TF::XlaCallModuleOp` with `_quantization_method = "static_range_ptq
+  // {}"`, representing a per-tensor static-range PTQ quantization.
+  constexpr absl::string_view
+      kXlaCallModuleOpWithPerTensorQuantizedConvolution = R"mlir(
+    func.func @main(%arg0: tensor<1x1x3xf32>, %arg1: tensor<3x4xf32>) -> tensor<1x1x4xf32> {
+      %0 = "tf.XlaCallModule"(%arg0, %arg1) <{Sout = [#tf_type.shape<1x1x4>], dim_args_spec = [], disabled_checks = [], function_list = [], has_token_input_output = false, module = "", platforms = ["CPU"], version = 9 : i64}>
+          {
+            _entry_function = @composite_conv_fn_1,
+            _stablehlo_version = "1.0.0",
+            _original_entry_function = "composite_conv_fn_1",
+            _quantization_method = "static_range_ptq {}",
+            _stablehlo_module_attrs = {jax.uses_shape_polymorphism = true},
+            _tfl_quant_trait = "fully_quantizable"
+          } : (tensor<1x1x3xf32>, tensor<3x4xf32>) -> tensor<1x1x4xf32>
+      return %0 : tensor<1x1x4xf32>
+    }
+  )mlir";
+
+  const OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kXlaCallModuleOpWithPerTensorQuantizedConvolution);
+  ASSERT_TRUE(module_op);
+
+  const FailureOr<TF::XlaCallModuleOp> xla_call_module_op =
+      FindFirstOpFromMainFunc<TF::XlaCallModuleOp>(*module_op);
+  ASSERT_TRUE(succeeded(xla_call_module_op));
+
+  const std::unique_ptr<OpQuantSpec> op_quant_spec =
+      GetStableHloOpQuantSpec(*xla_call_module_op);
+  ASSERT_THAT(op_quant_spec, NotNull());
+
+  EXPECT_THAT(op_quant_spec->coeff_op_quant_dim, IsEmpty());
+}
+
+TEST_F(GetStableHloOpQuantSpecTest,
+       EmptyCoeffOpQuantDimForPerChannelQuantizedConvolution) {
+  constexpr absl::string_view
+      kXlaCallModuleOpWithPerChannelQuantizedConvolution = R"mlir(
+    func.func @main(%arg0: tensor<1x1x3xf32>, %arg1: tensor<3x4xf32>) -> tensor<1x1x4xf32> {
+      %0 = "tf.XlaCallModule"(%arg0, %arg1) <{Sout = [#tf_type.shape<1x1x4>], dim_args_spec = [], disabled_checks = [], function_list = [], has_token_input_output = false, module = "", platforms = ["CPU"], version = 9 : i64}>
+          {
+            _entry_function = @composite_conv_fn_1,
+            _stablehlo_version = "1.0.0",
+            _original_entry_function = "composite_conv_fn_1",
+            _quantization_method = "static_range_ptq {input_quantized_types {key: 1, value {dimension_specs {dimension: 3}}}}",
+            _stablehlo_module_attrs = {jax.uses_shape_polymorphism = true},
+            _tfl_quant_trait = "fully_quantizable"
+          } : (tensor<1x1x3xf32>, tensor<3x4xf32>) -> tensor<1x1x4xf32>
+      return %0 : tensor<1x1x4xf32>
+    }
+  )mlir";
+
+  const OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kXlaCallModuleOpWithPerChannelQuantizedConvolution);
+  ASSERT_TRUE(module_op);
+
+  const FailureOr<TF::XlaCallModuleOp> xla_call_module_op =
+      FindFirstOpFromMainFunc<TF::XlaCallModuleOp>(*module_op);
+  ASSERT_TRUE(succeeded(xla_call_module_op));
+
+  const std::unique_ptr<OpQuantSpec> op_quant_spec =
+      GetStableHloOpQuantSpec(*xla_call_module_op);
+  ASSERT_THAT(op_quant_spec, NotNull());
+
+  EXPECT_THAT(op_quant_spec->coeff_op_quant_dim,
+              UnorderedElementsAre(Pair(1, 3)));
+}
+
+using GetStableHloQuantConstraintsTest = ::mlir::quant::QuantizationTestBase;
+
+TEST_F(GetStableHloQuantConstraintsTest,
+       HasSameOperandAndResultTypeRequirementSucceeds) {
+  // Quantizable ops: constants
+  // Non-quantizable ops: normal StableHLO ops and terminators
+  constexpr absl::string_view kModuleGather = R"mlir(
+    module {
+      func.func @main() -> (tensor<2x3x2x2xf32>) {
+        %0 = stablehlo.constant dense<1.0> : tensor<3x4x2xf32>
+        %1 = stablehlo.constant dense<2> : tensor<2x3x2xi64>
+        %2 = "stablehlo.gather"(%0, %1) {
+          dimension_numbers = #stablehlo.gather<
+            offset_dims = [2, 3],
+            collapsed_slice_dims = [0],
+            start_index_map = [1, 0],
+            index_vector_dim = 2>,
+          slice_sizes = array<i64: 1, 2, 2>,
+          indices_are_sorted = false
+        } : (tensor<3x4x2xf32>, tensor<2x3x2xi64>) -> tensor<2x3x2x2xf32>
+        func.return %2 : tensor<2x3x2x2xf32>
+      }
+    }
+  )mlir";
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleGather);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  Operation* gather_op = FindOperationOfType<GatherOp>(main_fn);
+  const auto spec = GetStableHloQuantConstraints(gather_op);
+
+  EXPECT_THAT(spec, NotNull());
+  EXPECT_THAT(spec->has_same_operand_and_result_type_requirement, IsTrue());
+}
+
+}  // namespace
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/convert_tf_quant_ops_to_mhlo.cc

@@ -0,0 +1,781 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <algorithm>
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Func/Extensions/AllExtensions.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/Quant.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/Dialect/Shape/IR/Shape.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Matchers.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "stablehlo/dialect/ChloOps.h"  // from @stablehlo
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/passes.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/utils/tf_type_utils.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/mangling_util.h"
+#include "tensorflow/compiler/mlir/tf2xla/transforms/utils.h"
+#include "xla/hlo/translate/hlo_to_mhlo/attribute_importer.h"
+#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/numeric_types.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor.pb.h"
+#include "tensorflow/core/util/quantization/uniform_quant_ops_attr.pb.h"
+#include "tensorflow/core/util/quantization/uniform_quant_ops_params.h"
+
+namespace mlir::quant::stablehlo {
+namespace {
+
+using quant::tensorflow::GetDenseAttrFromTensorProtoAttr;
+using quant::tensorflow::GetIntTypeFromTFQint;
+using quant::tensorflow::IsTFQintType;
+
+#define GEN_PASS_DEF_CONVERTTFQUANTOPSTOMHLO
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/passes.h.inc"
+
+template <typename UniformQuantizedOp>
+FailureOr<TensorType> GetUniformQuantizedType(
+    UniformQuantizedOp op, Type original_type,
+    TypedValue<TensorType> scales_value,
+    TypedValue<TensorType> zero_points_value, FloatType expressed_type,
+    int64_t storage_type_min, int64_t storage_type_max,
+    int64_t quantized_dimension, PatternRewriter &rewriter) {
+  // Check whether the scales operand has constant op.
+  DenseFPElementsAttr scales;
+  if (!matchPattern(scales_value, m_Constant(&scales))) {
+    return rewriter.notifyMatchFailure(op, "scales must be constant");
+  }
+
+  // Check whether the zero_points operand has constant op.
+  DenseIntElementsAttr zero_points;
+  if (!matchPattern(zero_points_value, m_Constant(&zero_points))) {
+    return rewriter.notifyMatchFailure(op, "zero_points must be constant");
+  }
+
+  auto original_element_type = getElementTypeOrSelf(original_type);
+  if (!mlir::isa<TF::Qint8Type, TF::Quint8Type, TF::Qint32Type>(
+          original_element_type)) {
+    return rewriter.notifyMatchFailure(
+        op, "Quantized type must be qint8, quint8 or qint32.");
+  }
+  auto storage_type = GetIntTypeFromTFQint(original_element_type);
+
+  const unsigned flags = mlir::isa<TF::Quint8Type>(original_element_type)
+                             ? 0
+                             : quant::QuantizationFlags::Signed;
+  Type elem_ty;
+  if (quantized_dimension == -1) {
+    elem_ty = quant::UniformQuantizedType::get(
+        flags, storage_type, expressed_type, scales.getValues<float>()[0],
+        zero_points.getValues<int32_t>()[0], storage_type_min,
+        storage_type_max);
+  } else {
+    SmallVector<double> scales_vec;
+    SmallVector<int64_t> zero_points_vec;
+    for (auto elem : scales.getValues<float>()) scales_vec.push_back(elem);
+    for (auto elem : zero_points.getValues<int32_t>())
+      zero_points_vec.push_back(elem);
+    elem_ty = quant::UniformQuantizedPerAxisType::get(
+        flags, storage_type, expressed_type, scales_vec, zero_points_vec,
+        quantized_dimension, storage_type_min, storage_type_max);
+  }
+
+  return mlir::cast<TensorType>(original_type).clone(elem_ty);
+}
+
+// If operand is TF const op, create MHLO constant op from the contents.
+// Otherwise convert the operand to the desired type.
+FailureOr<Value> CreateConstantOrConvertOp(Operation *op, Value operand,
+                                           TensorType new_operand_type,
+                                           PatternRewriter &rewriter) {
+  // Check whether the rhs operand has constant op.
+  TF::TensorProtoAttr tensor_proto_attr;
+  if (!matchPattern(operand, m_Constant(&tensor_proto_attr))) {
+    return Value(mhlo::BitcastConvertOp::create(rewriter, op->getLoc(),
+                                                new_operand_type, operand));
+  }
+
+  auto dense_attr_or = GetDenseAttrFromTensorProtoAttr(
+      tensor_proto_attr.getValue(), new_operand_type);
+  if (failed(dense_attr_or)) return failure();
+
+  return Value(mhlo::ConstantOp::create(rewriter, op->getLoc(),
+                                        new_operand_type, *dense_attr_or));
+}
+
+xla::ConvolutionDimensionNumbers ConvertConvolutionDimensionNumbers(
+    const ::tensorflow::UniformQuantizedConvolutionDimensionNumbersAttr
+        &dnums_input) {
+  xla::ConvolutionDimensionNumbers dnums;
+  dnums.set_input_batch_dimension(dnums_input.input_batch_dimension());
+  dnums.set_input_feature_dimension(dnums_input.input_feature_dimension());
+  for (auto value : dnums_input.input_spatial_dimensions()) {
+    dnums.add_input_spatial_dimensions(value);
+  }
+  dnums.set_kernel_input_feature_dimension(
+      dnums_input.kernel_input_feature_dimension());
+  dnums.set_kernel_output_feature_dimension(
+      dnums_input.kernel_output_feature_dimension());
+  for (auto value : dnums_input.kernel_spatial_dimensions()) {
+    dnums.add_kernel_spatial_dimensions(value);
+  }
+  dnums.set_output_batch_dimension(dnums_input.output_batch_dimension());
+  dnums.set_output_feature_dimension(dnums_input.output_feature_dimension());
+  for (auto value : dnums_input.output_spatial_dimensions()) {
+    dnums.add_output_spatial_dimensions(value);
+  }
+  return dnums;
+}
+
+DenseIntElementsAttr ConvertToDenseElementsAttr(ArrayAttr array_attr,
+                                                PatternRewriter &rewriter) {
+  SmallVector<int64_t> array;
+  array.reserve(array_attr.size());
+  for (auto elem : array_attr.getAsRange<IntegerAttr>()) {
+    array.push_back(elem.getInt());
+  }
+  return DenseIntElementsAttr::get(
+      RankedTensorType::get({static_cast<int64_t>(array_attr.size())},
+                            rewriter.getIntegerType(64)),
+      array);
+}
+
+template <typename UniformQuantizedConvolutionOp>
+FailureOr<ElementsAttr> ConvertPaddingAttr(
+    UniformQuantizedConvolutionOp op,
+    const xla::ConvolutionDimensionNumbers &dnums, PatternRewriter &rewriter) {
+  StringAttr conv_padding = op.getPaddingAttr();
+  SmallVector<int64_t> padding_nums;
+  ShapedType lhs_shape = mlir::cast<ShapedType>(op.getLhs().getType());
+  ShapedType rhs_shape = mlir::cast<ShapedType>(op.getRhs().getType());
+
+  // Handle only static shape cases.
+  // TODO(b/260284866): Handle dynamic shape cases.
+  if (!lhs_shape.hasStaticShape()) {
+    return op.emitError("lhs must have static shape.");
+  }
+  if (!rhs_shape.hasStaticShape()) {
+    return op.emitError("rhs must have static shape.");
+  }
+
+  const int64_t padding_nums_size = 2 * (rhs_shape.getRank() - 2);
+  padding_nums.reserve(padding_nums_size);
+  if (conv_padding.strref() == "EXPLICIT") {
+    for (auto padding_elem :
+         op.getExplicitPaddingAttr().template getAsRange<IntegerAttr>()) {
+      padding_nums.push_back(padding_elem.getInt());
+    }
+  } else if (conv_padding.strref() == "VALID") {
+    padding_nums.resize(padding_nums_size, 0);
+  } else {
+    padding_nums.resize(padding_nums_size);
+    for (int i = 0; i < dnums.input_spatial_dimensions_size(); ++i) {
+      const int64_t stride =
+          mlir::cast<IntegerAttr>(op.getWindowStridesAttr()[i]).getInt();
+      const int64_t lhs_size_dilated =
+          ::tensorflow::UniformQuantizedConvolutionParams::DilatedSize(
+              lhs_shape.getDimSize(dnums.input_spatial_dimensions(i)),
+              mlir::cast<IntegerAttr>(op.getLhsDilationAttr()[i]).getInt());
+      const int64_t rhs_size_dilated =
+          ::tensorflow::UniformQuantizedConvolutionParams::DilatedSize(
+              rhs_shape.getDimSize(dnums.kernel_spatial_dimensions(i)),
+              mlir::cast<IntegerAttr>(op.getRhsDilationAttr()[i]).getInt());
+
+      const int64_t output_size = (lhs_size_dilated + stride - 1) / stride;
+      const int64_t total_padding = std::max(
+          (output_size - 1) * stride + rhs_size_dilated - lhs_size_dilated,
+          static_cast<int64_t>(0));
+      const int64_t padding_begin = total_padding / 2;
+      const int64_t padding_end = total_padding - padding_begin;
+      padding_nums[2 * i] = padding_begin;
+      padding_nums[2 * i + 1] = padding_end;
+    }
+  }
+
+  ElementsAttr padding_attr = DenseIntElementsAttr::get(
+      RankedTensorType::get({static_cast<int32_t>(padding_nums.size() / 2), 2},
+                            rewriter.getIntegerType(64)),
+      padding_nums);
+  return padding_attr;
+}
+
+template <typename UniformQuantizedConvolutionOp>
+FailureOr<SmallVector<NamedAttribute>> ConvertToMhloConvolutionOpAttrs(
+    UniformQuantizedConvolutionOp op, PatternRewriter &rewriter) {
+  // TODO(b/261005147): Update the lowering logic after migration to mhlo
+  // ConvolutionDimensionNumbers.
+  ::tensorflow::UniformQuantizedConvolutionDimensionNumbersAttr dnums_input;
+  if (!dnums_input.ParseFromString(std::string(op.getDimensionNumbers()))) {
+    return op->emitError("Parse dimension_numbers failed.");
+  }
+  xla::ConvolutionDimensionNumbers dnums =
+      ConvertConvolutionDimensionNumbers(dnums_input);
+
+  SmallVector<NamedAttribute> converted_attrs;
+  for (auto attr : op->getAttrs()) {
+    if (attr.getName() == op.getFeatureGroupCountAttrName() ||
+        attr.getName() == op.getBatchGroupCountAttrName()) {
+      converted_attrs.push_back(attr);
+    } else if (attr.getName() == op.getDimensionNumbersAttrName()) {
+      attr.setValue(xla::ConvertConvDimensionNumbers(dnums, &rewriter));
+      converted_attrs.push_back(attr);
+    } else if (attr.getName() == op.getPaddingAttrName()) {
+      auto value_or = ConvertPaddingAttr(op, dnums, rewriter);
+      if (failed(value_or)) {
+        return failure();
+      }
+      attr.setValue(*value_or);
+      converted_attrs.push_back(attr);
+    } else if (attr.getName() == op.getWindowStridesAttrName() ||
+               attr.getName() == op.getLhsDilationAttrName() ||
+               attr.getName() == op.getRhsDilationAttrName()) {
+      attr.setValue(ConvertToDenseElementsAttr(
+          mlir::cast<ArrayAttr>(attr.getValue()), rewriter));
+      converted_attrs.push_back(attr);
+    }
+  }
+  return converted_attrs;
+}
+
+// TODO(hinsu): Move this pattern to legalize_tf after resolving the dependency
+// on the tensor proto.
+class ConvertUniformQuantizedDotHybridOp
+    : public OpConversionPattern<TF::UniformQuantizedDotHybridOp> {
+ public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      TF::UniformQuantizedDotHybridOp op,
+      TF::UniformQuantizedDotHybridOpAdaptor adaptor,
+      ConversionPatternRewriter &rewriter) const override {
+    // Uniform Quantized type for the rhs.
+    int64_t rhs_quantized_dimension = op.getRhsQuantizationAxis();
+    // Currently for dot, PTQ supports per-tensor quantization.
+    if (rhs_quantized_dimension != -1) {
+      return rewriter.notifyMatchFailure(
+          op, "Legalization supports only rhs_quantization_axis -1.");
+    }
+    auto rhs_type = GetUniformQuantizedType(
+        op, op.getRhs().getType(), op.getRhsScales(), op.getRhsZeroPoints(),
+        /*expressed_type=*/rewriter.getF32Type(), op.getRhsQuantizationMinVal(),
+        op.getRhsQuantizationMaxVal(), rhs_quantized_dimension, rewriter);
+    if (failed(rhs_type)) {
+      return failure();
+    }
+
+    auto rhs_or =
+        CreateConstantOrConvertOp(op, adaptor.getRhs(), *rhs_type, rewriter);
+    if (failed(rhs_or)) {
+      return failure();
+    }
+
+    rewriter.replaceOpWithNewOp<mhlo::DotOp>(op, op.getType(), op.getLhs(),
+                                             *rhs_or,
+                                             /*precision_config=*/nullptr);
+    return success();
+  }
+};
+
+class ConvertUniformQuantizedConvolutionHybridOp
+    : public OpConversionPattern<TF::UniformQuantizedConvolutionHybridOp> {
+ public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      TF::UniformQuantizedConvolutionHybridOp op,
+      TF::UniformQuantizedConvolutionHybridOpAdaptor adaptor,
+      ConversionPatternRewriter &rewriter) const override {
+    // Uniform Quantized type for the rhs.
+    auto rhs_type = GetUniformQuantizedType(
+        op, op.getRhs().getType(), op.getRhsScales(), op.getRhsZeroPoints(),
+        /*expressed_type=*/rewriter.getF32Type(), op.getRhsQuantizationMinVal(),
+        op.getRhsQuantizationMaxVal(), op.getRhsQuantizationAxis(), rewriter);
+    if (failed(rhs_type)) {
+      return failure();
+    }
+
+    auto rhs_or =
+        CreateConstantOrConvertOp(op, adaptor.getRhs(), *rhs_type, rewriter);
+    if (failed(rhs_or)) {
+      return failure();
+    }
+
+    auto converted_attrs_or = ConvertToMhloConvolutionOpAttrs(op, rewriter);
+    if (failed(converted_attrs_or)) {
+      return failure();
+    }
+    SmallVector<Value, 2> operands{op.getLhs(), *rhs_or};
+    rewriter.replaceOpWithNewOp<mhlo::ConvolutionOp>(op, op.getType(), operands,
+                                                     *converted_attrs_or);
+    return success();
+  }
+};
+
+class ConvertUniformQuantizeOp
+    : public OpRewritePattern<TF::UniformQuantizeOp> {
+ public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(TF::UniformQuantizeOp op,
+                                PatternRewriter &rewriter) const override {
+    auto output_type = GetUniformQuantizedType(
+        op, op.getOutput().getType(), op.getScales(), op.getZeroPoints(),
+        /*expressed_type=*/rewriter.getF32Type(), op.getQuantizationMinVal(),
+        op.getQuantizationMaxVal(), op.getQuantizationAxis(), rewriter);
+    if (failed(output_type)) {
+      return failure();
+    }
+
+    auto result = mhlo::UniformQuantizeOp::create(rewriter, op->getLoc(),
+                                                  *output_type, op.getInput());
+    rewriter.replaceOpWithNewOp<mhlo::BitcastConvertOp>(
+        op,
+        output_type->clone(
+            mlir::dyn_cast<quant::QuantizedType>(output_type->getElementType())
+                .getStorageType()),
+        result);
+
+    return success();
+  }
+};
+
+// UniformDequantizeOp takes TF quantized types as input which would have been
+// converted to the mhlo quantized types. Use OpConversionPattern in order to
+// retrieve the operand type *after* conversion, using OpAdaptor operand
+// accessor.
+// Same for other Uniform Quant Ops that take TF quantized types as input.
+class ConvertUniformDequantizeOp
+    : public OpConversionPattern<TF::UniformDequantizeOp> {
+ public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      TF::UniformDequantizeOp op, TF::UniformDequantizeOpAdaptor adaptor,
+      ConversionPatternRewriter &rewriter) const override {
+    Value input = adaptor.getInput();
+
+    auto input_quant_type = GetUniformQuantizedType(
+        op, op.getInput().getType(), op.getScales(), op.getZeroPoints(),
+        /*expressed_type=*/rewriter.getF32Type(), op.getQuantizationMinVal(),
+        op.getQuantizationMaxVal(), op.getQuantizationAxis(), rewriter);
+    if (failed(input_quant_type)) {
+      return failure();
+    }
+    input = mhlo::BitcastConvertOp::create(rewriter, op->getLoc(),
+                                           *input_quant_type, input);
+
+    rewriter.replaceOpWithNewOp<mhlo::UniformDequantizeOp>(
+        op, op.getOutput().getType(), input);
+    return success();
+  }
+};
+
+class ConvertUniformRequantizeOp
+    : public OpConversionPattern<TF::UniformRequantizeOp> {
+ public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      TF::UniformRequantizeOp op, TF::UniformRequantizeOpAdaptor adaptor,
+      ConversionPatternRewriter &rewriter) const override {
+    Value input = adaptor.getInput();
+
+    auto input_quant_type = GetUniformQuantizedType(
+        op, op.getInput().getType(), op.getInputScales(),
+        op.getInputZeroPoints(),
+        /*expressed_type=*/rewriter.getF32Type(),
+        op.getInputQuantizationMinVal(), op.getInputQuantizationMaxVal(),
+        op.getInputQuantizationAxis(), rewriter);
+    if (failed(input_quant_type)) {
+      return failure();
+    }
+
+    auto output_type = GetUniformQuantizedType(
+        op, op.getOutput().getType(), op.getOutputScales(),
+        op.getOutputZeroPoints(),
+        /*expressed_type=*/rewriter.getF32Type(),
+        op.getOutputQuantizationMinVal(), op.getOutputQuantizationMaxVal(),
+        op.getOutputQuantizationAxis(), rewriter);
+    if (failed(output_type)) {
+      return failure();
+    }
+
+    auto input_quant = mhlo::BitcastConvertOp::create(rewriter, op->getLoc(),
+                                                      *input_quant_type, input);
+    auto result = mhlo::UniformQuantizeOp::create(rewriter, op->getLoc(),
+                                                  *output_type, input_quant);
+    rewriter.replaceOpWithNewOp<mhlo::BitcastConvertOp>(
+        op,
+        output_type->clone(
+            mlir::dyn_cast<quant::QuantizedType>(output_type->getElementType())
+                .getStorageType()),
+        result);
+    return success();
+  }
+};
+
+class ConvertUniformQuantizedDotOp
+    : public OpConversionPattern<TF::UniformQuantizedDotOp> {
+ public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      TF::UniformQuantizedDotOp op, TF::UniformQuantizedDotOpAdaptor adaptor,
+      ConversionPatternRewriter &rewriter) const override {
+    Value lhs = adaptor.getLhs();
+
+    auto lhs_quant_type = GetUniformQuantizedType(
+        op, op.getLhs().getType(), op.getLhsScales(), op.getLhsZeroPoints(),
+        /*expressed_type=*/rewriter.getF32Type(), op.getLhsQuantizationMinVal(),
+        op.getLhsQuantizationMaxVal(), op.getLhsQuantizationAxis(), rewriter);
+    if (failed(lhs_quant_type)) {
+      return failure();
+    }
+    lhs = mhlo::BitcastConvertOp::create(rewriter, op->getLoc(),
+                                         *lhs_quant_type, adaptor.getLhs());
+
+    // Uniform Quantized type for the rhs.
+    int64_t rhs_quantized_dimension = op.getRhsQuantizationAxis();
+    // Currently for dot, PTQ supports per-tensor quantization.
+    if (rhs_quantized_dimension != -1) {
+      return rewriter.notifyMatchFailure(
+          op, "Legalization supports only rhs_quantization_axis -1.");
+    }
+    auto rhs_type = GetUniformQuantizedType(
+        op, op.getRhs().getType(), op.getRhsScales(), op.getRhsZeroPoints(),
+        /*expressed_type=*/rewriter.getF32Type(), op.getRhsQuantizationMinVal(),
+        op.getRhsQuantizationMaxVal(), rhs_quantized_dimension, rewriter);
+    if (failed(rhs_type)) {
+      return failure();
+    }
+
+    auto rhs_or =
+        CreateConstantOrConvertOp(op, adaptor.getRhs(), *rhs_type, rewriter);
+    if (failed(rhs_or)) {
+      return failure();
+    }
+
+    auto output_type = GetUniformQuantizedType(
+        op, op.getOutput().getType(), op.getOutputScales(),
+        op.getOutputZeroPoints(),
+        /*expressed_type=*/rewriter.getF32Type(),
+        op.getOutputQuantizationMinVal(), op.getOutputQuantizationMaxVal(),
+        op.getOutputQuantizationAxis(), rewriter);
+    if (failed(output_type)) {
+      return failure();
+    }
+
+    auto result =
+        mhlo::DotOp::create(rewriter, op->getLoc(), *output_type, lhs, *rhs_or,
+                            /*precision_config=*/nullptr);
+    rewriter.replaceOpWithNewOp<mhlo::BitcastConvertOp>(
+        op,
+        output_type->clone(
+            mlir::dyn_cast<quant::QuantizedType>(output_type->getElementType())
+                .getStorageType()),
+        result);
+    return success();
+  }
+};
+
+class ConvertUniformQuantizedConvolutionOp
+    : public OpConversionPattern<TF::UniformQuantizedConvolutionOp> {
+ public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      TF::UniformQuantizedConvolutionOp op,
+      TF::UniformQuantizedConvolutionOpAdaptor adaptor,
+      ConversionPatternRewriter &rewriter) const override {
+    Value lhs = adaptor.getLhs();
+
+    auto lhs_quant_type = GetUniformQuantizedType(
+        op, op.getLhs().getType(), op.getLhsScales(), op.getLhsZeroPoints(),
+        /*expressed_type=*/rewriter.getF32Type(), op.getLhsQuantizationMinVal(),
+        op.getLhsQuantizationMaxVal(), op.getLhsQuantizationAxis(), rewriter);
+    if (failed(lhs_quant_type)) {
+      return failure();
+    }
+    lhs = mhlo::BitcastConvertOp::create(rewriter, op->getLoc(),
+                                         *lhs_quant_type, adaptor.getLhs());
+
+    auto rhs_type = GetUniformQuantizedType(
+        op, op.getRhs().getType(), op.getRhsScales(), op.getRhsZeroPoints(),
+        /*expressed_type=*/rewriter.getF32Type(), op.getRhsQuantizationMinVal(),
+        op.getRhsQuantizationMaxVal(), op.getRhsQuantizationAxis(), rewriter);
+    if (failed(rhs_type)) {
+      return failure();
+    }
+
+    auto rhs_or =
+        CreateConstantOrConvertOp(op, adaptor.getRhs(), *rhs_type, rewriter);
+    if (failed(rhs_or)) {
+      return failure();
+    }
+
+    auto output_type = GetUniformQuantizedType(
+        op, op.getOutput().getType(), op.getOutputScales(),
+        op.getOutputZeroPoints(),
+        /*expressed_type=*/rewriter.getF32Type(),
+        op.getOutputQuantizationMinVal(), op.getOutputQuantizationMaxVal(),
+        op.getOutputQuantizationAxis(), rewriter);
+    if (failed(output_type)) {
+      return failure();
+    }
+
+    auto converted_attrs_or = ConvertToMhloConvolutionOpAttrs(op, rewriter);
+    if (failed(converted_attrs_or)) {
+      return failure();
+    }
+    SmallVector<Value, 2> operands{lhs, *rhs_or};
+    auto result = mhlo::ConvolutionOp::create(
+        rewriter, op->getLoc(), *output_type, operands, *converted_attrs_or);
+    rewriter.replaceOpWithNewOp<mhlo::BitcastConvertOp>(
+        op,
+        output_type->clone(
+            mlir::dyn_cast<quant::QuantizedType>(output_type->getElementType())
+                .getStorageType()),
+        result);
+    return success();
+  }
+};
+
+class ConvertUniformQuantizedAddOp
+    : public OpConversionPattern<TF::UniformQuantizedAddOp> {
+ public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      TF::UniformQuantizedAddOp op, TF::UniformQuantizedAddOpAdaptor adaptor,
+      ConversionPatternRewriter &rewriter) const override {
+    Value lhs = adaptor.getLhs();
+
+    auto lhs_type = mlir::cast<ShapedType>(lhs.getType());
+    if (!lhs_type.hasRank()) {
+      return rewriter.notifyMatchFailure(
+          op, "Legalization supports cases where only lhs rank known.");
+    }
+
+    auto lhs_quant_type = GetUniformQuantizedType(
+        op, op.getLhs().getType(), op.getLhsScales(), op.getLhsZeroPoints(),
+        /*expressed_type=*/rewriter.getF32Type(), op.getLhsQuantizationMinVal(),
+        op.getLhsQuantizationMaxVal(), op.getLhsQuantizationAxis(), rewriter);
+    if (failed(lhs_quant_type)) {
+      return failure();
+    }
+    lhs = mhlo::BitcastConvertOp::create(rewriter, op->getLoc(),
+                                         *lhs_quant_type, adaptor.getLhs());
+
+    // rhs (bias) is always 1D that broadcasts to the last dim of lhs.
+    auto broadcast_dims =
+        rewriter.getDenseI64ArrayAttr({lhs_type.getRank() - 1});
+
+    auto rhs_type = GetUniformQuantizedType(
+        op, op.getRhs().getType(), op.getRhsScales(), op.getRhsZeroPoints(),
+        /*expressed_type=*/rewriter.getF32Type(), op.getRhsQuantizationMinVal(),
+        op.getRhsQuantizationMaxVal(), op.getRhsQuantizationAxis(), rewriter);
+    if (failed(rhs_type)) {
+      return failure();
+    }
+
+    auto rhs_or =
+        CreateConstantOrConvertOp(op, adaptor.getRhs(), *rhs_type, rewriter);
+    if (failed(rhs_or)) {
+      return failure();
+    }
+
+    auto output_type = GetUniformQuantizedType(
+        op, op.getOutput().getType(), op.getOutputScales(),
+        op.getOutputZeroPoints(),
+        /*expressed_type=*/rewriter.getF32Type(),
+        op.getOutputQuantizationMinVal(), op.getOutputQuantizationMaxVal(),
+        op.getOutputQuantizationAxis(), rewriter);
+    if (failed(output_type)) {
+      return failure();
+    }
+
+    // lhs, rhs, output scales and zero_points are guaranteed (by the TF
+    // quantizer) to be identical, respectively.
+    auto result = chlo::BroadcastAddOp::create(
+        rewriter, op->getLoc(), *output_type, lhs, *rhs_or, broadcast_dims);
+    rewriter.replaceOpWithNewOp<mhlo::BitcastConvertOp>(
+        op,
+        output_type->clone(
+            mlir::dyn_cast<quant::QuantizedType>(output_type->getElementType())
+                .getStorageType()),
+        result);
+    return success();
+  }
+};
+
+class ConvertUniformQuantizedClipByValueOp
+    : public OpConversionPattern<TF::UniformQuantizedClipByValueOp> {
+ public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      TF::UniformQuantizedClipByValueOp op,
+      TF::UniformQuantizedClipByValueOpAdaptor adaptor,
+      ConversionPatternRewriter &rewriter) const override {
+    Value operand = adaptor.getOperand();
+
+    const int64_t quantization_axis = op.getQuantizationAxis();
+    llvm::SmallVector<int64_t> broadcast_dims_values = {};
+    if (quantization_axis >= 0) {
+      broadcast_dims_values.push_back(quantization_axis);
+    }
+    auto broadcast_dims = rewriter.getDenseI64ArrayAttr(broadcast_dims_values);
+
+    auto min_max_type = GetUniformQuantizedType(
+        op, op.getMin().getType(), op.getScales(), op.getZeroPoints(),
+        /*expressed_type=*/rewriter.getF32Type(), op.getQuantizationMinVal(),
+        op.getQuantizationMaxVal(), op.getQuantizationAxis(), rewriter);
+    if (failed(min_max_type)) {
+      return failure();
+    }
+    auto min_or = CreateConstantOrConvertOp(op, adaptor.getMin(), *min_max_type,
+                                            rewriter);
+    if (failed(min_or)) {
+      return failure();
+    }
+    auto max_or = CreateConstantOrConvertOp(op, adaptor.getMax(), *min_max_type,
+                                            rewriter);
+    if (failed(max_or)) {
+      return failure();
+    }
+
+    auto output_type = GetUniformQuantizedType(
+        op, op.getOutput().getType(), op.getScales(), op.getZeroPoints(),
+        /*expressed_type=*/rewriter.getF32Type(), op.getQuantizationMinVal(),
+        op.getQuantizationMaxVal(), op.getQuantizationAxis(), rewriter);
+    if (failed(output_type)) {
+      return failure();
+    }
+    operand = mhlo::BitcastConvertOp::create(rewriter, op->getLoc(),
+                                             *output_type, operand);
+
+    Value res_min_clipped = chlo::BroadcastMaxOp::create(
+        rewriter, op->getLoc(), *output_type, operand, *min_or, broadcast_dims);
+    Value res_max_clipped =
+        chlo::BroadcastMinOp::create(rewriter, op->getLoc(), *output_type,
+                                     res_min_clipped, *max_or, broadcast_dims);
+    rewriter.replaceOpWithNewOp<mhlo::BitcastConvertOp>(
+        op,
+        output_type->clone(
+            mlir::dyn_cast<quant::QuantizedType>(output_type->getElementType())
+                .getStorageType()),
+        res_max_clipped);
+    return success();
+  }
+};
+
+// This pattern converts qint <-> int CastOp to int -> int ConvertOps.
+// The former are introduced in ConvertTFQuantTypes pass. The resulting int <->
+// int ConvertOps are no-ops and can be removed later in a Canonicalizer pass.
+class ConvertTfCastOp : public OpConversionPattern<TF::CastOp> {
+ public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      TF::CastOp op, TF::CastOpAdaptor adaptor,
+      ConversionPatternRewriter &rewriter) const override {
+    Type output_type = op.getDstT();
+    if (!IsTFQintType(output_type) && !IsTFQintType(op.getSrcT())) {
+      // skip CastOps with no qint types.
+      return failure();
+    }
+    Value input = adaptor.getX();
+    rewriter.replaceOpWithNewOp<mhlo::ConvertOp>(
+        op, input, GetIntTypeFromTFQint(output_type));
+    return success();
+  }
+};
+
+class ConvertTFQuantOpsToMHLO
+    : public impl::ConvertTFQuantOpsToMHLOBase<ConvertTFQuantOpsToMHLO> {
+ public:
+  ConvertTFQuantOpsToMHLO() = default;
+  ConvertTFQuantOpsToMHLO(const ConvertTFQuantOpsToMHLO &) = default;
+
+  // Performs conversion of MHLO quant ops to primitive ops.
+  void runOnOperation() override;
+};
+
+void ConvertTFQuantOpsToMHLO::runOnOperation() {
+  MLIRContext *ctx = &getContext();
+  func::FuncOp func = getOperation();
+  ConversionTarget target(*ctx);
+  target.addLegalDialect<TF::TensorFlowDialect, mhlo::MhloDialect,
+                         chlo::ChloDialect>();
+  target.addIllegalOp<
+      TF::UniformQuantizeOp, TF::UniformRequantizeOp, TF::UniformDequantizeOp,
+      TF::UniformQuantizedDotOp, TF::UniformQuantizedDotHybridOp,
+      TF::UniformQuantizedConvolutionOp,
+      TF::UniformQuantizedConvolutionHybridOp, TF::UniformQuantizedAddOp,
+      TF::UniformQuantizedClipByValueOp>();
+  target.addDynamicallyLegalOp<TF::CastOp>([](Operation *op) {
+    auto cast_op = llvm::dyn_cast<TF::CastOp>(op);
+    return !IsTFQintType(cast_op.getSrcT()) && !IsTFQintType(cast_op.getDstT());
+  });
+
+  RewritePatternSet patterns(ctx);
+  PopulateLegalizeTfQuantizationPatterns(ctx, &patterns);
+  if (failed(applyPartialConversion(func, target, std::move(patterns)))) {
+    signalPassFailure();
+  }
+}
+
+}  // namespace
+
+void PopulateLegalizeTfQuantizationPatterns(MLIRContext *context,
+                                            RewritePatternSet *patterns) {
+  patterns
+      ->add<ConvertUniformQuantizedDotHybridOp,
+            ConvertUniformQuantizedConvolutionHybridOp,
+            ConvertUniformQuantizeOp, ConvertUniformRequantizeOp,
+            ConvertUniformDequantizeOp, ConvertUniformQuantizedDotOp,
+            ConvertUniformQuantizedConvolutionOp, ConvertUniformQuantizedAddOp,
+            ConvertUniformQuantizedClipByValueOp, ConvertTfCastOp>(context);
+}
+
+std::unique_ptr<OperationPass<func::FuncOp>>
+CreateConvertTFQuantOpsToMHLOPass() {
+  return std::make_unique<ConvertTFQuantOpsToMHLO>();
+}
+
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/convert_tf_quant_to_mhlo_int_test.cc

@@ -0,0 +1,813 @@
+/* Copyright 2023 The StableHLO Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cstdint>
+#include <cstring>
+#include <memory>
+#include <optional>
+#include <utility>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "absl/log/check.h"
+#include "absl/random/random.h"
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/Quant.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/DialectRegistry.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Parser/Parser.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "stablehlo/dialect/ChloOps.h"  // from @stablehlo
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/passes.h"
+#include "tensorflow/compiler/mlir/quantization/tensorflow/cc/constant_fold.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/convert_tensor.h"
+#include "tensorflow/compiler/tf2xla/shape_util.h"
+#include "tensorflow/compiler/tf2xla/type_util.h"
+#include "xla/error_spec.h"
+#include "xla/literal.h"
+#include "xla/literal_util.h"
+#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
+#include "xla/pjrt/maybe_owning_mlir_module.h"
+#include "xla/pjrt/pjrt_client.h"
+#include "xla/pjrt/pjrt_executable.h"
+#include "xla/pjrt/plugin/xla_cpu/cpu_client_options.h"
+#include "xla/pjrt/plugin/xla_cpu/xla_cpu_pjrt_client.h"
+#include "xla/shape.h"
+#include "xla/shape_util.h"
+#include "xla/tests/literal_test_util.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/statusor.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+
+namespace mlir::quant::stablehlo {
+namespace {
+
+using ::testing::Test;
+
+class ConvertTfQuantToMhloIntTest : public Test {
+ protected:
+  void SetUp() override {
+    DialectRegistry dialects;
+    dialects.insert<TF::TensorFlowDialect, func::FuncDialect, chlo::ChloDialect,
+                    mhlo::MhloDialect, quant::QuantDialect>();
+    ctx_ = std::make_unique<MLIRContext>(dialects);
+    ctx_->loadAllAvailableDialects();
+
+    // Create a CPU client with 1 device.
+    xla::CpuClientOptions options;
+    options.asynchronous = false;
+    options.cpu_device_count = 1;
+    TF_ASSERT_OK_AND_ASSIGN(pjrt_client_, xla::GetXlaPjrtCpuClient(options));
+    device_ = pjrt_client_->addressable_devices().front();
+    TF_ASSERT_OK_AND_ASSIGN(memory_space_, device_->default_memory_space());
+    CHECK(device_);
+  }
+
+  absl::StatusOr<OwningOpRef<ModuleOp>> ReplaceFuncArgsByConstant(
+      absl::string_view program,
+      absl::Span<const xla::Literal* const> arguments,
+      bool use_mhlo_const = false) {
+    auto module_op = parseSourceString<ModuleOp>(program, ctx_.get());
+    CHECK(module_op);
+    auto func_op = llvm::dyn_cast<func::FuncOp>(
+        *module_op->getBodyRegion().getOps().begin());
+    if (!func_op) {
+      return absl::InternalError("Input MLIR must have only 1 func");
+    }
+    if (arguments.size() != func_op.getNumArguments()) {
+      return absl::InternalError("Input argument has wrong size");
+    }
+
+    // Convert input xla::Literal arguments to constants, this allows using
+    // constant folding to evaluate function return value.
+    mlir::OpBuilder builder(ctx_.get());
+    for (int i = 0; i < arguments.size(); ++i) {
+      const xla::Literal* const xla_literal = arguments[i];
+      tensorflow::TensorShape shape;
+      TF_ASSIGN_OR_RETURN(auto data_type,
+                          tensorflow::EncodePrimitiveTypeAsDataType(
+                              xla_literal->shape().element_type()));
+      TF_RETURN_IF_ERROR(
+          tensorflow::XLAShapeToTensorShape(xla_literal->shape(), &shape));
+      tensorflow::Tensor tensor(data_type, shape);
+      std::memcpy(static_cast<char*>(tensor.data()),
+                  xla_literal->untyped_data(),
+                  xla::ShapeUtil::ByteSizeOfPrimitiveType(
+                      xla_literal->shape().element_type()) *
+                      xla_literal->element_count());
+      TF_ASSIGN_OR_RETURN(auto attrs,
+                          tensorflow::ConvertTensor(tensor, &builder));
+      builder.setInsertionPoint(
+          &func_op.getFunctionBody().getBlocks().front().front());
+      // Use mhlo.Constant when it is consumed by the lowering passes since they
+      // can't lower tf.Const.
+      Value cst;
+      if (use_mhlo_const) {
+        cst = mhlo::ConstantOp::create(builder, func_op->getLoc(), attrs);
+      } else {
+        cst = TF::ConstOp::create(builder, func_op->getLoc(), attrs);
+      }
+      func_op.getArgument(i).replaceAllUsesWith(cst);
+    }
+    return module_op;
+  }
+
+  // Evaluate return value of a function using TF kernel.
+  // This assumes that the module op has only 1 function and it has TF ops only.
+  absl::StatusOr<std::shared_ptr<xla::Literal>> EvaluateTfFunction(
+      absl::string_view program,
+      absl::Span<const xla::Literal* const> arguments) {
+    TF_ASSIGN_OR_RETURN(auto module_op,
+                        ReplaceFuncArgsByConstant(program, arguments));
+    // Constant fold the func.Return op's producer op to evaluate the return
+    // value. The evaluation will use TF kernels.
+    // This assumes that func.Return is the last op in the function and it
+    // returns only 1 value.
+    auto& return_op = llvm::dyn_cast<func::FuncOp>(
+                          *module_op->getBodyRegion().getOps().begin())
+                          .getFunctionBody()
+                          .getBlocks()
+                          .back()
+                          .back();
+    if (!llvm::isa<func::ReturnOp>(return_op) ||
+        return_op.getNumOperands() != 1) {
+      return absl::InternalError(
+          "Func must have ReturnOp as last op and must return 1 value");
+    }
+    auto def_op = return_op.getOperand(0).getDefiningOp();
+    auto fold_results = ConstantFoldOpIfPossible(def_op);
+    if (fold_results.size() != 1 ||
+        !llvm::isa<TF::ConstOp>(fold_results[0].getDefiningOp())) {
+      return absl::InternalError("Failed to evaluate TF ops");
+    }
+
+    // Convert output tensor back to xla::Literal.
+    tensorflow::Tensor tensor;
+    TF_RETURN_IF_ERROR(tensorflow::ConvertToTensor(
+        llvm::dyn_cast<TF::ConstOp>(fold_results[0].getDefiningOp()).getValue(),
+        &tensor));
+    xla::Shape xla_shape;
+    TF_RETURN_IF_ERROR(tensorflow::TensorShapeToXLAShape(
+        tensor.dtype(), tensor.shape(), &xla_shape));
+    xla::PjRtClient::HostBufferSemantics host_buffer_semantics =
+        xla::PjRtClient::HostBufferSemantics::kImmutableUntilTransferCompletes;
+    TF_ASSIGN_OR_RETURN(
+        auto buffer,
+        pjrt_client_->BufferFromHostBuffer(
+            tensor.data(), xla_shape.element_type(), xla_shape.dimensions(),
+            /*byte_strides=*/std::nullopt, host_buffer_semantics,
+            /*on_done_with_host_buffer=*/nullptr,
+            *device_->default_memory_space(), /*device_layout=*/nullptr));
+    return buffer->ToLiteral().Await();
+  }
+
+  absl::StatusOr<std::unique_ptr<xla::PjRtLoadedExecutable>> CompileProgram(
+      absl::string_view program,
+      absl::Span<const xla::Literal* const> arguments) {
+    // Replace args by mhlo.constant since the lowering passes can't lower
+    // tf.Const.
+    TF_ASSIGN_OR_RETURN(
+        auto module_op,
+        ReplaceFuncArgsByConstant(program, arguments, /*use_mhlo_const=*/true));
+
+    // Run the Convert TF Quant Types, TF Quant -> MHLO Quant and MHLO Quant ->
+    // MHLO int passes.
+    PassManager pm(module_op->getContext());
+    pm.addNestedPass<func::FuncOp>(CreateConvertTFQuantTypesPass());
+    AddQuantizationLoweringPasses(pm);
+    CHECK(succeeded(pm.run(module_op.get())));
+    // Compile the program.
+    return pjrt_client_->CompileAndLoad(
+        xla::MaybeOwningMlirModule(std::move(module_op)),
+        xla::CompileOptions{});
+  }
+
+  absl::StatusOr<std::shared_ptr<xla::Literal>>
+  ExecuteProgramAndReturnSingleResult(
+      xla::PjRtLoadedExecutable* executable,
+      absl::Span<const xla::Literal* const> arguments) {
+    // Process and buffer arguments.
+    std::vector<std::unique_ptr<xla::PjRtBuffer>> buffers;
+    std::vector<xla::PjRtBuffer*> buffer_ptrs;
+    buffers.reserve(arguments.size());
+    for (const xla::Literal* argument : arguments) {
+      TF_ASSIGN_OR_RETURN(auto buffer, pjrt_client_->BufferFromHostLiteral(
+                                           *argument, memory_space_));
+      buffer_ptrs.push_back(buffer.get());
+      buffers.push_back(std::move(buffer));
+    }
+    // Run the executable.
+    TF_ASSIGN_OR_RETURN(auto result,
+                        executable->Execute({buffer_ptrs}, /*options=*/{}));
+    CHECK(result.size() == 1 && result[0].size() == 1);
+    return result[0][0]->ToLiteral().Await();
+  }
+
+  void ExecuteAndCompareResultsWithTfKernel(
+      absl::string_view program,
+      absl::Span<const xla::Literal* const> arguments,
+      std::optional<absl::string_view> tf_program = std::nullopt,
+      double error_tolerance = 0.1) {
+    // Expected result is calculated by evaluating using TF kernels. In some
+    // cases, TF kernel behaves differently from lowered graph (e.g. Hybrid
+    // ops). So we optionally use a different graph to calculate the expected
+    // result.
+    TF_ASSERT_OK_AND_ASSIGN(
+        auto expected,
+        this->EvaluateTfFunction(
+            (tf_program.has_value() ? *tf_program : program), arguments));
+
+    TF_ASSERT_OK_AND_ASSIGN(auto executable,
+                            this->CompileProgram(program, arguments));
+    TF_ASSERT_OK_AND_ASSIGN(
+        auto result,
+        this->ExecuteProgramAndReturnSingleResult(executable.get(), arguments));
+
+    // Convert to double for comparison. This is needed for comparing integers
+    // since it LiteralTestUtil asserts different integers even if it is within
+    // error_spec.
+    TF_ASSERT_OK_AND_ASSIGN(auto expected_double, expected->Convert(xla::F64));
+    TF_ASSERT_OK_AND_ASSIGN(auto result_double, result->Convert(xla::F64));
+    EXPECT_TRUE(xla::LiteralTestUtil::Near(expected_double, result_double,
+                                           xla::ErrorSpec(error_tolerance)));
+  }
+
+  absl::StatusOr<xla::Literal> CreateRandomF32Literal(
+      absl::Span<const int64_t> dims, float min = -100, float max = 100) {
+    TF_ASSIGN_OR_RETURN(auto shape,
+                        xla::ShapeUtil::MakeValidatedShape(xla::F32, dims));
+    return xla::LiteralUtil::CreateLiteralWithGenerator<xla::F32, float>(
+        shape, [this, min, max](absl::Span<const int64_t> dims) -> float {
+          return absl::Uniform(bitgen_, min, max);
+        });
+  }
+
+  absl::StatusOr<xla::Literal> CreateRandomI8Literal(
+      absl::Span<const int64_t> dims, int8_t min = -128, int8_t max = 127) {
+    TF_ASSIGN_OR_RETURN(auto shape,
+                        xla::ShapeUtil::MakeValidatedShape(xla::S8, dims));
+    return xla::LiteralUtil::CreateLiteralWithGenerator<xla::S8, int8_t>(
+        shape, [this, min, max](absl::Span<const int64_t> dims) -> int8_t {
+          return absl::Uniform(bitgen_, min, max);
+        });
+  }
+
+  absl::StatusOr<xla::Literal> CreateRandomI32Literal(
+      absl::Span<const int64_t> dims, int32_t min = -128, int32_t max = 127) {
+    TF_ASSIGN_OR_RETURN(auto shape,
+                        xla::ShapeUtil::MakeValidatedShape(xla::S32, dims));
+    return xla::LiteralUtil::CreateLiteralWithGenerator<xla::S32, int32_t>(
+        shape, [this, min, max](absl::Span<const int64_t> dims) -> int32_t {
+          return absl::Uniform(bitgen_, min, max);
+        });
+  }
+
+  std::unique_ptr<MLIRContext> ctx_;
+  std::unique_ptr<xla::PjRtClient> pjrt_client_;
+  xla::PjRtDevice* device_;
+  xla::PjRtMemorySpace* memory_space_;
+  absl::BitGen bitgen_;
+};
+
+TEST_F(ConvertTfQuantToMhloIntTest, UniformQuantizeAndDequantizeToValidGraph) {
+  constexpr absl::string_view kProgram = R"mlir(
+func.func @main(%arg0: tensor<10xf32>) -> tensor<10xf32> {
+  %scale = "tf.Const"() { value = dense<0.347> : tensor<f32> } : () -> tensor<f32>
+  %zp = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
+  %0 = "tf.UniformQuantize"(%arg0, %scale, %zp) {
+    quantization_axis = -1 : i64,
+    quantization_min_val = -128 : i64,
+    quantization_max_val = 127 : i64
+  } : (tensor<10xf32>, tensor<f32>, tensor<i32>) -> tensor<10x!tf_type.qint8>
+  %1 = "tf.UniformDequantize"(%0, %scale, %zp) {
+    quantization_axis = -1 : i64,
+    quantization_min_val = -128 : i64,
+    quantization_max_val = 127 : i64
+  } : (tensor<10x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<10xf32>
+  return %1 : tensor<10xf32>
+})mlir";
+  TF_ASSERT_OK_AND_ASSIGN(auto arg0, CreateRandomF32Literal({10}));
+  // error_tolerance is set to be slightly > scale because different rounding
+  // implementations for UniformQuantize in TF kernel and the lowering passes
+  // may cause +/-1 differences.
+  ExecuteAndCompareResultsWithTfKernel(
+      kProgram, {&arg0}, /*tf_program=*/std::nullopt, /*error_tolerance=*/0.35);
+}
+
+TEST_F(ConvertTfQuantToMhloIntTest,
+       UniformQuantizeAndDequantizeToValidGraphQuint8) {
+  constexpr absl::string_view kProgram = R"mlir(
+func.func @main(%arg0: tensor<10xf32>) -> tensor<10xf32> {
+  %scale = "tf.Const"() { value = dense<0.347> : tensor<f32> } : () -> tensor<f32>
+  %zp = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
+  %0 = "tf.UniformQuantize"(%arg0, %scale, %zp) {
+    quantization_axis = -1 : i64,
+    quantization_min_val = 0 : i64,
+    quantization_max_val = 255 : i64
+  } : (tensor<10xf32>, tensor<f32>, tensor<i32>) -> tensor<10x!tf_type.quint8>
+  %1 = "tf.UniformDequantize"(%0, %scale, %zp) {
+    quantization_axis = -1 : i64,
+    quantization_min_val = 0 : i64,
+    quantization_max_val = 255 : i64
+  } : (tensor<10x!tf_type.quint8>, tensor<f32>, tensor<i32>) -> tensor<10xf32>
+  return %1 : tensor<10xf32>
+})mlir";
+  TF_ASSERT_OK_AND_ASSIGN(auto arg0, CreateRandomF32Literal({10}));
+  ExecuteAndCompareResultsWithTfKernel(kProgram, {&arg0},
+                                       /*tf_program=*/std::nullopt,
+                                       /*error_tolerance=*/0.35);
+}
+
+TEST_F(ConvertTfQuantToMhloIntTest,
+       UniformQuantizeAndDequantizePerChannelToValidGraphQuint8) {
+  constexpr absl::string_view kProgram = R"mlir(
+func.func @main(
+  %arg0: tensor<10x10xf32>, %scale: tensor<10xf32>, %zp: tensor<10xi32>
+  ) -> tensor<10x10xf32> {
+  %0 = "tf.UniformQuantize"(%arg0, %scale, %zp) {
+    quantization_axis = 1 : i64,
+    quantization_min_val = 0 : i64,
+    quantization_max_val = 255 : i64
+  } : (tensor<10x10xf32>, tensor<10xf32>, tensor<10xi32>) -> tensor<10x10x!tf_type.quint8>
+  %1 = "tf.UniformDequantize"(%0, %scale, %zp) {
+    quantization_axis = 1 : i64,
+    quantization_min_val = 0 : i64,
+    quantization_max_val = 255 : i64
+  } : (tensor<10x10x!tf_type.quint8>, tensor<10xf32>, tensor<10xi32>) -> tensor<10x10xf32>
+  return %1 : tensor<10x10xf32>
+})mlir";
+  TF_ASSERT_OK_AND_ASSIGN(auto arg0, CreateRandomF32Literal({10, 10}));
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto scale, CreateRandomF32Literal({10}, /*min=*/0.0001, /*max=*/2));
+  TF_ASSERT_OK_AND_ASSIGN(auto zp,
+                          CreateRandomI32Literal({10}, /*min=*/0, /*max=*/255));
+  ExecuteAndCompareResultsWithTfKernel(kProgram, {&arg0, &scale, &zp},
+                                       /*tf_program=*/std::nullopt,
+                                       /*error_tolerance=*/0.35);
+}
+
+TEST_F(ConvertTfQuantToMhloIntTest, UniformQuantizePerChannelToValidGraph) {
+  constexpr absl::string_view kProgram = R"mlir(
+func.func @main(
+    %arg0: tensor<10x10xf32>, %scale: tensor<10xf32>, %zp: tensor<10xi32>
+  ) -> tensor<10x10xi8> {
+  %0 = "tf.UniformQuantize"(%arg0, %scale, %zp) {
+    quantization_axis = 1 : i64,
+    quantization_min_val = -128 : i64,
+    quantization_max_val = 127 : i64
+  } : (tensor<10x10xf32>, tensor<10xf32>, tensor<10xi32>) -> tensor<10x10x!tf_type.qint8>
+  %1 = "tf.Cast"(%0) {} : (tensor<10x10x!tf_type.qint8>) -> tensor<10x10xi8>
+  return %1 : tensor<10x10xi8>
+})mlir";
+  TF_ASSERT_OK_AND_ASSIGN(auto arg0, CreateRandomF32Literal({10, 10}));
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto scale, CreateRandomF32Literal({10}, /*min=*/0.0001, /*max=*/2));
+  TF_ASSERT_OK_AND_ASSIGN(auto zp, CreateRandomI32Literal({10}));
+  // Different rounding implementations for UniformQuantize in TF kernel and the
+  // lowering passes may cause +/-1 differences.
+  ExecuteAndCompareResultsWithTfKernel(kProgram, {&arg0, &scale, &zp},
+                                       /*tf_program=*/std::nullopt,
+                                       /*error_tolerance=*/1.0);
+}
+
+TEST_F(ConvertTfQuantToMhloIntTest, UniformDequantizePerChannelToValidGraph) {
+  constexpr absl::string_view kProgram = R"mlir(
+func.func @main(
+    %arg0: tensor<10x10xi8>, %scale: tensor<10xf32>, %zp: tensor<10xi32>
+  ) -> tensor<10x10xf32> {
+  %0 = "tf.Cast"(%arg0) {} : (tensor<10x10xi8>) -> tensor<10x10x!tf_type.qint8>
+  %1 = "tf.UniformDequantize"(%0, %scale, %zp) {
+    quantization_axis = 1 : i64,
+    quantization_min_val = -128 : i64,
+    quantization_max_val = 127 : i64
+  } : (tensor<10x10x!tf_type.qint8>, tensor<10xf32>, tensor<10xi32>) -> tensor<10x10xf32>
+  return %1 : tensor<10x10xf32>
+})mlir";
+  TF_ASSERT_OK_AND_ASSIGN(auto arg0, CreateRandomI8Literal({10, 10}));
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto scale, CreateRandomF32Literal({10}, /*min=*/0.0001, /*max=*/2));
+  TF_ASSERT_OK_AND_ASSIGN(auto zp, CreateRandomI32Literal({10}));
+  ExecuteAndCompareResultsWithTfKernel(kProgram, {&arg0, &scale, &zp});
+}
+
+TEST_F(ConvertTfQuantToMhloIntTest, UniformQuantizeConvolutionToValidGraph) {
+  constexpr absl::string_view kProgram = R"mlir(
+func.func @main(%input: tensor<1x9x9x9xi8>, %filter: tensor<3x3x9x10xi8>) -> tensor<1x9x9x10xi32> {
+  %input_scale = "tf.Const"() { value = dense<2.0> : tensor<f32> } : () -> tensor<f32>
+  %input_zp = "tf.Const"() { value = dense<-10> : tensor<i32> } : () -> tensor<i32>
+  %filter_scale = "tf.Const"() { value = dense<0.5> : tensor<f32> } : () -> tensor<f32>
+  %filter_zp = "tf.Const"() { value = dense<0> : tensor<i32> } : () -> tensor<i32>
+  %accum_scale = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
+  %accum_zp = "tf.Const"() { value = dense<0> : tensor<i32> } : () -> tensor<i32>
+  %quant_input = "tf.Cast"(%input) {} : (tensor<1x9x9x9xi8>) ->
+    tensor<1x9x9x9x!tf_type.qint8>
+  %quant_filter = "tf.Cast"(%filter) {} : (tensor<3x3x9x10xi8>) ->
+    tensor<3x3x9x10x!tf_type.qint8>
+  %0 = "tf.UniformQuantizedConvolution"(
+    %quant_input, %quant_filter, %input_scale, %input_zp,
+    %filter_scale, %filter_zp, %accum_scale, %accum_zp
+  ) {
+    Tin = "tfdtype$DT_QINT8", Tout = "tfdtype$DT_QINT32",
+    attr_map = "", batch_group_count = 1 : i64,
+    dimension_numbers = "\10\03\1A\02\01\02 \02(\032\02\00\01@\03J\02\01\02",
+    explicit_padding = [], feature_group_count = 1 : i64, lhs_dilation = [1, 1],
+    lhs_quantization_axis = -1 : i64, lhs_quantization_max_val = 127 : i64,
+    lhs_quantization_min_val = -128 : i64, output_quantization_axis = -1 : i64,
+    output_quantization_max_val = 2147483647 : i64,
+    output_quantization_min_val = -2147483648 : i64, padding = "SAME",
+    rhs_dilation = [1, 1], rhs_quantization_axis = -1 : i64,
+    rhs_quantization_max_val = 127 : i64, rhs_quantization_min_val = -128 : i64,
+    window_strides = [1, 1]
+  } : (tensor<1x9x9x9x!tf_type.qint8>, tensor<3x3x9x10x!tf_type.qint8>,
+    tensor<f32>, tensor<i32>, tensor<f32>, tensor<i32>, tensor<f32>, tensor<i32>
+  ) -> tensor<1x9x9x10x!tf_type.qint32>
+  %output = "tf.Cast"(%0) {} : (tensor<1x9x9x10x!tf_type.qint32>) -> tensor<1x9x9x10xi32>
+  return %output : tensor<1x9x9x10xi32>
+})mlir";
+  TF_ASSERT_OK_AND_ASSIGN(auto input, CreateRandomI8Literal({1, 9, 9, 9}));
+  TF_ASSERT_OK_AND_ASSIGN(auto filter, CreateRandomI8Literal({3, 3, 9, 10}));
+  ExecuteAndCompareResultsWithTfKernel(kProgram, {&input, &filter});
+}
+
+TEST_F(ConvertTfQuantToMhloIntTest,
+       UniformQuantizeConvolutionPerChannelToValidGraph) {
+  constexpr absl::string_view kProgram = R"mlir(
+func.func @main(
+    %input: tensor<1x9x9x9xi8>, %filter: tensor<3x3x9x10xi8>, %scale: tensor<10xf32>
+  ) -> tensor<1x9x9x10xi32> {
+  %input_scale = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
+  %input_zp = "tf.Const"() { value = dense<-10> : tensor<i32> } : () -> tensor<i32>
+  %zp = "tf.Const"() { value = dense<0> : tensor<10xi32> } : () -> tensor<10xi32>
+  %quant_input = "tf.Cast"(%input) {} : (tensor<1x9x9x9xi8>) ->
+    tensor<1x9x9x9x!tf_type.qint8>
+  %quant_filter = "tf.Cast"(%filter) {} : (tensor<3x3x9x10xi8>) ->
+    tensor<3x3x9x10x!tf_type.qint8>
+  %0 = "tf.UniformQuantizedConvolution"(
+    %quant_input, %quant_filter, %input_scale, %input_zp, %scale, %zp, %scale, %zp
+  ) {
+    Tin = "tfdtype$DT_QINT8", Tout = "tfdtype$DT_QINT32",
+    attr_map = "", batch_group_count = 1 : i64,
+    dimension_numbers = "\10\03\1A\02\01\02 \02(\032\02\00\01@\03J\02\01\02",
+    explicit_padding = [], feature_group_count = 1 : i64, lhs_dilation = [1, 1],
+    lhs_quantization_axis = -1 : i64, lhs_quantization_max_val = 127 : i64,
+    lhs_quantization_min_val = -128 : i64, output_quantization_axis = 3 : i64,
+    output_quantization_max_val = 2147483647 : i64,
+    output_quantization_min_val = -2147483648 : i64, padding = "SAME",
+    rhs_dilation = [1, 1], rhs_quantization_axis = 3 : i64,
+    rhs_quantization_max_val = 127 : i64, rhs_quantization_min_val = -128 : i64,
+    window_strides = [1, 1]
+  } : (tensor<1x9x9x9x!tf_type.qint8>, tensor<3x3x9x10x!tf_type.qint8>,
+    tensor<f32>, tensor<i32>, tensor<10xf32>, tensor<10xi32>, tensor<10xf32>, tensor<10xi32>
+  ) -> tensor<1x9x9x10x!tf_type.qint32>
+  %output = "tf.Cast"(%0) {} : (tensor<1x9x9x10x!tf_type.qint32>) -> tensor<1x9x9x10xi32>
+  return %output : tensor<1x9x9x10xi32>
+})mlir";
+  TF_ASSERT_OK_AND_ASSIGN(auto input, CreateRandomI8Literal({1, 9, 9, 9}));
+  TF_ASSERT_OK_AND_ASSIGN(auto filter, CreateRandomI8Literal({3, 3, 9, 10}));
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto scale, CreateRandomF32Literal({10}, /*min=*/0.0001, /*max=*/2));
+  ExecuteAndCompareResultsWithTfKernel(kProgram, {&input, &filter, &scale});
+}
+
+TEST_F(ConvertTfQuantToMhloIntTest,
+       UniformQuantizeConvolutionHybridToValidGraph) {
+  constexpr absl::string_view kTfProgram = R"mlir(
+func.func @main(%input: tensor<2x10x10x10xf32>, %filter: tensor<3x3x10x20xi8>) -> tensor<2x10x10x20xf32> {
+  %filter_scale = "tf.Const"() { value = dense<0.047> : tensor<f32> } : () -> tensor<f32>
+  %filter_zp = "tf.Const"() { value = dense<0> : tensor<i32> } : () -> tensor<i32>
+  %quant_filter = "tf.Cast"(%filter) {} : (tensor<3x3x10x20xi8>) ->
+    tensor<3x3x10x20x!tf_type.qint8>
+  %filter_new = "tf.UniformDequantize"(%quant_filter, %filter_scale, %filter_zp) {
+    quantization_axis = -1 : i64, quantization_min_val = -128 : i64,
+    quantization_max_val = 127 : i64
+  } : (
+    tensor<3x3x10x20x!tf_type.qint8>, tensor<f32>, tensor<i32>
+  ) -> tensor<3x3x10x20xf32>
+  %0 = "tf.Conv2D"(%input, %filter_new) {
+    Tin = "tfdtype$DT_FLOAT", Tout = "tfdtype$DT_FLOAT",
+    attr_map = "", batch_group_count = 1 : i64,
+    explicit_padding = [], feature_group_count = 1 : i64, lhs_dilation = [1, 1],
+    padding = "SAME", rhs_dilation = [1, 1], strides = [1, 1, 1, 1]
+  } : (tensor<2x10x10x10xf32>, tensor<3x3x10x20xf32>) -> tensor<2x10x10x20xf32>
+  return %0 : tensor<2x10x10x20xf32>
+})mlir";
+  constexpr absl::string_view kProgram = R"mlir(
+func.func @main(%input: tensor<2x10x10x10xf32>, %filter: tensor<3x3x10x20xi8>) -> tensor<2x10x10x20xf32> {
+  %filter_scale = "tf.Const"() { value = dense<0.047> : tensor<f32> } : () -> tensor<f32>
+  %filter_zp = "tf.Const"() { value = dense<0> : tensor<i32> } : () -> tensor<i32>
+  %quant_filter = "tf.Cast"(%filter) {} : (tensor<3x3x10x20xi8>) -> tensor<3x3x10x20x!tf_type.qint8>
+  %0 = "tf.UniformQuantizedConvolutionHybrid"(
+    %input, %quant_filter, %filter_scale, %filter_zp
+  ) {
+    Tin = "tfdtype$DT_QINT8", Tout = "tfdtype$DT_FLOAT",
+    attr_map = "", batch_group_count = 1 : i64,
+    dimension_numbers = "\10\03\1A\02\01\02 \02(\032\02\00\01@\03J\02\01\02",
+    explicit_padding = [], feature_group_count = 1 : i64, lhs_dilation = [1, 1],
+    padding = "SAME", rhs_dilation = [1, 1], rhs_quantization_axis = -1 : i64,
+    rhs_quantization_max_val = 127 : i64, rhs_quantization_min_val = -128 : i64,
+    window_strides = [1, 1]
+  } : (tensor<2x10x10x10xf32>, tensor<3x3x10x20x!tf_type.qint8>,
+    tensor<f32>, tensor<i32>) -> tensor<2x10x10x20xf32>
+  return %0 : tensor<2x10x10x20xf32>
+})mlir";
+  TF_ASSERT_OK_AND_ASSIGN(auto input, CreateRandomF32Literal({2, 10, 10, 10}));
+  TF_ASSERT_OK_AND_ASSIGN(auto filter, CreateRandomI8Literal({3, 3, 10, 20}));
+  // TF kernels for UniformQuantizedConvolutionHybrid does DRQ. But StableHLO
+  // hybrid ops does weight-only. So we use a different TF graph for evaluating
+  // expected weight-only quantized results.
+  ExecuteAndCompareResultsWithTfKernel(kProgram, {&input, &filter}, kTfProgram);
+}
+
+TEST_F(ConvertTfQuantToMhloIntTest, UniformQuantizeDotToValidGraph) {
+  constexpr absl::string_view kProgram = R"mlir(
+func.func @main(%input: tensor<8x9xi8>, %filter: tensor<9x10xi8>) -> tensor<8x10xi32> {
+  %input_scale = "tf.Const"() { value = dense<0.588> : tensor<f32> } : () -> tensor<f32>
+  %input_zp = "tf.Const"() { value = dense<42> : tensor<i32> } : () -> tensor<i32>
+  %filter_scale = "tf.Const"() { value = dense<0.0235> : tensor<f32> } : () -> tensor<f32>
+  %filter_zp = "tf.Const"() { value = dense<0> : tensor<i32> } : () -> tensor<i32>
+  %accum_scale = "tf.Const"() { value = dense<0.013818> : tensor<f32> } : () -> tensor<f32>
+  %accum_zp = "tf.Const"() { value = dense<0> : tensor<i32> } : () -> tensor<i32>
+  %quant_input = "tf.Cast"(%input) {} : (tensor<8x9xi8>) -> tensor<8x9x!tf_type.qint8>
+  %quant_filter = "tf.Cast"(%filter) {} : (tensor<9x10xi8>) -> tensor<9x10x!tf_type.qint8>
+  %0 = "tf.UniformQuantizedDot"(
+    %quant_input, %quant_filter, %input_scale, %input_zp, %filter_scale,
+    %filter_zp, %accum_scale, %accum_zp
+  ) {
+    Tin = "tfdtype$DT_QINT8", Tout = "tfdtype$DT_QINT32", attr_map = "",
+    device = "", lhs_quantization_axis = -1 : i64,
+    lhs_quantization_max_val = 127 : i64,
+    lhs_quantization_min_val = -128 : i64,
+    output_quantization_axis = -1 : i64,
+    output_quantization_max_val = 2147483647 : i64,
+    output_quantization_min_val = -2147483648 : i64,
+    rhs_quantization_axis = -1 : i64,
+    rhs_quantization_max_val = 127 : i64,
+    rhs_quantization_min_val = -128 : i64
+  } : (
+    tensor<8x9x!tf_type.qint8>, tensor<9x10x!tf_type.qint8>, tensor<f32>,
+    tensor<i32>, tensor<f32>, tensor<i32>, tensor<f32>, tensor<i32>
+  ) -> tensor<8x10x!tf_type.qint32>
+  %output = "tf.Cast"(%0) {} : (tensor<8x10x!tf_type.qint32>) -> tensor<8x10xi32>
+  return %output : tensor<8x10xi32>
+})mlir";
+  TF_ASSERT_OK_AND_ASSIGN(auto input, CreateRandomI8Literal({8, 9}));
+  TF_ASSERT_OK_AND_ASSIGN(auto filter, CreateRandomI8Literal({9, 10}));
+  ExecuteAndCompareResultsWithTfKernel(kProgram, {&input, &filter});
+}
+
+TEST_F(ConvertTfQuantToMhloIntTest, UniformQuantizeDotHybridToValidGraph) {
+  constexpr absl::string_view kTfProgram = R"mlir(
+func.func @main(%input: tensor<8x9xf32>, %filter: tensor<9x10xi8>) -> tensor<8x10xf32> {
+  %filter_scale = "tf.Const"() { value = dense<0.0235> : tensor<f32> } : () -> tensor<f32>
+  %filter_zp = "tf.Const"() { value = dense<0> : tensor<i32> } : () -> tensor<i32>
+  %quant_filter = "tf.Cast"(%filter) {} : (tensor<9x10xi8>) -> tensor<9x10x!tf_type.qint8>
+  %filter_new = "tf.UniformDequantize"(%quant_filter, %filter_scale, %filter_zp) {
+    quantization_axis = -1 : i64, quantization_min_val = -128 : i64,
+    quantization_max_val = 127 : i64
+  } : (tensor<9x10x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<9x10xf32>
+  %0 = "tf.MatMul"(%input, %filter_new) {
+  } : (tensor<8x9xf32>, tensor<9x10xf32>) -> tensor<8x10xf32>
+  return %0 : tensor<8x10xf32>
+})mlir";
+  constexpr absl::string_view kProgram = R"mlir(
+func.func @main(%input: tensor<8x9xf32>, %filter: tensor<9x10xi8>) -> tensor<8x10xf32> {
+  %filter_scale = "tf.Const"() { value = dense<0.0235> : tensor<f32> } : ()
+  -> tensor<f32>
+  %filter_zp = "tf.Const"() { value = dense<0> : tensor<i32> } : () -> tensor<i32>
+  %quant_filter = "tf.Cast"(%filter) {} : (tensor<9x10xi8>) -> tensor<9x10x!tf_type.qint8>
+  %0 = "tf.UniformQuantizedDotHybrid"(
+    %input, %quant_filter, %filter_scale, %filter_zp
+  ) {
+    Tin = "tfdtype$DT_QINT8", Tout = "tfdtype$DT_FLOAT", attr_map = "",
+    device = "", rhs_quantization_axis = -1 : i64,
+    rhs_quantization_max_val = 127 : i64, rhs_quantization_min_val = -128 : i64
+  } : (tensor<8x9xf32>, tensor<9x10x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<8x10xf32>
+  return %0 : tensor<8x10xf32>
+})mlir";
+  TF_ASSERT_OK_AND_ASSIGN(auto input, CreateRandomF32Literal({8, 9}));
+  TF_ASSERT_OK_AND_ASSIGN(auto filter, CreateRandomI8Literal({9, 10}));
+  // TF kernels for UniformQuantizedDotHybrid does DRQ. But StableHLO hybrid ops
+  // does weight-only. So we use a different TF graph for evaluating expected
+  // weight-only quantized results.
+  ExecuteAndCompareResultsWithTfKernel(kProgram, {&input, &filter}, kTfProgram);
+}
+
+TEST_F(ConvertTfQuantToMhloIntTest, UniformRequantizeToValidGraph) {
+  constexpr absl::string_view kProgram = R"mlir(
+func.func @main(%input: tensor<10xi8>) -> tensor<10xi8> {
+  %input_scale = "tf.Const"() { value = dense<0.2235> : tensor<f32> } : () -> tensor<f32>
+  %input_zp = "tf.Const"() { value = dense<-2> : tensor<i32> } : () -> tensor<i32>
+  %output_scale = "tf.Const"() { value = dense<0.11> : tensor<f32> } : () -> tensor<f32>
+  %output_zp = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
+  %0 = "tf.Cast"(%input) {} : (tensor<10xi8>) -> tensor<10x!tf_type.qint8>
+  %1 = "tf.UniformRequantize"(
+    %0, %input_scale, %input_zp, %output_scale, %output_zp
+  ) {
+    Tin = "tfdtype$DT_QINT8", Tout = "tfdtype$DT_QINT8", attr_map = "",
+    device = "", input_quantization_axis = -1,
+    input_quantization_max_val = 127 : i64,
+    input_quantization_min_val = -128 : i64,
+    output_quantization_axis = -1 : i64,
+    output_quantization_max_val = 127 : i64,
+    output_quantization_min_val = -128 : i64
+  } : (
+    tensor<10x!tf_type.qint8>, tensor<f32>, tensor<i32>, tensor<f32>,
+    tensor<i32>
+  ) -> tensor<10x!tf_type.qint8>
+  %2 = "tf.Cast"(%1) {} : (tensor<10x!tf_type.qint8>) -> tensor<10xi8>
+  return %2 : tensor<10xi8>
+})mlir";
+  TF_ASSERT_OK_AND_ASSIGN(auto input, CreateRandomI8Literal({10}));
+  ExecuteAndCompareResultsWithTfKernel(kProgram, {&input});
+}
+
+TEST_F(ConvertTfQuantToMhloIntTest, UniformRequantizePerChannelToValidGraph) {
+  constexpr absl::string_view kProgram = R"mlir(
+func.func @main(
+    %input: tensor<10x10xi8>, %input_scale: tensor<10xf32>,
+    %input_zp: tensor<10xi32>, %output_scale: tensor<10xf32>,
+    %output_zp: tensor<10xi32>
+  ) -> tensor<10x10xi8> {
+  %0 = "tf.Cast"(%input) {} : (tensor<10x10xi8>) -> tensor<10x10x!tf_type.qint8>
+  %1 = "tf.UniformRequantize"(
+    %0, %input_scale, %input_zp, %output_scale, %output_zp
+  ) {
+    Tin = "tfdtype$DT_QINT8", Tout = "tfdtype$DT_QINT8", attr_map = "",
+    device = "", input_quantization_axis = 1,
+    input_quantization_max_val = 127 : i64,
+    input_quantization_min_val = -128 : i64,
+    output_quantization_axis = 1 : i64,
+    output_quantization_max_val = 127 : i64,
+    output_quantization_min_val = -128 : i64
+  } : (
+    tensor<10x10x!tf_type.qint8>, tensor<10xf32>, tensor<10xi32>,
+    tensor<10xf32>, tensor<10xi32>
+  ) -> tensor<10x10x!tf_type.qint8>
+  %2 = "tf.Cast"(%1) {} : (tensor<10x10x!tf_type.qint8>) -> tensor<10x10xi8>
+  return %2 : tensor<10x10xi8>
+})mlir";
+  TF_ASSERT_OK_AND_ASSIGN(auto input, CreateRandomI8Literal({10, 10}));
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto input_scale,
+      CreateRandomF32Literal({10}, /*min=*/0.0001, /*max=*/2));
+  TF_ASSERT_OK_AND_ASSIGN(auto input_zp, CreateRandomI32Literal({10}));
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto output_scale,
+      CreateRandomF32Literal({10}, /*min=*/0.0001, /*max=*/2));
+  TF_ASSERT_OK_AND_ASSIGN(auto output_zp, CreateRandomI32Literal({10}));
+  // error_tolerance is set to be 1 because different rounding implementations
+  // in TF kernel and the lowering passes may cause +/-1 differences.
+  ExecuteAndCompareResultsWithTfKernel(
+      kProgram, {&input, &input_scale, &input_zp, &output_scale, &output_zp},
+      /*tf_program=*/std::nullopt,
+      /*error_tolerance=*/1.0);
+}
+
+TEST_F(ConvertTfQuantToMhloIntTest,
+       UniformRequantizePerTensorToPerChannelToValidGraph) {
+  constexpr absl::string_view kProgram = R"mlir(
+func.func @main(
+    %input: tensor<10x10xi8>, %input_scale: tensor<f32>, %input_zp: tensor<i32>,
+    %output_scale: tensor<10xf32>, %output_zp: tensor<10xi32>
+  ) -> tensor<10x10xi8> {
+  %0 = "tf.Cast"(%input) {} : (tensor<10x10xi8>) -> tensor<10x10x!tf_type.qint8>
+  %1 = "tf.UniformRequantize"(
+    %0, %input_scale, %input_zp, %output_scale, %output_zp
+  ) {
+    Tin = "tfdtype$DT_QINT8", Tout = "tfdtype$DT_QINT8", attr_map = "",
+    device = "", input_quantization_axis = -1,
+    input_quantization_max_val = 127 : i64,
+    input_quantization_min_val = -128 : i64,
+    output_quantization_axis = 1 : i64,
+    output_quantization_max_val = 127 : i64,
+    output_quantization_min_val = -128 : i64
+  } : (
+    tensor<10x10x!tf_type.qint8>, tensor<f32>, tensor<i32>,
+    tensor<10xf32>, tensor<10xi32>
+  ) -> tensor<10x10x!tf_type.qint8>
+  %2 = "tf.Cast"(%1) {} : (tensor<10x10x!tf_type.qint8>) -> tensor<10x10xi8>
+  return %2 : tensor<10x10xi8>
+})mlir";
+  TF_ASSERT_OK_AND_ASSIGN(auto input, CreateRandomI8Literal({10, 10}));
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto input_scale, CreateRandomF32Literal({}, /*min=*/0.0001, /*max=*/2));
+  TF_ASSERT_OK_AND_ASSIGN(auto input_zp, CreateRandomI32Literal({}));
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto output_scale,
+      CreateRandomF32Literal({10}, /*min=*/0.0001, /*max=*/2));
+  TF_ASSERT_OK_AND_ASSIGN(auto output_zp, CreateRandomI32Literal({10}));
+  // error_tolerance is set to be 1 because different rounding implementations
+  // in TF kernel and the lowering passes may cause +/-1 differences.
+  ExecuteAndCompareResultsWithTfKernel(
+      kProgram, {&input, &input_scale, &input_zp, &output_scale, &output_zp},
+      /*tf_program=*/std::nullopt,
+      /*error_tolerance=*/1.0);
+}
+
+TEST_F(ConvertTfQuantToMhloIntTest,
+       UniformRequantizePerChannelToPerTensorToValidGraph) {
+  constexpr absl::string_view kProgram = R"mlir(
+func.func @main(
+    %input: tensor<10x10xi8>, %input_scale: tensor<10xf32>,
+    %input_zp: tensor<10xi32>, %output_scale: tensor<f32>, %output_zp: tensor<i32>
+  ) -> tensor<10x10xi8> {
+  %0 = "tf.Cast"(%input) {} : (tensor<10x10xi8>) -> tensor<10x10x!tf_type.qint8>
+  %1 = "tf.UniformRequantize"(
+    %0, %input_scale, %input_zp, %output_scale, %output_zp
+  ) {
+    Tin = "tfdtype$DT_QINT8", Tout = "tfdtype$DT_QINT8", attr_map = "",
+    device = "", input_quantization_axis = 1,
+    input_quantization_max_val = 127 : i64,
+    input_quantization_min_val = -128 : i64,
+    output_quantization_axis = -1 : i64,
+    output_quantization_max_val = 127 : i64,
+    output_quantization_min_val = -128 : i64
+  } : (
+    tensor<10x10x!tf_type.qint8>, tensor<10xf32>, tensor<10xi32>,
+    tensor<f32>, tensor<i32>
+  ) -> tensor<10x10x!tf_type.qint8>
+  %2 = "tf.Cast"(%1) {} : (tensor<10x10x!tf_type.qint8>) -> tensor<10x10xi8>
+  return %2 : tensor<10x10xi8>
+})mlir";
+  TF_ASSERT_OK_AND_ASSIGN(auto input, CreateRandomI8Literal({10, 10}));
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto input_scale,
+      CreateRandomF32Literal({10}, /*min=*/0.0001, /*max=*/2));
+  TF_ASSERT_OK_AND_ASSIGN(auto input_zp, CreateRandomI32Literal({10}));
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto output_scale, CreateRandomF32Literal({}, /*min=*/0.0001, /*max=*/2));
+  TF_ASSERT_OK_AND_ASSIGN(auto output_zp, CreateRandomI32Literal({}));
+  // error_tolerance is set to be 1 because different rounding implementations
+  // in TF kernel and the lowering passes may cause +/-1 differences.
+  ExecuteAndCompareResultsWithTfKernel(
+      kProgram, {&input, &input_scale, &input_zp, &output_scale, &output_zp},
+      /*tf_program=*/std::nullopt,
+      /*error_tolerance=*/1.0);
+}
+
+TEST_F(ConvertTfQuantToMhloIntTest, UniformQuantizeAddToValidGraph) {
+  constexpr absl::string_view kProgram = R"mlir(
+func.func @main(%lhs: tensor<10x10xi32>, %rhs: tensor<10x10xi32>) -> tensor<10x10xi32> {
+  %lhs_scale = "tf.Const"() { value = dense<0.518> : tensor<f32> } : () -> tensor<f32>
+  %lhs_zp = "tf.Const"() { value = dense<42> : tensor<i32> } : () -> tensor<i32>
+  %rhs_scale = "tf.Const"() { value = dense<0.0239> : tensor<f32> } : () -> tensor<f32>
+  %rhs_zp = "tf.Const"() { value = dense<0> : tensor<i32> } : () -> tensor<i32>
+  %accum_scale = "tf.Const"() { value = dense<0.013> : tensor<f32> } : () -> tensor<f32>
+  %accum_zp = "tf.Const"() { value = dense<0> : tensor<i32> } : () -> tensor<i32>
+  %quant_lhs = "tf.Cast"(%lhs) {} : (tensor<10x10xi32>) -> tensor<10x10x!tf_type.qint32>
+  %quant_rhs = "tf.Cast"(%rhs) {} : (tensor<10x10xi32>) -> tensor<10x10x!tf_type.qint32>
+  %0 = "tf.UniformQuantizedAdd"(
+    %quant_lhs, %quant_rhs, %lhs_scale, %lhs_zp, %rhs_scale,
+    %rhs_zp, %accum_scale, %accum_zp
+  ) {
+    Tin = "tfdtype$DT_QINT32", Tout = "tfdtype$DT_QINT32", attr_map = "",
+    device = "", lhs_quantization_axis = -1 : i64,
+    lhs_quantization_max_val = 2147483647 : i64,
+    lhs_quantization_min_val = -2147483648 : i64,
+    output_quantization_axis = -1 : i64,
+    output_quantization_max_val = 2147483647 : i64,
+    output_quantization_min_val = -2147483648 : i64,
+    rhs_quantization_axis = -1 : i64,
+    rhs_quantization_max_val = 2147483647 : i64,
+    rhs_quantization_min_val = -2147483648 : i64
+  } : (
+    tensor<10x10x!tf_type.qint32>, tensor<10x10x!tf_type.qint32>, tensor<f32>,
+    tensor<i32>, tensor<f32>, tensor<i32>, tensor<f32>, tensor<i32>
+  ) -> tensor<10x10x!tf_type.qint32>
+  %1 = "tf.Cast"(%0) {} : (tensor<10x10x!tf_type.qint32>) ->  tensor<10x10xi32>
+  return %1 : tensor<10x10xi32>
+})mlir";
+  TF_ASSERT_OK_AND_ASSIGN(auto lhs, CreateRandomI32Literal({10, 10}));
+  TF_ASSERT_OK_AND_ASSIGN(auto rhs, CreateRandomI32Literal({10, 10}));
+  // error_tolerance is set to be 1 because different rounding implementations
+  // in TF kernel and the lowering passes may cause +/-1 differences.
+  ExecuteAndCompareResultsWithTfKernel(kProgram, {&lhs, &rhs},
+                                       /*tf_program=*/std::nullopt,
+                                       /*error_tolerance=*/1.0);
+}
+
+}  // namespace
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/convert_tf_quant_types.cc

@@ -0,0 +1,331 @@
+/* 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.
+==============================================================================*/
+
+// The TF dialect uses some TF types that are illegal in the MHLO dialect and
+// some generic types that are legal in MHLO. This pass legalizes TF types into
+// types that are legal in MHLO. For example, TF::Qint8Type is converted to i8.
+// Rewrites here should run before TF to MHLO op legalizations are run.
+
+#include <memory>
+#include <string>
+#include <utility>
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Matchers.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/quantization/stablehlo/utils/tf_type_utils.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_attributes.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/core/lib/monitoring/counter.h"
+
+namespace mlir::quant::stablehlo {
+namespace {
+
+using quant::tensorflow::GetDenseAttrFromTensorProtoAttr;
+using quant::tensorflow::GetIntTypeFromTFQint;
+using quant::tensorflow::IsTFQintType;
+using quant::tensorflow::IsTFUniformQuantizedOp;
+
+#define GEN_PASS_DEF_CONVERTTFQUANTTYPES
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/passes.h.inc"
+
+// TODO: b/290366702 - Temporarily added metrics for debugging.
+auto *mlir_tf_quant_op_count = ::tensorflow::monitoring::Counter<1>::New(
+    "/tensorflow/core/tf2xla/tf_quant_op_count" /*metric_name*/,
+    "Counts the number of ops that has qint types" /*metric description*/,
+    "op_name" /*metric label*/);
+
+// Returns whether a type is illegal. Here we consider TF qint types illegal.
+// See pass description in passes.td for more info about how illegal types are
+// treated in this pass.
+bool IsIllegalType(Type type) {
+  return IsTFQintType(getElementTypeOrSelf(type));
+}
+
+// Get the corresponding int type from TF qint types.
+// If input is not TF qint types, returns the original type.
+Type ToLegalType(Type type) {
+  if (IsTFQintType(type)) return GetIntTypeFromTFQint(type);
+  if (auto shaped = mlir::dyn_cast<ShapedType>(type)) {
+    Type elem = shaped.getElementType();
+    if (IsTFQintType(elem)) return shaped.clone(ToLegalType(elem));
+  }
+  return type;
+}
+
+bool IsQintToIntCast(Operation *op) {
+  auto cast_op = llvm::dyn_cast<TF::CastOp>(op);
+  return cast_op && IsIllegalType(cast_op.getX().getType()) &&
+         ToLegalType(cast_op.getX().getType()) == cast_op.getY().getType();
+}
+
+bool IsIntToQintCast(Operation *op) {
+  auto cast_op = llvm::dyn_cast<TF::CastOp>(op);
+  return cast_op && IsIllegalType(cast_op.getY().getType()) &&
+         ToLegalType(cast_op.getY().getType()) == cast_op.getX().getType();
+}
+
+// Check if an op result value is consumed by qint -> int TF Cast OP.
+bool IsQintValueQintToIntCast(Value v) {
+  if (!IsIllegalType(v.getType())) {
+    return true;
+  }
+  if (v.getUsers().empty()) {
+    return false;
+  }
+  return llvm::all_of(v.getUsers(), [&](OpOperand operand) {
+    return IsQintToIntCast(operand.getOwner());
+  });
+}
+
+// Check if an op operand value is defined by int -> qint TF Cast OP.
+bool IsQintValueDefinedByIntToQintCast(Value v) {
+  if (!IsIllegalType(v.getType())) {
+    return true;
+  }
+  if (!v.getDefiningOp() || !llvm::isa<TF::CastOp>(v.getDefiningOp())) {
+    return false;
+  }
+  return IsIntToQintCast(v.getDefiningOp());
+}
+
+bool IsTFUniformQuantizedOpLegal(Operation *op) {
+  // UniformQuantized Ops are considered legal if its qint operands and
+  // results are connected to TF CastOp.
+  return op && llvm::all_of(op->getResults(), IsQintValueQintToIntCast) &&
+         llvm::all_of(op->getOperands(), IsQintValueDefinedByIntToQintCast);
+}
+
+bool IsCastOpLegal(TF::CastOp cast_op) {
+  // Consider qint <-> qint casts illegal.
+  if (IsIllegalType(cast_op.getSrcT()) && IsIllegalType(cast_op.getDstT())) {
+    return false;
+  }
+  // Consider CastOp illegal if either of its Src/Dst type is qint and is
+  // connected to a non-UQ op.
+  if (IsIllegalType(cast_op.getSrcT()) &&
+      !(cast_op.getX().getDefiningOp() &&
+        IsTFUniformQuantizedOp(cast_op.getX().getDefiningOp()))) {
+    return false;
+  }
+  if (IsIllegalType(cast_op.getDstT()) &&
+      !IsTFUniformQuantizedOp(*cast_op.getY().getUsers().begin())) {
+    return false;
+  }
+  return true;
+}
+
+class TFQuantTypeConverter : public TypeConverter {
+ public:
+  TFQuantTypeConverter() {
+    addConversion([](Type type) -> Type {
+      return IsIllegalType(type) ? ToLegalType(type) : type;
+    });
+  }
+};
+
+// An Op is illegal iff it is non-UQ op and it contains qint types.
+class TFQuantTypeConversionTarget : public ConversionTarget {
+ public:
+  explicit TFQuantTypeConversionTarget(MLIRContext &ctx,
+                                       TFQuantTypeConverter &converter)
+      : ConversionTarget(ctx), converter_(converter) {
+    markUnknownOpDynamicallyLegal([this](Operation *op) {
+      // Consider UQ op legal if it has a CastOp next to the qint input/output.
+      if (IsTFUniformQuantizedOp(op)) {
+        return IsTFUniformQuantizedOpLegal(op);
+      } else if (auto cast_op = llvm::dyn_cast<TF::CastOp>(op)) {
+        return IsCastOpLegal(cast_op);
+      } else if (auto const_op = llvm::dyn_cast<TF::ConstOp>(op)) {
+        return !IsIllegalType(const_op.getOutput().getType());
+      }
+      // The FuncOp type can contain types that the op's operand and result
+      // types do not contain.
+      if (auto func = dyn_cast<func::FuncOp>(op)) {
+        if (!converter_.isSignatureLegal(func.getFunctionType())) return false;
+      }
+      return converter_.isLegal(op);
+    });
+  }
+
+ private:
+  TFQuantTypeConverter &converter_;
+};
+
+class TFQuantTypePattern : public ConversionPattern {
+ public:
+  TFQuantTypePattern(MLIRContext *ctx, TypeConverter &converter)
+      : ConversionPattern(converter, MatchAnyOpTypeTag(), 1, ctx) {}
+
+  LogicalResult matchAndRewrite(
+      Operation *op, ArrayRef<Value> operands,
+      ConversionPatternRewriter &rewriter) const override {
+    // This pattern only handle non-UQ, non-const ops.
+    if (IsTFUniformQuantizedOp(op) || llvm::isa<TF::ConstOp>(op)) {
+      return failure();
+    }
+
+    // Update the results.
+    llvm::SmallVector<Type, 4> new_results;
+    if (failed(getTypeConverter()->convertTypes(op->getResultTypes(),
+                                                new_results)))
+      return failure();
+
+    // Update the regions. The dialect conversion framework wants new regions to
+    // be created and updated, rather than updating the old op. Thus we use an
+    // OperationState so we can add regions to the new up.
+    OperationState state(op->getLoc(), op->getName().getStringRef(), operands,
+                         new_results, op->getAttrs(), op->getSuccessors());
+    for (Region &region : op->getRegions()) {
+      auto new_region = std::make_unique<Region>(op);
+      rewriter.inlineRegionBefore(region, *new_region, new_region->begin());
+      if (failed(rewriter.convertRegionTypes(new_region.get(),
+                                             *getTypeConverter()))) {
+        return failure();
+      }
+      state.addRegion(std::move(new_region));
+    }
+    rewriter.replaceOp(op, rewriter.create(state)->getResults());
+
+    // TODO: b/290366702 - Temporarily added metrics for debugging.
+    mlir_tf_quant_op_count->GetCell(std::string(op->getName().getStringRef()))
+        ->IncrementBy(1);
+    return success();
+  }
+};
+
+// This pattern adds qint <-> int Cast to all qint operands and results for UQ
+// ops.
+class TFUniformQuantizedOpsPattern : public ConversionPattern {
+ public:
+  explicit TFUniformQuantizedOpsPattern(MLIRContext *ctx)
+      : ConversionPattern(MatchAnyOpTypeTag(), 1, ctx) {}
+
+  LogicalResult matchAndRewrite(
+      Operation *op, ArrayRef<Value> operands,
+      ConversionPatternRewriter &rewriter) const override {
+    // This pattern only handle UQ ops.
+    if (!IsTFUniformQuantizedOp(op)) {
+      return failure();
+    }
+
+    // Add CastOp int->qint before an input operand only when it original type
+    // is qint and its defining op is not already an int->qint CastOp.
+    llvm::SmallVector<Value, 4> new_operands;
+    for (int i = 0; i < operands.size(); ++i) {
+      Type orig_op_type = op->getOperandTypes()[i];
+      if (IsIllegalType(orig_op_type) &&
+          !IsQintValueDefinedByIntToQintCast(op->getOperand(i))) {
+        new_operands.push_back(TF::CastOp::create(rewriter, op->getLoc(),
+                                                  orig_op_type, operands[i]));
+      } else {
+        new_operands.push_back(operands[i]);
+      }
+    }
+
+    // Create a new UQ op.
+    OperationState state(op->getLoc(), op->getName().getStringRef(),
+                         new_operands, op->getResultTypes(), op->getAttrs(),
+                         op->getSuccessors());
+    Operation *new_op = rewriter.create(state);
+    llvm::SmallVector<Value, 4> new_results = new_op->getResults();
+
+    // Add qint->int CastOp after output result if its original type is qint and
+    // its users are not all qint->int CastOps.
+    for (int i = 0; i < new_results.size(); ++i) {
+      Value &result = new_results[i];
+      if (IsIllegalType(result.getType()) &&
+          !IsQintValueQintToIntCast(op->getResult(i))) {
+        result = TF::CastOp::create(rewriter, op->getLoc(),
+                                    ToLegalType(result.getType()), result);
+      }
+      // If the result is already consumed by qint->int CastOp, manually replace
+      // its use by the new UQ op. This is because such CastOp is already legal,
+      // it will not go through any conversion pattern later. Without this, that
+      // CastOp will still be consuming the original UQ op and cause errors.
+      op->getResult(i).replaceUsesWithIf(
+          new_op->getResult(i), [](OpOperand &operand) {
+            return IsQintToIntCast(operand.getOwner());
+          });
+    }
+    rewriter.replaceOp(op, new_results);
+    return success();
+  }
+};
+
+class TFConstOpQuantToIntPattern : public OpConversionPattern<TF::ConstOp> {
+ public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      TF::ConstOp op, TF::ConstOpAdaptor adaptor,
+      ConversionPatternRewriter &rewriter) const override {
+    if (!IsIllegalType(op.getOutput().getType())) return failure();
+    TF::TensorProtoAttr tensor_proto_attr;
+    if (!matchPattern(op.getOperation(), m_Constant(&tensor_proto_attr))) {
+      return rewriter.notifyMatchFailure(op, "operand must be constant.");
+    }
+    auto dense_attr_or = GetDenseAttrFromTensorProtoAttr(
+        tensor_proto_attr.getValue(),
+        mlir::dyn_cast<TensorType>(ToLegalType(op.getOutput().getType())));
+    if (failed(dense_attr_or)) {
+      op->emitError("failed to get DenseElementAttr.");
+      return failure();
+    }
+
+    rewriter.replaceOpWithNewOp<TF::ConstOp>(
+        op, ToLegalType(op.getOutput().getType()), *dense_attr_or);
+    return success();
+  }
+};
+
+struct ConvertTFQuantTypes
+    : public impl::ConvertTFQuantTypesBase<ConvertTFQuantTypes> {
+  void runOnOperation() override;
+};
+
+void ConvertTFQuantTypes::runOnOperation() {
+  TFQuantTypeConverter converter;
+  RewritePatternSet patterns(&getContext());
+  patterns.add<TFQuantTypePattern>(&getContext(), converter);
+  patterns.add<TFConstOpQuantToIntPattern, TFUniformQuantizedOpsPattern>(
+      &getContext());
+  populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(patterns,
+                                                                 converter);
+  TFQuantTypeConversionTarget target(getContext(), converter);
+  if (failed(applyFullConversion(getOperation(), target, std::move(patterns))))
+    return signalPassFailure();
+}
+
+}  // namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>> CreateConvertTFQuantTypesPass() {
+  return std::make_unique<ConvertTFQuantTypes>();
+}
+
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/convert_tf_quant_types_test.cc

@@ -0,0 +1,109 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cstdint>
+#include <memory>
+#include <string>
+
+#include <gtest/gtest.h>
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/passes.h"
+#include "tensorflow/compiler/mlir/register_common_dialects.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/deserialize_mlir_module_utils.h"
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "xla/tsl/platform/statusor.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/lib/monitoring/cell_reader.h"
+
+namespace mlir::quant::stablehlo {
+namespace {
+
+using ::mlir::DialectRegistry;
+using ::mlir::MLIRContext;
+using ::mlir::ModuleOp;
+using ::mlir::OwningOpRef;
+using ::tensorflow::monitoring::testing::CellReader;
+using ::testing::Test;
+
+static constexpr char kMetricsName[] =
+    "/tensorflow/core/tf2xla/tf_quant_op_count";
+
+class LegalizeTfTypesTest : public Test {
+ protected:
+  void CreateModule(const char* module_string) {
+    DialectRegistry mlir_registry;
+    RegisterCommonToolingDialects(mlir_registry);
+    context_.appendDialectRegistry(mlir_registry);
+    TF_ASSERT_OK(
+        tensorflow::DeserializeMlirModule(module_string, &context_, &module_));
+
+    pm_ = std::make_unique<mlir::PassManager>(&context_);
+    pm_->addNestedPass<mlir::func::FuncOp>(
+        quant::stablehlo::CreateConvertTFQuantTypesPass());
+  }
+  mlir::LogicalResult Run() { return pm_->run(module_.get()); }
+
+ private:
+  MLIRContext context_;
+  OwningOpRef<ModuleOp> module_;
+  std::unique_ptr<mlir::PassManager> pm_;
+};
+
+TEST_F(LegalizeTfTypesTest, RecordsStreamzQuantOps) {
+  static constexpr char kMlirModuleStr[] = R"(
+  module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
+    func.func @main(%arg0: tensor<3x3x!tf_type.qint8>, %arg1: tensor<3x3x!tf_type.qint8>) -> tensor<6x3x!tf_type.qint8> {
+      %axis = "tf.Const"() { value = dense<0> : tensor<i64> } : () -> tensor<i64>
+      %1 = "tf.ConcatV2"(%arg0, %arg1, %axis) : (tensor<3x3x!tf_type.qint8>, tensor<3x3x!tf_type.qint8>, tensor<i64>) -> tensor<6x3x!tf_type.qint8>
+      func.return %1 : tensor<6x3x!tf_type.qint8>
+    }
+  })";
+  CreateModule(kMlirModuleStr);
+  CellReader<int64_t> reader(kMetricsName);
+
+  auto result = Run();
+
+  EXPECT_TRUE(result.succeeded());
+  EXPECT_EQ(reader.Delta("tf.ConcatV2"), 1);
+  EXPECT_EQ(reader.Delta("func.return"), 1);
+  EXPECT_EQ(reader.Delta("func.func"), 0);
+}
+
+TEST_F(LegalizeTfTypesTest, RecordsStreamzNoQuantOps) {
+  static constexpr char kMlirModuleStr[] = R"(
+  module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
+    func.func @main(%arg0: tensor<3x3xf32>, %arg1: tensor<3x3xf32>) -> tensor<6x3xf32> {
+      %axis = "tf.Const"() { value = dense<0> : tensor<i64> } : () -> tensor<i64>
+      %1 = "tf.ConcatV2"(%arg0, %arg1, %axis) : (tensor<3x3xf32>, tensor<3x3xf32>, tensor<i64>) -> tensor<6x3xf32>
+      func.return %1 : tensor<6x3xf32>
+    }
+  })";
+  CreateModule(kMlirModuleStr);
+  CellReader<int64_t> reader(kMetricsName);
+
+  auto result = Run();
+
+  EXPECT_TRUE(result.succeeded());
+  EXPECT_EQ(reader.Delta("tf.ConcatV2"), 0);
+  EXPECT_EQ(reader.Delta("func.return"), 0);
+  EXPECT_EQ(reader.Delta("func.func"), 0);
+}
+
+}  // namespace
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/legalize_tf_quant_test.cc

@@ -0,0 +1,149 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <memory>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/StringRef.h"
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.h"
+#include "xla/client/client_library.h"
+#include "xla/shape.h"
+#include "xla/stream_executor/platform.h"
+#include "xla/stream_executor/platform_manager.h"
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/protobuf/config.pb.h"
+#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
+#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
+
+namespace mlir::quant::stablehlo {
+namespace {
+
+using ::testing::Test;
+
+class LegalizeTFQuantTest : public Test {
+ protected:
+  void TestBridgeLowering(llvm::StringRef mlir_module_string,
+                          llvm::ArrayRef<tensorflow::TensorShape> arg_shapes,
+                          tensorflow::DataType dtype) {
+    tensorflow::tpu::MlirToHloArgs mlir_to_hlo_args;
+    mlir_to_hlo_args.rollout_state =
+        tensorflow::ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED;
+    mlir_to_hlo_args.mlir_module = mlir_module_string;
+    tensorflow::se::Platform* platform =
+        tensorflow::se::PlatformManager::PlatformWithName("Host").value();
+    auto client =
+        xla::ClientLibrary::GetOrCreateCompileOnlyClient(platform).value();
+    tensorflow::tpu::TPUCompileMetadataProto metadata_proto;
+    // Set up an arg per arg_shape with the specified type.
+    for (int i = 0; i < arg_shapes.size(); ++i) {
+      auto metadata_arg = metadata_proto.add_args();
+      metadata_arg->set_kind(
+          tensorflow::tpu::TPUCompileMetadataProto::Arg::PARAMETER);
+      metadata_arg->set_dtype(dtype);
+    }
+    // Set up one dummy retval.
+    metadata_proto.add_retvals();
+    bool use_tuple_args = true;
+    std::vector<tensorflow::tpu::ShardingAndIndex> arg_core_mapping;
+    std::vector<std::vector<xla::Shape>> per_core_arg_shapes;
+    std::vector<std::unique_ptr<mlir::Pass>> custom_legalization_passes;
+
+    TF_EXPECT_OK(tensorflow::tf2xla::v2::LegalizeMlirToHlo(
+                     mlir_to_hlo_args, metadata_proto, use_tuple_args,
+                     /*device_type=*/"XLA_TPU_JIT", custom_legalization_passes,
+                     /*shape_determination_fns=*/{}, arg_shapes,
+                     &arg_core_mapping, &per_core_arg_shapes, client)
+                     .status());
+  }
+};
+
+TEST_F(LegalizeTFQuantTest, LegalizesModuleWithTFUniformQuantization) {
+  constexpr char mlir_module_string[] = R"mlir(
+  module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
+    func.func @main(%arg0 : tensor<1xf32>) -> tensor<1xf32> {
+      %scales = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
+      %zps = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
+
+      %0 = "tf.UniformQuantize"(%arg0, %scales, %zps) {
+        quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
+      } : (tensor<1xf32>, tensor<f32>, tensor<i32>) -> tensor<1x!tf_type.qint8>
+      %1 = "tf.UniformDequantize"(%0, %scales, %zps) {
+        quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
+      } : (tensor<1x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<1xf32>
+      func.return %1 : tensor<1xf32>
+    }
+  })mlir";
+
+  std::vector<tensorflow::TensorShape> arg_shapes = {{1}};
+
+  TestBridgeLowering(mlir_module_string, arg_shapes, tensorflow::DT_FLOAT);
+}
+
+TEST_F(LegalizeTFQuantTest, LegalizesModuleWithDequantize) {
+  constexpr char mlir_module_string[] = R"mlir(
+  module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
+    func.func @main(%arg0: tensor<1x!tf_type.qint8>) -> tensor<1xf32> {
+      %min_range = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
+      %max_range = "tf.Const"() { value = dense<5.0> : tensor<f32> } : () -> tensor<f32>
+      %0 = "tf.Dequantize"(%arg0, %min_range, %max_range) : (tensor<1x!tf_type.qint8>, tensor<f32>, tensor<f32>) -> tensor<1xf32>
+      func.return %0 : tensor<1xf32>
+    }
+  })mlir";
+  std::vector<tensorflow::TensorShape> arg_shapes = {{1}};
+
+  TestBridgeLowering(mlir_module_string, arg_shapes, tensorflow::DT_QINT8);
+}
+
+TEST_F(LegalizeTFQuantTest, LegalizesModuleWithClipByValue) {
+  constexpr char mlir_module_string[] = R"mlir(
+  module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
+    func.func @main(%arg0 : tensor<2x2xf32>) -> tensor<2x2xf32> {
+      %max = "tf.Const"() { value = dense<12.0> : tensor<f32> } : () -> tensor<f32>
+      %min = "tf.Const"() { value = dense<-25.0> : tensor<f32> } : () -> tensor<f32>
+      %scales = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
+      %zps = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
+
+      %0 = "tf.UniformQuantize"(%arg0, %scales, %zps) {
+        quantization_axis = -1 : i64, quantization_min_val = -2147483648 : i64, quantization_max_val = 2147483647 : i64
+      } : (tensor<2x2xf32>, tensor<f32>, tensor<i32>) -> tensor<2x2x!tf_type.qint32>
+      %qmax = "tf.UniformQuantize"(%max, %scales, %zps) {
+        quantization_axis = -1 : i64, quantization_min_val = -2147483648 : i64, quantization_max_val = 2147483647 : i64
+      } : (tensor<f32>, tensor<f32>, tensor<i32>) -> tensor<!tf_type.qint32>
+      %qmin = "tf.UniformQuantize"(%min, %scales, %zps) {
+        quantization_axis = -1 : i64, quantization_min_val = -2147483648 : i64, quantization_max_val = 2147483647 : i64
+      } : (tensor<f32>, tensor<f32>, tensor<i32>) -> tensor<!tf_type.qint32>
+
+      %1 = "tf.UniformQuantizedClipByValue"(%0, %qmin, %qmax, %scales, %zps) {
+        quantization_axis = -1 : i64, quantization_min_val = -2147483648 : i64, quantization_max_val = 2147483647 : i64
+      } : (tensor<2x2x!tf_type.qint32>, tensor<!tf_type.qint32>, tensor<!tf_type.qint32>, tensor<f32>, tensor<i32>) -> tensor<2x2x!tf_type.qint32>
+
+      %2 = "tf.UniformDequantize"(%1, %scales, %zps) {
+        quantization_axis = -1 : i64, quantization_min_val = -2147483648 : i64, quantization_max_val = 2147483647 : i64
+      } : (tensor<2x2x!tf_type.qint32>, tensor<f32>, tensor<i32>) -> tensor<2x2xf32>
+      func.return %2 : tensor<2x2xf32>
+    }
+  })mlir";
+  std::vector<tensorflow::TensorShape> arg_shapes = {{2, 2}};
+
+  TestBridgeLowering(mlir_module_string, arg_shapes, tensorflow::DT_FLOAT);
+}
+
+}  // namespace
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/optimize.cc

@@ -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.
+==============================================================================*/
+
+#include <memory>
+#include <utility>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "stablehlo/dialect/ChloOps.h"  // from @stablehlo  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/passes.h"
+#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"  // IWYU pragma: keep
+
+namespace mlir::quant::stablehlo {
+namespace {
+
+#define GEN_PASS_DEF_OPTIMIZEINTGRAPH
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/passes.h.inc"
+
+class OptimizeIntGraph : public impl::OptimizeIntGraphBase<OptimizeIntGraph> {
+ public:
+  OptimizeIntGraph() = default;
+  OptimizeIntGraph(const OptimizeIntGraph &) = default;
+  void runOnOperation() override;
+};
+
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/optimize.inc"
+
+void OptimizeIntGraph::runOnOperation() {
+  RewritePatternSet patterns(&getContext());
+  populateWithGenerated(patterns);
+  auto func = getOperation();
+  if (failed(applyPatternsGreedily(func, std::move(patterns)))) {
+    signalPassFailure();
+  }
+}
+
+}  // namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>> CreateOptimizeIntGraphPass() {
+  return std::make_unique<OptimizeIntGraph>();
+}
+
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/optimize.td

@@ -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.
+==============================================================================*/
+
+include "mlir/IR/OpBase.td"
+include "mlir/IR/PatternBase.td"
+include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"
+include "mhlo/IR/hlo_ops.td"
+include "stablehlo/dialect/ChloOps.td"
+
+def IsDefinedByConvOrDotGeneralOp : Constraint<Or<[
+  CPred<"$0.getDefiningOp<mhlo::ConvolutionOp>()">,
+  CPred<"$0.getDefiningOp<mhlo::DotGeneralOp>()">]>>;
+
+def IsNull : Constraint<CPred<"!$0">>;
+
+// This pattern optimizes:
+//   conv/dot_general + a + b -> conv/dot_general + (a + b)
+//   conv/dot_general - a - b -> conv/dot_general - (a + b)
+//   conv/dot_general + a - b -> conv/dot_general + (a - b)
+//   conv/dot_general - a + b -> conv/dot_general - (a - b)
+foreach OpsTuple = [
+    [CHLO_BroadcastAddOp, CHLO_BroadcastAddOp, CHLO_BroadcastAddOp],
+    [CHLO_BroadcastSubOp, CHLO_BroadcastSubOp, CHLO_BroadcastAddOp],
+    [CHLO_BroadcastAddOp, CHLO_BroadcastSubOp, CHLO_BroadcastSubOp],
+    [CHLO_BroadcastSubOp, CHLO_BroadcastAddOp, CHLO_BroadcastSubOp]] in {
+  def optimizeConsecutiveConv#OpsTuple[0]#OpsTuple[1] : Pat<
+    (OpsTuple[1]
+      (OpsTuple[0] $input, $zp_offset, $broadcast_dims_1),
+      $bias, $broadcast_dims_2),
+    (OpsTuple[0]
+      $input,
+      (OpsTuple[2] $zp_offset, $bias, $broadcast_dims_2), $broadcast_dims_1),
+    [
+      (IsNull $broadcast_dims_1),
+      (IsNull $broadcast_dims_2),
+      (TensorOf<[AnyInteger]> $input),
+      (IsDefinedByConvOrDotGeneralOp $input)]>;
+}

tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/passes.cc

@@ -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.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/passes.h"
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Pass/PassOptions.h"  // from @llvm-project
+#include "mlir/Transforms/Passes.h"  // from @llvm-project
+#include "stablehlo/transforms/Passes.h"  // from @stablehlo
+#include "xla/mlir_hlo/mhlo/transforms/passes.h"
+
+namespace mlir::quant::stablehlo {
+
+void AddQuantizationLoweringPasses(mlir::OpPassManager& pm) {
+  // These passes are grouped together and must run in this specific order.
+  pm.addNestedPass<mlir::func::FuncOp>(CreateConvertTFQuantOpsToMHLOPass());
+  pm.addNestedPass<mlir::func::FuncOp>(mhlo::createChloLegalizeToHloPass());
+  pm.addNestedPass<mlir::func::FuncOp>(mlir::createCanonicalizerPass());
+  pm.addPass(mhlo::createHloLegalizeToStablehloPass());
+  pm.addNestedPass<mlir::func::FuncOp>(
+      mlir::stablehlo::createStablehloLegalizeQuantToMathPass());
+  pm.addPass(mhlo::createStablehloLegalizeToHloPass());
+  pm.addNestedPass<mlir::func::FuncOp>(mlir::createCanonicalizerPass());
+  pm.addNestedPass<mlir::func::FuncOp>(CreateVerifyQuantLegalizationPass());
+}
+
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/passes.h

@@ -0,0 +1,62 @@
+/* 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_PASSES_BRIDGE_PASSES_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_PASSES_BRIDGE_PASSES_H_
+
+#include <memory>
+
+#define GEN_PASS_DECL
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+
+namespace mlir::quant::stablehlo {
+
+// Creates an instance of the ConvertTFQuantOpsToMHLOPass pass, which will
+// convert TF uniform quantized ops to the corresponding quantized MHLO ops.
+std::unique_ptr<OperationPass<func::FuncOp>>
+CreateConvertTFQuantOpsToMHLOPass();
+
+// TODO(b/288094093): Migrate uniform quantization legalization in a separate
+// pass.
+void PopulateLegalizeTfQuantizationPatterns(MLIRContext *context,
+                                            RewritePatternSet *patterns);
+
+// Creates an instance of the ConvertTFQuantTypes pass, which will convert TF
+// qint types to int types and surround TF UniformQuantized ops with qint <->
+// int casts.
+std::unique_ptr<OperationPass<func::FuncOp>> CreateConvertTFQuantTypesPass();
+
+// Creates an instance of the VerifyQuantLegalization pass, which verifies all
+// quant ops and types are lowered.
+std::unique_ptr<OperationPass<func::FuncOp>>
+CreateVerifyQuantLegalizationPass();
+
+// Add all passes for lowering TF quant ops and types to MHLO int.
+void AddQuantizationLoweringPasses(mlir::OpPassManager &pm);
+
+// Creates an instance of OptimizeIntGraphPass, which optimizes the int graph
+// lowered from the quantized graph.
+std::unique_ptr<OperationPass<func::FuncOp>> CreateOptimizeIntGraphPass();
+
+#define GEN_PASS_REGISTRATION
+#define GEN_PASS_DECL_CONVERTTFQUANTOPSTOMHLO
+#define GEN_PASS_DECL_CONVERTTFQUANTTYPES
+#define GEN_PASS_DECL_VERIFYQUANTLEGALIZATION
+#define GEN_PASS_DECL_OPTIMIZEINTGRAPH
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/passes.h.inc"
+}  // namespace mlir::quant::stablehlo
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_PASSES_BRIDGE_PASSES_H_

tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/passes.td

@@ -0,0 +1,69 @@
+/* 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.
+==============================================================================*/
+
+// Declare bridge passes that perform legalization and/or lowering.
+
+include "mlir/Pass/PassBase.td"
+
+def ConvertTFQuantOpsToMHLO : Pass<"quant-convert-tf-quant-ops-to-mhlo", "mlir::func::FuncOp"> {
+  let summary = "Convert TF Quant ops to MHLO quantizated ops.";
+
+  let description = [{
+    Convert TF Quant ops to MHLO quant ops.
+  }];
+
+  let constructor = "mlir::quant::stablehlo::CreateConvertTFQuantOpsToMHLOPass()";
+  let dependentDialects = ["TF::TensorFlowDialect", "chlo::ChloDialect",
+                           "mhlo::MhloDialect", "tf_type::TFTypeDialect",
+                           "quant::QuantDialect"];
+}
+
+def ConvertTFQuantTypes : Pass<"convert-tf-quant-types", "mlir::func::FuncOp"> {
+  let summary = "Replace TensorFlow qint types with int types.";
+
+  let description = [{
+    Converts TF ops with qint types to int types. Some UniformQuantized ops
+    argument/result allow qint type only. For such cases, add qint <-> int
+    tf.Cast around the ops so that they are still valid.
+  }];
+
+  let constructor = "mlir::quant::stablehlo::CreateConvertTFQuantTypesPass()";
+  let dependentDialects = ["TF::TensorFlowDialect", "tf_type::TFTypeDialect"];
+}
+
+def VerifyQuantLegalization : Pass<"verify-quant-legalization", "mlir::func::FuncOp"> {
+  let summary = "Verifies that all TF quant ops and types have been legalized.";
+
+  let description = [{
+    Ensures that all TF quant ops and types have been legalized to HLO
+    and reports an error about which op failed to legalize. This pass
+    does not transform any ops and is checking.}];
+
+  let constructor = "mlir::quant::stablehlo::CreateVerifyQuantLegalizationPass()";
+  let dependentDialects =  ["tf_type::TFTypeDialect",
+                            "quant::QuantDialect"];
+
+}
+
+def OptimizeIntGraph : Pass<"optimize-int-graph", "mlir::func::FuncOp"> {
+  let summary = "Optimization patterns for quantized integer graph";
+
+  let description = [{
+    This includes patterns for merging addition of zp offset and bias.
+  }];
+
+  let constructor = "mlir::quant::stablehlo::CreateOptimizeIntGraphPass()";
+  let dependentDialects = ["mhlo::MhloDialect"];
+}

tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/verify_quant_legalization.cc

@@ -0,0 +1,93 @@
+/* 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.
+==============================================================================*/
+
+// The TF dialect uses some TF types that are illegal in the MHLO dialect and
+// some generic types that are legal in MHLO. This pass legalizes TF types into
+// types that are legal in MHLO. For example, TF::Qint8Type is converted to i8.
+// Rewrites here should run before TF to MHLO op legalizations are run.
+
+#include <memory>
+
+#include "absl/log/log.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/Quant.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/IR/Visitors.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/quantization/stablehlo/utils/tf_type_utils.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
+
+namespace mlir::quant::stablehlo {
+namespace {
+
+using quant::tensorflow::IsTFQintType;
+using quant::tensorflow::IsTFUniformQuantizedOp;
+
+#define GEN_PASS_DEF_VERIFYQUANTLEGALIZATION
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/passes.h.inc"
+
+class VerifyQuantLegalization
+    : public impl::VerifyQuantLegalizationBase<VerifyQuantLegalization> {
+ public:
+  void runOnOperation() override;
+};
+
+bool IsQuantType(Type type) {
+  auto element_type = getElementTypeOrSelf(type);
+  return mlir::isa<quant::UniformQuantizedType>(element_type) ||
+         IsTFQintType(element_type);
+}
+
+bool IsMhloUniformQuantizedOp(Operation& op) {
+  return llvm::isa<mhlo::UniformQuantizeOp, mhlo::UniformDequantizeOp>(op);
+}
+
+void VerifyQuantLegalization::runOnOperation() {
+  Operation* func_op = getOperation();
+
+  auto walk_result = func_op->walk([&](Operation* op) {
+    // Verify all uq and qint types are lowered.
+    if (llvm::any_of(op->getOperandTypes(), IsQuantType) ||
+        llvm::any_of(op->getResultTypes(), IsQuantType) ||
+        IsTFUniformQuantizedOp(op) || IsMhloUniformQuantizedOp(*op)) {
+      op->emitOpError("is illegal as it is a UQ op or contains uq/qint types");
+      LOG(ERROR) << "Found illegal op containing uq/qint type: "
+                 << op->getName().getStringRef().str();
+      return WalkResult::interrupt();
+    }
+    return WalkResult::advance();
+  });
+
+  if (walk_result.wasInterrupted()) {
+    signalPassFailure();
+  }
+}
+
+}  // namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>>
+CreateVerifyQuantLegalizationPass() {
+  return std::make_unique<VerifyQuantLegalization>();
+}
+
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/passes/convert_func_to_bfloat16.cc

@@ -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 <memory>
+#include <utility>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/TypeRange.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.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 "mlir/Support/TypeID.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h"  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/quantization/stablehlo/utils/bfloat16_type.h"
+#include "tensorflow/core/platform/bfloat16.h"
+
+namespace mlir::quant::stablehlo {
+namespace {
+
+class BFloat16TypeConverter : public TypeConverter {
+ public:
+  BFloat16TypeConverter() {
+    addConversion([](const Type type) -> Type {
+      return quant::stablehlo::IsLargeFloatType(type)
+                 ? quant::stablehlo::ToBfloat16Type(type)
+                 : type;
+    });
+  }
+};
+
+// This helper function makes legality check easier. Both convert ops in the
+// patterns below are considered legal:
+//  - `BitcastConvertOp` (i32 -> f32) + `ConvertOp` (f32 -> bf16)
+//  - `ConvertOp` (bf16 -> f32) -> `BitcastConvertOp` (f32 -> i32)
+template <typename ConvertOp, typename OtherConvertOp>
+bool IsConvertOpLegal(ConvertOp convert_op, BFloat16TypeConverter& converter) {
+  if (!converter.isLegal(convert_op.getOperand().getType())) {
+    auto other_convert_op = dyn_cast_or_null<OtherConvertOp>(
+        convert_op.getOperand().getDefiningOp());
+    return other_convert_op &&
+           converter.isLegal(other_convert_op.getOperand().getType());
+  } else if (!converter.isLegal(convert_op.getResult().getType())) {
+    if (!convert_op.getResult().hasOneUse()) {
+      return false;
+    }
+    auto other_convert_op = dyn_cast_or_null<OtherConvertOp>(
+        *convert_op.getResult().getUsers().begin());
+    return other_convert_op &&
+           converter.isLegal(other_convert_op.getResult().getType());
+  }
+  return true;
+}
+
+class BFloat16TypeConversionTarget : public ConversionTarget {
+ public:
+  explicit BFloat16TypeConversionTarget(MLIRContext& ctx,
+                                        BFloat16TypeConverter& converter)
+      : ConversionTarget(ctx), converter_(converter) {
+    markUnknownOpDynamicallyLegal([this](Operation* op) {
+      // The FuncOp type can contain types that the op's operand and result
+      // types do not contain.
+      if (auto func = dyn_cast<func::FuncOp>(op)) {
+        if (!converter_.isSignatureLegal(func.getFunctionType())) return false;
+      } else if (auto bitcast_convert_op =
+                     dyn_cast<mlir::stablehlo::BitcastConvertOp>(op)) {
+        return IsConvertOpLegal<mlir::stablehlo::BitcastConvertOp,
+                                mlir::stablehlo::ConvertOp>(bitcast_convert_op,
+                                                            converter_);
+      } else if (auto convert_op = dyn_cast<mlir::stablehlo::ConvertOp>(op)) {
+        return IsConvertOpLegal<mlir::stablehlo::ConvertOp,
+                                mlir::stablehlo::BitcastConvertOp>(convert_op,
+                                                                   converter_);
+      }
+      return converter_.isLegal(op);
+    });
+  }
+
+ private:
+  BFloat16TypeConverter& converter_;
+};
+
+class BFloat16TypePattern : public ConversionPattern {
+ public:
+  BFloat16TypePattern(TypeConverter& converter, MLIRContext* ctx)
+      : ConversionPattern(converter, MatchAnyOpTypeTag(), /*benefit=*/1, ctx) {}
+
+  LogicalResult matchAndRewrite(
+      Operation* op, const ArrayRef<Value> operands,
+      ConversionPatternRewriter& rewriter) const override {
+    if (getTypeConverter()->isLegal(op)) {
+      return failure();
+    }
+    if (isa<mlir::stablehlo::BitcastConvertOp>(op)) {
+      // Skip `BitcastConvertOp`, which is handled by the other pattern.
+      return failure();
+    }
+
+    // Update the results.
+    SmallVector<Type, 4> new_results;
+    if (failed(getTypeConverter()->convertTypes(op->getResultTypes(),
+                                                new_results)))
+      return failure();
+
+    // Update the regions. The dialect conversion framework wants new regions to
+    // be created and updated, rather than updating the old op. Thus we use an
+    // OperationState so we can add regions to the new op.
+    OperationState state(op->getLoc(), op->getName().getStringRef(), operands,
+                         new_results, op->getAttrs(), op->getSuccessors());
+    for (Region& region : op->getRegions()) {
+      auto new_region = std::make_unique<Region>(op);
+      rewriter.inlineRegionBefore(region, *new_region, new_region->begin());
+      if (failed(rewriter.convertRegionTypes(new_region.get(),
+                                             *getTypeConverter()))) {
+        return failure();
+      }
+      state.addRegion(std::move(new_region));
+    }
+
+    // Convert value of ConstantOp to bfloat16.
+    if (auto const_op = dyn_cast<mlir::stablehlo::ConstantOp>(op)) {
+      const auto values = const_op.getValue().tryGetValues<float>();
+      if (!values.has_value()) {
+        return failure();
+      }
+      const SmallVector<tensorflow::bfloat16> bfloat16_values(values->begin(),
+                                                              values->end());
+      state.attributes.set(
+          const_op.getValueAttrName(),
+          DenseFPElementsAttr::get(
+              mlir::dyn_cast<ShapedType>(const_op.getValue().getType())
+                  .clone(rewriter.getBF16Type()),
+              bfloat16_values));
+    }
+
+    rewriter.replaceOp(op, rewriter.create(state)->getResults());
+
+    return success();
+  }
+};
+
+class BitcastConvertOpPattern
+    : public OpConversionPattern<mlir::stablehlo::BitcastConvertOp> {
+ public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      mlir::stablehlo::BitcastConvertOp op,
+      mlir::stablehlo::BitcastConvertOpAdaptor adaptor,
+      ConversionPatternRewriter& rewriter) const override {
+    const bool is_input_legal =
+        getTypeConverter()->isLegal(op.getOperand().getType());
+    const bool is_output_legal =
+        getTypeConverter()->isLegal(op.getResult().getType());
+    if (is_input_legal && is_output_legal) {
+      return failure();
+    } else if (is_input_legal) {
+      // output is f32, we bitcast_convert to f32 and then convert to bf16.
+      const Value output = mlir::stablehlo::BitcastConvertOp::create(
+          rewriter, op->getLoc(), op.getResult().getType(),
+          adaptor.getOperand());
+      rewriter.replaceOpWithNewOp<mlir::stablehlo::ConvertOp>(
+          op, getTypeConverter()->convertType(op.getResult().getType()),
+          output);
+    } else if (is_output_legal) {
+      // input is f32, we convert from bf16 and then bitcast_convert.
+      const Value output = mlir::stablehlo::ConvertOp::create(
+          rewriter, op->getLoc(), op.getOperand().getType(),
+          adaptor.getOperand());
+      rewriter.replaceOpWithNewOp<mlir::stablehlo::BitcastConvertOp>(
+          op, op.getResult().getType(), output);
+    } else {
+      // Both input/output are f32. Convert to no-op.
+      rewriter.replaceOp(op, adaptor.getOperand());
+    }
+    return success();
+  }
+};
+}  // namespace
+
+#define GEN_PASS_DEF_CONVERTFUNCTOBFLOAT16PASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"
+namespace {
+class ConvertFuncToBfloat16Pass
+    : public impl::ConvertFuncToBfloat16PassBase<ConvertFuncToBfloat16Pass> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ConvertFuncToBfloat16Pass)
+
+  explicit ConvertFuncToBfloat16Pass() = default;
+
+ private:
+  void runOnOperation() override;
+};
+
+void ConvertFuncToBfloat16Pass::runOnOperation() {
+  func::FuncOp func_op = getOperation();
+  MLIRContext* context = func_op.getContext();
+  RewritePatternSet patterns(context);
+
+  BFloat16TypeConverter converter;
+  patterns.add<BFloat16TypePattern, BitcastConvertOpPattern>(converter,
+                                                             context);
+  populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(patterns,
+                                                                 converter);
+  BFloat16TypeConversionTarget target(*context, converter);
+  if (failed(applyPartialConversion(func_op.getOperation(), target,
+                                    std::move(patterns)))) {
+    return signalPassFailure();
+  }
+}
+
+}  // namespace
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/passes/convert_shape_constraint_to_assert.cc

@@ -0,0 +1,218 @@
+/* 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 <algorithm>
+#include <cstdint>
+#include <utility>
+
+#include "llvm/ADT/STLExtras.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Shape/IR/Shape.h"  // from @llvm-project
+#include "mlir/Dialect/Tensor/IR/Tensor.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/TypeRange.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.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 "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo
+#include "stablehlo/transforms/Passes.h"  // from @stablehlo
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h"  // IWYU pragma: keep
+
+namespace mlir::quant::stablehlo {
+
+#define GEN_PASS_DEF_CONVERTSHAPETOSTABLEHLOWITHCONSTRAINTSPASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"
+
+namespace {
+using ::mlir::stablehlo::AndOp;
+using ::mlir::stablehlo::CompareOp;
+using ::mlir::stablehlo::ComparisonDirection;
+using ::mlir::stablehlo::ConcatenateOp;
+using ::mlir::stablehlo::ConstantOp;
+using ::mlir::stablehlo::CustomCallOp;
+using ::mlir::stablehlo::OrOp;
+using ::mlir::stablehlo::ReshapeOp;
+using ::mlir::stablehlo::SliceOp;
+
+// Cast from index-based shape representation used in the Shape dialect to the
+// i32-based representation used in HLO:
+//   * index => tensor<i32>.
+//   * tensor<Nxindex> => tensor<Nxi32>.
+//   * All i32-based types from above => themselves.
+// There is no convenient op that can express this, so we're using
+// unrealized_conversion_cast (with the idea that all these casts will
+// annihilate at the end of the pass).
+Value castToI32(PatternRewriter& rewriter, Location loc, Value value) {
+  Type resultType;
+  if (value.getType().isIndex())
+    resultType = RankedTensorType::get({}, rewriter.getI32Type());
+  if (auto valueType = mlir::dyn_cast<ShapedType>(value.getType())) {
+    if (!valueType.hasStaticShape()) return {};
+    if (valueType.getElementType().isInteger(32)) return value;
+    if (valueType.getElementType().isIndex())
+      resultType =
+          RankedTensorType::get(valueType.getShape(), rewriter.getI32Type());
+  }
+  if (!resultType) return {};
+  auto cast =
+      UnrealizedConversionCastOp::create(rewriter, loc, resultType, value);
+  return cast.getResult(0);
+}
+
+// Pads input tensor<N x i32> by X ones from the left. The number X is
+// determined by input pad. Result is tensor<(X+N) x i32>, where the first X
+// elements are ones.
+Value padFromLeft(PatternRewriter& rewriter, Location loc, Value input,
+                  int64_t pad) {
+  Value padI32 = ConstantOp::create(
+      rewriter, loc,
+      DenseIntElementsAttr::get<int32_t>(
+          RankedTensorType::get({pad}, rewriter.getI32Type()), 1));
+  return ConcatenateOp::create(rewriter, loc, ValueRange{padI32, input},
+                               /*dimension=*/0);
+}
+
+void insertShapeAssertionCustomCall(OpBuilder builder, Location loc,
+                                    Value assert) {
+  auto customCall =
+      CustomCallOp::create(builder, loc, TypeRange{}, ValueRange{assert});
+  customCall.setCallTargetName("shape_assertion");
+  customCall.setHasSideEffect(true);
+  customCall->setAttr("error_message",
+                      builder.getStringAttr("Shape assertion failed"));
+}
+
+struct ConvertCstrBroadcastableOp
+    : public OpRewritePattern<shape::CstrBroadcastableOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(shape::CstrBroadcastableOp op,
+                                PatternRewriter& rewriter) const override {
+    // As defined, op inputs must be 1D tensor or !shape.shape.
+    // We only support inputs of two 1D tensors.
+    if (op.getShapes().size() != 2) return failure();
+    auto shape1 = castToI32(rewriter, op.getLoc(), op.getShapes().front());
+    auto shape2 = castToI32(rewriter, op.getLoc(), op.getShapes().back());
+    if (!shape1 || !shape2) return failure();
+    auto tensorType1 = mlir::dyn_cast<RankedTensorType>(shape1.getType());
+    auto tensorType2 = mlir::dyn_cast<RankedTensorType>(shape2.getType());
+    if (!tensorType1 || !tensorType2) return failure();
+
+    // If the two operand shapes are of different sizes, the smaller one is
+    // padded with 1's from the left.
+    int32_t rank =
+        std::max(tensorType1.getDimSize(0), tensorType2.getDimSize(0));
+    if (tensorType1.getDimSize(0) < tensorType2.getDimSize(0)) {
+      shape1 =
+          padFromLeft(rewriter, op.getLoc(), shape1,
+                      tensorType2.getDimSize(0) - tensorType1.getDimSize(0));
+    } else if (tensorType1.getDimSize(0) > tensorType2.getDimSize(0)) {
+      shape2 =
+          padFromLeft(rewriter, op.getLoc(), shape2,
+                      tensorType1.getDimSize(0) - tensorType2.getDimSize(0));
+    }
+
+    // Compute if each dim is broadcastable. A dim is broadcastable iff
+    // dimSize1 == dimSize2 or dimSize1 == 1 or dimSize2 == 1
+    auto allOne = ConstantOp::create(
+        rewriter, op.getLoc(),
+        DenseIntElementsAttr::get<int32_t>(
+            RankedTensorType::get({rank}, rewriter.getI32Type()),
+            static_cast<int32_t>(1)));
+    Value dimSize1Is1 = CompareOp::create(rewriter, op.getLoc(), shape1, allOne,
+                                          ComparisonDirection::EQ);
+    Value dimSize2Is1 = CompareOp::create(rewriter, op.getLoc(), shape2, allOne,
+                                          ComparisonDirection::EQ);
+    Value eitherDimSizeIs1 =
+        OrOp::create(rewriter, op.getLoc(), dimSize1Is1, dimSize2Is1);
+    Value dimSizeEq = CompareOp::create(rewriter, op.getLoc(), shape1, shape2,
+                                        ComparisonDirection::EQ);
+    Value dimBroadcastable =
+        OrOp::create(rewriter, op.getLoc(), eitherDimSizeIs1, dimSizeEq);
+
+    // Iterate over each dim to check that all dims are broadcastable.
+    auto boolType = RankedTensorType::get({1}, rewriter.getI1Type());
+    Value allBroadcastable = ConstantOp::create(
+        rewriter, op.getLoc(), DenseIntElementsAttr::get<bool>(boolType, true));
+    for (auto i = 0; i < rank; ++i) {
+      Value broadcastable =
+          SliceOp::create(rewriter, op.getLoc(), dimBroadcastable,
+                          rewriter.getDenseI64ArrayAttr(i),
+                          rewriter.getDenseI64ArrayAttr(i + 1),
+                          rewriter.getDenseI64ArrayAttr(1));
+      allBroadcastable =
+          AndOp::create(rewriter, op.getLoc(), allBroadcastable, broadcastable);
+    }
+    Value allBroadcastableScalar = ReshapeOp::create(
+        rewriter, op.getLoc(), RankedTensorType::get({}, rewriter.getI1Type()),
+        allBroadcastable);
+
+    // Add CustomCallOp and replace Cstr op with const witness, which is useful
+    // for canonicalizer to remove the shape.assuming region.
+    insertShapeAssertionCustomCall(rewriter, op->getLoc(),
+                                   allBroadcastableScalar);
+    rewriter.replaceOpWithNewOp<shape::ConstWitnessOp>(op.getOperation(), true);
+    return success();
+  }
+};
+
+bool hasIndexStyle(Value value) {
+  if (value.getType().isIndex()) return true;
+  auto type = mlir::dyn_cast<ShapedType>(value.getType());
+  return type && type.getElementType().isIndex();
+}
+
+struct ConvertShapeToStablehloWithConstraintsPass
+    : public impl::ConvertShapeToStablehloWithConstraintsPassBase<
+          ConvertShapeToStablehloWithConstraintsPass> {
+  void runOnOperation() override {
+    ConversionTarget target(getContext());
+    target.addIllegalDialect<shape::ShapeDialect>();
+    target.addIllegalDialect<tensor::TensorDialect>();
+    target.addIllegalOp<arith::IndexCastOp>();
+    target.addIllegalOp<arith::MulIOp>();
+    target.addDynamicallyLegalDialect<::mlir::stablehlo::StablehloDialect>(
+        [](Operation* op) {
+          return !llvm::any_of(op->getOperands(), hasIndexStyle);
+        });
+    target.addLegalOp<tensor::CastOp>();
+    target.addLegalOp<UnrealizedConversionCastOp>();
+    target.addLegalOp<shape::ConstWitnessOp, shape::AssumingOp,
+                      shape::AssumingYieldOp>();
+
+    RewritePatternSet patterns(&getContext());
+    ::mlir::stablehlo::populateShapeToStablehloPatterns(&getContext(),
+                                                        &patterns);
+
+    patterns.add<ConvertCstrBroadcastableOp>(&getContext());
+    if (failed(applyPartialConversion(getOperation(), target,
+                                      std::move(patterns))))
+      return signalPassFailure();
+  }
+};
+
+}  // namespace
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/passes/convert_xla_call_module_op_to_bfloat16.cc

@@ -0,0 +1,149 @@
+/* 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 <string>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "llvm/Support/raw_ostream.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/Quant.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Dialect/Shape/IR/Shape.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/IR/Visitors.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 "mlir/Support/TypeID.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "shardy/dialect/sdy/ir/register.h"  // from @shardy
+#include "stablehlo/dialect/Serialization.h"  // from @stablehlo
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h"  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/quantization/stablehlo/utils/bfloat16_type.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+
+namespace mlir::quant::stablehlo {
+
+absl::StatusOr<std::string> ConvertSerializedStableHloModuleToBfloat16(
+    const StringRef serialized_stablehlo_module) {
+  // StableHLO module is empty often because the XlaCallModuleOp is already
+  // deserialized, e.g. after invoking XlaCallModuleDeserializationPass. We
+  // don't handle this situation.
+  if (serialized_stablehlo_module.empty()) {
+    return absl::InvalidArgumentError("StableHLO module is empty.");
+  }
+
+  MLIRContext context;
+  mlir::sdy::loadAllRequiredDialects(&context);
+  OwningOpRef<ModuleOp> stablehlo_module_op =
+      mlir::stablehlo::deserializePortableArtifact(serialized_stablehlo_module,
+                                                   &context);
+  auto version =
+      mlir::stablehlo::getPortableArtifactVersion(serialized_stablehlo_module);
+  if (failed(version)) {
+    return absl::InternalError(
+        "Failed to get the deserialized StableHLO version, XlaCallModuleOp "
+        "must have a valid StableHLO module serialized using "
+        "stablehlo::serializePortableArtifact APIs.");
+  }
+
+  // Convert the StableHLO module to bfloat16.
+  PassManager pm(&context);
+  pm.addNestedPass<func::FuncOp>(createConvertFuncToBfloat16Pass());
+  if (failed(pm.run(stablehlo_module_op.get()))) {
+    return absl::InternalError(
+        "Failed to convert StableHLO module to bfloat16.");
+  }
+
+  std::string bytecode;
+  llvm::raw_string_ostream os(bytecode);
+  if (failed(mlir::stablehlo::serializePortableArtifact(
+          stablehlo_module_op.get(), version.value().toString(), os,
+          /*allowOtherDialects=*/true))) {
+    return absl::InternalError("Failed to serialize StableHLO module.");
+  }
+  return bytecode;
+}
+
+#define GEN_PASS_DEF_CONVERTXLACALLMODULEOPTOBFLOAT16PASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"
+
+namespace {
+class ConvertXlaCallModuleOpToBfloat16Pass
+    : public impl::ConvertXlaCallModuleOpToBfloat16PassBase<
+          ConvertXlaCallModuleOpToBfloat16Pass> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
+      ConvertXlaCallModuleOpToBfloat16Pass)
+
+  explicit ConvertXlaCallModuleOpToBfloat16Pass() = default;
+
+ private:
+  void runOnOperation() override;
+};
+
+void ConvertXlaCallModuleOpToBfloat16Pass::runOnOperation() {
+  Operation* func_op = getOperation();
+  SymbolTableCollection symbol_table;
+  OpBuilder builder(&getContext());
+
+  auto result = func_op->walk([&](TF::XlaCallModuleOp op) {
+    // Converts the serialized StableHLO module to bfloat16.
+    auto result =
+        ConvertSerializedStableHloModuleToBfloat16(op.getModuleAttr());
+    if (!result.ok()) {
+      llvm::errs() << "Failed to convert StableHLO module to bfloat16: "
+                   << result.status().message();
+      return WalkResult::interrupt();
+    }
+    op.setModuleAttr(StringAttr::get(&getContext(), *result));
+
+    // Convert the `tf.XlaCallModuleOp` to bfloat16 and add casts around it.
+    builder.setInsertionPoint(op);
+    for (auto& op_operand : op->getOpOperands()) {
+      if (quant::stablehlo::IsLargeFloatType(op_operand.get().getType())) {
+        op_operand.set(TF::CastOp::create(
+            builder, op->getLoc(),
+            quant::stablehlo::ToBfloat16Type(op_operand.get().getType()),
+            op_operand.get()));
+      }
+    }
+    builder.setInsertionPointAfter(op);
+    for (auto op_result : op->getOpResults()) {
+      if (quant::stablehlo::IsLargeFloatType(op_result.getType())) {
+        const Type original_type = op_result.getType();
+        op_result.setType(quant::stablehlo::ToBfloat16Type(original_type));
+        const Value cast =
+            TF::CastOp::create(builder, op->getLoc(), original_type, op_result);
+        op_result.replaceAllUsesExcept(cast, cast.getDefiningOp());
+      }
+    }
+    return WalkResult::advance();
+  });
+
+  if (result.wasInterrupted()) return signalPassFailure();
+}
+
+}  // namespace
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/passes/defer_activation_transpose.cc

@@ -0,0 +1,293 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include <cstdint>
+#include <optional>
+#include <utility>
+
+#include "absl/base/nullability.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/IRMapping.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo
+#include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/permutation.h"
+
+namespace mlir::quant::stablehlo {
+
+#define GEN_PASS_DEF_DEFERACTIVATIONTRANSPOSEPASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"
+
+namespace {
+
+using ::mlir::stablehlo::AddOp;
+using ::mlir::stablehlo::BroadcastInDimOp;
+using ::mlir::stablehlo::MaxOp;
+using ::mlir::stablehlo::TransposeOp;
+
+// Returns `success()` if `op` is a `TransposeOp` with permutation attribute
+// equivalent to `permuation`.
+LogicalResult IsTransposeOpWithPermuation(Operation* absl_nullable op,
+                                          const ArrayRef<int64_t> permutation) {
+  auto transpose_op = dyn_cast_or_null<TransposeOp>(op);
+  return success(transpose_op != nullptr && transpose_op.getPermutation() ==
+                                                ArrayRef<int64_t>(permutation));
+}
+
+// Convenience function to create a `TransposeOp` with a given `permutation`.
+// The Location is set as `input`'s loc.
+TransposeOp CreateTransposeOp(Value input, const ArrayRef<int64_t> permutation,
+                              PatternRewriter& rewriter) {
+  return TransposeOp::create(rewriter, input.getLoc(), input,
+                             rewriter.getDenseI64ArrayAttr(permutation));
+}
+
+// Defers the transpose of the left-hand side (LHS) to the right-hand side and
+// the result of a binary operation. In detail, this rewrites the
+// `op(transpose(%rhs), %lhs)` to `transpose(op(%rhs, transpose(%lhs)))`. The
+// LHS transpose permutation must be a NCHW->NHWC permutation.
+template <typename OpT>
+void DeferRhsTransposeForBinaryOp(OpT op, PatternRewriter& rewriter) {
+  auto transpose_op = cast<TransposeOp>(op.getOperand(0).getDefiningOp());
+  Value lhs_pre_transpose = transpose_op.getOperand();
+
+  // NCHW -> NHWC for the right-hand side, to match the operand's shape.
+  Value rhs = op.getOperand(1);
+  TransposeOp rhs_transpose_op = CreateTransposeOp(
+      /*input=*/rhs, kNchwToNhwcPermutation, rewriter);
+
+  auto new_binary_op =
+      OpT::create(rewriter, op.getLoc(), lhs_pre_transpose, rhs_transpose_op);
+
+  // NHWC -> NCHW for the output, to match the shapes of `op`'s users.
+  TransposeOp output_transpose_op = CreateTransposeOp(
+      /*input=*/new_binary_op, kNhwcToNchwPermutation, rewriter);
+
+  rewriter.replaceAllUsesWith(op.getResult(), output_transpose_op);
+}
+
+// "Climbs up" the `op` if `op` is a `BraodcastInDimOp` and returns the defining
+// op of its operand. Returns `op` otherwise. May return `nullptr` when the
+// `BroadcastInDimOp`'s operand is a block argument.
+Operation* absl_nullable SkipUpwardsOptionalBroadcastInDimOp(
+    Operation* absl_nonnull op) {
+  if (auto broadcast_in_dim_op = dyn_cast_or_null<BroadcastInDimOp>(op);
+      broadcast_in_dim_op != nullptr) {
+    return broadcast_in_dim_op.getOperand().getDefiningOp();
+  }
+  return op;
+}
+
+class DeferActivationTransposeForAddOp : public OpRewritePattern<AddOp> {
+ public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(AddOp op,
+                                PatternRewriter& rewriter) const override {
+    // Only supports the case for 2D convolution.
+    const Value lhs = op.getOperand(0);
+    if (!HasRankOf(lhs, /*rank=*/4)) return failure();
+
+    const Value rhs = op.getOperand(1);
+    Operation* rhs_op = rhs.getDefiningOp();
+    if (rhs_op == nullptr) return failure();
+
+    // Ignore the optional `BroadcastInDimOp` in between the constant and RHS.
+    rhs_op = SkipUpwardsOptionalBroadcastInDimOp(rhs_op);
+
+    if (rhs_op == nullptr || !rhs_op->hasTrait<OpTrait::ConstantLike>()) {
+      return failure();
+    }
+
+    // Match LHS permutation that converts: NHWC -> NCHW.
+    if (IsTransposeOpWithPermuation(lhs.getDefiningOp(), kNhwcToNchwPermutation)
+            .failed()) {
+      return failure();
+    }
+
+    DeferRhsTransposeForBinaryOp(op, rewriter);
+    return success();
+  }
+};
+
+// Rewrites the `reduce_window(transpose(%activation), %init_value)` patterns to
+// `transpose(reduce_window(%activation), %init_value)`, deferring the transpose
+// to the result. The reduce function should be equivalent to
+// `stablehlo.maximum`, representing max pooling.
+class DeferActivationTransposeForMaxPoolReduceWindowOp
+    : public OpRewritePattern<mlir::stablehlo::ReduceWindowOp> {
+ public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mlir::stablehlo::ReduceWindowOp op,
+                                PatternRewriter& rewriter) const override {
+    if (failed(MatchMaxPoolReduceWindowOp(op))) return failure();
+
+    // Match only when the lhs is connected to a transpose.
+    // Only supports the case commonly appearing for 2D convolutions.
+    Value lhs = op.getOperand(0);
+    if (!HasRankOf(lhs, /*rank=*/4)) return failure();
+
+    // Match input permutation that converts: NHWC -> NCHW.
+    if (IsTransposeOpWithPermuation(lhs.getDefiningOp(), kNhwcToNchwPermutation)
+            .failed()) {
+      return failure();
+    }
+
+    // Pushes the transpose op at the input to the result.
+    auto transpose_op = cast<TransposeOp>(op.getOperand(0).getDefiningOp());
+
+    const auto result_type = mlir::cast<TensorType>(op.getResult(0).getType());
+    const SmallVector<int64_t> new_result_shape =
+        quant::Permute<int64_t>(result_type.getShape(), kNchwToNhwcPermutation);
+
+    const TensorType new_result_type =
+        result_type.cloneWith(new_result_shape, result_type.getElementType());
+
+    // Create a new `stablehlo.reduce_window` with all relevant attributes
+    // permutated to match the new operand & result type.
+    auto new_reduce_window_op = mlir::stablehlo::ReduceWindowOp::create(
+        rewriter, op.getLoc(), new_result_type, transpose_op.getOperand(),
+        /*init_value=*/op.getOperand(1),
+        /*window_dimensions=*/
+        PermuteI64ArrayAttr(rewriter, op.getWindowDimensions(),
+                            kNchwToNhwcPermutation),
+        /*window_strides=*/
+        PermuteI64ArrayAttr(rewriter, op.getWindowStrides(),
+                            kNchwToNhwcPermutation),
+        /*base_dilations=*/
+        PermuteI64ArrayAttr(rewriter, op.getBaseDilations(),
+                            kNchwToNhwcPermutation),
+        /*window_dilations=*/
+        PermuteI64ArrayAttr(rewriter, op.getWindowDilations(),
+                            kNchwToNhwcPermutation),
+        /*padding=*/DenseIntElementsAttr(nullptr));
+
+    // Clone the reduce body. It is not affected by the permutation.
+    IRMapping mapping;
+    op.getBody().cloneInto(&new_reduce_window_op.getBody(), mapping);
+
+    // Introduce a transpose to the result to match the shapes of `op`'s uses.
+    TransposeOp result_transpose_op = CreateTransposeOp(
+        /*input=*/new_reduce_window_op.getResult(0), kNhwcToNchwPermutation,
+        rewriter);
+
+    rewriter.replaceAllUsesWith(op.getResult(0), result_transpose_op);
+    return success();
+  }
+
+ private:
+  // Permutes `array_attr` with `permutation`. The number of elements in
+  // `array_attr` and `permutation` must be equal. Returns a null attribute
+  // if `array_attr` is null.
+  DenseI64ArrayAttr PermuteI64ArrayAttr(
+      PatternRewriter& rewriter,
+      const std::optional<ArrayRef<int64_t>> array_attr,
+      const ArrayRef<int64_t> permutation) const {
+    if (!array_attr.has_value()) return DenseI64ArrayAttr(nullptr);
+
+    return rewriter.getDenseI64ArrayAttr(
+        quant::Permute<int64_t>(array_attr.value(), permutation));
+  }
+
+  LogicalResult MatchMaxPoolReduceWindowOp(
+      mlir::stablehlo::ReduceWindowOp op) const {
+    // TODO: b/321099943 - Support explicit padding.
+    if (HasPadding(op)) return failure();
+
+    // Check that the reduce-window body is a max operation.
+    return success(IsMaxFunction(op.getBody().front()));
+  }
+
+  // Whether `block` semantically corresponds to a `stablehlo.maximum` op.
+  bool IsMaxFunction(Block& block) const {
+    if (block.getNumArguments() != 2) return false;
+
+    auto return_op = cast<mlir::stablehlo::ReturnOp>(block.getTerminator());
+    if (return_op.getNumOperands() != 1) return false;
+
+    auto max_op = dyn_cast_or_null<MaxOp>(
+        return_op.getOperands().front().getDefiningOp());
+    if (!max_op) return false;
+
+    return (max_op.getLhs() == block.getArgument(0)) &&
+           (max_op.getRhs() == block.getArgument(1));
+  }
+
+  // Whether `op` has the `padding` attribute (which is optional).
+  bool HasPadding(mlir::stablehlo::ReduceWindowOp op) const {
+    return op.getPadding() != std::nullopt;
+  }
+};
+
+// Rewrites `maximum(transpose(%rhs), %lhs)` patterns to
+// `transpose(maximum(%rhs, transpose(%lhs)))`.
+class DeferActivationTransposeForMaxOp : public OpRewritePattern<MaxOp> {
+ public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(MaxOp op,
+                                PatternRewriter& rewriter) const override {
+    Value input = op.getOperand(0);
+    if (!HasRankOf(input, /*rank=*/4)) return failure();
+
+    const Value max_value = op.getOperand(1);
+    Operation* max_value_op = max_value.getDefiningOp();
+    if (max_value_op == nullptr ||
+        !max_value_op->hasTrait<OpTrait::ConstantLike>()) {
+      return failure();
+    }
+
+    if (IsTransposeOpWithPermuation(input.getDefiningOp(),
+                                    kNhwcToNchwPermutation)
+            .failed()) {
+      return failure();
+    }
+    DeferRhsTransposeForBinaryOp(op, rewriter);
+    return success();
+  }
+};
+
+}  // namespace
+
+class DeferActivationTransposePass
+    : public impl::DeferActivationTransposePassBase<
+          DeferActivationTransposePass> {
+ private:
+  void runOnOperation() override;
+};
+
+void DeferActivationTransposePass::runOnOperation() {
+  func::FuncOp func_op = getOperation();
+  MLIRContext& ctx = getContext();
+
+  RewritePatternSet patterns(&ctx);
+  patterns.add<DeferActivationTransposeForAddOp,
+               DeferActivationTransposeForMaxPoolReduceWindowOp,
+               DeferActivationTransposeForMaxOp>(&ctx);
+  if (failed(applyPatternsGreedily(func_op, std::move(patterns)))) {
+    func_op->emitWarning() << "Failed to converge patterns: " << getArgument();
+  }
+}
+
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/passes/fold_constant_transpose.cc

@@ -0,0 +1,195 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include <cstdint>
+#include <utility>
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/permutation.h"
+
+namespace mlir::quant::stablehlo {
+
+#define GEN_PASS_DEF_FOLDCONSTANTTRANSPOSEPASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"
+
+namespace {
+
+// Returns contiguous offset (address) of the position represented by `indices`
+// in a `shape` shaped tensor. Assumes row-major order. `indices` and `shape`
+// should have the same size.
+// Example: Index (2, 3) of a (4, 5)-shaped tensor has the contiguous offset of
+// 2 * 5 + 3 = 13.
+int64_t GetContiguousOffset(const ArrayRef<int64_t> indices,
+                            const ArrayRef<int64_t> shape) {
+  int64_t contiguous_offset = 0;
+  int64_t base_offset = 1;
+  for (auto [i, dimension] : llvm::reverse(llvm::zip_equal(indices, shape))) {
+    contiguous_offset += base_offset * i;
+    base_offset *= dimension;
+  }
+
+  return contiguous_offset;
+}
+
+// Performs transposition of a tensor represented as a contiguous element array.
+// Assumes row-major order. The shape of the input tensor and the desired
+// permutation is registered during construction, and calling `TransposeValues`
+// returns the transposed tensor values.
+class DenseElementsTransposer {
+ public:
+  DenseElementsTransposer(const ArrayRef<int64_t> original_shape,
+                          const ArrayRef<int64_t> permutation)
+      : rank_(original_shape.size()),
+        original_shape_(original_shape),
+        target_shape_(quant::Permute<int64_t>(original_shape, permutation)),
+        permutation_(permutation) {}
+
+  // Transposes `values` with the permutation. Returns the transposed values.
+  SmallVector<float> TransposeValues(const ArrayRef<float> values) const {
+    SmallVector<float> transposed_values(values.size());
+    SmallVector<int64_t> current_indices = {};
+    TransposeRecursively(values, transposed_values, current_indices);
+
+    return transposed_values;
+  }
+
+  // Returns the shape after permutation.
+  llvm::ArrayRef<int64_t> GetTargetShape() const { return target_shape_; }
+
+ private:
+  // Helper function that performs transposition recursively by mapping each set
+  // of indices from the original values to the target values.
+  void TransposeRecursively(const ArrayRef<float> original_values,
+                            const MutableArrayRef<float> target_values,
+                            SmallVector<int64_t>& current_indices) const {
+    // Map an element from `original_values` to `target_values` when a set of
+    // indices is formed.
+    if (current_indices.size() == rank_) {
+      const int64_t original_index =
+          GetContiguousOffset(current_indices, original_shape_);
+
+      const SmallVector<int64_t> target_indices =
+          quant::Permute<int64_t>(current_indices, permutation_);
+      const int64_t target_index =
+          GetContiguousOffset(target_indices, target_shape_);
+
+      target_values[target_index] = original_values[original_index];
+      return;
+    }
+
+    // Recursively iterate by selecting the index of the next dimension.
+    const int64_t next_shape_idx = current_indices.size();
+    for (int64_t i = 0; i < original_shape_[next_shape_idx]; ++i) {
+      current_indices.push_back(i);
+      TransposeRecursively(original_values, target_values, current_indices);
+      current_indices.pop_back();
+    }
+  }
+
+  int64_t rank_;                         // Rank of the input values.
+  SmallVector<int64_t> original_shape_;  // Shape of the original tensor.
+  SmallVector<int64_t> target_shape_;    // Shape of the target tensor.
+  SmallVector<int64_t> permutation_;
+};
+
+class FoldTransposedConstantOp
+    : public OpRewritePattern<mlir::stablehlo::TransposeOp> {
+ public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mlir::stablehlo::TransposeOp op,
+                                PatternRewriter& rewriter) const override {
+    Value operand = op.getOperand();
+    auto const_op =
+        dyn_cast_or_null<mlir::stablehlo::ConstantOp>(operand.getDefiningOp());
+    if (!const_op) return failure();
+
+    // Only support float tensors.
+    auto tensor_type = mlir::dyn_cast<TensorType>(const_op.getType());
+    if (!tensor_type || !tensor_type.getElementType().isF32()) {
+      return failure();
+    }
+
+    const auto value_attr =
+        mlir::dyn_cast_or_null<DenseFPElementsAttr>(const_op.getValue());
+    if (!value_attr) {
+      return failure();
+    }
+    const ArrayRef<int64_t> original_shape =
+        value_attr.getShapedType().getShape();
+
+    const SmallVector<float> original_values =
+        llvm::to_vector(value_attr.getValues<float>());
+
+    // Fold the constant value by transposing the values according to the
+    // `TransposeOp`'s permutation attribute.
+    const DenseElementsTransposer transposer(original_shape,
+                                             op.getPermutation());
+    SmallVector<float> transposed_values =
+        transposer.TransposeValues(original_values);
+
+    // Create a new constant op with the transposed values.
+    const Location combined_loc =
+        rewriter.getFusedLoc({const_op.getLoc(), op.getLoc()});
+    RankedTensorType new_value_type =
+        RankedTensorType::getChecked(combined_loc, transposer.GetTargetShape(),
+                                     /*elementType=*/rewriter.getF32Type());
+    DenseFPElementsAttr new_value_attr =
+        DenseFPElementsAttr::get(new_value_type, std::move(transposed_values));
+    mlir::stablehlo::ConstantOp new_const_op =
+        mlir::stablehlo::ConstantOp::create(rewriter, combined_loc,
+                                            new_value_attr);
+
+    rewriter.replaceOp(op, new_const_op);
+    return success();
+  }
+};
+
+}  // namespace
+
+class FoldConstantTransposePass
+    : public impl::FoldConstantTransposePassBase<FoldConstantTransposePass> {
+ public:
+  using impl::FoldConstantTransposePassBase<
+      FoldConstantTransposePass>::FoldConstantTransposePassBase;
+
+ private:
+  void runOnOperation() override;
+};
+
+void FoldConstantTransposePass::runOnOperation() {
+  func::FuncOp func_op = getOperation();
+  MLIRContext& ctx = getContext();
+
+  RewritePatternSet patterns(&ctx);
+  patterns.add<FoldTransposedConstantOp>(&ctx);
+  if (failed(applyPatternsGreedily(func_op, std::move(patterns)))) {
+    func_op.emitError("Failed to fold constant->transpose pattern.");
+    signalPassFailure();
+  }
+}
+
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/passes/insert_calibration_statistics_saver.cc

@@ -0,0 +1,190 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <unordered_set>
+#include <vector>
+
+#include "absl/strings/string_view.h"
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/tf_quant_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"  // IWYU pragma: keep
+#include "tsl/platform/path.h"
+
+namespace mlir::quant::stablehlo {
+namespace {
+
+std::string GetOutputFilePath(absl::string_view calibration_data_dir,
+                              absl::string_view func_name,
+                              int32_t output_file_idx) {
+  return tsl::io::JoinPath(calibration_data_dir,
+                           llvm::Twine(func_name)
+                               .concat("_")
+                               .concat(std::to_string(output_file_idx))
+                               .concat(".pb")
+                               .str());
+}
+
+// Finds `CustomAggregator` ops and collects their outputs and attributes.
+void FindCustomAggregatorOps(
+    Region& region,
+    const std::unordered_set<std::string>& aggregator_ops_to_ignore,
+    SmallVector<Value>& statistics_outputs, SmallVector<StringRef>& ids,
+    SmallVector<int32_t>& calibration_methods) {
+  for (auto op : region.getOps<TF::CustomAggregatorOp>()) {
+    if (aggregator_ops_to_ignore.count(op.getId().str())) continue;
+
+    ids.push_back(op.getId());
+    calibration_methods.push_back(op.getCalibrationMethod());
+    statistics_outputs.push_back(op.getMin());
+    statistics_outputs.push_back(op.getMax());
+    statistics_outputs.push_back(op.getHistogram());
+  }
+}
+
+// Inserts a `CalibrationStatisticsSaverOp` to the end of the region.
+LogicalResult InsertCalibrationStatisticsSaverOp(
+    Region& region, MLIRContext& ctx, absl::string_view output_file_path,
+    const std::unordered_set<std::string>& aggregator_ops_to_ignore) {
+  SmallVector<Value> statistics_outputs;
+  SmallVector<StringRef> ids;
+  SmallVector<int32_t> calibration_methods;
+  FindCustomAggregatorOps(region, aggregator_ops_to_ignore, statistics_outputs,
+                          ids, calibration_methods);
+  if (statistics_outputs.empty()) return failure();
+
+  OpBuilder builder(&ctx);
+  // Set the insertion point right before the return op.
+  builder.setInsertionPoint(&region.back().back());
+
+  StringAttr output_file_path_attr = builder.getStringAttr(output_file_path);
+  ArrayAttr ids_attr = builder.getStrArrayAttr(ids);
+  ArrayAttr calibration_methods_attr =
+      builder.getI32ArrayAttr(calibration_methods);
+  TF::CalibrationStatisticsSaverOp::create(
+      builder, region.getLoc(), statistics_outputs, output_file_path_attr,
+      ids_attr, calibration_methods_attr);
+  return success();
+}
+
+// Returns true if the op contains a `CalibrationStatisticsSaverOp`.
+bool ContainCalibrationStatisticsSaverOp(Operation* op) {
+  // Check the region for CaseRegionOp, IfRegionOp and WhileRegionOp.
+  for (Region& region : op->getRegions()) {
+    if (!region.getOps<TF::CalibrationStatisticsSaverOp>().empty()) {
+      return true;
+    }
+  }
+
+  SymbolTable symbol_table(op->getParentOfType<ModuleOp>());
+  // Check the functions associated to CaseOp, IfOp and WhileOp.
+  for (const NamedAttribute& attr : op->getAttrs()) {
+    FlatSymbolRefAttr symbol_attr =
+        dyn_cast_or_null<FlatSymbolRefAttr>(attr.getValue());
+    if (!symbol_attr) continue;
+
+    func::FuncOp target_func = dyn_cast_or_null<func::FuncOp>(
+        symbol_table.lookup(symbol_attr.getValue()));
+    if (!target_func) continue;
+
+    if (!target_func.getBody()
+             .getOps<TF::CalibrationStatisticsSaverOp>()
+             .empty()) {
+      return true;
+    }
+  }
+  return false;
+}
+
+}  // namespace
+
+#define GEN_PASS_DECL_INSERTCALIBRATIONSTATISTICSSAVERPASS
+#define GEN_PASS_DEF_INSERTCALIBRATIONSTATISTICSSAVERPASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"
+
+class InsertCalibrationStatisticsSaverPass
+    : public impl::InsertCalibrationStatisticsSaverPassBase<
+          InsertCalibrationStatisticsSaverPass> {
+ public:
+  using impl::InsertCalibrationStatisticsSaverPassBase<
+      InsertCalibrationStatisticsSaverPass>::
+      InsertCalibrationStatisticsSaverPassBase;
+
+ private:
+  void runOnOperation() override;
+};
+
+void InsertCalibrationStatisticsSaverPass::runOnOperation() {
+  ModuleOp module_op = getOperation();
+  MLIRContext& ctx = getContext();
+
+  std::unordered_set<std::string> aggregator_ops_to_ignore(
+      aggregator_ops_to_ignore_.begin(), aggregator_ops_to_ignore_.end());
+
+  // Insert CalibrationStatisticsSaverOp to the end of each region.
+  for (auto func_op : module_op.getOps<func::FuncOp>()) {
+    int32_t output_file_idx = 0;
+    StringRef func_name = func_op.getSymName();
+
+    func_op.walk([&output_file_idx, &ctx, &func_name, &aggregator_ops_to_ignore,
+                  this](Operation* op) {
+      for (Region& region : op->getRegions()) {
+        if (succeeded(InsertCalibrationStatisticsSaverOp(
+                region, ctx,
+                GetOutputFilePath(calibration_data_dir_, func_name,
+                                  output_file_idx),
+                aggregator_ops_to_ignore))) {
+          ++output_file_idx;
+        };
+      }
+    });
+  }
+
+  // Control flow ops that contains CalibrationStatisticsSaver ops must be set
+  // to stateful, otherwise the op will not be executed.
+  OpBuilder builder(&ctx);
+  module_op.walk([&builder](Operation* op) {
+    if (op->hasAttrOfType<BoolAttr>("is_stateless") &&
+        ContainCalibrationStatisticsSaverOp(op)) {
+      op->setAttr("is_stateless", builder.getBoolAttr(false));
+    }
+  });
+}
+
+std::unique_ptr<OperationPass<ModuleOp>>
+CreateInsertCalibrationStatisticsSaverPass(
+    StringRef calibration_data_dir,
+    const std::vector<std::string>& aggregator_ops_to_ignore) {
+  InsertCalibrationStatisticsSaverPassOptions options = {
+      .aggregator_ops_to_ignore_ = llvm::to_vector(aggregator_ops_to_ignore),
+      .calibration_data_dir_ = calibration_data_dir.str(),
+  };
+  return std::make_unique<InsertCalibrationStatisticsSaverPass>(options);
+}
+
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/passes/insert_weight_param.cc

@@ -0,0 +1,250 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cstdint>
+#include <utility>
+
+#include "llvm/ADT/STLExtras.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/Quant.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/Dialect/Shape/IR/Shape.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Matchers.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/IR/Visitors.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project  // IWYU pragma: keep
+#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 "mlir/Support/TypeID.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.h"
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"
+#include "tensorflow/compiler/mlir/quantization/common/lift_as_function_call.h"
+#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h"  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+
+namespace mlir::quant::stablehlo {
+
+
+#define GEN_PASS_DEF_INSERTWEIGHTPARAMPASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"
+
+namespace {
+
+using ::stablehlo::quantization::Method;
+using ::stablehlo::quantization::QuantizedType;
+using ::stablehlo::quantization::WeightOnlyPtq;
+
+// Inserts quantization parameters of weights for weight-only quantization and
+// dynamic range quantization of `stablehlo.convolution` and
+// `stablehlo.dot_general`.
+class InsertWeightParamPass
+    : public impl::InsertWeightParamPassBase<InsertWeightParamPass> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(InsertWeightParamPass)
+
+  using impl::InsertWeightParamPassBase<
+      InsertWeightParamPass>::InsertWeightParamPassBase;
+
+ private:
+  void runOnOperation() override;
+};
+
+// Inserts quantization parameters for weights for hybrid quantization of
+// `stablehlo.convolution` and `stablehlo.dot_general`.
+class InsertWeightParamPattern
+    : public OpTraitRewritePattern<OpTrait::ConstantLike> {
+ public:
+  explicit InsertWeightParamPattern(MLIRContext* context)
+      : OpTraitRewritePattern(context) {}
+
+  LogicalResult matchAndRewrite(Operation* op,
+                                PatternRewriter& rewriter) const override {
+    if (op->getNumResults() != 1) {
+      return failure();
+    }
+    auto type = mlir::cast<TensorType>(op->getResult(0).getType());
+    if (!type || !type.getElementType().isF32()) {
+      return failure();
+    }
+    if (!op->hasOneUse() ||
+        !IsWeightQuantizableFunction(*op->getUses().begin(), type.getRank())) {
+      return failure();
+    }
+
+    Operation* quantizable_op = *op->getUsers().begin();
+    DenseFPElementsAttr attr;
+    matchPattern(op->getResult(0), m_Constant(&attr));
+
+    Method method = GetQuantizationMethodOrDefault(quantizable_op);
+    const WeightOnlyPtq& weight_only_ptq = method.weight_only_ptq();
+
+    Type weight_type;
+    if (IsPerTensor(weight_only_ptq)) {
+      weight_type =
+          dyn_cast<quant::QuantizedType>(GetUniformQuantizedTypeForWeight(
+              attr, /*symmetric=*/true, /*num_bits=*/8, /*is_signed=*/true,
+              /*narrow_range=*/true, /*legacy_float_scale=*/false));
+    } else {
+      int quantization_dimension = GetQuantizationDimension(
+          weight_only_ptq, cast<TF::XlaCallModuleOp>(quantizable_op));
+      weight_type = GetUniformQuantizedPerAxisTypeForWeight(
+          attr, quantization_dimension, /*symmetric=*/true, /*num_bits=*/8,
+          /*is_signed=*/true,
+          /*narrow_range=*/true, /*legacy_float_scale=*/false);
+    }
+
+    auto quant_type = dyn_cast<quant::QuantizedType>(weight_type);
+    if (!quant_type) {
+      op->emitError(
+          "Failed to get weight quantization parameters for weight-only "
+          "quantization.");
+      return failure();
+    }
+
+    const Type expressed_type = op->getResult(0).getType();
+    const Type quantized_type =
+        quant_type.castFromExpressedType(expressed_type);
+
+    rewriter.setInsertionPointAfter(op);
+    auto q = mlir::quant::ir::QuantizeCastOp::create(
+        rewriter, op->getLoc(), quantized_type, op->getResult(0));
+    auto dq = mlir::quant::ir::DequantizeCastOp::create(rewriter, op->getLoc(),
+                                                        expressed_type, q);
+    quantizable_op->setOperand(1, dq.getResult());
+    return success();
+  }
+
+  // Checks if the operand is second operand of `tf.XlaCallModule` op for
+  // `stablehlo.convolution` or `stablehlo.dot_general` with fully_quantizable
+  // trait.
+  static bool IsWeightQuantizableFunction(OpOperand& operand, int64_t rank) {
+    if (operand.getOperandNumber() != 1) {
+      return false;
+    }
+    Operation* user = operand.getOwner();
+    if (!IsWeightOnlyQuantizableOp(*user)) {
+      return false;
+    }
+    Method method = GetQuantizationMethodOrDefault(user);
+    return HasValidWeightOnlyPtqMethod(method.weight_only_ptq(), rank);
+  }
+
+ private:
+  static bool HasValidWeightOnlyPtqMethod(const WeightOnlyPtq& weight_only_ptq,
+                                          int64_t rank) {
+    const auto& input_quantized_types = weight_only_ptq.input_quantized_types();
+    if (IsPerTensor(weight_only_ptq)) {
+      return true;
+    }
+    // `input_quantized_types` should contain spec for quantization type of the
+    // second operand, which is weight.
+    const QuantizedType& quantized_type = input_quantized_types.at(1);
+    if (const auto& specs = quantized_type.dimension_specs();
+        specs.has_dimension()) {
+      return specs.dimension() >= 0 && specs.dimension() < rank;
+    }
+    return true;
+  }
+
+  static bool IsPerTensor(const WeightOnlyPtq& weight_only_ptq) {
+    const auto& input_quantized_types = weight_only_ptq.input_quantized_types();
+    if (input_quantized_types.empty()) {
+      return true;
+    }
+    auto weight_type = input_quantized_types.find(1);
+    if (weight_type == input_quantized_types.end()) {
+      return true;
+    }
+    return weight_type->second.has_per_tensor();
+  }
+
+  static int GetQuantizationDimension(const WeightOnlyPtq& weight_only_ptq,
+                                      TF::XlaCallModuleOp op) {
+    const QuantizedType& quantized_type =
+        weight_only_ptq.input_quantized_types().at(1);
+    if (quantized_type.dimension_specs().has_dimension()) {
+      return quantized_type.dimension_specs().dimension();
+    }
+    return GetDefaultQuantizationDimension(op);
+  }
+
+  // Determines quantization dimension of weights for given `tf.XlaCallModule`
+  // op. For convolution, returns output feature dimension of the kernel. For
+  // dot_general, returns the first non-contracting dimension, non-batching
+  // dimension. If such dimension does not exists, returns the last dimension of
+  // rhs.
+  static int64_t GetDefaultQuantizationDimension(TF::XlaCallModuleOp op) {
+    const StringRef function_name = GetEntryFunctionName(op);
+    const auto module_op = op->getParentOfType<ModuleOp>();
+    const SymbolTable symbol_table(module_op);
+    func::FuncOp func = symbol_table.lookup<func::FuncOp>(function_name);
+
+    if (function_name.contains("conv")) {
+      return (*(func.getOps<mlir::stablehlo::ConvolutionOp>().begin()))
+          .getDimensionNumbers()
+          .getKernelOutputFeatureDimension();
+    } else if (function_name.contains("dot_general")) {
+      auto dot = *(func.getOps<mlir::stablehlo::DotGeneralOp>().begin());
+      const ::mlir::stablehlo::DotDimensionNumbersAttr dimension_numbers =
+          dot.getDotDimensionNumbers();
+      ArrayRef<int64_t> rhs_contracting_dims =
+          dimension_numbers.getRhsContractingDimensions();
+      ArrayRef<int64_t> rhs_batching_dims =
+          dimension_numbers.getRhsBatchingDimensions();
+      int64_t rank = cast<TensorType>(dot.getRhs().getType()).getRank();
+      for (int i = 0; i < rank; ++i) {
+        // Return the first non-contracting, non-batching dimension of rhs.
+        if (llvm::find(rhs_contracting_dims, i) == rhs_contracting_dims.end() &&
+            llvm::find(rhs_batching_dims, i) == rhs_batching_dims.end()) {
+          return i;
+        }
+      }
+    }
+    return cast<TensorType>(op.getOperand(1).getType()).getRank() - 1;
+  }
+};
+
+void InsertWeightParamPass::runOnOperation() {
+  func::FuncOp func = getOperation();
+  MLIRContext* context = func.getContext();
+  RewritePatternSet patterns(context);
+
+  patterns.add<InsertWeightParamPattern>(context);
+
+  if (failed(applyPatternsGreedily(func, std::move(patterns)))) {
+    signalPassFailure();
+  }
+}
+
+}  // namespace
+
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/passes/lift_quantizable_spots_as_functions.cc

@@ -0,0 +1,243 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include <memory>
+#include <string>
+#include <utility>
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/Debug.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Rewrite/FrozenRewritePatternSet.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.h"
+#include "tensorflow/compiler/mlir/quantization/common/lift_as_function_call.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
+#include "tsl/platform/regexp.h"    // IWYU pragma: keep
+
+#define DEBUG_TYPE "lift_quantizable_spots_as_functions"
+
+namespace mlir::quant::stablehlo {
+
+#define GEN_PASS_DEF_LIFTQUANTIZABLESPOTSASFUNCTIONSPASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"
+
+namespace {
+
+using ::stablehlo::quantization::FunctionNameMatcherSpec;
+using ::stablehlo::quantization::Method;
+using ::stablehlo::quantization::QuantizationSpec;
+using ::stablehlo::quantization::QuantizationSpecs;
+using ::tsl::protobuf::TextFormat;
+
+// TODO - b/303543789: Move the helper functions below to a separate util.
+// Fetches the default or null attribute, used for pattern matching.
+Attribute DefaultOrNullAttr(OpBuilder& builder, const Attribute& attr) {
+  if (attr) return attr;
+  return builder.getStringAttr(kNullAttributeValue);
+}
+
+// Checks whether the value of a constant equals the given float, regardless
+// of the tensor dimension.
+bool FloatValueEquals(const Attribute& attr, const double value) {
+  const auto fp_attr = mlir::dyn_cast_or_null<DenseFPElementsAttr>(attr);
+  if (!fp_attr) return false;
+
+  if (fp_attr.isSplat()) {
+    return fp_attr.getSplatValue<APFloat>().isExactlyValue(value);
+  }
+  return llvm::all_of(fp_attr.getValues<APFloat>(), [value](const APFloat& f) {
+    return f.isExactlyValue(value);
+  });
+}
+
+inline void TrimTrailingWhitespaces(std::string& str) {
+  while (!str.empty() && str.back() == ' ') {
+    str.pop_back();
+  }
+}
+
+// Lifts quantizable units as separate functions, thereby identifying the
+// boundaries of quantizable subgraphs. `QuantizationSpecs` influences how
+// quantizable units are lifted.
+//
+// FileCheck test cases using various `QuantizationSpecs` can be seen at
+// `TestLiftQuantizableSpotsAsFunctionsWithQuantizationSpecsPass`.
+class LiftQuantizableSpotsAsFunctionsPass
+    : public impl::LiftQuantizableSpotsAsFunctionsPassBase<
+          LiftQuantizableSpotsAsFunctionsPass> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
+      LiftQuantizableSpotsAsFunctionsPass)
+
+  LiftQuantizableSpotsAsFunctionsPass() = default;
+
+  // Constructor with explicit user-provided `QuantizationSpecs`.
+  explicit LiftQuantizableSpotsAsFunctionsPass(
+      QuantizationSpecs quantization_specs)
+      : quantization_specs_(std::move(quantization_specs)) {}
+
+ private:
+  void runOnOperation() override;
+
+  // No explicit quantization spec is specified by default. Implicitly this
+  // means that all quantizable units will be identified and lifted.
+  QuantizationSpecs quantization_specs_{};
+};
+
+namespace simple_patterns {
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/lift_quantizable_spots_as_functions_simple.inc"
+}
+
+namespace fusion_patterns {
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/lift_quantizable_spots_as_functions_fusion.inc"
+}
+
+// Quantizable Unit matcher that uses lifted function's name for matching.
+class FunctionNameMatcher {
+ public:
+  explicit FunctionNameMatcher(const FunctionNameMatcherSpec& spec)
+      : match_regex_(GetMatchRegex(spec)) {}
+
+  // Returns `true` when matched with the entry function of
+  // `xla_call_module_op`.
+  bool Match(TF::XlaCallModuleOp xla_call_module_op) const {
+    if (match_regex_ == nullptr) return false;
+
+    const std::string lifted_func_name =
+        xla_call_module_op->getAttrOfType<FlatSymbolRefAttr>("_entry_function")
+            .getValue()
+            .str();
+
+    return RE2::FullMatch(lifted_func_name, *match_regex_);  // NOLINT
+  }
+
+ private:
+  // Returns an owned `RE2` object that corresponds to the `spec`. Returns
+  // `nullptr` if the `spec` is invalid.
+  // NOLINTNEXTLINE - RE2 included via TSL regexp.h
+  std::unique_ptr<RE2> GetMatchRegex(const FunctionNameMatcherSpec& spec) {
+    const std::string& regex = spec.regex();
+    if (regex.empty()) return nullptr;
+
+    return std::make_unique<RE2>(regex);  // NOLINT
+  }
+
+  // Regex object used for matching against a lifted function's name.
+  std::unique_ptr<RE2> match_regex_;  // NOLINT
+};
+
+// Converts `Method` to a single-line textproto representation. Returns
+// `failure()` when converting to textproto failed.
+FailureOr<std::string> QuantizationMethodToTextProto(const Method& method) {
+  TextFormat::Printer printer;
+  printer.SetSingleLineMode(true);
+
+  std::string method_txtpb;
+  if (!printer.PrintToString(method, &method_txtpb)) {
+    LLVM_DEBUG(llvm::dbgs() << "Failed to convert Method to textproto\n.");
+    return failure();
+  }
+
+  // Single line mode might have an extra space at the end, due to the internal
+  // details of `Printer`.
+  TrimTrailingWhitespaces(method_txtpb);
+
+  return method_txtpb;
+}
+
+// Applies quantization spec to all matched lifted functions. At this point only
+// denylisting (`NoQuantization`) will be applied if specs is nonempty.
+// TODO: b/307620778 - Support more advanced selective quantization methods.
+LogicalResult ApplyQuantizationSpec(const QuantizationSpec& spec,
+                                    ModuleOp module_op) {
+  const Method& quantization_method = spec.method();
+
+  FailureOr<std::string> quantization_method_txtpb =
+      QuantizationMethodToTextProto(quantization_method);
+  if (failed(quantization_method_txtpb)) return failure();
+
+  const FunctionNameMatcher matcher(spec.matcher().function_name());
+  // Iterate over all XlaCallModuleOp in all FuncOps.
+  for (auto func : module_op.getOps<func::FuncOp>()) {
+    for (auto xla_call_module_op : func.getOps<TF::XlaCallModuleOp>()) {
+      if (!matcher.Match(xla_call_module_op)) continue;
+
+      // Set the text representation of `Method` to matched
+      // `TF::XlaCallModuleOp`.
+      xla_call_module_op->setAttr(
+          kQuantizationMethodAttr,
+          StringAttr::get(module_op.getContext(),
+                          std::move(*quantization_method_txtpb)));
+    }
+  }
+  return success();
+}
+
+void LiftQuantizableSpotsAsFunctionsPass::runOnOperation() {
+  MLIRContext* ctx = &getContext();
+  RewritePatternSet patterns(ctx);
+  ModuleOp module_op = getOperation();
+
+  simple_patterns::populateWithGenerated(patterns);
+  fusion_patterns::populateWithGenerated(patterns);
+  FrozenRewritePatternSet frozen_patterns(std::move(patterns));
+
+  // Iterate over the sorted list of functions to keep order deterministic.
+  for (func::FuncOp func : GetSortedFunctions(module_op)) {
+    if (failed(applyPatternsGreedily(func, frozen_patterns))) {
+      func.emitError()
+          << "quant-stablehlo-lift-quantizable-spots-as-functions failed.";
+      signalPassFailure();
+    }
+  }
+
+  // Remove all attr_map attributes.
+  module_op.walk([](Operation* op) { op->removeAttr(kAttrMapAttribute); });
+
+  // Perform selective quantization. Iterates over the quantization specs and
+  // applies quantization methods to each matched lifted function.
+  for (const QuantizationSpec& spec : quantization_specs_.specs()) {
+    if (failed(ApplyQuantizationSpec(spec, module_op))) {
+      signalPassFailure();
+      return;
+    }
+  }
+}
+
+}  // namespace
+
+// Creates `LiftQuantizableSpotsAsFunctionsPass` with user-defined
+// `QuantizationSpecs`.
+std::unique_ptr<OperationPass<ModuleOp>>
+CreateLiftQuantizableSpotsAsFunctionsPass(
+    const QuantizationSpecs& quantization_specs) {
+  return std::make_unique<LiftQuantizableSpotsAsFunctionsPass>(
+      quantization_specs);
+}
+
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/passes/lift_quantizable_spots_as_functions_fusion.td

@@ -0,0 +1,492 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+include "mlir/Dialect/Arith/IR/ArithOps.td"
+include "mlir/Dialect/Func/IR/FuncOps.td"
+include "mlir/Dialect/Shape/IR/ShapeOps.td"
+include "mlir/IR/OpBase.td"
+include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"
+include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"
+include "stablehlo/dialect/StablehloOps.td"
+include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.td"
+include "tensorflow/compiler/mlir/quantization/common/lift_as_function_call.td"
+
+//===----------------------------------------------------------------------===//
+// Pattern rules for lifting ops with bias as functions
+//===----------------------------------------------------------------------===//
+
+def LiftDotGeneralWithBiasSameShape : Pat<
+  (StableHLO_AddOp:$res
+    (StableHLO_DotGeneralOp
+        $lhs, $rhs, $dot_dimension_numbers, $precision_config, $algorithm),
+    $bias),
+  (LiftAsTFXlaCallModule<"composite_dot_general_with_bias_same_shape_fn">
+    (ArgumentList $lhs, $rhs, $bias),
+    (ResultList $res),
+    (NamedAttributeList
+      (NamedAttr<"dot_dimension_numbers"> $dot_dimension_numbers),
+      (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)),
+      (NamedAttr<"algorithm"> (DefaultOrNullAttr $algorithm)))),
+  [(IsNotInLiftedFunc $res), (IsStableHLOConstantOp $bias)], [], (addBenefit 5)>;
+
+def LiftConvWithBiasSameShape : Pat<
+  (StableHLO_AddOp:$res
+    (StableHLO_ConvolutionOp $lhs, $rhs, $window_strides, $padding,
+        $lhs_dilation, $rhs_dilation, $window_reversal, $dimension_numbers,
+        $feature_group_count, $batch_group_count, $precision_config),
+    $bias),
+  (LiftAsTFXlaCallModule<"composite_conv_with_bias_same_shape_fn">
+    (ArgumentList $lhs, $rhs, $bias),
+    (ResultList $res),
+    (NamedAttributeList
+      (NamedAttr<"window_strides"> (DefaultOrNullAttr $window_strides)),
+      (NamedAttr<"padding"> (DefaultOrNullAttr $padding)),
+      (NamedAttr<"lhs_dilation"> (DefaultOrNullAttr $lhs_dilation)),
+      (NamedAttr<"rhs_dilation"> (DefaultOrNullAttr $rhs_dilation)),
+      (NamedAttr<"window_reversal"> (DefaultOrNullAttr $window_reversal)),
+      (NamedAttr<"dimension_numbers"> $dimension_numbers),
+      (NamedAttr<"feature_group_count"> $feature_group_count),
+      (NamedAttr<"batch_group_count"> $batch_group_count),
+      (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)))),
+  [(IsNotInLiftedFunc $res), (IsStableHLOConstantOp $bias)], [], (addBenefit 5)>;
+
+
+def LiftConvWithBias : Pat<
+  (StableHLO_AddOp:$res
+    (StableHLO_ConvolutionOp $lhs, $rhs, $window_strides, $padding,
+        $lhs_dilation, $rhs_dilation, $window_reversal, $dimension_numbers,
+        $feature_group_count, $batch_group_count, $precision_config),
+    (StableHLO_BroadcastInDimOp $bias, $dims)),
+  (LiftAsTFXlaCallModule<"composite_conv_with_bias_fn">
+    (ArgumentList $lhs, $rhs, $bias),
+    (ResultList $res),
+    (NamedAttributeList
+      (NamedAttr<"window_strides"> (DefaultOrNullAttr $window_strides)),
+      (NamedAttr<"padding"> (DefaultOrNullAttr $padding)),
+      (NamedAttr<"lhs_dilation"> (DefaultOrNullAttr $lhs_dilation)),
+      (NamedAttr<"rhs_dilation"> (DefaultOrNullAttr $rhs_dilation)),
+      (NamedAttr<"window_reversal"> (DefaultOrNullAttr $window_reversal)),
+      (NamedAttr<"dimension_numbers"> $dimension_numbers),
+      (NamedAttr<"feature_group_count"> $feature_group_count),
+      (NamedAttr<"batch_group_count"> $batch_group_count),
+      (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)))),
+  [(IsNotInLiftedFunc $res), (IsStableHLOConstantOp $bias)], [], (addBenefit 5)>;
+
+def LiftDotGeneralWithBias : Pat<
+  (StableHLO_AddOp:$res
+    (StableHLO_DotGeneralOp
+        $lhs, $rhs, $dot_dimension_numbers, $precision_config, $algorithm),
+    (StableHLO_BroadcastInDimOp $bias, $dims)),
+  (LiftAsTFXlaCallModule<"composite_dot_general_with_bias_fn">
+    (ArgumentList $lhs, $rhs, $bias),
+    (ResultList $res),
+    (NamedAttributeList
+      (NamedAttr<"dot_dimension_numbers"> $dot_dimension_numbers),
+      (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)),
+      (NamedAttr<"algorithm"> (DefaultOrNullAttr $algorithm)))),
+  [(IsNotInLiftedFunc $res), (IsStableHLOConstantOp $bias)], [], (addBenefit 5)>;
+
+def LiftConvWithBiasDynamic : Pat<
+  (StableHLO_AddOp:$res
+    (StableHLO_ConvolutionOp:$conv_0 $lhs, $rhs, $window_strides, $padding,
+        $lhs_dilation, $rhs_dilation, $window_reversal, $dimension_numbers,
+        $feature_group_count, $batch_group_count, $precision_config),
+    (StableHLO_DynamicBroadcastInDimOp
+      $bias,
+      (Shape_ShapeOfOp $conv_1), $_, $_, $_)),
+  (LiftAsTFXlaCallModule<"composite_conv_with_bias_dynamic_fn">
+    (ArgumentList $lhs, $rhs, $bias),
+    (ResultList $res),
+    (NamedAttributeList
+      (NamedAttr<"window_strides"> (DefaultOrNullAttr $window_strides)),
+      (NamedAttr<"padding"> (DefaultOrNullAttr $padding)),
+      (NamedAttr<"lhs_dilation"> (DefaultOrNullAttr $lhs_dilation)),
+      (NamedAttr<"rhs_dilation"> (DefaultOrNullAttr $rhs_dilation)),
+      (NamedAttr<"window_reversal"> (DefaultOrNullAttr $window_reversal)),
+      (NamedAttr<"dimension_numbers"> $dimension_numbers),
+      (NamedAttr<"feature_group_count"> $feature_group_count),
+      (NamedAttr<"batch_group_count"> $batch_group_count),
+      (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)))),
+  [(IsNotInLiftedFunc $res), (IsStableHLOConstantOp $bias), (AreTheSameValue $conv_0, $conv_1)], [], (addBenefit 10)>;
+
+def LiftDotGeneralWithBiasDynamic : Pat<
+  (StableHLO_AddOp:$res
+    (StableHLO_DotGeneralOp:$dot_general_0 $lhs, $rhs, $dot_dimension_numbers, $precision_config, $algorithm),
+    (StableHLO_DynamicBroadcastInDimOp
+      $bias,
+      (Shape_ShapeOfOp $dot_general_1), $_, $_, $_)),
+  (LiftAsTFXlaCallModule<"composite_dot_general_with_bias_dynamic_fn">
+    (ArgumentList $lhs, $rhs, $bias),
+    (ResultList $res),
+    (NamedAttributeList
+      (NamedAttr<"dot_dimension_numbers"> $dot_dimension_numbers),
+      (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)),
+      (NamedAttr<"algorithm"> (DefaultOrNullAttr $algorithm)))),
+  [(IsNotInLiftedFunc $res), (IsStableHLOConstantOp $bias), (AreTheSameValue $dot_general_0, $dot_general_1)], [], (addBenefit 10)>;
+
+//===----------------------------------------------------------------------===//
+// Pattern rules for lifting ops with activation as functions
+//===----------------------------------------------------------------------===//
+
+def LiftConvWithRelu : Pat<
+  (StableHLO_MaxOp:$res
+    (StableHLO_ConvolutionOp $lhs, $rhs, $window_strides, $padding,
+        $lhs_dilation, $rhs_dilation, $window_reversal, $dimension_numbers,
+        $feature_group_count, $batch_group_count, $precision_config),
+    (StableHLO_ConstantOp $cst)),
+  (LiftAsTFXlaCallModule<"composite_conv_with_relu_fn">
+    (ArgumentList $lhs, $rhs),
+    (ResultList $res),
+    (NamedAttributeList
+      (NamedAttr<"window_strides"> (DefaultOrNullAttr $window_strides)),
+      (NamedAttr<"padding"> (DefaultOrNullAttr $padding)),
+      (NamedAttr<"lhs_dilation"> (DefaultOrNullAttr $lhs_dilation)),
+      (NamedAttr<"rhs_dilation"> (DefaultOrNullAttr $rhs_dilation)),
+      (NamedAttr<"window_reversal"> (DefaultOrNullAttr $window_reversal)),
+      (NamedAttr<"dimension_numbers"> $dimension_numbers),
+      (NamedAttr<"feature_group_count"> $feature_group_count),
+      (NamedAttr<"batch_group_count"> $batch_group_count),
+      (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)))),
+  [(IsNotInLiftedFunc $res), (FloatValueEquals<"0"> $cst)], [], (addBenefit 10)>;
+
+def LiftDotGeneralWithRelu : Pat<
+  (StableHLO_MaxOp:$res
+    (StableHLO_DotGeneralOp
+        $lhs, $rhs, $dot_dimension_numbers, $precision_config, $algorithm),
+    (StableHLO_ConstantOp $cst)),
+  (LiftAsTFXlaCallModule<"composite_dot_general_with_relu_fn">
+    (ArgumentList $lhs, $rhs),
+    (ResultList $res),
+    (NamedAttributeList
+      (NamedAttr<"dot_dimension_numbers"> $dot_dimension_numbers),
+      (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)),
+      (NamedAttr<"algorithm"> (DefaultOrNullAttr $algorithm)))),
+  [(IsNotInLiftedFunc $res),
+   (FloatValueEquals<"0"> $cst)], [], (addBenefit 10)>;
+
+def LiftConvWithReluDynamic : Pat<
+  (StableHLO_MaxOp:$res
+    (StableHLO_ConvolutionOp:$conv_0 $lhs, $rhs, $window_strides, $padding,
+        $lhs_dilation, $rhs_dilation, $window_reversal, $dimension_numbers,
+        $feature_group_count, $batch_group_count, $precision_config),
+    (StableHLO_DynamicBroadcastInDimOp
+      (StableHLO_ConstantOp $cst),
+      (Shape_ShapeOfOp $conv_1), $_, $_, $_)),
+  (LiftAsTFXlaCallModule<"composite_conv_with_relu_dynamic_fn">
+    (ArgumentList $lhs, $rhs),
+    (ResultList $res),
+    (NamedAttributeList
+      (NamedAttr<"window_strides"> (DefaultOrNullAttr $window_strides)),
+      (NamedAttr<"padding"> (DefaultOrNullAttr $padding)),
+      (NamedAttr<"lhs_dilation"> (DefaultOrNullAttr $lhs_dilation)),
+      (NamedAttr<"rhs_dilation"> (DefaultOrNullAttr $rhs_dilation)),
+      (NamedAttr<"window_reversal"> (DefaultOrNullAttr $window_reversal)),
+      (NamedAttr<"dimension_numbers"> $dimension_numbers),
+      (NamedAttr<"feature_group_count"> $feature_group_count),
+      (NamedAttr<"batch_group_count"> $batch_group_count),
+      (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)))),
+  [(IsNotInLiftedFunc $res),
+   (FloatValueEquals<"0"> $cst), (AreTheSameValue $conv_0, $conv_1)], [], (addBenefit 15)>;
+
+def LiftDotGeneralWithReluDynamic : Pat<
+  (StableHLO_MaxOp:$res
+    (StableHLO_DotGeneralOp:$dot_general_0 $lhs, $rhs, $dot_dimension_numbers, $precision_config, $algorithm),
+    (StableHLO_DynamicBroadcastInDimOp
+      (StableHLO_ConstantOp $cst),
+      (Shape_ShapeOfOp $dot_general_1), $_, $_, $_)),
+  (LiftAsTFXlaCallModule<"composite_dot_general_with_relu_dynamic_fn">
+    (ArgumentList $lhs, $rhs),
+    (ResultList $res),
+    (NamedAttributeList
+      (NamedAttr<"dot_dimension_numbers"> $dot_dimension_numbers),
+      (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)),
+      (NamedAttr<"algorithm"> (DefaultOrNullAttr $algorithm)))),
+  [(IsNotInLiftedFunc $res),
+   (FloatValueEquals<"0"> $cst), (AreTheSameValue $dot_general_0, $dot_general_1)], [], (addBenefit 15)>;
+
+def LiftConvWithRelu6 : Pat<
+  (StableHLO_ClampOp:$res
+    (StableHLO_ConstantOp $cst_0),
+    (StableHLO_ConvolutionOp $lhs, $rhs, $window_strides, $padding,
+        $lhs_dilation, $rhs_dilation, $window_reversal, $dimension_numbers,
+        $feature_group_count, $batch_group_count, $precision_config),
+    (StableHLO_ConstantOp $cst_1)),
+  (LiftAsTFXlaCallModule<"composite_conv_with_relu6_fn">
+    (ArgumentList $lhs, $rhs),
+    (ResultList $res),
+    (NamedAttributeList
+      (NamedAttr<"window_strides"> (DefaultOrNullAttr $window_strides)),
+      (NamedAttr<"padding"> (DefaultOrNullAttr $padding)),
+      (NamedAttr<"lhs_dilation"> (DefaultOrNullAttr $lhs_dilation)),
+      (NamedAttr<"rhs_dilation"> (DefaultOrNullAttr $rhs_dilation)),
+      (NamedAttr<"window_reversal"> (DefaultOrNullAttr $window_reversal)),
+      (NamedAttr<"dimension_numbers"> $dimension_numbers),
+      (NamedAttr<"feature_group_count"> $feature_group_count),
+      (NamedAttr<"batch_group_count"> $batch_group_count),
+      (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)))),
+  [(IsNotInLiftedFunc $res), (FloatValueEquals<"0"> $cst_0), (FloatValueEquals<"6"> $cst_1)], [], (addBenefit 10)>;
+
+def LiftDotGeneralWithRelu6 : Pat<
+  (StableHLO_ClampOp:$res
+    (StableHLO_ConstantOp $cst_0),
+    (StableHLO_DotGeneralOp
+        $lhs, $rhs, $dot_dimension_numbers, $precision_config, $algorithm),
+    (StableHLO_ConstantOp $cst_1)),
+  (LiftAsTFXlaCallModule<"composite_dot_general_with_relu6_fn">
+    (ArgumentList $lhs, $rhs),
+    (ResultList $res),
+    (NamedAttributeList
+      (NamedAttr<"dot_dimension_numbers"> $dot_dimension_numbers),
+      (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)),
+      (NamedAttr<"algorithm"> (DefaultOrNullAttr $algorithm)))),
+  [(IsNotInLiftedFunc $res), (FloatValueEquals<"0"> $cst_0), (FloatValueEquals<"6"> $cst_1)], [], (addBenefit 10)>;
+
+//===----------------------------------------------------------------------===//
+// Pattern rules for lifting ops with bias and activation as functions
+//===----------------------------------------------------------------------===//
+
+def LiftDotGeneralWithBiasSameShapeAndRelu : Pat<
+  (StableHLO_MaxOp:$res
+    (StableHLO_AddOp
+      (StableHLO_DotGeneralOp
+          $lhs, $rhs, $dot_dimension_numbers, $precision_config, $algorithm),
+      $bias),
+    (StableHLO_ConstantOp $cst)),
+  (LiftAsTFXlaCallModule<"composite_dot_general_with_bias_same_shape_and_relu_fn">
+    (ArgumentList $lhs, $rhs, $bias),
+    (ResultList $res),
+    (NamedAttributeList
+      (NamedAttr<"dot_dimension_numbers"> $dot_dimension_numbers),
+      (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)),
+      (NamedAttr<"algorithm"> (DefaultOrNullAttr $algorithm)))),
+  [(IsNotInLiftedFunc $res),
+   (FloatValueEquals<"0"> $cst), (IsStableHLOConstantOp $bias)], [], (addBenefit 10)>;
+
+def LiftConvWithBiasSameShapeAndRelu : Pat<
+  (StableHLO_MaxOp:$res
+    (StableHLO_AddOp
+      (StableHLO_ConvolutionOp $lhs, $rhs, $window_strides, $padding,
+          $lhs_dilation, $rhs_dilation, $window_reversal, $dimension_numbers,
+          $feature_group_count, $batch_group_count, $precision_config),
+      $bias),
+    (StableHLO_ConstantOp $cst)),
+  (LiftAsTFXlaCallModule<"composite_conv_with_bias_same_shape_and_relu_fn">
+    (ArgumentList $lhs, $rhs, $bias),
+    (ResultList $res),
+    (NamedAttributeList
+      (NamedAttr<"window_strides"> (DefaultOrNullAttr $window_strides)),
+      (NamedAttr<"padding"> (DefaultOrNullAttr $padding)),
+      (NamedAttr<"lhs_dilation"> (DefaultOrNullAttr $lhs_dilation)),
+      (NamedAttr<"rhs_dilation"> (DefaultOrNullAttr $rhs_dilation)),
+      (NamedAttr<"window_reversal"> (DefaultOrNullAttr $window_reversal)),
+      (NamedAttr<"dimension_numbers"> $dimension_numbers),
+      (NamedAttr<"feature_group_count"> $feature_group_count),
+      (NamedAttr<"batch_group_count"> $batch_group_count),
+      (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)))),
+  [(IsNotInLiftedFunc $res),
+   (FloatValueEquals<"0"> $cst), (IsStableHLOConstantOp $bias)], [], (addBenefit 10)>;
+
+
+def LiftConvWithBiasAndRelu : Pat<
+  (StableHLO_MaxOp:$res
+    (StableHLO_AddOp
+      (StableHLO_ConvolutionOp $lhs, $rhs, $window_strides, $padding,
+          $lhs_dilation, $rhs_dilation, $window_reversal, $dimension_numbers,
+          $feature_group_count, $batch_group_count, $precision_config),
+      (StableHLO_BroadcastInDimOp $bias, $dims)),
+    (StableHLO_ConstantOp $cst)),
+  (LiftAsTFXlaCallModule<"composite_conv_with_bias_and_relu_fn">
+    (ArgumentList $lhs, $rhs, $bias),
+    (ResultList $res),
+    (NamedAttributeList
+      (NamedAttr<"window_strides"> (DefaultOrNullAttr $window_strides)),
+      (NamedAttr<"padding"> (DefaultOrNullAttr $padding)),
+      (NamedAttr<"lhs_dilation"> (DefaultOrNullAttr $lhs_dilation)),
+      (NamedAttr<"rhs_dilation"> (DefaultOrNullAttr $rhs_dilation)),
+      (NamedAttr<"window_reversal"> (DefaultOrNullAttr $window_reversal)),
+      (NamedAttr<"dimension_numbers"> $dimension_numbers),
+      (NamedAttr<"feature_group_count"> $feature_group_count),
+      (NamedAttr<"batch_group_count"> $batch_group_count),
+      (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)))),
+  [(IsNotInLiftedFunc $res),
+   (FloatValueEquals<"0"> $cst), (IsStableHLOConstantOp $bias)], [], (addBenefit 10)>;
+
+def LiftDotGeneralWithBiasAndRelu : Pat<
+  (StableHLO_MaxOp:$res
+    (StableHLO_AddOp
+      (StableHLO_DotGeneralOp
+          $lhs, $rhs, $dot_dimension_numbers, $precision_config, $algorithm),
+      (StableHLO_BroadcastInDimOp $bias, $dims)),
+    (StableHLO_ConstantOp $cst)),
+  (LiftAsTFXlaCallModule<"composite_dot_general_with_bias_and_relu_fn">
+    (ArgumentList $lhs, $rhs, $bias),
+    (ResultList $res),
+    (NamedAttributeList
+      (NamedAttr<"dot_dimension_numbers"> $dot_dimension_numbers),
+      (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)),
+      (NamedAttr<"algorithm"> (DefaultOrNullAttr $algorithm)))),
+  [(IsNotInLiftedFunc $res),
+   (FloatValueEquals<"0"> $cst), (IsStableHLOConstantOp $bias)], [], (addBenefit 10)>;
+
+def LiftConvWithBiasAndReluDynamic : Pat<
+  (StableHLO_MaxOp:$res
+    (StableHLO_AddOp:$add_0
+      (StableHLO_ConvolutionOp:$conv_0 $lhs, $rhs, $window_strides, $padding,
+          $lhs_dilation, $rhs_dilation, $window_reversal, $dimension_numbers,
+          $feature_group_count, $batch_group_count, $precision_config),
+      (StableHLO_DynamicBroadcastInDimOp
+        $bias,
+        (Shape_ShapeOfOp $conv_1), $_, $_, $_)),
+    (StableHLO_DynamicBroadcastInDimOp
+      (StableHLO_ConstantOp $cst),
+      (Shape_ShapeOfOp $add_1), $_, $_, $_)),
+  (LiftAsTFXlaCallModule<"composite_conv_with_bias_and_relu_dynamic_fn">
+    (ArgumentList $lhs, $rhs, $bias),
+    (ResultList $res),
+    (NamedAttributeList
+      (NamedAttr<"window_strides"> (DefaultOrNullAttr $window_strides)),
+      (NamedAttr<"padding"> (DefaultOrNullAttr $padding)),
+      (NamedAttr<"lhs_dilation"> (DefaultOrNullAttr $lhs_dilation)),
+      (NamedAttr<"rhs_dilation"> (DefaultOrNullAttr $rhs_dilation)),
+      (NamedAttr<"window_reversal"> (DefaultOrNullAttr $window_reversal)),
+      (NamedAttr<"dimension_numbers"> $dimension_numbers),
+      (NamedAttr<"feature_group_count"> $feature_group_count),
+      (NamedAttr<"batch_group_count"> $batch_group_count),
+      (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)))),
+  [(IsNotInLiftedFunc $res),
+   (FloatValueEquals<"0"> $cst), (IsStableHLOConstantOp $bias), (AreTheSameValue $conv_0, $conv_1), (AreTheSameValue $add_0, $add_1)], [], (addBenefit 15)>;
+
+def LiftDotGeneralWithBiasAndReluDynamic : Pat<
+  (StableHLO_MaxOp:$res
+    (StableHLO_AddOp:$add_0
+      (StableHLO_DotGeneralOp:$dot_general_0 $lhs, $rhs, $dot_dimension_numbers, $precision_config, $algorithm),
+      (StableHLO_DynamicBroadcastInDimOp
+        $bias,
+        (Shape_ShapeOfOp $dot_general_1), $_, $_, $_)),
+    (StableHLO_DynamicBroadcastInDimOp
+      (StableHLO_ConstantOp $cst),
+      (Shape_ShapeOfOp $add_1), $_, $_, $_)),
+  (LiftAsTFXlaCallModule<"composite_dot_general_with_bias_and_relu_dynamic_fn">
+    (ArgumentList $lhs, $rhs, $bias),
+    (ResultList $res),
+    (NamedAttributeList
+      (NamedAttr<"dot_dimension_numbers"> $dot_dimension_numbers),
+      (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)),
+      (NamedAttr<"algorithm"> (DefaultOrNullAttr $algorithm)))),
+  [(IsNotInLiftedFunc $res),
+   (FloatValueEquals<"0"> $cst), (IsStableHLOConstantOp $bias), (AreTheSameValue $dot_general_0, $dot_general_1), (AreTheSameValue $add_0, $add_1)], [], (addBenefit 15)>;
+
+def LiftDotGeneralWithBiasSameShapeAndRelu6 : Pat<
+  (StableHLO_ClampOp:$res
+    (StableHLO_ConstantOp $cst_0),
+    (StableHLO_AddOp
+      (StableHLO_DotGeneralOp
+          $lhs, $rhs, $dot_dimension_numbers, $precision_config, $algorithm),
+      $bias),
+    (StableHLO_ConstantOp $cst_1)),
+  (LiftAsTFXlaCallModule<"composite_dot_general_with_bias_same_shape_and_relu6_fn">
+    (ArgumentList $lhs, $rhs, $bias),
+    (ResultList $res),
+    (NamedAttributeList
+      (NamedAttr<"dot_dimension_numbers"> $dot_dimension_numbers),
+      (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)),
+      (NamedAttr<"algorithm"> (DefaultOrNullAttr $algorithm)))),
+  [(IsNotInLiftedFunc $res), (IsStableHLOConstantOp $bias), (FloatValueEquals<"0"> $cst_0), (FloatValueEquals<"6"> $cst_1)], [], (addBenefit 10)>;
+
+def LiftConvWithBiasAndRelu6 : Pat<
+  (StableHLO_ClampOp:$res
+    (StableHLO_ConstantOp $cst_0),
+    (StableHLO_AddOp
+      (StableHLO_ConvolutionOp $lhs, $rhs, $window_strides, $padding,
+          $lhs_dilation, $rhs_dilation, $window_reversal, $dimension_numbers,
+          $feature_group_count, $batch_group_count, $precision_config),
+      (StableHLO_BroadcastInDimOp $bias, $dims)),
+    (StableHLO_ConstantOp $cst_1)),
+  (LiftAsTFXlaCallModule<"composite_conv_with_bias_and_relu6_fn">
+    (ArgumentList $lhs, $rhs, $bias),
+    (ResultList $res),
+    (NamedAttributeList
+      (NamedAttr<"window_strides"> (DefaultOrNullAttr $window_strides)),
+      (NamedAttr<"padding"> (DefaultOrNullAttr $padding)),
+      (NamedAttr<"lhs_dilation"> (DefaultOrNullAttr $lhs_dilation)),
+      (NamedAttr<"rhs_dilation"> (DefaultOrNullAttr $rhs_dilation)),
+      (NamedAttr<"window_reversal"> (DefaultOrNullAttr $window_reversal)),
+      (NamedAttr<"dimension_numbers"> $dimension_numbers),
+      (NamedAttr<"feature_group_count"> $feature_group_count),
+      (NamedAttr<"batch_group_count"> $batch_group_count),
+      (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)))),
+  [(IsNotInLiftedFunc $res), (IsStableHLOConstantOp $bias), (FloatValueEquals<"0"> $cst_0), (FloatValueEquals<"6"> $cst_1)], [], (addBenefit 10)>;
+
+def LiftDotGeneralWithBiasAndRelu6 : Pat<
+  (StableHLO_ClampOp:$res
+    (StableHLO_ConstantOp $cst_0),
+    (StableHLO_AddOp
+      (StableHLO_DotGeneralOp
+          $lhs, $rhs, $dot_dimension_numbers, $precision_config, $algorithm),
+      (StableHLO_BroadcastInDimOp $bias, $dims)),
+    (StableHLO_ConstantOp $cst_1)),
+  (LiftAsTFXlaCallModule<"composite_dot_general_with_bias_and_relu6_fn">
+    (ArgumentList $lhs, $rhs, $bias),
+    (ResultList $res),
+    (NamedAttributeList
+      (NamedAttr<"dot_dimension_numbers"> $dot_dimension_numbers),
+      (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)),
+      (NamedAttr<"algorithm"> (DefaultOrNullAttr $algorithm)))),
+  [(IsNotInLiftedFunc $res), (IsStableHLOConstantOp $bias), (FloatValueEquals<"0"> $cst_0), (FloatValueEquals<"6"> $cst_1)], [], (addBenefit 10)>;
+
+def LiftConvWithBiasAndRelu6Dynamic : Pat<
+  (StableHLO_ClampOp:$res
+    (StableHLO_ConstantOp $cst_0),
+    (StableHLO_AddOp
+      (StableHLO_ConvolutionOp:$conv_0 $lhs, $rhs, $window_strides, $padding,
+          $lhs_dilation, $rhs_dilation, $window_reversal, $dimension_numbers,
+          $feature_group_count, $batch_group_count, $precision_config),
+      (StableHLO_DynamicBroadcastInDimOp
+        $bias,
+        (Shape_ShapeOfOp $conv_1), $_, $_, $_)),
+    (StableHLO_ConstantOp $cst_1)),
+  (LiftAsTFXlaCallModule<"composite_conv_with_bias_and_relu6_dynamic_fn">
+    (ArgumentList $lhs, $rhs, $bias),
+    (ResultList $res),
+    (NamedAttributeList
+      (NamedAttr<"window_strides"> (DefaultOrNullAttr $window_strides)),
+      (NamedAttr<"padding"> (DefaultOrNullAttr $padding)),
+      (NamedAttr<"lhs_dilation"> (DefaultOrNullAttr $lhs_dilation)),
+      (NamedAttr<"rhs_dilation"> (DefaultOrNullAttr $rhs_dilation)),
+      (NamedAttr<"window_reversal"> (DefaultOrNullAttr $window_reversal)),
+      (NamedAttr<"dimension_numbers"> $dimension_numbers),
+      (NamedAttr<"feature_group_count"> $feature_group_count),
+      (NamedAttr<"batch_group_count"> $batch_group_count),
+      (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)))),
+  [(IsNotInLiftedFunc $res), (IsStableHLOConstantOp $bias), (FloatValueEquals<"0"> $cst_0), (FloatValueEquals<"6"> $cst_1), (AreTheSameValue $conv_0, $conv_1)], [], (addBenefit 15)>;
+
+def LiftDotGeneralWithBiasAndRelu6Dynamic : Pat<
+  (StableHLO_ClampOp:$res
+    (StableHLO_ConstantOp $cst_0),
+    (StableHLO_AddOp
+      (StableHLO_DotGeneralOp:$dot_general_0
+          $lhs, $rhs, $dot_dimension_numbers, $precision_config, $algorithm),
+      (StableHLO_DynamicBroadcastInDimOp
+        $bias,
+        (Shape_ShapeOfOp $dot_general_1), $_, $_, $_)),
+    (StableHLO_ConstantOp $cst_1)),
+  (LiftAsTFXlaCallModule<"composite_dot_general_with_bias_and_relu6_dynamic_fn">
+    (ArgumentList $lhs, $rhs, $bias),
+    (ResultList $res),
+    (NamedAttributeList
+      (NamedAttr<"dot_dimension_numbers"> $dot_dimension_numbers),
+      (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)),
+      (NamedAttr<"algorithm"> (DefaultOrNullAttr $algorithm)))),
+  [(IsNotInLiftedFunc $res), (IsStableHLOConstantOp $bias), (FloatValueEquals<"0"> $cst_0), (FloatValueEquals<"6"> $cst_1), (AreTheSameValue $dot_general_0, $dot_general_1)], [], (addBenefit 15)>;

tensorflow/compiler/mlir/quantization/stablehlo/passes/lift_quantizable_spots_as_functions_simple.td

@@ -0,0 +1,78 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+include "mlir/IR/OpBase.td"
+include "mlir/Dialect/Func/IR/FuncOps.td"
+include "mlir/Dialect/Arith/IR/ArithOps.td"
+include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"
+include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"
+include "stablehlo/dialect/StablehloOps.td"
+include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.td"
+include "tensorflow/compiler/mlir/quantization/common/lift_as_function_call.td"
+
+//===----------------------------------------------------------------------===//
+// Pattern rules for lifting ops as functions
+//===----------------------------------------------------------------------===//
+
+def LiftConv : Pat<
+  (StableHLO_ConvolutionOp:$res $lhs, $rhs, $window_strides, $padding,
+      $lhs_dilation, $rhs_dilation, $window_reversal, $dimension_numbers,
+      $feature_group_count, $batch_group_count, $precision_config),
+  (LiftAsTFXlaCallModule<"composite_conv_fn">
+    (ArgumentList $lhs, $rhs),
+    (ResultList $res),
+    (NamedAttributeList
+      (NamedAttr<"window_strides"> (DefaultOrNullAttr $window_strides)),
+      (NamedAttr<"padding"> (DefaultOrNullAttr $padding)),
+      (NamedAttr<"lhs_dilation"> (DefaultOrNullAttr $lhs_dilation)),
+      (NamedAttr<"rhs_dilation"> (DefaultOrNullAttr $rhs_dilation)),
+      (NamedAttr<"window_reversal"> (DefaultOrNullAttr $window_reversal)),
+      (NamedAttr<"dimension_numbers"> $dimension_numbers),
+      (NamedAttr<"feature_group_count"> $feature_group_count),
+      (NamedAttr<"batch_group_count"> $batch_group_count),
+      (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)))),
+  [(IsNotInLiftedFunc $res)], [], (addBenefit 1)>;
+
+def LiftDotGeneral : Pat<
+  (StableHLO_DotGeneralOp:$res
+      $lhs, $rhs, $dot_dimension_numbers, $precision_config, $algorithm),
+  (LiftAsTFXlaCallModule<"composite_dot_general_fn">
+    (ArgumentList $lhs, $rhs),
+    (ResultList $res),
+    (NamedAttributeList
+      (NamedAttr<"dot_dimension_numbers"> $dot_dimension_numbers),
+      (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)),
+      (NamedAttr<"algorithm"> (DefaultOrNullAttr $algorithm)))),
+  [(IsNotInLiftedFunc $res)], [], (addBenefit 1)>;
+
+def LiftGather : Pat<
+  (StableHLO_GatherOp:$res
+      $operand, $start_indices, $dimension_numbers, $slice_sizes, $indices_are_sorted),
+  (LiftAsTFXlaCallModule<"composite_gather_fn">
+    (ArgumentList $operand, $start_indices),
+    (ResultList $res),
+    (NamedAttributeList
+      (NamedAttr<"dimension_numbers"> $dimension_numbers),
+      (NamedAttr<"slice_sizes"> $slice_sizes),
+      (NamedAttr<"indices_are_sorted"> (DefaultOrNullAttr $indices_are_sorted)))),
+  [(IsNotInLiftedFunc $res), (IsStableHLOConstantOp $operand)], [], (addBenefit 1)>;
+
+def LiftAdd : Pat<
+  (StableHLO_AddOp:$res
+      $lhs, $rhs),
+  (LiftAsTFXlaCallModule<"composite_add_fn">
+    (ArgumentList $lhs, $rhs),
+    (ResultList $res)),
+  [(IsNotInLiftedFunc $res), (IsNotInStableHloOpRegion $res)], [], (addBenefit 1)>;

tensorflow/compiler/mlir/quantization/stablehlo/passes/merge_fusion_with_dequantize.cc

@@ -0,0 +1,151 @@
+/* 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 <utility>
+
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "stablehlo/dialect/ChloOps.h"  // from @stablehlo
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo
+
+namespace mlir::quant::stablehlo {
+
+#define GEN_PASS_DEF_MERGEFUSIONWITHDEQUANTIZEPASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"
+
+namespace {
+
+class MergeFusionWithDequantizePass
+    : public impl::MergeFusionWithDequantizePassBase<
+          MergeFusionWithDequantizePass> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MergeFusionWithDequantizePass)
+
+  explicit MergeFusionWithDequantizePass() = default;
+
+ private:
+  void runOnOperation() override;
+};
+
+class MergeFusionWithUniformDequantizePattern
+    : public OpRewritePattern<func::CallOp> {
+ public:
+  explicit MergeFusionWithUniformDequantizePattern(MLIRContext* context)
+      : OpRewritePattern<func::CallOp>(context) {}
+  LogicalResult matchAndRewrite(func::CallOp call_op,
+                                PatternRewriter& rewriter) const override {
+    if (call_op.getNumResults() != 1) return failure();
+    auto users = call_op->getUsers();
+    for (auto user : users) {
+      if (!llvm::isa<mlir::stablehlo::UniformDequantizeOp>(user)) {
+        return failure();
+      }
+    }
+    auto func_name = call_op.getCallee();
+    if (!func_name.starts_with("quantized_")) return failure();
+    if (call_op->getNumResults() != 1) return failure();
+    if (!mlir::isa<quant::UniformQuantizedType>(
+            getElementTypeOrSelf(call_op->getResult(0).getType())))
+      return failure();
+
+    // Fetch the callee function.
+    SymbolTable symbol_table(call_op->getParentOfType<ModuleOp>());
+    auto func_op =
+        dyn_cast_or_null<func::FuncOp>(symbol_table.lookup(func_name));
+    if (!func_op) return failure();
+    // The quantized fusion should have requantize and return ops at the end.
+    auto return_op = dyn_cast_or_null<func::ReturnOp>(
+        func_op.getRegion().getBlocks().front().getTerminator());
+    if (!return_op) return failure();
+    auto req_op = llvm::dyn_cast_or_null<mlir::stablehlo::UniformQuantizeOp>(
+        return_op.getOperands()[0].getDefiningOp());
+    if (!req_op) return failure();
+
+    // Create a new func.call op with f32 output.
+    auto new_call_op = call_op.clone();
+    new_call_op->getResult(0).setType(
+        mlir::cast<ShapedType>(call_op.getResult(0).getType())
+            .clone(rewriter.getF32Type()));
+    rewriter.setInsertionPoint(call_op);
+    rewriter.insert(new_call_op);
+
+    // Remove the dequantize ops and replace uses by the new func.call op.
+    SmallVector<Operation*> users_to_erase;
+    for (auto user : users) {
+      llvm::dyn_cast<mlir::stablehlo::UniformDequantizeOp>(user)
+          .replaceAllUsesWith(new_call_op.getResult(0));
+      users_to_erase.push_back(user);
+    }
+    for (auto user : users_to_erase) rewriter.eraseOp(user);
+    rewriter.eraseOp(call_op);
+    if (failed(func_op.eraseResult(0))) {
+      return failure();
+    }
+    if (failed(func_op.insertResult(0, new_call_op.getResult(0).getType(),
+                                    /*resultAttrs=*/nullptr))) {
+      return failure();
+    }
+
+    // Modify the quantized fused function to do dequantize+relu(6).
+    rewriter.setInsertionPoint(req_op);
+    Value new_result = mlir::stablehlo::UniformDequantizeOp::create(
+        rewriter, req_op.getLoc(), func_op.getResultTypes()[0],
+        req_op.getOperand());
+    if (func_name.contains("_relu6_")) {
+      auto min = mlir::stablehlo::ConstantOp::create(
+          rewriter, req_op.getLoc(), rewriter.getF32FloatAttr(0));
+      auto max = mlir::stablehlo::ConstantOp::create(
+          rewriter, req_op.getLoc(), rewriter.getF32FloatAttr(6));
+      new_result = mlir::stablehlo::ClampOp::create(rewriter, req_op.getLoc(),
+                                                    min, new_result, max);
+    } else if (func_name.contains("_relu_")) {
+      auto min = mlir::stablehlo::ConstantOp::create(
+          rewriter, req_op.getLoc(), rewriter.getF32FloatAttr(0));
+      new_result = mlir::chlo::BroadcastMaxOp::create(rewriter, req_op.getLoc(),
+                                                      min, new_result, nullptr);
+    }
+    return_op->setOperand(0, new_result);
+    rewriter.eraseOp(req_op);
+
+    return success();
+  }
+};
+
+void MergeFusionWithDequantizePass::runOnOperation() {
+  ModuleOp module_op = getOperation();
+  MLIRContext* ctx = module_op.getContext();
+  RewritePatternSet patterns(ctx);
+  patterns.add<MergeFusionWithUniformDequantizePattern>(ctx);
+  if (failed(applyPatternsGreedily(module_op, std::move(patterns)))) {
+    signalPassFailure();
+  }
+}
+
+}  // namespace
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/passes/nchw_convolution_to_nhwc.cc

@@ -0,0 +1,194 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include <cstdint>
+#include <utility>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo
+#include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.h"
+#include "tensorflow/compiler/mlir/quantization/common/uniform_quantized_types.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/permutation.h"
+
+namespace mlir::quant::stablehlo {
+
+#define GEN_PASS_DEF_NCHWCONVOLUTIONTONHWCPASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"
+
+namespace {
+
+using ::mlir::stablehlo::ConvDimensionNumbersAttr;
+
+class NchwConvolutionToNhwcPass
+    : public impl::NchwConvolutionToNhwcPassBase<NchwConvolutionToNhwcPass> {
+ private:
+  void runOnOperation() override;
+};
+
+// Rewrites NCHW convolution to NHWC.
+// * Src dimension numbers: [b, f, 0, 1]x[o, i, 0, 1]->[b, f, 0, 1]
+// * Dst dimension numbers: [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]
+class RewriteNchwConvolutionToNhwc
+    : public OpRewritePattern<mlir::stablehlo::ConvolutionOp> {
+ public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mlir::stablehlo::ConvolutionOp op,
+                                PatternRewriter& rewriter) const override {
+    // Handles 2D convolutions only.
+    if (!HasRankOf(op.getOperand(0), /*rank=*/4) ||
+        !HasRankOf(op.getOperand(1), /*rank=*/4)) {
+      return failure();
+    }
+
+    if (!quant::IsOpNotQuantized(op)) return failure();
+
+    const ConvDimensionNumbersAttr dimension_nums = op.getDimensionNumbers();
+    const bool dimension_nums_matched =
+        MatchInputDimensionNumbers(dimension_nums) &&
+        MatchKernelDimensionNumbers(dimension_nums) &&
+        MatchOutputDimensionNumbers(dimension_nums);
+    if (!dimension_nums_matched) {
+      return failure();
+    }
+
+    // Transpose the input tensor: [b, f, 0, 1] => [b, 0, 1, f]
+    Value input = op->getOperand(0);
+    const TensorType new_input_tensor_type = GetTransposedTensorType(
+        mlir::cast<TensorType>(input.getType()), kNchwToNhwcPermutation);
+
+    auto input_transpose_op = mlir::stablehlo::TransposeOp::create(
+        rewriter, op.getLoc(), /*resultType0=*/new_input_tensor_type,
+        /*operand=*/input,
+        rewriter.getDenseI64ArrayAttr(kNchwToNhwcPermutation));
+
+    // Transpose the filter tensor: [o, i, 0, 1] => [0, 1, i, o]
+    Value filter = op->getOperand(1);
+    const TensorType new_filter_tensor_type = GetTransposedTensorType(
+        mlir::cast<TensorType>(filter.getType()), kOihwToHwioPermutation);
+
+    auto filter_transpose_op = mlir::stablehlo::TransposeOp::create(
+        rewriter, op.getLoc(), /*resultType0=*/new_filter_tensor_type,
+        /*operand=*/filter,
+        rewriter.getDenseI64ArrayAttr(kOihwToHwioPermutation));
+
+    // [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]
+    const auto new_dimension_nums = rewriter.getAttr<ConvDimensionNumbersAttr>(
+        /*inputBatchDimension=*/0, /*inputFeatureDimension=*/3,
+        /*inputSpatialDimensions=*/SmallVector<int64_t>{1, 2},
+        /*kernelInputFeatureDimension=*/2, /*kernelOutputFeatureDimension=*/3,
+        /*kernelSpatialDimensions=*/SmallVector<int64_t>{0, 1},
+        /*outputBatchDimension=*/0, /*outputFeatureDimension=*/3,
+        /*outputSpatialDimensions=*/SmallVector<int64_t>{1, 2});
+
+    // Determine the shape of the output tensor: [b, f, 0, 1] => [b, 0, 1, f]
+    auto output_tensor_type =
+        mlir::cast<TensorType>(op->getResult(0).getType());
+    const TensorType new_conv_output_tensor_type =
+        GetTransposedTensorType(output_tensor_type, kNchwToNhwcPermutation);
+
+    // window_strides, padding, lhs_dilation, rhs_dilation, window_reversal are
+    // reused without modification because the ordering of spatial dimensions
+    // is not modified (i.e. before: [b, f, 0, 1], after: [b, 0, 1, f] => the
+    // spatial dimension is still ordered as {0, 1}).
+    auto new_convolution_op = mlir::stablehlo::ConvolutionOp::create(
+        rewriter, op.getLoc(), /*resultType0=*/new_conv_output_tensor_type,
+        /*lhs=*/input_transpose_op,
+        /*rhs=*/filter_transpose_op,
+        /*window_strides=*/op.getWindowStridesAttr(),
+        /*padding=*/op.getPaddingAttr(),
+        /*lhs_dilation=*/op.getLhsDilationAttr(),
+        /*rhs_dilation=*/op.getRhsDilationAttr(),
+        /*window_reversal=*/op.getWindowReversalAttr(),
+        /*dimension_numbers=*/new_dimension_nums,
+        /*feature_group_count=*/op.getFeatureGroupCountAttr(),
+        /*batch_group_count=*/op.getBatchGroupCountAttr(),
+        /*precision_config=*/op.getPrecisionConfigAttr());
+
+    // Transpose the output of the `ConvolutionOp` back to the original op's
+    // output shape so that users' shapes match.
+    // [b, 0, 1, f] => [b, f, 0, 1]
+    auto output_transpose_op = mlir::stablehlo::TransposeOp::create(
+        rewriter, new_convolution_op.getLoc(),
+        /*resultType0=*/output_tensor_type,
+        /*operand=*/new_convolution_op,
+        rewriter.getDenseI64ArrayAttr(kNhwcToNchwPermutation));
+
+    rewriter.replaceAllUsesWith(op, output_transpose_op);
+    return success();
+  }
+
+ private:
+  // Matches input dimensions corresponding to: [b, f, 0, 1].
+  bool MatchInputDimensionNumbers(
+      const ConvDimensionNumbersAttr dimension_numbers) const {
+    return dimension_numbers.getInputBatchDimension() == 0 &&
+           dimension_numbers.getInputFeatureDimension() == 1 &&
+           dimension_numbers.getInputSpatialDimensions() ==
+               ArrayRef<int64_t>{2, 3};
+  }
+
+  // Matches kernel dimensions corresponding to: [o, i, 0, 1].
+  bool MatchKernelDimensionNumbers(
+      const ConvDimensionNumbersAttr dimension_numbers) const {
+    return dimension_numbers.getKernelInputFeatureDimension() == 1 &&
+           dimension_numbers.getKernelOutputFeatureDimension() == 0 &&
+           dimension_numbers.getKernelSpatialDimensions() ==
+               ArrayRef<int64_t>{2, 3};
+  }
+
+  // Matches output dimensions corresponding to: [b, f, 0, 1].
+  bool MatchOutputDimensionNumbers(
+      const ConvDimensionNumbersAttr dimension_numbers) const {
+    return dimension_numbers.getOutputBatchDimension() == 0 &&
+           dimension_numbers.getOutputFeatureDimension() == 1 &&
+           dimension_numbers.getOutputSpatialDimensions() ==
+               ArrayRef<int64_t>{2, 3};
+  }
+
+  // Returns a new tensor type with the shape transposed according to the
+  // permutation. The rank of `type` and the size of `permutation` must be
+  // equal.
+  TensorType GetTransposedTensorType(
+      const TensorType type, const ArrayRef<int64_t> permutation) const {
+    const SmallVector<int64_t> after_shape =
+        quant::Permute<int64_t>(type.getShape(), permutation);
+    return type.cloneWith(after_shape, type.getElementType());
+  }
+};
+
+}  // namespace
+
+void NchwConvolutionToNhwcPass::runOnOperation() {
+  func::FuncOp func_op = getOperation();
+  MLIRContext& ctx = getContext();
+
+  RewritePatternSet patterns(&ctx);
+  patterns.add<RewriteNchwConvolutionToNhwc>(&ctx);
+
+  if (failed(applyPatternsGreedily(func_op, std::move(patterns)))) {
+    func_op.emitError() << "Failed to run NchwConvolutionToNhwcPass.";
+    signalPassFailure();
+  }
+}
+
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/passes/optimize_graph.cc

@@ -0,0 +1,55 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <utility>
+
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h"  // IWYU pragma: keep
+
+namespace mlir::quant::stablehlo {
+
+#define GEN_PASS_DEF_OPTIMIZEGRAPHPASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"
+
+namespace {
+
+class OptimizeGraphPass
+    : public impl::OptimizeGraphPassBase<OptimizeGraphPass> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(OptimizeGraphPass)
+
+  explicit OptimizeGraphPass() = default;
+
+ private:
+  void runOnOperation() override;
+};
+
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/optimize_graph.inc"
+
+void OptimizeGraphPass::runOnOperation() {
+  RewritePatternSet patterns(&getContext());
+  populateWithGenerated(patterns);
+  auto func = getOperation();
+  if (failed(applyPatternsGreedily(func, std::move(patterns)))) {
+    signalPassFailure();
+  }
+}
+}  // namespace
+
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/passes/optimize_graph.td

@@ -0,0 +1,24 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+include "stablehlo/dialect/StablehloOps.td"
+include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.td"
+
+// Merge requantization followed by dequantization.
+def MergeRequantizationFollowedByDequantization : Pat<
+  (StableHLO_UniformDequantizeOp:$res
+    (StableHLO_UniformQuantizeOp $input)),
+  (StableHLO_UniformDequantizeOp $input),
+  [(IsUniformQuantizedType $input)]>;

tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h

@@ -0,0 +1,61 @@
+/* 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_PASSES_PASSES_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_PASSES_PASSES_H_
+
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "absl/status/statusor.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_options.pb.h"
+
+namespace mlir::quant::stablehlo {
+
+// Creates a pass that quantizes weight component of StableHLO graph.
+std::unique_ptr<OperationPass<func::FuncOp>> CreateQuantizeWeightPass(
+    const ::stablehlo::quantization::QuantizationComponentSpec&
+        quantization_component_spec = {});
+
+// Converts a serialized StableHLO module to bfloat16 and output serialized
+// module.
+absl::StatusOr<std::string> ConvertSerializedStableHloModuleToBfloat16(
+    StringRef serialized_stablehlo_module);
+
+std::unique_ptr<OperationPass<ModuleOp>>
+CreateLiftQuantizableSpotsAsFunctionsPass(
+    const ::stablehlo::quantization::QuantizationSpecs& quantization_specs);
+
+// Creates a pass that inserts CalibrationStatisticsSaverOp.
+std::unique_ptr<OperationPass<ModuleOp>>
+CreateInsertCalibrationStatisticsSaverPass(
+    StringRef calibration_data_dir,
+    const std::vector<std::string>& aggregator_ops_to_ignore);
+
+// Adds generated pass default constructors or options definitions.
+#define GEN_PASS_DECL
+// Adds generated pass registration functions.
+#define GEN_PASS_REGISTRATION
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"
+
+}  // namespace mlir::quant::stablehlo
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_PASSES_PASSES_H_

tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.td

@@ -0,0 +1,248 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+include "mlir/Pass/PassBase.td"
+
+def QuantizeWeightPass : Pass<"tf-stablehlo-quantize-weight", "mlir::func::FuncOp"> {
+  let summary = "Quantizes the weight component of StableHLO graph.";
+  let dependentDialects = ["mlir::stablehlo::StablehloDialect"];
+  let constructor = "mlir::quant::stablehlo::CreateQuantizeWeightPass()";
+}
+
+def UnfuseMhloBatchNormPass : Pass<"tf-stablehlo-unfuse-mhlo-batch-norm", "mlir::func::FuncOp"> {
+  let summary = "Unfuses batch normalization into arithmetic ops.";
+}
+
+def LiftQuantizableSpotsAsFunctionsPass : Pass<"tf-stablehlo-lift-quantizable-spots-as-functions", "mlir::ModuleOp"> {
+  let summary = "Replace quantization candidates with composite functions into the module.";
+  let description = [{
+    Mark frequent fusible patterns as functions for quantization targets.
+    In addition to brining performance benefits by reducing q/dq op overhead in non-full quantization,
+    this brings higher accuracy by keeping a smaller range when quantizing ops
+    that disperse values. (ex: convolution, dot_general)
+  }];
+  let dependentDialects = [
+      "mlir::func::FuncDialect",
+      "mlir::stablehlo::StablehloDialect",
+      "TF::TensorFlowDialect",
+  ];
+}
+
+def ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPass : Pass<"tf-stablehlo-replace-stablehlo-ops-in-main-function-with-xla-call-module-ops", "mlir::ModuleOp"> {
+  let summary = "Replaces the StableHLO ops with a separate XlaCallModuleOps.";
+  let description = [{
+     Replaces the StableHLO ops in the main function block with
+     tf.XlaCallModuleOps as separate subgraphs. Wires them back to the main
+     function block to be compatible with SavedModel structure.
+  }];
+}
+
+def RestoreFunctionNamePass : Pass<"tf-stablehlo-restore-function-name", "ModuleOp"> {
+  let summary = "Restores function name from XlaCallModule op.";
+}
+
+def QuantizeCompositeFunctionsPass : Pass<"tf-stablehlo-quantize-composite-functions", "ModuleOp"> {
+  let summary = "Quantize composite functions with QDQ input / outputs.";
+  let options = [
+    Option<"enable_per_channel_quantized_weight_",
+        "enable-per-channel-quantized-weight",
+        "bool", /*default=*/"true",
+        "Whether to enable per-channel quantized weights.">,
+    Option<"mlir_dump_file_name_", "mlir-dump-file-name",
+        "std::optional<std::string>", /*default=*/"std::nullopt",
+        "MLIR dump file name.">,
+    Option<"merge_fusion_with_dequantize_",
+        "merge-fusion-with-dequantize",
+        "bool", /*default=*/"false",
+        "Whether to merge quantized conv/dot_general fusion with subsequent dequantize.">,
+  ];
+  let dependentDialects = [
+    "mlir::arith::ArithDialect",
+    "mlir::stablehlo::StablehloDialect",
+    "mlir::quant::QuantDialect",
+    "mlir::quant::ir::TFQuantDialect",
+    "TF::TensorFlowDialect",
+  ];
+}
+
+def PrepareQuantizePass : Pass<"tf-stablehlo-prepare-quantize", "mlir::ModuleOp"> {
+  let summary = "Prepare StableHLO dialect for static range quantization by converting quantfork.stats into quantfork.qcast and dcast ops.";
+  let options = [
+    Option<"enable_per_channel_quantized_weight_",
+        "enable-per-channel-quantized-weight",
+        "bool", /*default=*/"true",
+        "Whether to enable per-channel quantized weights.">,
+    Option<"bit_width_", "bit-width", "int", /*default=*/"8",
+        "Bitwidth of quantized integer">
+    ];
+  let dependentDialects = [
+      "mlir::stablehlo::StablehloDialect",
+      "mlir::quant::QuantDialect",
+      "mlir::quant::ir::TFQuantDialect",
+      "mlir::arith::ArithDialect",
+  ];
+}
+
+def QuantizePass : Pass<"tf-stablehlo-quantize", "mlir::ModuleOp"> {
+  let summary = "Applies static-range quantization on ops by converting quantfork.qcast, quantfork.dcast, and float op into uniform quantized ops .";
+  let options = [
+    Option<"enable_per_channel_quantized_weight_",
+        "enable-per-channel-quantized-weight",
+        "bool", /*default=*/"true",
+        "Whether to enable per-channel quantized weights.">,
+  ];
+  let dependentDialects = [
+    "mlir::stablehlo::StablehloDialect",
+    "mlir::quant::QuantDialect",
+    "mlir::quant::ir::TFQuantDialect",
+  ];
+}
+
+def PostQuantizePass : Pass<"tf-stablehlo-post-quantize", "mlir::func::FuncOp"> {
+  let summary = "Apply clean-up after quantization.";
+  let dependentDialects = [
+    "mlir::stablehlo::StablehloDialect",
+    "mlir::quant::ir::TFQuantDialect",
+  ];
+}
+
+def XlaCallModuleToCallPass : Pass<"tf-stablehlo-xla-call-module-to-call", "ModuleOp"> {
+  let summary = "Convert XlaCallModuleOp to func.call op";
+  let dependentDialects = [
+    "TF::TensorFlowDialect",
+  ];
+}
+
+def MergeFusionWithDequantizePass : Pass<"tf-stablehlo-merge-fusion-with-dequantize", "mlir::ModuleOp"> {
+  let summary = "Merge quantized conv/dot_general fusion with subsequent dequantize.";
+  let dependentDialects = [
+    "chlo::ChloDialect",
+    "mlir::stablehlo::StablehloDialect",
+  ];
+}
+
+def UnwrapXlaCallModuleOpPass : Pass<"tf-stablehlo-unwrap-xla-call-module-op", "ModuleOp"> {
+  let summary = "Unwrap XlaCallModuleOps into inline functions if not used for quantizing fused patterns.";
+  let dependentDialects = ["TF::TensorFlowDialect"];
+}
+
+def ConvertFuncToBfloat16Pass : Pass<"tf-stablehlo-convert-func-to-bfloat16", "mlir::func::FuncOp"> {
+  let summary = "Convert a StableHLO function to bfloat16";
+  let dependentDialects = ["mlir::stablehlo::StablehloDialect"];
+}
+
+def ConvertXlaCallModuleOpToBfloat16Pass : Pass<"tf-stablehlo-convert-xla-call-module-op-to-bfloat16", "mlir::func::FuncOp"> {
+  let summary = "Convert serialized XlaCallModuleOp to bfloat16";
+  let dependentDialects = [
+    "TF::TensorFlowDialect",
+    "mlir::quant::QuantDialect",
+    "mlir::shape::ShapeDialect",
+    "mlir::stablehlo::StablehloDialect",
+  ];
+}
+
+def ConvertShapeToStablehloWithConstraintsPass : Pass<"tf-stablehlo-convert-shape-to-stablehlo-with-constraints", "mlir::func::FuncOp"> {
+  let summary = "Convert shape.cstr_broadcastable to stablehlo.custom_call @shape_assertion";
+  let dependentDialects = [
+    "mlir::shape::ShapeDialect",
+    "mlir::tensor::TensorDialect",
+    "mlir::stablehlo::StablehloDialect",
+  ];
+}
+
+def OptimizeGraphPass : Pass<"tf-optimize-graph", "ModuleOp"> {
+  let summary = "Optimize the sub-optimal patterns after quantization.";
+  let dependentDialects = ["mlir::stablehlo::StablehloDialect",];
+}
+
+def NchwConvolutionToNhwcPass : Pass<"tf-stablehlo-nchw-convolution-to-nhwc", "mlir::func::FuncOp"> {
+  let summary = "Converts stablehlo.convolution op of NCHW format to -> NHWC.";
+  let description = [{
+    Matches `ConvolutionOp`s with NCHW format and converts it to NHWC
+    format by inserting `TransposeOp`s to input, filter, and output tensors.
+    In terms of dimension numbers, this matches
+    `[b, f, 0, 1]x[o, i, 0, 1]->[b, f, 0, 1]` format and converts it to
+    `[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]` format.
+
+    This pass is useful to convert models that conventionally use the NCHW
+    format to target hardwares that are more NHWC-friendly.
+  }];
+  let dependentDialects = ["mlir::stablehlo::StablehloDialect"];
+}
+
+def DeferActivationTransposePass : Pass<"tf-stablehlo-defer-activation-transpose", "mlir::func::FuncOp"> {
+  let summary = "Merges stablehlo.transpose for activations.";
+  let description = [{
+    Defers activation transposes (e.g. LHS of `stablehlo.add`) to the output and
+    optionally inserts `stablehlo.transpose`s to match the shape of operands.
+    This is useful when recursively pushing down the extra `stablehlo.transpose`
+    inserted to activation tensors after running `NchwConvolutionToNhwcPass`.
+
+    Currently only converts limited cases that appear in NCHW->NHWC 2D
+    convolution conversion, to avoid introducing unwanted pessimizations.
+  }];
+  let dependentDialects = ["mlir::stablehlo::StablehloDialect"];
+}
+
+def InsertWeightParamPass : Pass<"tf-stablehlo-insert-weight-param", "mlir::func::FuncOp"> {
+  let summary = "Insert quantization parameters of weights for weight-only quantization and dynamic range quantization.";
+  let dependentDialects = [
+      "mlir::stablehlo::StablehloDialect",
+      "TF::TensorFlowDialect",
+      "mlir::quant::QuantDialect",
+      "mlir::quant::ir::TFQuantDialect",
+  ];
+}
+
+def FoldConstantTransposePass : Pass<"tf-stablehlo-fold-constant-transpose", "mlir::func::FuncOp"> {
+  let summary = "Folds stablehlo.constant -> stablehlo.transpose patterns.";
+  let description = [{
+    Finds patterns where a `stablehlo.constant` is directly followed by a
+    `stablehlo.transpose` and folds them into a single `stablehlo.constant`.
+    This is considered an aggressive optimization, but it is useful to eliminate
+    `stablehlo.constant`->`stablehlo.transpose` patterns which are often
+    by-products of other shape conversion optimizations, such as NCHW->NHWC
+    convolution conversion.
+  }];
+  let dependentDialects = ["mlir::stablehlo::StablehloDialect"];
+}
+
+def RemoveShardingCustomCallPass : Pass<"tf-stablehlo-remove-sharding-custom-call", "mlir::func::FuncOp"> {
+  let summary = "Removes `stablehlo.custom_call @Sharding`";
+  let description = [{
+    Finds `stablehlo.custom_call @Sharding` and removes all instances of them,
+    replacing the usages by its operand. This is used where sharding doesn't
+    make much sense or sharding custom calls are incompatible, e.g. on-device
+    targets.
+  }];
+  let dependentDialects = ["mlir::stablehlo::StablehloDialect"];
+}
+
+def InsertCalibrationStatisticsSaverPass : Pass<"tf-stablehlo-insert-calibration-statistics-saver", "ModuleOp"> {
+  let summary = "Inserts `CalibrationStatisticsSaver` op to collect and save calibration statistics.";
+  let description = [{
+    Finds all `CustomAggregator` ops in the each function and add a single
+    `CalibrationStatisticsSaver` op at the end of the function to collect their
+    statistics.
+  }];
+  let options = [
+    ListOption<"aggregator_ops_to_ignore_", "aggregator-ops-to-ignore", "std::string",
+               "Ops to ignore when inserting CalibrationStatisticsSaver.">,
+    Option<"calibration_data_dir_", "calibration-data-dir",
+        "std::string", /*default=*/"",
+        "The directory to save calibration data.">,
+  ];
+  let dependentDialects = ["TF::TensorFlowDialect"];
+}

tensorflow/compiler/mlir/quantization/stablehlo/passes/post_quantize.cc

@@ -0,0 +1,160 @@
+/* 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 <utility>
+
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Matchers.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"
+#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h"  // IWYU pragma: keep
+
+namespace mlir::quant::stablehlo {
+
+#define GEN_PASS_DEF_POSTQUANTIZEPASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"
+
+namespace {
+
+// Applies clean-up patterns after quantization.
+class PostQuantizePass : public impl::PostQuantizePassBase<PostQuantizePass> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PostQuantizePass)
+
+  explicit PostQuantizePass() = default;
+
+ private:
+  void runOnOperation() override;
+};
+
+// TODO: b/305815328 - Consider preserving leading and trailing QDQs for
+// ModifyIONodesPass in TFLite use cases.
+// Removes the back-to-back quantize and dequantize ops with volatile attribute.
+class RemoveVolatileQdqPattern
+    : public OpRewritePattern<mlir::quant::ir::DequantizeCastOp> {
+ public:
+  explicit RemoveVolatileQdqPattern(MLIRContext* context)
+      : OpRewritePattern<mlir::quant::ir::DequantizeCastOp>(context) {}
+
+  LogicalResult matchAndRewrite(mlir::quant::ir::DequantizeCastOp op,
+                                PatternRewriter& rewriter) const override {
+    auto input_op = op.getArg().getDefiningOp();
+    if (auto q =
+            llvm::dyn_cast_or_null<mlir::quant::ir::QuantizeCastOp>(input_op)) {
+      if (!q->getAttr(kVolatileOpAttrName)) return failure();
+
+      // If the quantize op is a requantize op, it is being used in other scale
+      // adjustments and should be kept. Instead, move dequantize op before the
+      // requantize op to remove the unnecessary requantize op.
+      if (const QuantizedType qtype =
+              QuantizedType::getQuantizedElementType(q.getArg().getType())) {
+        rewriter.setInsertionPoint(op);
+        rewriter.replaceOpWithNewOp<mlir::quant::ir::DequantizeCastOp>(
+            op, op.getResult().getType(), q.getArg());
+        return success();
+      }
+
+      op.replaceAllUsesWith(q.getArg());
+      return success();
+    }
+    return failure();
+  }
+};
+
+// Replaces constant and uniform_quantize ops with single quantized constant op.
+class QuantizeConstPattern
+    : public OpRewritePattern<mlir::stablehlo::UniformQuantizeOp> {
+ public:
+  explicit QuantizeConstPattern(MLIRContext* context)
+      : OpRewritePattern<mlir::stablehlo::UniformQuantizeOp>(context) {}
+
+  LogicalResult matchAndRewrite(mlir::stablehlo::UniformQuantizeOp op,
+                                PatternRewriter& rewriter) const override {
+    DenseFPElementsAttr attr;
+    if (matchPattern(op.getOperand(), m_Constant(&attr))) {
+      const Type qtype = op.getResult().getType();
+      ElementsAttr quantized_attr = Quantize(attr, qtype);
+      if (quantized_attr) {
+        rewriter.replaceOpWithNewOp<mlir::stablehlo::ConstantOp>(
+            op, qtype, quantized_attr);
+        return success();
+      }
+    }
+    return failure();
+  }
+};
+
+// Replaces quantfork.dcast with stablehlo.uniform_dequantize.
+class ConvertDequantizeCastToUniformDequantizePattern
+    : public OpRewritePattern<mlir::quant::ir::DequantizeCastOp> {
+ public:
+  explicit ConvertDequantizeCastToUniformDequantizePattern(MLIRContext* context)
+      : OpRewritePattern<mlir::quant::ir::DequantizeCastOp>(context) {}
+  LogicalResult matchAndRewrite(mlir::quant::ir::DequantizeCastOp dq_op,
+                                PatternRewriter& rewriter) const override {
+    rewriter.replaceOpWithNewOp<mlir::stablehlo::UniformDequantizeOp>(
+        dq_op, dq_op.getResult().getType(), dq_op.getArg());
+    return success();
+  }
+};
+
+// Replaces quantfork.qcast with stablehlo.uniform_quantize.
+class ConvertQuantizeCastToUniformQuantizePattern
+    : public OpRewritePattern<mlir::quant::ir::QuantizeCastOp> {
+ public:
+  explicit ConvertQuantizeCastToUniformQuantizePattern(MLIRContext* context)
+      : OpRewritePattern<mlir::quant::ir::QuantizeCastOp>(context) {}
+  LogicalResult matchAndRewrite(mlir::quant::ir::QuantizeCastOp q_op,
+                                PatternRewriter& rewriter) const override {
+    rewriter.replaceOpWithNewOp<mlir::stablehlo::UniformQuantizeOp>(
+        q_op, q_op.getResult().getType(), q_op.getArg());
+    return success();
+  }
+};
+
+void PostQuantizePass::runOnOperation() {
+  RewritePatternSet patterns(&getContext());
+  func::FuncOp func = getOperation();
+  MLIRContext* ctx = func.getContext();
+  // TODO: b/307463853 - Consider splitting passes for each pattern set.
+  patterns.add<FoldTrivalRequantizeOp<mlir::quant::ir::QuantizeCastOp>,
+               RemoveVolatileQdqPattern>(ctx);
+  if (failed(applyPatternsGreedily(func, std::move(patterns)))) {
+    signalPassFailure();
+  }
+
+  RewritePatternSet patterns_2(&getContext());
+  patterns_2
+      .add<QuantizeConstPattern, ConvertQuantizeCastToUniformQuantizePattern,
+           ConvertDequantizeCastToUniformDequantizePattern>(ctx);
+  if (failed(applyPatternsGreedily(func, std::move(patterns_2)))) {
+    signalPassFailure();
+  }
+}
+
+}  // namespace
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/passes/prepare_quantize.cc

@@ -0,0 +1,200 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include <memory>
+#include <utility>
+
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/Quant.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"
+#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_driver.h"
+#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/ops/stablehlo_op_quant_spec.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h"  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+
+namespace mlir {
+namespace quant {
+namespace stablehlo {
+
+#define GEN_PASS_DEF_PREPAREQUANTIZEPASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"
+
+namespace {
+
+// Applies prepare quantization on the model in TF dialect. This pass runs
+// before the quantization pass and propagate the quantization parameters
+// across ops. This step is necessary for post-training quantization and also
+// making the quantization rule for some operations in the quantization-aware
+// training quantization simpler.
+class PrepareQuantizePass
+    : public impl::PrepareQuantizePassBase<PrepareQuantizePass> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PrepareQuantizePass)
+
+  using impl::PrepareQuantizePassBase<
+      PrepareQuantizePass>::PrepareQuantizePassBase;
+
+  explicit PrepareQuantizePass(const bool enable_per_channel_quantized_weight,
+                               const int bit_width) {
+    enable_per_channel_quantized_weight_ = enable_per_channel_quantized_weight;
+    bit_width_ = bit_width;
+  }
+
+  void runOnOperation() override;
+};
+
+// Merges consecutive QuantizeCast ops. See b/246655213 for details.
+// For example, the following case:
+// %1 = quantfork.QuantizeCastOp(%0) : f32 -> qtype1
+// %2 = quantfork.QuantizeCastOp(%1) : qtype1 -> qtype2
+// %3 = quantfork.QuantizedOp1(%1)
+// %4 = quantfork.QuantizedOp2(%2)
+// will be tranformed to:
+// %1 = quantfork.QuantizeCastOp(%0) : f32 -> qtype1
+// %2 = quantfork.QuantizeCastOp(%0) : f32 -> qtype2
+// %3 = quantfork.QuantizedOp1(%1)
+// %4 = quantfork.QuantizedOp2(%2)
+// Converting from f32 -> qtype1 -> qtype2 will add unexpected quantization
+// lost for %2. This pattern avoids that by converting from f32 -> qtype2
+// directly.
+class MergeConsecutiveQuantizeCast
+    : public mlir::OpRewritePattern<mlir::quant::ir::QuantizeCastOp> {
+ public:
+  explicit MergeConsecutiveQuantizeCast(MLIRContext* context)
+      : OpRewritePattern<mlir::quant::ir::QuantizeCastOp>(context) {}
+
+ private:
+  LogicalResult matchAndRewrite(mlir::quant::ir::QuantizeCastOp q_op,
+                                PatternRewriter& rewriter) const override {
+    auto preceding_qcast =
+        q_op.getArg().getDefiningOp<mlir::quant::ir::QuantizeCastOp>();
+    if (!preceding_qcast) return failure();
+
+    auto new_qcast = mlir::quant::ir::QuantizeCastOp::create(
+        rewriter, q_op.getLoc(), q_op.getType(), preceding_qcast.getArg());
+    new_qcast->setAttr(kVolatileOpAttrName, rewriter.getUnitAttr());
+    q_op->replaceAllUsesWith(new_qcast);
+    return success();
+  }
+};
+
+class ConvertTFConstOpToArithConstOp : public OpRewritePattern<TF::ConstOp> {
+ public:
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(TF::ConstOp op,
+                                PatternRewriter& rewriter) const override {
+    rewriter.replaceOpWithNewOp<arith::ConstantOp>(op, op.getValue());
+    return success();
+  }
+};
+
+class ConvertStablehloConstToArithConstOp
+    : public OpRewritePattern<mlir::stablehlo::ConstantOp> {
+ public:
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(mlir::stablehlo::ConstantOp op,
+                                PatternRewriter& rewriter) const override {
+    rewriter.replaceOpWithNewOp<arith::ConstantOp>(op, op.getValue());
+    return success();
+  }
+};
+
+class ConvertArithConstToStablehloConstOp
+    : public OpRewritePattern<arith::ConstantOp> {
+ public:
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(arith::ConstantOp op,
+                                PatternRewriter& rewriter) const override {
+    rewriter.replaceOpWithNewOp<mlir::stablehlo::ConstantOp>(op, op.getValue());
+    return success();
+  }
+};
+
+void PrepareQuantizePass::runOnOperation() {
+  ModuleOp module_op = getOperation();
+  MLIRContext* ctx = module_op.getContext();
+
+  auto func_op_quant_spec = GetStableHloOpQuantSpec;
+  auto func_op_quant_scale_spec = GetStableHloQuantConstraints;
+
+  for (auto func_op : module_op.getOps<func::FuncOp>()) {
+    // The function might contain more stats ops than required, and it will
+    // introduce requantize if the calibration stats have conflicts. This tries
+    // to remove all the redundant stats ops.
+    RemoveRedundantStatsOps(func_op, func_op_quant_spec,
+                            func_op_quant_scale_spec);
+
+    RewritePatternSet patterns(ctx);
+    // Convert quant stats to int8 quantization parameters.
+    // Currently, only activation stats are imported, so narrow_range = false.
+    patterns.add<ConvertStatsToQDQs<mlir::quant::ir::QuantizeCastOp,
+                                    mlir::quant::ir::DequantizeCastOp>>(
+        bit_width_,
+        /*narrow_range=*/false,
+        /*is_signed=*/true,
+        /*legacy_float_scale=*/false, ctx);
+    // Convert all constants to arith::ConstantOp as quantization driver can
+    // deal with the arith::ConstantOp instances.
+    patterns.add<ConvertTFConstOpToArithConstOp>(ctx);
+    patterns.add<ConvertStablehloConstToArithConstOp>(ctx);
+    if (failed(applyPatternsGreedily(func_op, std::move(patterns)))) {
+      signalPassFailure();
+    }
+
+    // Finally, the quantization parameters can be propagated to the rest of the
+    // values (tensors).
+    ApplyQuantizationParamsPropagation(
+        func_op, /*is_signed=*/true, bit_width_,
+        !enable_per_channel_quantized_weight_, func_op_quant_spec,
+        func_op_quant_scale_spec,
+        /*infer_tensor_ranges=*/true, /*legacy_float_scale=*/false,
+        /*is_qdq_conversion=*/false);
+
+    // Restore constants as stablehlo::ConstantOp.
+    RewritePatternSet patterns_2(ctx);
+    patterns_2
+        .add<MergeConsecutiveQuantizeCast, ConvertArithConstToStablehloConstOp>(
+            ctx);
+    if (failed(applyPatternsGreedily(func_op, std::move(patterns_2)))) {
+      signalPassFailure();
+    }
+  }
+}
+
+}  // namespace
+
+// Creates an instance of the TensorFlow dialect PrepareQuantize pass.
+std::unique_ptr<OperationPass<ModuleOp>> CreatePrepareQuantizePass(
+    const bool enable_per_channel_quantized_weight, const int bit_width) {
+  return std::make_unique<PrepareQuantizePass>(
+      enable_per_channel_quantized_weight, bit_width);
+}
+
+}  // namespace stablehlo
+}  // namespace quant
+}  // namespace mlir

tensorflow/compiler/mlir/quantization/stablehlo/passes/quantization_patterns.h

@@ -0,0 +1,254 @@
+/* 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_PASSES_QUANTIZATION_PATTERNS_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_PASSES_QUANTIZATION_PATTERNS_H_
+
+#include <type_traits>
+#include <utility>
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/IRMapping.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo
+#include "tensorflow/compiler/mlir/quantization/common/lift_as_function_call.h"
+#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/ops/stablehlo_op_quant_spec.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace mlir::quant::stablehlo {
+
+// Checks whether an op is connected with a quantized composite function. If
+// not, the same-scale op will not be quantized. This decision is based on the
+// current assumption that the performance gain of the same-scale op itself
+// could not beat the overhead of the quantize and dequantize routines need to
+// be added around that op. When the assumption changes, this policy might
+// change as well.
+bool IsConnectedWithQuantizedCompsiteFunction(Operation* same_scale_op);
+
+// A base rewrite pattern which matches any N-in-M-out operations with
+// quantization parameters propagated to at least one of its operands. The
+// quantization parameters are annotated by the QuantizeOp/DequantizeOp pairs.
+// Each matched pattern are rewritten by its quantized alternatives.
+//
+// Quantization method is determined by the `_quantization_method` attributes
+// attached to each quantizable units.
+//
+// Template constraints are imposed as follows:
+//
+// * `QuantizeOpT` should have only one operand.
+// * `DequantizeOpT` should have only one result.
+template <typename ConcreteT, typename QuantizeOpT, typename DequantizeOpT,
+          typename VerifierT, typename RootOpT = DequantizeOpT,
+          typename = std::enable_if_t<
+              QuantizeOpT::template hasTrait<OpTrait::OneOperand>() &&
+              DequantizeOpT::template hasTrait<OpTrait::OneResult>()>>
+class StableHloQuantizationPattern : public OpRewritePattern<RootOpT> {
+ public:
+  explicit StableHloQuantizationPattern(MLIRContext* context)
+      // Set the benefit to a large number so that it is always preferred.
+      : OpRewritePattern<RootOpT>(context, /*benefit=*/300) {}
+
+ private:
+  // Collects all candidate ops for quantization, which are the
+  // `dequantize_op`'s users.
+  FailureOr<SmallVector<Operation*>> CollectCandidateOps(
+      DequantizeOpT dequantize_op) const {
+    auto users = dequantize_op->getResult(0).getUsers();
+    return SmallVector<Operation*>(users.begin(), users.end());
+  }
+
+  // Collects all candidate ops for quantization, which is the operand of
+  // `quantize_op`. If successful, this always returns one element which is the
+  // operand of `quantize_op`.
+  FailureOr<SmallVector<Operation*>> CollectCandidateOps(
+      QuantizeOpT quantize_op) const {
+    Value operand = quantize_op->getOperand(0);
+    if (QuantizedType::getQuantizedElementType(operand.getType())) {
+      // The input of the quantize op has already been quantized, i.e.
+      // rescale.
+      return failure();
+    }
+
+    Operation* operand_op = operand.getDefiningOp();
+    if (operand_op == nullptr) {
+      // When `QuantizeOpT`'s operand does not have a defining op, it means it
+      // is a `BlockArgument`. The pattern does not match if there is no op to
+      // quantize.
+      return failure();
+    }
+
+    if (operand_op->hasTrait<OpTrait::ConstantLike>()) {
+      // Const-> QuantizeOp pattern will be handled separately.
+      return failure();
+    }
+
+    return SmallVector<Operation*>{operand_op};
+  }
+
+  LogicalResult matchAndRewrite(RootOpT op,
+                                PatternRewriter& rewriter) const override {
+    // Collect all the candidate ops for quantization.
+    FailureOr<SmallVector<Operation*>> candidate_ops = CollectCandidateOps(op);
+    // Safeguard check to ensure that there is at least one quantizable op.
+    if (failed(candidate_ops) || candidate_ops->empty()) return failure();
+
+    // Rewrite the floating-point ops to the quantized version, by fusing
+    // preceding dequantize ops and succeding quantize ops.
+    for (Operation* candidate_op : *candidate_ops) {
+      // If it is requantize op, we shouldn't rewrite this op.
+      if (isa<QuantizeOpT, DequantizeOpT>(candidate_op)) {
+        return failure();
+      }
+
+      // If the op is terminator, we shouldn't rewrite.
+      if (candidate_op->hasTrait<OpTrait::IsTerminator>()) {
+        return failure();
+      }
+
+      if (!IsOpQuantizableStableHlo(candidate_op)) {
+        return failure();
+      }
+
+      if (GetStableHloQuantConstraints(candidate_op)
+              ->has_same_scale_requirement &&
+          !IsConnectedWithQuantizedCompsiteFunction(candidate_op)) {
+        return failure();
+      }
+
+      // Ops with regions will be quantized in a separate pattern.
+      if (isa<mlir::stablehlo::ReduceWindowOp>(candidate_op)) {
+        return failure();
+      }
+
+      const bool weight_only_quantizable =
+          IsWeightOnlyQuantizableOp(*candidate_op);
+
+      // Collect all the quantized inputs and "clone" the matched op by these
+      // inputs.
+      SmallVector<Value, 4> inputs;
+      inputs.reserve(candidate_op->getNumOperands());
+      for (auto operand : candidate_op->getOperands()) {
+        Type operand_type = operand.getType();
+        if (mlir::isa<NoneType>(operand_type)) {
+          inputs.push_back(operand);
+          continue;
+        }
+
+        auto ele_type =
+            mlir::cast<TensorType>(operand.getType()).getElementType();
+        if (auto dq_op =
+                dyn_cast_or_null<DequantizeOpT>(operand.getDefiningOp())) {
+          inputs.push_back(dq_op.getOperand());
+        } else if (!ele_type.isF32()) {
+          // If the operand is an integer tensor, then it doesn't require the
+          // DequantizeOp in the pattern.
+          inputs.push_back(operand);
+        } else if (weight_only_quantizable) {
+          inputs.push_back(operand);
+        } else {
+          return failure();
+        }
+      }
+
+      // Collect all the quantized outputs and replace them by the results of
+      // the new quantized op.
+      llvm::SmallDenseMap<Value, int> outputs_replaced;
+      SmallVector<Type, 4> output_types;
+      output_types.reserve(candidate_op->getNumResults());
+      for (const auto& enumerated_result :
+           llvm::enumerate(candidate_op->getResults())) {
+        Value result = enumerated_result.value();
+        Type result_type = result.getType();
+        // Add this to the test coverage once we create test ops with none type
+        // results.
+        if (mlir::isa<NoneType>(result_type)) {
+          outputs_replaced.insert({result, enumerated_result.index()});
+          output_types.push_back(result_type);
+          continue;
+        }
+        Type result_ele_type =
+            mlir::cast<TensorType>(result.getType()).getElementType();
+        // If the user is the QuantizeOp, it must be the only user.
+        if (result.hasOneUse() && isa<QuantizeOpT>(*result.user_begin())) {
+          auto user = cast<QuantizeOpT>(*result.user_begin());
+          outputs_replaced.insert(
+              {user.getResult(), enumerated_result.index()});
+          output_types.push_back(user.getType());
+        } else if (!result_ele_type.isF32()) {
+          // If the result is an integer tensor, then it doesn't require the
+          // D op in the pattern.
+          outputs_replaced.insert({result, enumerated_result.index()});
+          output_types.push_back(result.getType());
+        } else if (weight_only_quantizable) {
+          outputs_replaced.insert({result, enumerated_result.index()});
+          output_types.push_back(result.getType());
+        } else {
+          return failure();
+        }
+      }
+
+      rewriter.setInsertionPointAfter(candidate_op);
+      OperationState new_state(candidate_op->getLoc(),
+                               candidate_op->getName().getStringRef(), inputs,
+                               output_types, candidate_op->getAttrs());
+      for (int i = 0; i < candidate_op->getNumRegions(); ++i) {
+        new_state.addRegion();
+      }
+      Operation* quantized_op = rewriter.create(new_state);
+      if (candidate_op->getNumRegions() != 0) {
+        for (const auto& indexed_regions :
+             llvm::enumerate(candidate_op->getRegions())) {
+          Region& target_region =
+              quantized_op->getRegion(indexed_regions.index());
+          IRMapping mapping;
+          indexed_regions.value().cloneInto(&target_region, mapping);
+        }
+      }
+      for (auto output : outputs_replaced) {
+        output.getFirst().replaceAllUsesWith(
+            quantized_op->getResult(output.getSecond()));
+      }
+    }
+    return success();
+  }
+};
+
+// Populates common patterns that are usually compute heavy or memory bound.
+void PopulateCommonQuantizationPatterns(
+    MLIRContext& ctx, RewritePatternSet& patterns,
+    bool enable_per_channel_quantized_weight);
+
+// Populates conversion patterns for all quantizable ops, including
+// ops that are not compute-heavy and data movement ops.
+void PopulateAllQuantizablePatterns(MLIRContext& ctx,
+                                    RewritePatternSet& patterns);
+
+}  // namespace mlir::quant::stablehlo
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_PASSES_QUANTIZATION_PATTERNS_H_

tensorflow/compiler/mlir/quantization/stablehlo/passes/quantize.cc

@@ -0,0 +1,111 @@
+/* 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 <utility>
+
+#include "mlir/Dialect/Quant/IR/Quant.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h"  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/quantization_patterns.h"
+
+namespace mlir::quant::stablehlo {
+
+#define GEN_PASS_DEF_QUANTIZEPASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"
+
+namespace {
+
+// Base struct for quantization.
+template <typename ConcreteT,
+          typename RootOpT = mlir::quant::ir::DequantizeCastOp>
+struct StableHloQuantizationBase
+    : public StableHloQuantizationPattern<ConcreteT,
+                                          mlir::quant::ir::QuantizeCastOp,
+                                          mlir::quant::ir::DequantizeCastOp,
+                                          /*VerifierT=*/void, RootOpT> {
+  explicit StableHloQuantizationBase(MLIRContext* ctx)
+      : StableHloQuantizationPattern<ConcreteT, mlir::quant::ir::QuantizeCastOp,
+                                     mlir::quant::ir::DequantizeCastOp,
+                                     /*VerifierT=*/void, RootOpT>(ctx) {}
+
+  static bool AllowWeightOnlyQuantization(Operation& op) { return false; }
+};
+
+// Quantization rewrite pattern using DQ as the root op.
+struct StableHloQuantization
+    : public StableHloQuantizationBase<StableHloQuantization> {
+  explicit StableHloQuantization(MLIRContext* ctx)
+      : StableHloQuantizationBase<StableHloQuantization>(ctx) {}
+};
+
+// Quantization rewrite pattern using Q as the root op. This is for the
+// quantizable ops without floating-point operands.
+struct StableHloQuantizationReverse
+    : public StableHloQuantizationBase<StableHloQuantizationReverse,
+                                       mlir::quant::ir::QuantizeCastOp> {
+  explicit StableHloQuantizationReverse(MLIRContext* ctx)
+      : StableHloQuantizationBase<StableHloQuantizationReverse,
+                                  mlir::quant::ir::QuantizeCastOp>(ctx) {}
+};
+
+class QuantizePass : public impl::QuantizePassBase<QuantizePass> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(QuantizePass)
+
+  using impl::QuantizePassBase<QuantizePass>::QuantizePassBase;
+
+  explicit QuantizePass(const bool enable_per_channel_quantized_weight) {
+    enable_per_channel_quantized_weight_ = enable_per_channel_quantized_weight;
+  }
+
+ private:
+  void runOnOperation() override;
+};
+
+void QuantizePass::runOnOperation() {
+  ModuleOp module_op = getOperation();
+  MLIRContext& ctx = getContext();
+
+  RewritePatternSet patterns(&ctx);
+  patterns.add<StableHloQuantization, StableHloQuantizationReverse>(&ctx);
+
+  PopulateCommonQuantizationPatterns(ctx, patterns,
+                                     enable_per_channel_quantized_weight_);
+
+  // Quantize all quantizable ops, including ops that are not compute-heavy.
+  PopulateAllQuantizablePatterns(ctx, patterns);
+
+  if (failed(applyPatternsGreedily(module_op, std::move(patterns)))) {
+    // There are cases where no rewrites happen even if a pattern matches,
+    // causing this to result in a convergence failure. Consider this as a
+    // best-effort.
+    module_op.emitWarning("Failed to converge pattern at QuantizePass.");
+  }
+}
+
+}  // namespace
+
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/passes/quantize_composite_functions.cc

@@ -0,0 +1,114 @@
+/* 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 "absl/log/log.h"
+#include "absl/status/status.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/Quant.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"  // IWYU pragma: keep
+#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/cc/run_passes.h"
+#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"  // IWYU pragma: keep
+
+#define DEBUG_TYPE "quantize-composite-functions"
+
+namespace mlir::quant::stablehlo {
+
+#define GEN_PASS_DEF_QUANTIZECOMPOSITEFUNCTIONSPASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"
+
+namespace {
+
+using ::tensorflow::quantization::RunPassesOnModuleOp;
+
+class QuantizeCompositeFunctionsPass
+    : public impl::QuantizeCompositeFunctionsPassBase<
+          QuantizeCompositeFunctionsPass> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(QuantizeCompositeFunctionsPass)
+
+  using impl::QuantizeCompositeFunctionsPassBase<
+      QuantizeCompositeFunctionsPass>::QuantizeCompositeFunctionsPassBase;
+
+  explicit QuantizeCompositeFunctionsPass(
+      const bool enable_per_channel_quantized_weight) {
+    enable_per_channel_quantized_weight_ = enable_per_channel_quantized_weight;
+  }
+
+ private:
+  void runOnOperation() override;
+};
+
+void QuantizeCompositeFunctionsPass::runOnOperation() {
+  MLIRContext& ctx = getContext();
+
+  PassManager pm(&ctx);
+  // Intermediate output from QuantizePass will have quantized ops
+  // (XlaCallModuleOps) with quantized input and output types, which are not
+  // allowed in the TF dialect.
+  pm.enableVerifier(false);
+
+  PrepareQuantizePassOptions options;
+  options.enable_per_channel_quantized_weight_ =
+      enable_per_channel_quantized_weight_;
+  // Change this to user-given bit width once we have custom configuration.
+  options.bit_width_ = 8;
+
+  // Insert quantization parameters for weights for ops with `weight_only_ptq`
+  // attribute.
+  pm.addNestedPass<func::FuncOp>(createInsertWeightParamPass());
+
+  // PrepareQuantizePass uses SymbolTable to fetch relevant GEMM ops for
+  // determining quantization attributes. This requires module-level context.
+  pm.addPass(createPrepareQuantizePass(options));
+
+  QuantizePassOptions quantize_options;
+  quantize_options.enable_per_channel_quantized_weight_ =
+      enable_per_channel_quantized_weight_;
+
+  // QuantizePass modifies FuncOps referenced outside of its given scope
+  // and therefore requires a module-level context.
+  pm.addPass(createQuantizePass(quantize_options));
+  pm.addNestedPass<func::FuncOp>(createPostQuantizePass());
+
+  // Convert XlaCallModuleOps lifted but not quantized to func.call op.
+  // The reasons these ops are not quantized may be:
+  // 1. Disabled due to selective quantization.
+  // 2. Not supported, e.g. add op for server.
+  pm.addPass(createXlaCallModuleToCallPass());
+
+  // TODO: b/321729008 - move this implementation to quantization_patterns.cc.
+  if (merge_fusion_with_dequantize_) {
+    pm.addPass(createMergeFusionWithDequantizePass());
+  }
+
+  ModuleOp module_op = getOperation();
+  if (const absl::Status pm_run_status =
+          RunPassesOnModuleOp(mlir_dump_file_name_, pm, module_op);
+      !pm_run_status.ok()) {
+    signalPassFailure();
+  }
+}
+
+}  // namespace
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/passes/quantize_weight.cc

@@ -0,0 +1,244 @@
+/* Copyright 2023 The StableHLO Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <algorithm>
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "Eigen/Core"  // from @eigen_archive
+#include "llvm/ADT/SetVector.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Dialect.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Rewrite/FrozenRewritePatternSet.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_options.pb.h"
+
+// NOLINTNEXTLINE
+//===----------------------------------------------------------------------===//
+// The Quantization Pass for Weight.
+//===----------------------------------------------------------------------===//
+
+namespace mlir::quant::stablehlo {
+
+// Put the definitions inside the ::mlir::tf_quant::stablehlo namespace, to
+// match the declarations in passes.h.
+#define GEN_PASS_DEF_QUANTIZEWEIGHTPASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"
+
+namespace {
+
+using QuantizationUnits = llvm::SetVector<std::pair<Operation*, int>>;
+using mlir::stablehlo::ConstantOp;
+using mlir::stablehlo::ConvertOp;
+using ::stablehlo::quantization::QuantizationComponentSpec;
+
+// Min/Max values used for creating ConstantOp.
+constexpr float kMaxFloat16Value = 65504.f;
+constexpr float kMinFloat16Value = -65504.f;
+
+class QuantizeWeightPass
+    : public impl::QuantizeWeightPassBase<QuantizeWeightPass> {
+ public:
+  explicit QuantizeWeightPass(
+      QuantizationComponentSpec quantization_component_spec)
+      : quantization_component_spec_(quantization_component_spec) {}
+
+ private:
+  void runOnOperation() override;
+  QuantizationComponentSpec quantization_component_spec_;
+};
+
+// Collects quantizable target ops, then insert Q-DQ quantization patterns.
+class QuantizeWeight : public OpRewritePattern<ConstantOp> {
+ public:
+  explicit QuantizeWeight(
+      MLIRContext* context,
+      const QuantizationComponentSpec& quantization_component_spec)
+      : OpRewritePattern<ConstantOp>(context),
+        quantization_component_spec_(quantization_component_spec) {}
+
+  LogicalResult matchAndRewrite(ConstantOp op,
+                                PatternRewriter& rewriter) const override {
+    // 1. Collect quantizable ops.
+    QuantizationUnits quantizable_ops = GetQuantizableOps(op);
+    if (quantizable_ops.empty()) {
+      return failure();
+    }
+
+    // 2. Quantize collected ops.
+    if (!QuantizeOps(rewriter, op, quantizable_ops)) {
+      return failure();
+    }
+
+    // 3. Complete the Q-DQ pair for each inference type.
+    if (!ConvertToFloat16Constant(rewriter, op)) {
+      return failure();
+    }
+    return success();
+  }
+
+ private:
+  const QuantizationComponentSpec quantization_component_spec_;
+  // Marks users that are applicable for quantization where the criteria for
+  // determining quantizable ops differs by the inference type.
+  QuantizationUnits GetQuantizableOps(ConstantOp op) const {
+    // Non-float tensors do not need quantization.
+    QuantizationUnits quantizable_ops;
+    const ShapedType type = mlir::dyn_cast<ShapedType>(op.getType());
+    if (!type || !type.getElementType().isF32()) return quantizable_ops;
+
+    const Value value = op.getResult();
+
+    for (OpOperand& use : value.getUses()) {
+      Operation* user = use.getOwner();
+      const int operand_num = use.getOperandNumber();
+      quantizable_ops.insert({user, operand_num});
+    }
+    return quantizable_ops;
+  }
+
+  // Returns whether quantization is applied to filtered users.
+  bool QuantizeOps(PatternRewriter& rewriter, ConstantOp op,
+                   const QuantizationUnits& quantizable_ops) const {
+    for (const std::pair<Operation*, int>& quant_op : quantizable_ops) {
+      // For f16 quantization, quantize all constant ops as float16.
+      QuantizeOpAsFloat16(rewriter, op, quant_op);
+    }
+    // TODO: b/264218457 - Return a value that accurately captures result
+    // status.
+    return true;
+  }
+
+  // Inserts ConvertOp which is used for converting float32 ConstantOp into
+  // float16 quantization. If there is an existing ConvertOp connected to the
+  // ConstantOp, the quantizable_op will be rewired to the existing ConvertOp.
+  // This guarantees at most one ConvertOp is created for float32 to float16
+  // conversion.
+  void QuantizeOpAsFloat16(PatternRewriter& rewriter, ConstantOp op,
+                           const std::pair<Operation*, int> quant_op) const {
+    const auto [quantizable_op, quantize_operand_num] = quant_op;
+    // If the constant is an output tensor, do nothing.
+    if (isa<func::ReturnOp>(quantizable_op)) {
+      return;
+    }
+
+    TensorType old_result_type =
+        mlir::dyn_cast<TensorType>(op.getResult().getType());
+    const FloatType quantized_type = Float16Type::get(op.getContext());
+    const ShapedType new_result_type = old_result_type.clone(quantized_type);
+
+    // Insert ConvertOp if it does not exist yet. Otherwise, just rewire without
+    // creating a ConvertOp.
+    for (const OpOperand& connected_op : op.getResult().getUses()) {
+      ConvertOp convert_op =
+          dyn_cast_or_null<ConvertOp>(connected_op.getOwner());
+      // ConvertOp already exists. Rewire the existing convert op into f16.
+      if (convert_op && convert_op.getType() == new_result_type) {
+        quantizable_op->setOperand(quantize_operand_num, convert_op);
+        return;
+      }
+    }
+    rewriter.setInsertionPointAfter(op);
+    ConvertOp new_convert_op = ConvertOp::create(
+        rewriter, op->getLoc(), new_result_type, op.getResult());
+    quantizable_op->setOperand(quantize_operand_num,
+                               new_convert_op.getResult());
+  }
+
+  // Returns whether a ConvertOp-Operation sequence can be converted into new
+  // ConstantOp-Convert-Operation. The new ConstantOp has float16 data type.
+  bool ConvertToFloat16Constant(PatternRewriter& rewriter,
+                                ConstantOp op) const {
+    for (Operation* connected_op : op.getResult().getUsers()) {
+      ConvertOp convert_op = dyn_cast_or_null<ConvertOp>(connected_op);
+      // Skip if no convert op exists.
+      if (!convert_op || convert_op.getResult().use_empty()) continue;
+
+      // Get types.
+      const Type old_result_type = op.getResult().getType();
+      const ShapedType new_result_type =
+          mlir::dyn_cast<ShapedType>(convert_op.getType());
+
+      // Proceeds only if the converting is to float16.
+      if (!new_result_type.getElementType().isF16()) continue;
+
+      // Convert values.
+      std::vector<Eigen::half> new_values;
+      const DenseFPElementsAttr value_attr =
+          mlir::cast<DenseFPElementsAttr>(op.getValue());
+      new_values.reserve(value_attr.getNumElements());
+
+      for (const float value : value_attr.getValues<float>()) {
+        new_values.push_back(Eigen::half(
+            std::min(std::max(value, kMinFloat16Value), kMaxFloat16Value)));
+      }
+      const DenseElementsAttr new_value_attr = DenseFPElementsAttr::get(
+          new_result_type, ArrayRef<Eigen::half>(new_values));
+      // Create new ConstantOp-ConvertOp-Operation sequences. At this moment,
+      // old ConstantOp is guaranteed to have one F32->F16 convert op regardless
+      // of its number of users.
+      rewriter.setInsertionPointAfter(op);
+      // create new F16 constant op in that location
+      ConstantOp new_const = ConstantOp::create(
+          rewriter, op->getLoc(), new_result_type, new_value_attr);
+      ConvertOp dcast =
+          ConvertOp::create(rewriter, op->getLoc(), old_result_type, new_const);
+      // replace all convert ops with dq op.
+      convert_op->replaceAllUsesWith(dcast);
+      // Return without scanning for the next ConvertOp as only one ConvertOp is
+      // connected to all quantizable ops.
+      return true;
+    }
+    return false;
+  }
+};
+
+// TODO: b/264218457 - Refactors the current file to parse preset quantization
+// options and allow modular control of quantization specs.
+void QuantizeWeightPass::runOnOperation() {
+  func::FuncOp func = getOperation();
+  MLIRContext* ctx = func.getContext();
+  RewritePatternSet patterns(ctx);
+
+  patterns.add<QuantizeWeight>(ctx, quantization_component_spec_);
+
+  FrozenRewritePatternSet frozen_patterns(std::move(patterns));
+
+  if (failed(applyPatternsGreedily(func, frozen_patterns))) {
+    signalPassFailure();
+  }
+}
+
+}  // namespace
+
+// Creates an instance of the StableHLO dialect Quantize Weight pass.
+std::unique_ptr<OperationPass<func::FuncOp>> CreateQuantizeWeightPass(
+    const QuantizationComponentSpec& quantization_component_spec) {
+  return std::make_unique<QuantizeWeightPass>(quantization_component_spec);
+}
+
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/passes/remove_sharding_custom_call.cc

@@ -0,0 +1,59 @@
+/* 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 <utility>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Rewrite/FrozenRewritePatternSet.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo  // IWYU pragma: keep
+
+namespace mlir::quant::stablehlo {
+
+#define GEN_PASS_DEF_REMOVESHARDINGCUSTOMCALLPASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"
+
+// Include patterns generated from `remove_sharding_custom_call.td`.
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/remove_sharding_custom_call.inc"
+
+class RemoveShardingCustomCallPass
+    : public impl::RemoveShardingCustomCallPassBase<
+          RemoveShardingCustomCallPass> {
+ public:
+  using impl::RemoveShardingCustomCallPassBase<
+      RemoveShardingCustomCallPass>::RemoveShardingCustomCallPassBase;
+
+ private:
+  void runOnOperation() override;
+};
+
+void RemoveShardingCustomCallPass::runOnOperation() {
+  func::FuncOp func_op = getOperation();
+  MLIRContext& ctx = getContext();
+
+  RewritePatternSet patterns(&ctx);
+  populateWithGenerated(patterns);
+
+  FrozenRewritePatternSet frozen_patterns(std::move(patterns));
+  if (failed(applyPatternsGreedily(func_op, frozen_patterns))) {
+    func_op.emitWarning() << "Failed to converge "
+                          << RemoveShardingCustomCallPass::getArgumentName();
+  }
+}
+
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/passes/remove_sharding_custom_call.td

@@ -0,0 +1,27 @@
+/* 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 "stablehlo/dialect/StablehloOps.td"
+
+class IsStringAttrOf<string value> : Constraint<
+  CPred<"::llvm::isa_and_nonnull<StringAttr>($_self) && llvm::cast<StringAttr>($_self).getValue() == \"" # value # "\"">,
+  "Is a string attribute whose value is \"" # value # "\""
+>;
+
+// Removes `stablehlo.custom_call @Sharding`. Assumes this call always accepts
+// one input. Other attributes are ignored.
+def RemoveShardingCustomCall : Pat<
+  (StableHLO_CustomCallOp (variadic $input0), $call_target_name, $_, $_, $_, $_, $_, $_, $_, $_),
+  (replaceWithValue $input0),
+  [(IsStringAttrOf<"Sharding">:$call_target_name)]>;

tensorflow/compiler/mlir/quantization/stablehlo/passes/replace_stablehlo_ops_in_main_function_with_xla_call_module_ops.cc

@@ -0,0 +1,545 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include <cstdint>
+#include <string>
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/IRMapping.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/IR/ValueRange.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo
+#include "stablehlo/dialect/Version.h"  // from @stablehlo
+#include "tensorflow/compiler/mlir/quantization/common/func.h"
+#include "tensorflow/compiler/mlir/quantization/common/lift_as_function_call.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/utils/stablehlo_type_utils.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/xla_call_module_attrs.h"
+#include "tensorflow/core/ir/types/dialect.h"
+
+namespace mlir::quant::stablehlo {
+
+#define GEN_PASS_DEF_REPLACESTABLEHLOOPSINMAINFUNCTIONWITHXLACALLMODULEOPSPASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"
+
+namespace {
+
+constexpr StringRef kStablehloModuleAttrsAttrName = "_stablehlo_module_attrs";
+constexpr StringRef kUsesShapePolymorphismAttr = "jax.uses_shape_polymorphism";
+constexpr StringRef kNoXlaCallModuleAttrName = "_no_xla_call_module";
+
+// Default version number for native serialization.
+constexpr int64_t kDefaultVersion = 9;
+// Platforms for XlaCallModuleOp.
+constexpr StringRef kPlatformCpu = "CPU";
+constexpr StringRef kPlatformTpu = "TPU";
+
+class ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPass
+    : public impl::
+          ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPassBase<
+              ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPass> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
+      ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPass)
+
+  ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPass() = default;
+
+  ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPass(
+      const ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPass& other) =
+      default;
+
+ private:
+  void runOnOperation() override;
+};
+
+// Creates a unique stablehlo function name based on op order.
+std::string CreateStablehloFunctionName(const int id) {
+  return Twine("_stablehlo_main_").concat(std::to_string(id)).str();
+}
+
+// Follows the structure of Live-variable analysis. It is a form of
+// CFG (Control Flow Graph) analysis, often used in compilers.
+//
+// A variable is live if it holds a value that may be used in the future.
+// It is live-in at node n if it is live on any of the node's in-edges.
+// It is live-out at node n if it is live on any of the node's out-edges.
+// def[n] refers to values that are defined at node n.
+// use[n] refers to values that are used at node n.
+//
+// Given a node n, variables' liveliness is defined like the following:
+// live_in[n] = use[n] U (live_out[n] - def[n])
+// live_out[n] = U {live_in[s] | s ε succ[n]}
+//
+// Consider a sequence of op:
+//
+// ```
+// node 1: %0 = stablehlo.constant
+// node 2: %1 = stablehlo.constant
+// node 3: %2 = stablehlo.add %0, %1
+// node 4: %3 = stablehlo.multiply %2, %1
+// node 5: return %3
+// ```
+//
+// In Backward Liveliness analysis, the liveliness for each node above becomes:
+// live_in[5] = use[5]   U (live_out[5] - def[5])
+//            = {%3}     U ({∅} - {∅})            = {%3}
+// live_in[4] = use[4]   U (live_out[4] - def[4])
+//            = {%1, %2} U ({%3} - {%3})          = {%1, %2}
+// live_in[3] = use[3]   U (live_out[3] - def[3])
+//            = {%0, %1} U ({%1, %2} - {%2})      = {%0, %1}
+// live_in[2] = use[2]   U (live_out[2] - def[2])
+//            = {∅}      U ({%0, %1} - {%1})      = {%0}
+// live_in[1] = use[1]   U (live_out[1] - def[1])
+//            = {∅}      U ({%0} - {%0})          = {∅}
+//
+// This analogy is used throughout this pass to ensure only live edges form
+// proper subgraphs.
+class LiveOuts {
+ public:
+  LiveOuts() = default;
+
+  explicit LiveOuts(OperandRange range)
+      : liveouts_(range.begin(), range.end()), prev_liveouts_(liveouts_) {}
+
+  // Delete the current op from liveouts and moves on to the parent ops.
+  void update(Operation& op) {
+    for (Value result_value : op.getResults()) {
+      liveouts_.remove(result_value);
+    }
+    for (Value operand : op.getOperands()) {
+      liveouts_.insert(operand);
+    }
+  }
+
+  // Snapshot the current live values to previous live values.
+  void snapshot_previous_state() { prev_liveouts_ = liveouts_; }
+
+  // Return the current live values.
+  const SetVector<Value>& get() const { return liveouts_; }
+
+  // Return the previous live values.
+  const SetVector<Value>& get_previous() const { return prev_liveouts_; }
+
+ private:
+  // Use SerVector to ensure deterministic traversal order.
+  SetVector<Value> liveouts_;
+  SetVector<Value> prev_liveouts_;
+};
+
+// Creates the tf.XlaCallModuleOp from attributes.
+void CreateXlaCallModuleOp(ValueRange inputs, ValueRange outputs,
+                           const TypeRange result_types,
+                           const SetVector<Operation*>& reverse_subgraph,
+                           const func::FuncOp stablehlo_func_op,
+                           ModuleOp module_op) {
+  MLIRContext* ctx = module_op.getContext();
+  OpBuilder builder(ctx);
+  Operation* last_subgraph_op = reverse_subgraph.front();
+  builder.setInsertionPointAfter(last_subgraph_op);
+
+  // Create attributes used for creating an XlaCallModuleOp.
+  SmallVector<Attribute> shape_attrs;
+  for (const Type result_type : result_types) {
+    shape_attrs.push_back(
+        tf_type::ShapeAttr::get(ctx, mlir::cast<ShapedType>(result_type)));
+  }
+  const auto empty_array_attr = ArrayAttr::get(ctx, {});
+  // TODO: b/310291615 - find a better way for platform support.
+  const auto platforms = ArrayAttr::get(
+      ctx,
+      {StringAttr::get(ctx, kPlatformCpu), StringAttr::get(ctx, kPlatformTpu)});
+
+  auto xla_call_module_op = TF::XlaCallModuleOp::create(
+      builder, module_op.getLoc(), /*output=*/result_types,
+      /*args=*/inputs,
+      /*version=*/kDefaultVersion, /*module=*/"",
+      /*Sout=*/ArrayAttr::get(ctx, shape_attrs),
+      /*dim_args_spec=*/empty_array_attr, platforms,
+      /*function_list=*/empty_array_attr,
+      /*has_token_input_output=*/false,
+      /*disabled_checks=*/empty_array_attr);
+  xla_call_module_op->setAttr(TF::kStablehloEntryFunctionAttrName,
+                              SymbolRefAttr::get(stablehlo_func_op));
+  std::string target_version =
+      mlir::vhlo::Version::fromCompatibilityRequirement(
+          vhlo::Version::CompatibilityRequirement::WEEK_4)
+          .toString();
+  xla_call_module_op->setAttr(TF::kStablehloVersionAttrName,
+                              builder.getStringAttr(target_version));
+  // Set jax.uses_shape_polymorphism=true to enable shape refinement at runtime.
+  // This is needed for native serialization version >= 8.
+  xla_call_module_op->setAttr(
+      kStablehloModuleAttrsAttrName,
+      builder.getDictionaryAttr(builder.getNamedAttr(
+          kUsesShapePolymorphismAttr, builder.getBoolAttr(true))));
+
+  for (auto [original_output_value, xla_call_module_op_result_value] :
+       llvm::zip_equal(outputs, xla_call_module_op->getResults())) {
+    original_output_value.replaceAllUsesExcept(xla_call_module_op_result_value,
+                                               /*exceptedUser=*/nullptr);
+  }
+}
+
+// Replaces the StableHLO ops with a separate XlaCallModuleOp, then wires it
+// back into the main graph.
+void ReplaceStablehloOpsWithXlaCallModuleOp(
+    const ArrayRef<Value> inputs, const ArrayRef<Value> outputs,
+    const SetVector<Operation*>& reverse_subgraph, const int stablehlo_func_id,
+    ModuleOp module_op) {
+  MLIRContext* ctx = module_op.getContext();
+  OpBuilder builder(ctx);
+
+  // Identify arg types & arg locs.
+  SmallVector<Type> arg_types;
+  SmallVector<Location> arg_locs;
+
+  // Add an argument for platform_index. This allows for multiple platforms.
+  // TODO: b/310291615 - find a better way for platform support.
+  arg_types.push_back(RankedTensorType::get({}, builder.getI32Type()));
+  arg_locs.push_back(module_op.getLoc());
+  for (const Value input_value : inputs) {
+    arg_types.push_back(input_value.getType());
+    arg_locs.push_back(input_value.getLoc());
+  }
+
+  // Identify result types.
+  SmallVector<Type> result_types;
+  for (const Value output_value : outputs) {
+    result_types.push_back(output_value.getType());
+  }
+
+  // 1) Create FuncOp for the StableHLO ops. They will be separate subgraphs.
+  builder.setInsertionPoint(&*module_op.begin());
+  auto stablehlo_func_op =
+      func::FuncOp::create(builder, module_op.getLoc(),
+                           CreateStablehloFunctionName(stablehlo_func_id),
+                           FunctionType::get(ctx, arg_types, result_types));
+  stablehlo_func_op.setVisibility(SymbolTable::Visibility::Private);
+  stablehlo_func_op->setAttr(TF::kFromXlaCallModuleAttrName,
+                             builder.getUnitAttr());
+
+  builder.createBlock(&stablehlo_func_op.getBody(), stablehlo_func_op.begin(),
+                      arg_types, arg_locs);
+
+  IRMapping mapper;
+  // stablehlo_func_op has 1 extra arg for platform index.
+  for (auto [input, stablehlo_func_arg] : llvm::zip_equal(
+           inputs, stablehlo_func_op.getArguments().take_back(inputs.size()))) {
+    mapper.map(input, stablehlo_func_arg);
+  }
+
+  for (Operation* subgraph_op : llvm::reverse(reverse_subgraph)) {
+    // Create a deep copy of the subgraph ops' operands to the func op.
+    stablehlo_func_op.getBody().begin()->push_back(subgraph_op->clone(mapper));
+  }
+
+  SmallVector<Value> result_values;
+  for (const Value original_output_value : outputs) {
+    // Use the mapped values in the newly created function that correspond to
+    // outputs in the original function.
+    result_values.push_back(mapper.lookup(original_output_value));
+  }
+  func::ReturnOp::create(builder, module_op.getLoc(), result_values);
+
+  // 2) Create XlaCallModuleOp (with ops mapped).
+  CreateXlaCallModuleOp(inputs, outputs, result_types, reverse_subgraph,
+                        stablehlo_func_op, module_op);
+
+  // 3) Erase the replaced ops.
+  for (Operation* subgraph_op : reverse_subgraph) {
+    subgraph_op->erase();
+  }
+}
+
+// Contains the actual logic for updating states and replacing StableHLO ops
+// with tf.XlaCallModuleOps.
+void UpdateStatesAndReplaceStablehloOps(
+    const SetVector<Value>& operands, const SetVector<Value>& defined_values,
+    const LiveOuts& liveouts, ModuleOp module_op,
+    const SetVector<Operation*>& reverse_subgraph, const int stablehlo_func_id,
+    func::FuncOp main_func, const bool is_last_subgraph = false) {
+  SetVector<Value> inputs = operands;
+  for (Value defined_value : defined_values) {
+    inputs.remove(defined_value);
+  }
+
+  SetVector<Value> outputs = liveouts.get_previous();
+  for (const Value live_value : liveouts.get()) {
+    outputs.remove(live_value);
+  }
+
+  if (is_last_subgraph) {
+    // Additionally remove arguments from the outputs, as it provides liveness
+    // throughout (functions as an invisible op above the very first op that
+    // returns the arguments).
+    for (const BlockArgument arg : main_func.getArguments()) {
+      outputs.remove(arg);
+    }
+  }
+
+  ReplaceStablehloOpsWithXlaCallModuleOp(
+      SmallVector<Value>(inputs.begin(), inputs.end()),
+      SmallVector<Value>(outputs.begin(), outputs.end()), reverse_subgraph,
+      stablehlo_func_id, module_op);
+}
+
+// Check if the op should be added to the subgraph.
+// The op should be added to the subgraph if all of its users match one
+// of following two conditions:
+// 1: The user is already in the current subgraph.
+// 2: The user will reach a dead end.
+//
+// If the op should be added to the subgraph and there are users who
+// will reach the dead end, add the ops on the dead end to the subgraph as well.
+bool ShouldAddOpToSubgraph(Operation* op,
+                           const SetVector<Operation*>& reverse_subgraph,
+                           const SetVector<Operation*>& ops_to_add,
+                           SmallVector<Operation*>& all_descendants) {
+  if (!op) {
+    return false;
+  }
+
+  SmallVector<Operation*> current_layer_descendants;
+  SmallVector<Operation*> next_layer_descendants;
+  int current_depth = 0;
+  current_layer_descendants.push_back(op);
+  // BFS downstream ops for current user.
+  // If any one of the descendants meet one of the three conditions, we return
+  // false for the current value:
+  // 1: The descendant is not in the ops_to_add.
+  // 2: The descendant is not a stablehlo op.
+  // 3: The depth of the descendant is larger than 5, we don't want to search
+  // too deep, max depth is arbitrarily chosen.
+  while (!current_layer_descendants.empty()) {
+    if (current_depth > 5) {
+      all_descendants.clear();
+      return false;
+    }
+    current_depth++;
+
+    for (Operation* descendant : current_layer_descendants) {
+      if (!quant::stablehlo::IsStablehloOp(descendant) ||
+          !ops_to_add.contains(descendant)) {
+        all_descendants.clear();
+        return false;
+      }
+      for (Operation* next_descendant : descendant->getUsers()) {
+        if (reverse_subgraph.contains(next_descendant)) {
+          continue;
+        }
+        next_layer_descendants.push_back(next_descendant);
+      }
+      all_descendants.push_back(descendant);
+    }
+
+    current_layer_descendants = next_layer_descendants;
+    next_layer_descendants.clear();
+  }
+
+  return true;
+}
+
+// Replaces the StableHLO ops in the main function block with
+// tf.XlaCallModuleOps as separate subgraphs. Wires them back to the main
+// function block to be compatible with SavedModel structure.
+void ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOps(
+    ModuleOp module_op, func::FuncOp main_func, int& stablehlo_func_id) {
+  Block& main_func_block = main_func.getBody().front();
+
+  // LiveOuts keeps track of live values at the output of some op. The updates
+  // must be made in a reverse, bottom-up manner.
+  const auto result_values = main_func_block.getTerminator()->getOperands();
+  LiveOuts liveouts(result_values);
+
+  // Copy ops to iterate because we will be modifying the block during
+  // iteration. The ordering should be reversed because liveness analysis is a
+  // bottom-up analysis. The terminator is not included because the return
+  // statement is not included in any subgraph (e.g. XlaCallModuleOp) and is
+  // untouched.
+  SmallVector<Operation*> reverse_main_func_block_ops;
+  SetVector<Operation*> ops_to_add;
+  for (Operation& main_func_block_op :
+       llvm::reverse(main_func_block.without_terminator())) {
+    reverse_main_func_block_ops.push_back(&main_func_block_op);
+    ops_to_add.insert(&main_func_block_op);
+  }
+
+  // Create a separate subgraph invoked with XlaCallModuleOp per each
+  // set of StableHLO ops in the main func block.
+  SetVector<Operation*> reverse_subgraph;
+  SetVector<Value> operands;
+  SetVector<Value> defined_values;
+
+  // Add op to the subgraph.
+  const auto add_to_subgraph = [&](Operation* op) {
+    // Move on to the parent ops.
+    liveouts.update(*op);
+    ops_to_add.remove(op);
+
+    if (!quant::stablehlo::IsStablehloOp(op)) {
+      // Always update the liveouts when the subgraph isn't being continued.
+      liveouts.snapshot_previous_state();
+      return;
+    }
+
+    reverse_subgraph.insert(op);
+    defined_values.insert(op->getResults().begin(), op->getResults().end());
+    operands.insert(op->getOperands().begin(), op->getOperands().end());
+  };
+
+  for (Operation* op : reverse_main_func_block_ops) {
+    if (!ops_to_add.contains(op)) continue;
+    // When hitting a non-StableHLO op, i.e. tf.CustomAggregatorOp, start
+    // recursively tracing defining ops of the current subgraph's operands. This
+    // makes sure that all dependencies needed for shape inference are included
+    // in the subgraph. We only trace StableHLO ops that have all users inside
+    // the current subgraph.
+    // TODO: b/311239049 - Consider rewrite this using BFS.
+    if (!quant::stablehlo::IsStablehloOp(op)) {
+      bool should_add_op = true;
+      while (should_add_op) {
+        should_add_op = false;
+        SmallVector<Operation*> all_descendants;
+        for (Value v : operands) {
+          if (defined_values.contains(v)) continue;
+          if (ShouldAddOpToSubgraph(v.getDefiningOp(), reverse_subgraph,
+                                    ops_to_add, all_descendants)) {
+            should_add_op = true;
+            break;
+          }
+        }
+        if (should_add_op) {
+          for (auto descendant : llvm::reverse(all_descendants)) {
+            add_to_subgraph(descendant);
+          }
+        }
+      }
+      // Create an XlaCallModuleOp if reverse_subgraph isn't empty.
+      if (!reverse_subgraph.empty()) {
+        UpdateStatesAndReplaceStablehloOps(operands, defined_values, liveouts,
+                                           module_op, reverse_subgraph,
+                                           ++stablehlo_func_id, main_func);
+
+        // Reset states and start a new subgraph.
+        reverse_subgraph.clear();
+        operands.clear();
+        defined_values.clear();
+      }
+    }
+    add_to_subgraph(op);
+  }
+
+  // Create the last subgraph if it isn't empty.
+  if (!reverse_subgraph.empty()) {
+    UpdateStatesAndReplaceStablehloOps(
+        operands, defined_values, liveouts, module_op, reverse_subgraph,
+        ++stablehlo_func_id, main_func, /*is_last_subgraph=*/true);
+  }
+}
+
+// Duplicates small constants for each use.
+//
+// In the subsequent graph partitioning, constants for shape inference need to
+// be in the same subgraph. But graph partitioning stops at ops with multiple
+// uses. So here we duplicate small constants for each use so that if a
+// constant is useful for shape inference for multiple subgraphs, they can be
+// included in each subgraphs. If duplicate constants are accidentally created
+// in the same subgraph, they can be easily removed with a canonicalizer pass.
+//
+// We set a size limit since constants needed for shape inference are no
+// larger than tensor rank. This avoids duplicating large constants.
+void DuplicateSmallConstantOps(ModuleOp module_op, func::FuncOp main_func) {
+  OpBuilder builder(main_func.getContext());
+  for (auto constant_op :
+       main_func.getBody().getOps<mlir::stablehlo::ConstantOp>()) {
+    builder.setInsertionPointAfter(constant_op);
+    if (constant_op.getResult().use_empty() ||
+        constant_op.getResult().hasOneUse())
+      continue;
+    // Do not duplicate constant op if the size is too large.
+    // 32 is chosen to be larger than all constants useful for shape references,
+    // while not too large to possibly significantly increase model size.
+    if (constant_op.getValue().getNumElements() > 32) continue;
+    while (!constant_op.getResult().hasOneUse()) {
+      auto new_constant_op = builder.clone(*constant_op.getOperation());
+      constant_op.getResult().getUses().begin()->assign(
+          dyn_cast<mlir::stablehlo::ConstantOp>(new_constant_op));
+    }
+  }
+}
+
+void ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPass::
+    runOnOperation() {
+  ModuleOp module_op = getOperation();
+
+  func::FuncOp main_func = quant::FindMainFuncOp(module_op);
+  if (!main_func) return;
+
+  // In case the model has tf.StatefulPartitionedCallOp or tf.PartitionedCallOp,
+  // we recursively find called functions and process StableHLO ops in them.
+  SmallVector<func::FuncOp> func_ops;
+  func_ops.push_back(main_func);
+  int stablehlo_func_id = -1;
+  while (!func_ops.empty()) {
+    auto func_op = func_ops.back();
+    func_ops.pop_back();
+    if (!func_op) continue;
+
+    // If the function has the _no_xla_call_module attribute, we don't need to
+    // replace StableHLO ops in this function with XlaCallModuleOps (including
+    // the functions called by this function).
+    if (func_op->getAttrOfType<mlir::UnitAttr>(kNoXlaCallModuleAttrName)) {
+      continue;
+    }
+
+    SymbolTable symbol_table(module_op);
+    for (auto call_op : func_op.getOps<TF::PartitionedCallOp>()) {
+      func_ops.push_back(dyn_cast_or_null<func::FuncOp>(symbol_table.lookup(
+          mlir::cast<FlatSymbolRefAttr>(call_op.getFAttr()).getValue())));
+    }
+    for (auto call_op : func_op.getOps<TF::StatefulPartitionedCallOp>()) {
+      func_ops.push_back(
+          dyn_cast_or_null<func::FuncOp>(symbol_table.lookup(call_op.getF())));
+    }
+
+    DuplicateSmallConstantOps(module_op, func_op);
+    ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOps(module_op, func_op,
+                                                          stablehlo_func_id);
+  }
+
+  // TODO - b/298966126: Currently quantizable functions are identified in TF
+  // Quantizer via the tf_quant.composite_function UnitAttr attached to
+  // func ops. We remove this attribute as this interferes with VHLO conversion.
+  // Remove this temporary hack.
+  for (auto func_op : module_op.getOps<func::FuncOp>()) {
+    func_op->removeAttr(kFusedFunctionAttr);
+  }
+}
+
+}  // namespace
+
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/passes/restore_function_name.cc

@@ -0,0 +1,94 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/quantization/common/lift_as_function_call.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/xla_call_module_attrs.h"
+
+//===----------------------------------------------------------------------===//
+// The stablehlo-restore-function-name Pass.
+//===----------------------------------------------------------------------===//
+
+namespace mlir::quant::stablehlo {
+
+#define GEN_PASS_DEF_RESTOREFUNCTIONNAMEPASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"
+
+namespace {
+
+// Restores entry function name from XlaCallModuleOp attribute.
+// This restoration is required because StableHLO functions are renamed during
+// the XlaCallModuleSerialization.
+class RestoreFunctionNamePass
+    : public impl::RestoreFunctionNamePassBase<RestoreFunctionNamePass> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(RestoreFunctionNamePass)
+
+  explicit RestoreFunctionNamePass() = default;
+
+  void runOnOperation() override;
+};
+
+void RestoreFunctionNameFromXlaCallModuleOp(TF::XlaCallModuleOp& call_op,
+                                            SymbolTable& symbol_table) {
+  if (!call_op->hasAttr(kOriginalStablehloEntryFunctionAttrName)) {
+    return;
+  }
+
+  const auto original_function_name = call_op->getAttrOfType<StringAttr>(
+      kOriginalStablehloEntryFunctionAttrName);
+  const auto current_function_name = call_op->getAttrOfType<FlatSymbolRefAttr>(
+      TF::kStablehloEntryFunctionAttrName);
+
+  if (!original_function_name || !current_function_name) {
+    return;
+  }
+
+  auto function =
+      symbol_table.lookup<func::FuncOp>(current_function_name.getValue());
+  if (function) {
+    function.setName(original_function_name);
+  }
+
+  call_op->setAttr(TF::kStablehloEntryFunctionAttrName,
+                   FlatSymbolRefAttr::get(original_function_name));
+}
+
+void RestoreFunctionNamePass::runOnOperation() {
+  ModuleOp module_op = getOperation();
+
+  MLIRContext* ctx = module_op.getContext();
+  OpBuilder builder(ctx);
+  SymbolTable symbol_table(module_op);
+
+  // TODO - b/298966126: Improve this logic if needed.
+  module_op.walk([&](TF::XlaCallModuleOp call_op) {
+    RestoreFunctionNameFromXlaCallModuleOp(call_op, symbol_table);
+  });
+}
+}  // namespace
+
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/passes/testing/passes.h

@@ -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_PASSES_TESTING_PASSES_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_PASSES_TESTING_PASSES_H_
+
+#include "mlir/Pass/Pass.h"  // from @llvm-project  // IWYU pragma: keep
+
+namespace mlir::quant::stablehlo::testing {
+
+// Identifies predefined `QuantizationSpecs` for
+// `TestLiftQuantizableSpotsAsFunctionsWithQuantizationSpecsPass`. The pass
+// option argument is specified in line comments for each enum value.
+enum class TestQuantizationSpecs {
+  kEmpty,                         // empty
+  kDisableAllDotGeneral,          // disable-all-dot-general
+  kStaticRangePtqToAll,           // static-range-ptq-to-all
+  kStaticRangePtqToComputeHeavy,  // static-range-ptq-to-compute-heavy
+};
+
+// Adds generated pass default constructors or options definitions.
+#define GEN_PASS_DECL
+// Adds generated pass registration functions.
+#define GEN_PASS_REGISTRATION
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/testing/passes.h.inc"
+
+}  // namespace mlir::quant::stablehlo::testing
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_PASSES_TESTING_PASSES_H_

tensorflow/compiler/mlir/quantization/stablehlo/passes/testing/passes.td

@@ -0,0 +1,94 @@
+/* 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.
+==============================================================================*/
+
+// Passes only used for testing purposes.
+
+include "mlir/Pass/PassBase.td"
+
+def TestPreCalibrationComponentPass : Pass<"stablehlo-test-pre-calibration-component", "mlir::ModuleOp"> {
+  let summary = "Test-only pass to test the PreCalibrationComponent.";
+  let description = [{
+    Runs the pre calibration passes for post-training quantization with default
+    configuration.
+  }];
+  let dependentDialects = [
+    "mlir::stablehlo::StablehloDialect", "mlir::TF::TensorFlowDialect",
+    "mlir::func::FuncDialect", "mlir::tf_executor::TensorFlowExecutorDialect",
+    "mlir::mhlo::MhloDialect", "mlir::vhlo::VhloDialect",
+  ];
+}
+
+def TestPostCalibrationComponentPass : Pass<"stablehlo-test-post-calibration-component", "mlir::ModuleOp"> {
+  let summary = "Test-only pass to test the PostCalibrationComponent.";
+  let description = [{
+    Runs the post-calibration passes for post-training quantization.
+  }];
+  let options = [
+    Option<"unpack_quantized_types_", "unpack-quantized-types", "bool",
+      /*default=*/"true", "Unpacks ops with uniform quantized types into "
+      "operations without uniform quantized types (mostly i8 or i32).">
+  ];
+  let dependentDialects = [
+    "mlir::stablehlo::StablehloDialect", "mlir::TF::TensorFlowDialect",
+    "mlir::func::FuncDialect", "mlir::mhlo::MhloDialect",
+    "mlir::quant::QuantDialect", "mlir::chlo::ChloDialect",
+    "mlir::vhlo::VhloDialect", "mlir::shape::ShapeDialect",
+    "mlir::quant::ir::TFQuantDialect",
+  ];
+}
+
+def TestTFToStablehloPass : Pass<"stablehlo-test-tf-to-stablehlo", "mlir::ModuleOp"> {
+  let summary = "Test-only pass to test TFToStablehloPasses.";
+  let description = [{
+    Runs the TFToStablehloPasses.
+  }];
+  let dependentDialects = [
+    "mlir::stablehlo::StablehloDialect", "mlir::TF::TensorFlowDialect",
+    "mlir::chlo::ChloDialect", "mlir::quant::QuantDialect",
+    "mlir::mhlo::MhloDialect", "mlir::shape::ShapeDialect",
+    "mlir::sparse_tensor::SparseTensorDialect", "mlir::ub::UBDialect",
+    "mlir::vhlo::VhloDialect",
+  ];
+}
+
+def TestLiftQuantizableSpotsAsFunctionsWithQuantizationSpecsPass :
+    Pass<"stablehlo-test-lift-quantizable-spots-as-functions-with-quantization-specs", "mlir::ModuleOp"> {
+  let summary = "Test-only pass for testing the LiftQuantizableSpotsAsFunctionsPass with a predefined QuantizationSpecs.";
+  let description = [{
+    This test-only pass is the same as `LiftQuantizableSpotsAsFunctionsPass` but
+    has predefined `QuantizationSpecs` to make FileCheck testing easier.
+  }];
+  let options = [
+    Option<"quantization_specs_", "quantization-specs",
+      "mlir::quant::stablehlo::testing::TestQuantizationSpecs",
+      /*default=*/"mlir::quant::stablehlo::testing::TestQuantizationSpecs::kEmpty",
+      "Sets one of the predefined `QuantizationSpecs` for testing.",
+      [{llvm::cl::values(
+        clEnumValN(mlir::quant::stablehlo::testing::TestQuantizationSpecs::kEmpty,
+          "empty", "Uses empty (default) QuantizationSpecs."),
+        clEnumValN(mlir::quant::stablehlo::testing::TestQuantizationSpecs::kDisableAllDotGeneral,
+          "disable-all-dot-general", "Disables all dot_general ops by matching lifted function names"),
+        clEnumValN(mlir::quant::stablehlo::testing::TestQuantizationSpecs::kStaticRangePtqToAll,
+          "static-range-ptq-to-all", "Applies `StaticRangePtq` to all quantizable units."),
+        clEnumValN(mlir::quant::stablehlo::testing::TestQuantizationSpecs::kStaticRangePtqToComputeHeavy,
+          "static-range-ptq-to-compute-heavy", "Applies `StaticRangePtq` to only compute heavy units.")
+      )}]>
+  ];
+  let dependentDialects = [
+      "mlir::func::FuncDialect",
+      "mlir::stablehlo::StablehloDialect",
+      "TF::TensorFlowDialect",
+  ];
+}

tensorflow/compiler/mlir/quantization/stablehlo/passes/testing/test_lift_quantizable_spots_as_functions_with_quantization_specs.cc

@@ -0,0 +1,140 @@
+/* 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 "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Pass/Pass.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/testing/passes.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"  // IWYU pragma: keep
+#include "tsl/platform/protobuf.h"  // IWYU pragma: keep
+
+namespace mlir::quant::stablehlo::testing {
+
+// NOLINTNEXTLINE - Automatically generated.
+#define GEN_PASS_DEF_TESTLIFTQUANTIZABLESPOTSASFUNCTIONSWITHQUANTIZATIONSPECSPASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/testing/passes.h.inc"
+
+namespace {
+
+using ::stablehlo::quantization::QuantizationSpecs;
+using ::tsl::protobuf::TextFormat;
+// NOLINTNEXTLINE(misc-include-cleaner) - Required for OSS.
+using ::tsl::protobuf::io::ArrayInputStream;
+
+// Empty (default) `QuantizationSpecs` proto.
+constexpr absl::string_view kSpecsEmpty = R"pb(specs
+                                               [])pb";
+
+// Configure `QuantizationSpecs` to disable quantization for all dot_general
+// quantizable units.
+constexpr absl::string_view kSpecsDisableAllDotGeneral =
+    R"pb(specs
+         [ {
+           matcher { function_name { regex: "composite_dot_general_.*" } }
+           method { no_quantization {} }
+         }])pb";
+
+// Configure `QuantizationSpecs` to apply `StaticRangePtq` to all quantizable
+// units.
+constexpr absl::string_view kSpecsStaticRangePtqToAll =
+    R"pb(specs
+         [ {
+           matcher { function_name { regex: ".*" } }
+           method { static_range_ptq {} }
+         }])pb";
+
+// Configure `QuantizationSpecs` to apply `StaticRangePtq` to compute heavy
+// units.
+constexpr absl::string_view kSpecsStaticRangePtqToComputeHeavy =
+    R"pb(specs
+         [ {
+           matcher { function_name { regex: "^.*(conv|dot|gather).*" } }
+           method { static_range_ptq {} }
+         }])pb";
+
+class TestLiftQuantizableSpotsAsFunctionsWithQuantizationSpecsPass
+    : public impl::
+          TestLiftQuantizableSpotsAsFunctionsWithQuantizationSpecsPassBase<
+              TestLiftQuantizableSpotsAsFunctionsWithQuantizationSpecsPass> {
+ public:
+  using impl::TestLiftQuantizableSpotsAsFunctionsWithQuantizationSpecsPassBase<
+      TestLiftQuantizableSpotsAsFunctionsWithQuantizationSpecsPass>::
+      TestLiftQuantizableSpotsAsFunctionsWithQuantizationSpecsPassBase;
+
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
+      TestLiftQuantizableSpotsAsFunctionsWithQuantizationSpecsPass)
+
+ private:
+  void runOnOperation() override;
+};
+
+// `TestQuantizationSpecs` -> predefined `QuantizationSpecs` textproto.
+absl::string_view GetQuantizationSpecsTextProto(
+    const TestQuantizationSpecs test_specs) {
+  switch (test_specs) {
+    case TestQuantizationSpecs::kEmpty:
+      return kSpecsEmpty;
+    case TestQuantizationSpecs::kDisableAllDotGeneral:
+      return kSpecsDisableAllDotGeneral;
+    case TestQuantizationSpecs::kStaticRangePtqToAll:
+      return kSpecsStaticRangePtqToAll;
+    case TestQuantizationSpecs::kStaticRangePtqToComputeHeavy:
+      return kSpecsStaticRangePtqToComputeHeavy;
+  }
+}
+
+// Parses a text proto into a `QuantizationSpecs` proto. Returns
+// `InvalidArgumentError` if `text_proto` is invalid.
+absl::StatusOr<QuantizationSpecs> ParseTextProto(
+    const absl::string_view text_proto) {
+  QuantizationSpecs quantization_specs;
+  TextFormat::Parser parser;
+  ArrayInputStream input_stream(text_proto.data(), text_proto.size());
+  if (parser.Parse(&input_stream, &quantization_specs)) {
+    return quantization_specs;
+  }
+  return absl::InvalidArgumentError("Could not parse text proto.");
+}
+
+void TestLiftQuantizableSpotsAsFunctionsWithQuantizationSpecsPass::
+    runOnOperation() {
+  PassManager pass_manager{&getContext()};
+
+  // Construct `QuantizationSpecs` from the pass option `quantization-specs`.
+  const absl::StatusOr<QuantizationSpecs> quantization_specs =
+      ParseTextProto(GetQuantizationSpecsTextProto(quantization_specs_));
+  if (!quantization_specs.ok()) {
+    signalPassFailure();
+    return;
+  }
+
+  pass_manager.addPass(
+      quant::stablehlo::CreateLiftQuantizableSpotsAsFunctionsPass(
+          *quantization_specs));
+
+  if (failed(pass_manager.run(getOperation()))) {
+    signalPassFailure();
+  }
+}
+
+}  // namespace
+}  // namespace mlir::quant::stablehlo::testing

tensorflow/compiler/mlir/quantization/stablehlo/passes/testing/test_post_calibration_component.cc

@@ -0,0 +1,83 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Dialect/Quant/IR/Quant.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "stablehlo/dialect/ChloOps.h"  // from @stablehlo  // IWYU pragma: keep
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo  // IWYU pragma: keep
+#include "stablehlo/dialect/VhloOps.h"  // from @stablehlo  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/config.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/post_calibration.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/testing/passes.h"  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"  // IWYU pragma: keep
+#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"  // IWYU pragma: keep
+
+namespace mlir::quant::stablehlo::testing {
+
+#define GEN_PASS_DEF_TESTPOSTCALIBRATIONCOMPONENTPASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/testing/passes.h.inc"
+
+namespace {
+
+using ::stablehlo::quantization::ExpandPresets;
+using ::stablehlo::quantization::PipelineConfig;
+using ::stablehlo::quantization::QuantizationConfig;
+
+class TestPostCalibrationComponentPass
+    : public impl::TestPostCalibrationComponentPassBase<
+          TestPostCalibrationComponentPass> {
+ public:
+  using impl::TestPostCalibrationComponentPassBase<
+      TestPostCalibrationComponentPass>::TestPostCalibrationComponentPassBase;
+
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestPostCalibrationComponentPass)
+
+ private:
+  void runOnOperation() override;
+};
+
+void TestPostCalibrationComponentPass::runOnOperation() {
+  ModuleOp module_op = getOperation();
+  MLIRContext& ctx = getContext();
+
+  OpPassManager pm(ModuleOp::getOperationName());
+
+  QuantizationConfig config = QuantizationConfig::default_instance();
+  config.mutable_static_range_ptq_preset();
+
+  const QuantizationConfig new_config = ExpandPresets(config);
+
+  PipelineConfig pipeline_config;
+  pipeline_config.set_unpack_quantized_types(unpack_quantized_types_);
+
+  quant::stablehlo::PostCalibrationComponent component(&ctx);
+  component.AddPasses(pm, new_config.specs(), pipeline_config);
+
+  if (failed(runPipeline(pm, module_op))) {
+    signalPassFailure();
+  }
+}
+
+}  // namespace
+
+}  // namespace mlir::quant::stablehlo::testing

tensorflow/compiler/mlir/quantization/stablehlo/passes/testing/test_pre_calibration_component.cc

@@ -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 "absl/status/statusor.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo  // IWYU pragma: keep
+#include "stablehlo/dialect/VhloOps.h"  // from @stablehlo  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/config.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/pre_calibration.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
+#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"  // IWYU pragma: keep
+#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"  // IWYU pragma: keep
+
+namespace mlir::quant::stablehlo::testing {
+
+#define GEN_PASS_DEF_TESTPRECALIBRATIONCOMPONENTPASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/testing/passes.h.inc"
+
+namespace {
+
+using ::stablehlo::quantization::ExpandPresets;
+using ::stablehlo::quantization::PopulateDefaults;
+using ::stablehlo::quantization::QuantizationConfig;
+
+class TestPreCalibrationComponentPass
+    : public impl::TestPreCalibrationComponentPassBase<
+          TestPreCalibrationComponentPass> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestPreCalibrationComponentPass)
+
+ private:
+  void runOnOperation() override;
+};
+
+void TestPreCalibrationComponentPass::runOnOperation() {
+  ModuleOp module_op = getOperation();
+  MLIRContext& ctx = getContext();
+
+  // Simply runs the PreCalibrationComponent with a default configuration.
+  PreCalibrationComponent component(&ctx);
+  QuantizationConfig quantization_config{};
+  quantization_config.mutable_static_range_ptq_preset();
+  quantization_config = ExpandPresets(PopulateDefaults(quantization_config));
+  if (!component.Run(module_op, quantization_config).ok()) {
+    signalPassFailure();
+  }
+}
+
+}  // namespace
+}  // namespace mlir::quant::stablehlo::testing