chore: import upstream snapshot with attribution

tensorflow/compiler/mlir/quantization/common/BUILD

@@ -0,0 +1,242 @@
+load("@llvm-project//mlir:tblgen.bzl", "td_library")
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("//tensorflow:tensorflow.bzl", "tf_cc_test")
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    # By default, these targets should only be used within the quantization library.
+    default_visibility = [
+        "//learning/brain/mlir/quantization:__subpackages__",
+        "//tensorflow/compiler/mlir/lite:__subpackages__",
+        "//tensorflow/compiler/mlir/quantization:__subpackages__",
+        "//tensorflow/compiler/mlir/stablehlo:__subpackages__",
+    ],
+    licenses = ["notice"],
+)
+
+td_library(
+    name = "lift_as_function_call_td_files",
+    srcs = [
+        "lift_as_function_call.td",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        "@llvm-project//mlir:FuncTdFiles",
+    ],
+)
+
+cc_library(
+    name = "lift_as_function_call",
+    srcs = ["lift_as_function_call.cc"],
+    hdrs = ["lift_as_function_call.h"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        ":attrs_and_constraints",
+        "//tensorflow/compiler/mlir/quantization/common/quantization_lib",
+        "//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
+        "//tensorflow/compiler/mlir/quantization/stablehlo:stablehlo_type_utils",
+        "//tensorflow/compiler/mlir/quantization/tensorflow/cc:quantization_unit_loc",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_ops",
+        "//tensorflow/compiler/mlir/tensorflow:xla_call_module_attrs",
+        "//tensorflow/core:framework_lite",
+        "//tensorflow/core/ir/types:Dialect",
+        "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/base:nullability",
+        "@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",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@stablehlo//:version",
+    ],
+)
+
+tf_cc_test(
+    name = "lift_as_function_call_test",
+    srcs = ["lift_as_function_call_test.cc"],
+    deps = [
+        ":attrs_and_constraints",
+        ":func",
+        ":lift_as_function_call",
+        ":test_base",
+        "//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "@com_google_absl//absl/algorithm:container",
+        "@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:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@stablehlo//:stablehlo_ops",
+        "@tsl//tsl/platform:protobuf",
+    ],
+)
+
+cc_library(
+    name = "func",
+    srcs = ["func.cc"],
+    hdrs = ["func.h"],
+    compatible_with = get_compatible_with_portable(),
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/cc/saved_model:signature_constants",
+        "//tensorflow/compiler/mlir/tensorflow:import_model",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
+tf_cc_test(
+    name = "func_test",
+    srcs = ["func_test.cc"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        ":func",
+        ":test_base",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+    ],
+)
+
+cc_library(
+    name = "test_base",
+    testonly = 1,
+    srcs = [],
+    hdrs = ["test_base.h"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        ":func",
+        "//tensorflow/compiler/mlir/quantization/common/ir:QuantOps",
+        "//tensorflow/compiler/mlir/quantization/stablehlo/cc:context",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
+        "//tensorflow/core:test",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Parser",
+        "@llvm-project//mlir:QuantOps",
+        "@llvm-project//mlir:Support",
+        "@stablehlo//:stablehlo_ops",
+    ],
+)
+
+cc_library(
+    name = "attrs_and_constraints",
+    srcs = [
+        "attrs_and_constraints.cc",
+    ],
+    hdrs = [
+        "attrs_and_constraints.h",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    visibility = ["//visibility:public"],
+    deps = [
+        ":tf_uniform_quantized_types",
+        "//tensorflow/compiler/mlir/quantization/common/quantization_lib",
+        "//tensorflow/compiler/mlir/quantization/tensorflow:quantization_options_proto_cc",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow:xla_call_module_attrs",
+        "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/status",
+        "@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",
+    ],
+)
+
+tf_cc_test(
+    name = "attrs_and_constraints_test",
+    srcs = ["attrs_and_constraints_test.cc"],
+    deps = [
+        ":attrs_and_constraints",
+        ":func",
+        ":test_base",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@stablehlo//:stablehlo_ops",
+    ],
+)
+
+td_library(
+    name = "quant_td_files",
+    srcs = [
+        "attrs_and_constraints.td",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        ":lift_as_function_call_td_files",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_ops_td_files",
+        "@llvm-project//mlir:ArithOpsTdFiles",
+        "@llvm-project//mlir:FuncTdFiles",
+    ],
+)
+
+cc_library(
+    name = "tf_uniform_quantized_types",
+    srcs = ["tf_uniform_quantized_types.cc"],
+    hdrs = ["tf_uniform_quantized_types.h"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:QuantOps",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
+cc_library(
+    name = "uniform_quantized_types",
+    srcs = ["uniform_quantized_types.cc"],
+    hdrs = ["uniform_quantized_types.h"],
+    compatible_with = get_compatible_with_portable(),
+    visibility = ["//visibility:public"],
+    deps = [
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:QuantOps",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
+tf_cc_test(
+    name = "uniform_quantized_types_test",
+    srcs = ["uniform_quantized_types_test.cc"],
+    deps = [
+        ":test_base",
+        ":uniform_quantized_types",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:QuantOps",
+        "@llvm-project//mlir:Support",
+        "@stablehlo//:stablehlo_ops",
+    ],
+)

tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.cc

@@ -0,0 +1,185 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.h"
+
+#include <cstdint>
+#include <optional>
+
+#include "absl/algorithm/container.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/MathExtras.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/IRMapping.h"  // from @llvm-project
+#include "mlir/IR/Operation.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  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/quantization/common/tf_uniform_quantized_types.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/xla_call_module_attrs.h"
+
+namespace mlir::quant {
+
+using ::mlir::stablehlo::DotGeneralOp;
+using tf_quant::IsQuantizedTensorType;
+
+bool HasStaticShape(Value value) {
+  auto shaped_type = mlir::dyn_cast<ShapedType>(value.getType());
+  if (!shaped_type) return false;
+
+  return shaped_type.hasStaticShape();
+}
+
+bool HasStaticShapeAtDims(Value value, const ArrayRef<int> dims) {
+  auto shaped_type = mlir::dyn_cast<ShapedType>(value.getType());
+  if (!shaped_type || !shaped_type.hasRank()) return false;
+
+  for (auto dim : dims) {
+    if (shaped_type.isDynamicDim(dim)) return false;
+  }
+  return true;
+}
+
+Type CloneTypeWithNewElementType(Type old_type, Type element_type) {
+  if (!mlir::isa<ShapedType>(old_type)) return {};
+
+  return mlir::cast<ShapedType>(old_type).clone(element_type);
+}
+
+SmallVector<Value> CloneOpWithReplacedOperands(
+    OpBuilder& builder, Operation* op, const ArrayRef<Value> new_operands) {
+  IRMapping mapping;
+  for (const auto& arg : enumerate(new_operands)) {
+    mapping.map(op->getOperand(arg.index()), arg.value());
+  }
+  return builder.clone(*op, mapping)->getResults();
+}
+
+FailureOr<int32_t> CastI64ToI32(const int64_t value) {
+  if (!llvm::isInt<32>(value)) {
+    DEBUG_WITH_TYPE(
+        "mlir-quant-attrs-and-constraints",
+        llvm::dbgs()
+            << "Tried to cast " << value
+            << "from int64 to int32, but lies out of range of int32.\n");
+    return failure();
+  }
+  return static_cast<int32_t>(value);
+}
+
+FailureOr<SmallVector<int32_t>> CastI64ArrayToI32(
+    const ArrayRef<int64_t> int64_array) {
+  SmallVector<int32_t> int32_array{};
+  int32_array.reserve(int64_array.size());
+
+  for (const int64_t i64 : int64_array) {
+    FailureOr<int32_t> cast_i32 = CastI64ToI32(i64);
+    if (failed(cast_i32)) return failure();
+
+    int32_array.push_back(*cast_i32);
+  }
+  return int32_array;
+}
+
+StringRef GetEntryFunctionName(TF::XlaCallModuleOp op) {
+  if (!op->hasAttrOfType<FlatSymbolRefAttr>(
+          TF::kStablehloEntryFunctionAttrName)) {
+    return StringRef();
+  }
+  return op
+      ->getAttrOfType<FlatSymbolRefAttr>(TF::kStablehloEntryFunctionAttrName)
+      .getValue();
+}
+
+bool IsHybridQuantizedOp(Operation* op) {
+  if ((op->getNumOperands() != 2 && op->getNumOperands() != 3) ||
+      op->getResultTypes().size() != 1) {
+    return false;
+  }
+  Type lhs_type = op->getOperand(0).getType();
+  Type rhs_type = op->getOperand(1).getType();
+  Type result_type = op->getResult(0).getType();
+  return !IsQuantizedTensorType(lhs_type) && IsQuantizedTensorType(rhs_type) &&
+         !IsQuantizedTensorType(result_type);
+}
+
+absl::StatusOr<bool> IsDotGeneralFullyConnected(DotGeneralOp dot_general_op) {
+  if (dot_general_op == nullptr)
+    return absl::InvalidArgumentError(
+        "Given dot_general op cannot be null when checking "
+        "`IsDotGeneralBatchMatmul`.");
+  const ::mlir::stablehlo::DotDimensionNumbersAttr dot_dimension_numbers =
+      dot_general_op.getDotDimensionNumbers();
+  const ArrayRef<int64_t> lhs_contracting_dims =
+      dot_dimension_numbers.getLhsContractingDimensions();
+  const ArrayRef<int64_t> rhs_contracting_dims =
+      dot_dimension_numbers.getRhsContractingDimensions();
+  const int64_t input_rank =
+      mlir::dyn_cast<ShapedType>(dot_general_op.getOperand(0).getType())
+          .getRank();
+  const int64_t filter_rank =
+      mlir::dyn_cast<ShapedType>(dot_general_op.getOperand(1).getType())
+          .getRank();
+  // The following conditions are such requirements:
+  //   - rank(lhs) is 1 or 2
+  //   - rank(rhs) = 2
+  //   - size(lhs_contracting_dimensions) = 1
+  //   - size(rhs_contracting_dimensions) = 1
+  //   - lhs_contracting_dimension = last dimension of lhs.
+  //   - `stablehlo.dot_general` should not have `lhs_batching_dim`.
+  //   - quantization_dimension(rhs) should not be in
+  //     `rhs_contracting_dimensions`.
+  // https://github.com/openxla/stablehlo/blob/main/docs/spec.md#dot_general
+  const bool has_proper_rank =
+      (input_rank == 1 || input_rank == 2) && filter_rank == 2;
+  const bool has_proper_contracting_dim =
+      lhs_contracting_dims.size() == 1 && rhs_contracting_dims.size() == 1 &&
+      lhs_contracting_dims[0] == input_rank - 1;
+  const bool is_not_batch_op =
+      dot_dimension_numbers.getLhsBatchingDimensions().empty();
+  const bool has_proper_quantization_dimension =
+      absl::c_find(rhs_contracting_dims, filter_rank) ==
+      rhs_contracting_dims.end();
+  return has_proper_rank && has_proper_contracting_dim && is_not_batch_op &&
+         has_proper_quantization_dimension;
+}
+
+std::optional<int64_t> GetDotGeneralQuantizationDim(
+    DotGeneralOp dot_general_op) {
+  if (dot_general_op == nullptr) return std::nullopt;
+  const int64_t filter_rank =
+      mlir::dyn_cast<ShapedType>(dot_general_op.getOperand(1).getType())
+          .getRank();
+
+  // To quantize rhs per-channel, we currently only consider the case where
+  // `stablehlo.dot_general` is legalizable to `tfl.fully_connected`.
+  const bool is_per_axis_quantizable =
+      IsDotGeneralFullyConnected(dot_general_op).value();
+  if (!is_per_axis_quantizable) return std::nullopt;
+  return filter_rank - 1;
+}
+
+bool ContainsConvOrDot(StringRef str) {
+  return str.contains("_conv") || str.contains("_dot_general");
+}
+
+}  // namespace mlir::quant

tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.h

@@ -0,0 +1,260 @@
+/* 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_COMMON_ATTRS_AND_CONSTRAINTS_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_ATTRS_AND_CONSTRAINTS_H_
+
+#include <array>
+#include <cstdint>
+#include <optional>
+#include <type_traits>
+
+#include "absl/status/statusor.h"
+#include "llvm/Support/Debug.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/Matchers.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.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/quantization_lib/quantization_utils.h"  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/xla_call_module_attrs.h"
+
+namespace mlir::quant {
+
+constexpr char kAttrMapAttribute[] = "attr_map";
+
+// Name of the string attribute attached to `XlaCallModuleOp`, which is the
+// textproto representation of `Method`.
+inline constexpr StringRef kQuantizationMethodAttr = "_quantization_method";
+
+// Permutation from the NHWC tensor format to NCHW. This is an inverse
+// permutation of `kNchwToNhwcPermutation`.
+inline constexpr std::array<int64_t, 4> kNhwcToNchwPermutation = {0, 3, 1, 2};
+
+// Permutation from the NCHW tensor format to NHWC. This is an inverse
+// permutation of `kNchwToNhwcPermutation`.
+inline constexpr std::array<int64_t, 4> kNchwToNhwcPermutation = {0, 2, 3, 1};
+
+// Permutation from the OIHW (== (output features, input features, height,
+// width)) tensor format to HWIO. This is commonly used to transpose convolution
+// weights represented as OIHW format to HWIO, which is more desirable for
+// certain downstream optimization passes (e.g. XLA).
+inline constexpr std::array<int64_t, 4> kOihwToHwioPermutation = {2, 3, 1, 0};
+
+// Returns true if the value has static shape.
+bool HasStaticShape(Value value);
+
+// Returns true if the value has static shape at given dims.
+bool HasStaticShapeAtDims(Value value, ArrayRef<int> dims);
+
+// Whether `value` has known rank of `rank`. Returns false when it is not a
+// `ShapedType` or its rank is unknown.
+inline bool HasRankOf(Value value, const int64_t rank) {
+  auto shaped_type = mlir::dyn_cast_or_null<ShapedType>(value.getType());
+  return shaped_type && shaped_type.hasRank() && shaped_type.getRank() == rank;
+}
+
+// Creates a new type that has the shape from the `old_type` and the element
+// type from the `element_type`.
+Type CloneTypeWithNewElementType(Type old_type, Type element_type);
+
+// Creates an array with integer/float type.
+template <typename T,
+          typename = std::enable_if_t<
+              (std::is_integral_v<T> || std::is_same_v<T, float>), void>>
+Value CreateConstValue(OpBuilder& builder, const Location loc,
+                       const SmallVector<int64_t>& shape,
+                       const SmallVector<T>& values) {
+  if constexpr (std::is_integral_v<T>) {
+    auto shape_type =
+        RankedTensorType::get(shape, builder.getIntegerType(sizeof(T) * 8));
+
+    const auto attr = DenseIntElementsAttr::get(shape_type, values);
+    return TF::ConstOp::create(builder, loc, attr);
+  }
+
+  const auto type = RankedTensorType::get(shape, builder.getF32Type());
+  const auto value_attr = DenseFPElementsAttr::get(type, values);
+  return TF::ConstOp::create(builder, loc, value_attr);
+}
+
+// Creates a 1D array with integer/float type.
+template <typename T>
+Value Create1DConstValue(OpBuilder& builder, const Location loc,
+                         const SmallVector<T>& values) {
+  return CreateConstValue<T>(builder, loc,
+                             {static_cast<int64_t>(values.size())}, values);
+}
+
+// Creates a scalar with integer / float type.
+template <typename T>
+Value CreateScalarConstValue(OpBuilder& builder, const Location loc,
+                             const T value) {
+  return CreateConstValue<T>(builder, loc, /*shape=*/{}, {value});
+}
+
+// Checks if the value is a constant and return its splat value.
+template <typename T,
+          typename = std::enable_if_t<
+              (std::is_integral_v<T> || std::is_same_v<T, float>), void>>
+bool GetSplatValue(Value value, T& splat_value) {
+  if constexpr (std::is_integral_v<T>) {
+    DenseIntElementsAttr value_attr;
+    if (!matchPattern(value, m_Constant(&value_attr)) ||
+        !value_attr.isSplat()) {
+      return false;
+    }
+    splat_value = value_attr.getSplatValue<T>();
+    return true;
+  }
+
+  DenseFPElementsAttr value_attr;
+  if (!matchPattern(value, m_Constant(&value_attr)) || !value_attr.isSplat()) {
+    return false;
+  }
+  splat_value = value_attr.getSplatValue<T>();
+  return true;
+}
+
+// Checks if the value is a constant and its splat value is equal to x.
+template <typename T>
+bool IsSplatValueEqual(Value value, const T x) {
+  T splat_value;
+  if (!GetSplatValue(value, splat_value)) return false;
+
+  return splat_value == x;
+}
+
+// Checks if two values are constants and their splat values are equal.
+template <typename T>
+bool AreSplatValuesEqual(Value x, Value y) {
+  T splat_x, splat_y;
+  if (!GetSplatValue(x, splat_x) || !GetSplatValue(y, splat_y)) {
+    return false;
+  }
+
+  return splat_x == splat_y;
+}
+
+// Clones an operation with new operands while keeping attributes.
+SmallVector<Value> CloneOpWithReplacedOperands(OpBuilder& builder,
+                                               Operation* op,
+                                               ArrayRef<Value> new_operands);
+
+// Tries casting `op` with a concrete op type `T`. If the cast fails or `op` is
+// a `nullptr`, returns `failure` and prints a debugging message identifying
+// the cast attempt as `name`.
+template <typename T>
+FailureOr<T> TryCast(Operation* op, const StringRef name) {
+  auto cast_op = dyn_cast_or_null<T>(op);
+  if (cast_op) {
+    return cast_op;
+  } else {
+    DEBUG_WITH_TYPE("mlir-quant-attrs-and-constraints",
+                    llvm::dbgs() << "Failed to match " << name << " ("
+                                 << T::getOperationName() << ").\n");
+    return failure();
+  }
+}
+
+FailureOr<int32_t> CastI64ToI32(int64_t value);
+
+// Tries to cast an array of int64 to int32. If any of the element in the
+// array is not in the range of int32, returns failure().
+FailureOr<SmallVector<int32_t>> CastI64ArrayToI32(
+    ArrayRef<int64_t> int64_array);
+
+// Returns the first operation with the given type in the function.
+template <typename OpType>
+OpType FindOperationOfType(func::FuncOp function) {
+  for (auto op : function.getBody().getOps<OpType>()) {
+    return op;
+  }
+  return nullptr;
+}
+
+// Returns the first user of the given operation, optionally of the given
+// type if provided. If there is no user or user of type, return nullptr.
+template <typename T = Operation*>
+Operation* FindUserOfType(Operation* op) {
+  for (Operation* user : op->getUsers()) {
+    if (isa<T>(user)) {
+      return user;
+    }
+  }
+  return nullptr;
+}
+
+// Returns the first user of the given operation, optionally of the given
+// type if provided. If there is no user or user of type, return nullptr.
+template <typename T = Operation*>
+Operation* FindOperandOfType(Operation* op) {
+  for (Value operand_value : op->getOperands()) {
+    if (isa<T>(operand_value.getDefiningOp())) {
+      return operand_value.getDefiningOp();
+    }
+  }
+  return nullptr;
+}
+
+// Returns the function attribute for the given call op which is lifted for
+// quantization.
+inline FlatSymbolRefAttr GetFuncAttr(TF::PartitionedCallOp call_op) {
+  return mlir::dyn_cast<FlatSymbolRefAttr>(call_op.getFAttr());
+}
+
+inline FlatSymbolRefAttr GetFuncAttr(TF::XlaCallModuleOp call_op) {
+  return call_op->getAttrOfType<FlatSymbolRefAttr>(
+      TF::kStablehloEntryFunctionAttrName);
+}
+
+// Returns the entry function name for the given tf.XlaCallModule op. Returns
+// empty string if such attribute does not exist.
+StringRef GetEntryFunctionName(TF::XlaCallModuleOp op);
+
+// Checks whether the given op contains QuantizationTrait::FullyQuantizable.
+inline bool HasQuantizableTrait(Operation* op) {
+  return op->hasAttrOfType<StringAttr>(kQuantTraitAttrName) &&
+         op->getAttrOfType<StringAttr>(kQuantTraitAttrName).getValue().str() ==
+             QuantTraitValues[QuantizationTrait::FullyQuantizable];
+}
+
+// Returns true if `op` has two operands and one result and only second operand
+// is quantized.
+bool IsHybridQuantizedOp(Operation* op);
+
+// Returns whether a given `stablehlo.dot_general` can be legalizable to
+// `tfl.fully_connected`.
+absl::StatusOr<bool> IsDotGeneralFullyConnected(
+    ::mlir::stablehlo::DotGeneralOp dot_general_op);
+
+// Returns the quantization dimension for a given `stablehlo.dot_general` op,
+// or `std::nullopt` if the given op is not per-channel quantizable.
+std::optional<int64_t> GetDotGeneralQuantizationDim(
+    ::mlir::stablehlo::DotGeneralOp dot_general_op);
+
+// Checks if a `StringRef` contains 'conv' or 'dot_general'.
+bool ContainsConvOrDot(StringRef str);
+
+}  // namespace mlir::quant
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_ATTRS_AND_CONSTRAINTS_H_

tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.td

@@ -0,0 +1,169 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+include "mlir/IR/OpBase.td"
+include "mlir/IR/PatternBase.td"
+include "tensorflow/compiler/mlir/tensorflow/ir/tf_op_base.td"
+
+def DenseElementsAttr : ElementsAttrBase<
+  CPred<"llvm::isa<DenseElementsAttr>($_self)">,
+  "non-opaque constant tensor">;
+
+// Checks if the data format is "NHWC".
+def IsDataFormatNHWC : ConstantAttr<TF_ConvnetDataFormatAttr, "\"NHWC\"">;
+
+// Checks if the data format is "NDHWC".
+def IsDataFormatNDHWC : ConstantAttr<TF_ConvnetDataFormatAttr, "\"NDHWC\"">;
+
+// Checks if the op is constant op.
+def IsConstTensor :  Constraint<CPred<"dyn_cast_or_null<TF::ConstOp>($0.getDefiningOp())">>;
+
+// Checks if the element value has a float type.
+def IsFloatElementsAttr : ElementsAttrBase<
+  CPred<"llvm::isa<ElementsAttr>($_self) && "
+        "llvm::isa<FloatType>(getElementTypeOrSelf(llvm::cast<ElementsAttr>($_self).getType()))">,
+        "float constant tensor">;
+
+// Checks if the boolean value is false.
+def IsFalseBoolAttr : AttrConstraint<
+  CPred<"!llvm::cast<BoolAttr>($_self).getValue()">>;
+
+// Checks if the value has only one user.
+def HasOneUse : Constraint<CPred<"$0.hasOneUse()">>;
+
+// Gets the type of a value.
+def GetValueType : NativeCodeCall<"$0.getType()">;
+
+// Checks if the value has the type of int8.
+def IsInt8ElementType : Constraint<
+  CPred<"getElementTypeOrSelf($0).isInteger(8)">>;
+
+// Checks if the value has the type of int32.
+def IsInt32ElementType : Constraint<
+  CPred<"getElementTypeOrSelf($0).isInteger(32)">>;
+
+// Checks if the value has the type of float32.
+def IsF32ElementType : Constraint<
+  CPred<"getElementTypeOrSelf($0).isF32()">>;
+
+// Checks if the value has the type of bfloat16.
+def IsBF16ElementType : Constraint<
+  CPred<"getElementTypeOrSelf($0).isBF16()">>;
+
+// Checks if the value has the type of UniformQuantizedType.
+def IsUniformQuantizedType : Constraint<
+  CPred<"llvm::isa<mlir::quant::UniformQuantizedType>(getElementTypeOrSelf($0))">>;
+
+// Checks if the given two values have the same type.
+def AreTheSameElementType : Constraint<
+  CPred<"$0.getType() == $1.getType()">>;
+
+// Checks if the given two values are the same.
+def AreTheSameValue : Constraint<
+  CPred<"$0 == $1">>;
+
+// Checks if the value has rank.
+def HasRank : Constraint<
+  CPred<"llvm::cast<ShapedType>($0.getType()).hasRank()">>;
+
+// Checks if the value has rank of `n`.
+class HasRankOf<int n> : Constraint<
+  CPred<"llvm::cast<ShapedType>($0.getType()).hasRank() && "
+        "llvm::cast<ShapedType>($0.getType()).getRank() == " # n>,
+  "Checks if the value has rank of 'n'.">;
+
+// Checks if the value has static shape.
+def HasStaticShapeConstraint : Constraint<CPred<"HasStaticShape($0)">>;
+
+// Checks if the value has static shape at given dims.
+class HasStaticShapeAtDimsConstraint<string dims> : Constraint<
+  CPred<"HasStaticShapeAtDims($0, {"# dims #"})">>;
+
+// The rewrite rule cannot replace a value with itself, so we work around
+// by cloning the root op to replicate that value. The old op will get folded.
+def CloningOpResult : NativeCodeCall<
+  "$_builder.clone(*op0)->getOpResult(0)">;
+
+// Same as CloningOpResult but is used for ops with multiple results.
+class CloningOpResults<int returns> : NativeCodeCall<
+  "$_builder.clone(*op0)->getOpResults()", returns>;
+
+// Creates an 1D array const with float values.
+class Create1DConst<string values> : NativeCodeCall<
+  "Create1DConstValue<float>($_builder, $_loc, "# values #")">;
+
+// Creates a scalar const with float value.
+class CreateScalarConst<string value> : NativeCodeCall<
+  "CreateScalarConstValue<float>($_builder, $_loc, "# value #")">;
+
+// Creates an 1D array const with integer values.
+// TODO: b/239490133 - Make the rule name and function name consistent.
+class Create1DIntegerConst<string type, string values> : NativeCodeCall<
+  "Create1DConstValue<"# type #">($_builder, $_loc, "# values #")">;
+
+// Creates a scalar const with integer value.
+class CreateScalarIntegerConst<string type, string value> : NativeCodeCall<
+  "CreateScalarConstValue<"# type #">($_builder, $_loc, "# value #")">;
+
+// Creates an I64 array attribute with given values.
+class CreateI64ArrayAttr<string values> : NativeCodeCall<
+  "$_builder.getI64ArrayAttr("# values #")">;
+
+// Creates a string attribute with given values.
+class CreateStringAttr<string values> : NativeCodeCall<
+  "$_builder.getStringAttr(\""# values #"\")">;
+
+// Creates a new F32 type with the same shape as the given value.
+def CloneTypeWithF32ElementType : NativeCodeCall<
+  "CloneTypeWithNewElementType($0.getType(), $_builder.getF32Type())">;
+
+// Creates a new I32 type with the same shape as the given value.
+def CloneTypeWithI32ElementType : NativeCodeCall<
+  "CloneTypeWithNewElementType($0.getType(), $_builder.getI32Type())">;
+
+// Checks if the value is a float constant and its splat value is equal to `x`.
+class IsSplatValueEqual<string x> : Constraint<CPred<
+  "IsSplatValueEqual<float>($0, "# x #")">>;
+
+// Checks if two values are float constants and their values are equal.
+def AreSplatValuesEqual : Constraint<CPred<
+  "AreSplatValuesEqual<float>($0, $1)">>;
+
+// Checks if the value is an integer constant and its splat value is equal to x.
+class IsIntSplatValueEqual<string type, string x> : Constraint<CPred<
+  "IsSplatValueEqual<"# type #">($0, "# x #")">>;
+
+// Checks if two values are integer constants and their values are equal.
+class AreIntSplatValuesEqual<string type> : Constraint<CPred<
+  "AreSplatValuesEqual<"# type #">($0, $1)">>;
+
+// Returns defining op of this value.
+def GetDefiningOp :  NativeCodeCall<"$0.getDefiningOp()">;
+
+// Clones an operation with new operands while keeping attributes.
+def CloneOpWithReplacedOperands : NativeCodeCall<
+  "CloneOpWithReplacedOperands("
+    "$_builder, $0, llvm::SmallVector<Value>{$1...}).front()">;
+
+// Checks whether the value of a constant equals the given float, regardless
+// of the tensor dimension.
+class FloatValueEquals<string val> : Constraint<CPred<
+  "FloatValueEquals($0, " # val # ")">>;
+
+// Fetches the default or null attribute, used for pattern matching.
+def DefaultOrNullAttr : NativeCodeCall<"DefaultOrNullAttr($_builder, $0)">;
+
+// Returns true if the given op is a StableHLO constant op.
+def IsStableHLOConstantOp : Constraint<CPred<"dyn_cast_or_null<::mlir::stablehlo::ConstantOp>($0.getDefiningOp())">>;
+

tensorflow/compiler/mlir/quantization/common/attrs_and_constraints_test.cc

@@ -0,0 +1,570 @@
+/* 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/common/attrs_and_constraints.h"
+
+#include <cstdint>
+#include <optional>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "absl/status/status_matchers.h"
+#include "absl/strings/string_view.h"
+#include "llvm/Support/MathExtras.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/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/func.h"
+#include "tensorflow/compiler/mlir/quantization/common/test_base.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+
+namespace mlir::quant {
+namespace {
+
+using ::mlir::stablehlo::AddOp;
+using ::mlir::stablehlo::ConstantOp;
+using ::mlir::stablehlo::ConvolutionOp;
+using ::mlir::stablehlo::DotGeneralOp;
+using ::mlir::stablehlo::SubtractOp;
+using ::testing::ElementsAreArray;
+using ::testing::Eq;
+using ::testing::IsEmpty;
+using ::testing::IsNull;
+using ::testing::NotNull;
+using ::testing::Optional;
+
+using AttrsAndConstraintsTest = ::mlir::quant::QuantizationTestBase;
+
+constexpr absl::string_view kModuleStatic = R"mlir(
+  module {
+    func.func @main(%arg0: tensor<1x1024xf32>, %arg1: tensor<1024x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
+      %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], precision = [] : (tensor<1x1024xf32>, tensor<1024x3xf32>) -> tensor<1x3xf32>
+      return %0 : tensor<1x3xf32>
+    }
+  }
+)mlir";
+
+constexpr absl::string_view kModuleDynamic = R"mlir(
+  module {
+    func.func @main(%arg0: tensor<?x1024xf32>, %arg1: tensor<1024x3xf32>) -> tensor<?x3xf32> attributes {_from_xla_call_module} {
+      %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], precision = [] : (tensor<?x1024xf32>, tensor<1024x3xf32>) -> tensor<?x3xf32>
+      return %0 : tensor<?x3xf32>
+    }
+  }
+)mlir";
+
+constexpr absl::string_view kModuleMultipleUses = R"mlir(
+  module {
+    func.func @main(%arg0: tensor<1x1024xf32>, %arg1: tensor<1024x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
+      %cst = stablehlo.constant dense<1.0> : tensor<1x3xf32>
+      %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], precision = [] : (tensor<1x1024xf32>, tensor<1024x3xf32>) -> tensor<1x3xf32>
+      %1 = stablehlo.subtract %cst, %0 : tensor<1x3xf32>
+      %2 = stablehlo.add %0, %cst : tensor<1x3xf32>
+      return %2 : tensor<1x3xf32>
+    }
+  }
+)mlir";
+
+constexpr absl::string_view kModuleXlaCallModule = R"mlir(
+  module {
+    func.func @main(%arg0: tensor<?x2xf32> {tf_saved_model.index_path = ["input_tensor"]}) -> (tensor<?x2xf32>) {
+      %0 = stablehlo.constant dense<[-0.211145893, -0.708605706]> : tensor<2xf32>
+      %1 = stablehlo.constant dense<[[-0.630731344, 0.54962182], [0.180364341, -0.764542698]]> : tensor<2x2xf32>
+      %2 = "tf.XlaCallModule"(%arg0, %1, %0) <{Sout = [#tf_type.shape<?x2>], module = "", version = 9 : i64}> {_entry_function = @composite_fn_1, _stablehlo_version = "1.0.0", _original_entry_function = "composite_fn_1", _tfl_quant_trait = "fully_quantizable"} : (tensor<?x2xf32>, tensor<2x2xf32>, tensor<2xf32>) -> tensor<?x2xf32>
+      return %2 : tensor<?x2xf32>
+    }
+    func.func private @composite_fn_1(%arg0: tensor<?x2xf32>, %arg1: tensor<2x2xf32>, %arg2: tensor<2xf32>) -> tensor<?x2xf32> attributes {_from_xla_call_module, tf_quant.composite_function} {
+      return %arg0 : tensor<?x2xf32>
+    }
+  }
+)mlir";
+
+constexpr absl::string_view kModuleDotWeightOnlyPtq = R"mlir(
+  module {
+    func.func @main(%arg0: tensor<?x2xf32> {tf_saved_model.index_path = ["input_tensor"]}) -> (tensor<?x2xf32>) {
+      %0 = stablehlo.constant dense<[-0.211145893, -0.708605706]> : tensor<2xf32>
+      %1 = stablehlo.constant dense<[[-0.630731344, 0.54962182], [0.180364341, -0.764542698]]> : tensor<2x2xf32>
+      %2 = "tf.XlaCallModule"(%arg0, %1, %0) <{Sout = [#tf_type.shape<?x2>], module = "", version = 9 : i64}> {_entry_function = @composite_dot_general_fn_1, _stablehlo_version = "1.0.0", _original_entry_function = "composite_dot_general_fn_1", _tfl_quant_trait = "fully_quantizable", _quantization_method = "weight_only_ptq { }"} : (tensor<?x2xf32>, tensor<2x2xf32>, tensor<2xf32>) -> tensor<?x2xf32>
+      return %2 : tensor<?x2xf32>
+    }
+    func.func private @composite_dot_general_fn_1(%arg0: tensor<?x2xf32>, %arg1: tensor<2x2xf32>, %arg2: tensor<2xf32>) -> tensor<?x2xf32> attributes {_from_xla_call_module, tf_quant.composite_function} {
+      %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<?x2xf32>, tensor<2x2xf32>) -> tensor<?x2xf32>
+      return %0 : tensor<?x2xf32>
+    }
+  }
+)mlir";
+
+constexpr absl::string_view kModuleXlaCallModuleNoEntryNoQuantTrait = R"mlir(
+  module {
+    func.func @main(%arg0: tensor<?x2xf32> {tf_saved_model.index_path = ["input_tensor"]}) -> (tensor<?x2xf32>) {
+      %0 = stablehlo.constant dense<[-0.211145893, -0.708605706]> : tensor<2xf32>
+      %1 = stablehlo.constant dense<[[-0.630731344, 0.54962182], [0.180364341, -0.764542698]]> : tensor<2x2xf32>
+      %2 = "tf.XlaCallModule"(%arg0, %1, %0) <{Sout = [#tf_type.shape<?x2>], module = "", version = 9 : i64}> {_original_entry_function = "composite_fn_1"} : (tensor<?x2xf32>, tensor<2x2xf32>, tensor<2xf32>) -> tensor<?x2xf32>
+      return %2 : tensor<?x2xf32>
+    }
+    func.func private @composite_fn_1(%arg0: tensor<?x2xf32>, %arg1: tensor<2x2xf32>, %arg2: tensor<2xf32>) -> tensor<?x2xf32> attributes {_from_xla_call_module, tf_quant.composite_function} {
+      return %arg0 : tensor<?x2xf32>
+    }
+  }
+)mlir";
+
+constexpr absl::string_view kModulePartitionedCall = R"mlir(
+  module {
+    func.func @main(%arg0: tensor<2x2xf32> {tf_saved_model.index_path = ["input_tensor"]}) -> (tensor<2x2xf32>) {
+      %cst = "tf.Const"() {device = "", value = dense<[[-0.630731344, 0.54962182], [0.180364341, -0.764542698]]> : tensor<2x2xf32>} : () -> tensor<2x2xf32>
+      %0 = "tf.PartitionedCall"(%arg0, %cst) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_fn_1} : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32> loc(callsite("test@main"("MatMul") at "QuantizationUnit(\12\06MatMul\1a\07main)"))
+      return %0 : tensor<2x2xf32>
+    }
+    func.func private @composite_fn_1(%arg0: tensor<2x2xf32>, %arg1: tensor<2x2xf32>) -> tensor<2x2xf32> attributes {tf_quant.composite_function} {
+      %0 = "tf.MatMul"(%arg0, %arg1) {attr_map = "0:transpose_a,1:transpose_b", device = "", transpose_a = false, transpose_b = false} : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
+      return %0 : tensor<2x2xf32>
+    }
+  }
+)mlir";
+
+constexpr absl::string_view kModuleHybridQuantized = R"mlir(
+  module {
+    func.func @main(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3x!quant.uniform<i8:f32, 6.000000e-03:0>> {tf_saved_model.index_path = ["input_tensor"]}) -> (tensor<1x3xf32>) {
+      %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3x!quant.uniform<i8:f32, 6.000000e-03:0>>) -> tensor<1x3xf32>
+      return %0 : tensor<1x3xf32>
+    }
+  }
+)mlir";
+
+TEST_F(AttrsAndConstraintsTest, HasStaticShapeSucceedsWithStaticShapes) {
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleStatic);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  Value dot_general_result =
+      FindOperationOfType<DotGeneralOp>(main_fn)->getResult(0);
+  EXPECT_TRUE(HasStaticShape(dot_general_result));
+  EXPECT_TRUE(HasStaticShapeAtDims(dot_general_result, /*dims=*/{0}));
+  EXPECT_TRUE(HasStaticShapeAtDims(dot_general_result, /*dims=*/{1}));
+}
+
+TEST_F(AttrsAndConstraintsTest, HasStaticShapeFailsWithDynamicShapes) {
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleDynamic);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  Value dot_general_result =
+      FindOperationOfType<DotGeneralOp>(main_fn)->getResult(0);
+  EXPECT_FALSE(HasStaticShape(dot_general_result));
+  EXPECT_FALSE(HasStaticShapeAtDims(dot_general_result, /*dims=*/{0}));
+  EXPECT_TRUE(HasStaticShapeAtDims(dot_general_result, /*dims=*/{1}));
+}
+
+TEST_F(AttrsAndConstraintsTest, HasRankOfReturnsTrueForMatchingRank) {
+  constexpr absl::string_view kConstantOpWithRankFour =
+      R"mlir(%0 = stablehlo.constant dense<0> : tensor<1x1x1x1xi8>)mlir";
+  OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kConstantOpWithRankFour);
+  ASSERT_TRUE(module_op);
+
+  ASSERT_FALSE(module_op->getBodyRegion().empty());
+  ASSERT_FALSE(module_op->getBodyRegion().front().empty());
+
+  auto constant_op = dyn_cast_or_null<mlir::stablehlo::ConstantOp>(
+      module_op->getBodyRegion().front().front());
+  ASSERT_THAT(constant_op, NotNull());
+
+  EXPECT_TRUE(HasRankOf(constant_op, /*rank=*/4));
+}
+
+TEST_F(AttrsAndConstraintsTest, HasRankOfReturnsFalseForNonMatchingRank) {
+  constexpr absl::string_view kConstantOpWithRankFour =
+      R"mlir(%0 = stablehlo.constant dense<0> : tensor<1x1x1x1xi8>)mlir";
+  OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kConstantOpWithRankFour);
+  ASSERT_TRUE(module_op);
+
+  ASSERT_FALSE(module_op->getBodyRegion().empty());
+  ASSERT_FALSE(module_op->getBodyRegion().front().empty());
+
+  auto constant_op = dyn_cast_or_null<mlir::stablehlo::ConstantOp>(
+      module_op->getBodyRegion().front().front());
+  ASSERT_THAT(constant_op, NotNull());
+
+  EXPECT_FALSE(HasRankOf(constant_op, /*rank=*/3));
+}
+
+TEST_F(AttrsAndConstraintsTest,
+       HasRankOfReturnsTrueForMatchingRankWithUnknownDimensions) {
+  // Argument has rank 2, but its dimensions are unknown.
+  constexpr absl::string_view kArgumentWithUnknownDims = R"mlir(
+    func.func @unknown_dims_arg(%arg: tensor<?x?xi8>) -> tensor<?x?xi8> {
+      return %arg : tensor<?x?xi8>
+    }
+  )mlir";
+
+  OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kArgumentWithUnknownDims);
+  ASSERT_TRUE(module_op);
+
+  auto func_op = module_op->lookupSymbol<func::FuncOp>("unknown_dims_arg");
+  ASSERT_THAT(func_op, NotNull());
+  ASSERT_THAT(func_op.getNumArguments(), Eq(1));
+
+  EXPECT_TRUE(HasRankOf(func_op.getArgument(0), /*rank=*/2));
+}
+
+TEST_F(AttrsAndConstraintsTest, HasRankOfReturnsFalseForUnknownRank) {
+  constexpr absl::string_view kArgumentWithUnknownRank = R"mlir(
+    func.func @unknown_rank_arg(%arg: tensor<*xi8>) -> tensor<*xi8> {
+      return %arg : tensor<*xi8>
+    }
+  )mlir";
+
+  OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kArgumentWithUnknownRank);
+  ASSERT_TRUE(module_op);
+
+  auto func_op = module_op->lookupSymbol<func::FuncOp>("unknown_rank_arg");
+  ASSERT_THAT(func_op, NotNull());
+  ASSERT_THAT(func_op.getNumArguments(), Eq(1));
+
+  EXPECT_FALSE(HasRankOf(func_op.getArgument(0), /*rank=*/1));
+}
+
+TEST_F(AttrsAndConstraintsTest, TryCastSucceeds) {
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleStatic);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  auto dot_general_op = FindOperationOfType<DotGeneralOp>(main_fn);
+  ASSERT_THAT(dot_general_op, NotNull());
+
+  EXPECT_TRUE(succeeded(
+      TryCast<DotGeneralOp>(dot_general_op, /*name=*/"dot_general_op")));
+}
+
+TEST_F(AttrsAndConstraintsTest, TryCastFailsOnWrongType) {
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleStatic);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  auto dot_general_op = FindOperationOfType<DotGeneralOp>(main_fn);
+  ASSERT_THAT(dot_general_op, NotNull());
+
+  EXPECT_TRUE(
+      failed(TryCast<AddOp>(dot_general_op, /*name=*/"dot_general_op")));
+}
+
+TEST_F(AttrsAndConstraintsTest, TryCastFailsOnNullPtr) {
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleStatic);
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  auto op_nullptr =
+      FindOperationOfType<DotGeneralOp>(main_fn)->getNextNode()->getNextNode();
+  // getNextNode() returns a nullptr if at the very last node.
+  EXPECT_THAT(op_nullptr, IsNull());
+
+  EXPECT_TRUE(failed(TryCast<DotGeneralOp>(op_nullptr, /*name=*/"op_nullptr")));
+  EXPECT_TRUE(failed(TryCast<DotGeneralOp>(nullptr, /*name=*/"nullptr")));
+}
+
+TEST_F(AttrsAndConstraintsTest, I64ValueInI32RangeAreCastedCorrectly) {
+  EXPECT_TRUE(succeeded(CastI64ToI32(llvm::minIntN(32))));
+  EXPECT_TRUE(succeeded(CastI64ToI32(llvm::maxIntN(32))));
+}
+
+TEST_F(AttrsAndConstraintsTest, CastingFailsForI64ValueOutOfI32Range) {
+  EXPECT_TRUE(failed(CastI64ToI32(llvm::minIntN(32) - 10)));
+  EXPECT_TRUE(failed(CastI64ToI32(llvm::maxIntN(32) + 10)));
+}
+
+TEST_F(AttrsAndConstraintsTest, I64ArrayInI32RangeAreCastedCorrectly) {
+  const SmallVector<int64_t> array_i64 = {llvm::minIntN(32), -2, -1, 0, 1, 2,
+                                          llvm::maxIntN(32)};
+
+  FailureOr<SmallVector<int32_t>> array_i32 = CastI64ArrayToI32(array_i64);
+  EXPECT_TRUE(succeeded(array_i32));
+  EXPECT_THAT(
+      *array_i32,
+      ElementsAreArray({static_cast<int32_t>(llvm::minIntN(32)), -2, -1, 0, 1,
+                        2, static_cast<int32_t>(llvm::maxIntN(32))}));
+}
+
+TEST_F(AttrsAndConstraintsTest, CastingFailsForI64ArrayUnderI32Range) {
+  const int64_t under_min_i32 = -2147483658;
+  ArrayRef<int64_t> array_i64(under_min_i32);
+
+  EXPECT_EQ(under_min_i32, llvm::minIntN(32) - 10);
+  EXPECT_TRUE(failed(CastI64ArrayToI32(array_i64)));
+}
+
+TEST_F(AttrsAndConstraintsTest, CastingFailsForI64ArrayAboveI32Range) {
+  const int64_t below_max_i32 = 2147483657;
+  ArrayRef<int64_t> array_i64(below_max_i32);
+
+  EXPECT_EQ(below_max_i32, llvm::maxIntN(32) + 10);
+  EXPECT_TRUE(failed(CastI64ArrayToI32(array_i64)));
+}
+
+TEST_F(AttrsAndConstraintsTest, FindUserOfDifferentTypes) {
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleMultipleUses);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  auto dot_general_op = FindOperationOfType<DotGeneralOp>(main_fn);
+  ASSERT_THAT(dot_general_op, NotNull());
+
+  EXPECT_THAT(FindUserOfType<AddOp>(dot_general_op), NotNull());
+  EXPECT_THAT(FindUserOfType<SubtractOp>(dot_general_op), NotNull());
+  EXPECT_THAT(FindUserOfType<>(dot_general_op), NotNull());
+  EXPECT_THAT(FindUserOfType<ConvolutionOp>(dot_general_op), IsNull());
+}
+
+TEST_F(AttrsAndConstraintsTest, FindOperandOfDifferentTypes) {
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleMultipleUses);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  auto subtract_op = FindOperationOfType<SubtractOp>(main_fn);
+  ASSERT_THAT(subtract_op, NotNull());
+
+  EXPECT_THAT(FindOperandOfType<DotGeneralOp>(subtract_op), NotNull());
+  EXPECT_THAT(FindOperandOfType<ConstantOp>(subtract_op), NotNull());
+  EXPECT_THAT(FindOperandOfType<>(subtract_op), NotNull());
+  EXPECT_THAT(FindOperandOfType<AddOp>(subtract_op), IsNull());
+}
+
+TEST_F(AttrsAndConstraintsTest, XlaCallModuleOpGetFuncAttr) {
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleXlaCallModule);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  auto xla_call_module_op = FindOperationOfType<TF::XlaCallModuleOp>(main_fn);
+  ASSERT_THAT(xla_call_module_op, NotNull());
+
+  FlatSymbolRefAttr xla_call_op_attr = GetFuncAttr(xla_call_module_op);
+  EXPECT_EQ(xla_call_op_attr.getValue(), "composite_fn_1");
+}
+
+TEST_F(AttrsAndConstraintsTest, PartitionedCallGetFuncAttr) {
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModulePartitionedCall);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  auto partitioned_call_op =
+      FindOperationOfType<TF::PartitionedCallOp>(main_fn);
+  ASSERT_THAT(partitioned_call_op, NotNull());
+
+  FlatSymbolRefAttr partitioned_call_op_attr = GetFuncAttr(partitioned_call_op);
+  EXPECT_EQ(partitioned_call_op_attr.getValue(), "composite_fn_1");
+}
+
+TEST_F(AttrsAndConstraintsTest, GetEntryFunctionNameCorrectly) {
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleXlaCallModule);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  auto xla_call_module_op = FindOperationOfType<TF::XlaCallModuleOp>(main_fn);
+  ASSERT_THAT(xla_call_module_op, NotNull());
+
+  EXPECT_EQ(GetEntryFunctionName(xla_call_module_op),
+            StringRef("composite_fn_1"));
+}
+
+TEST_F(AttrsAndConstraintsTest, GetEntryFunctionNameWhenNotSet) {
+  OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kModuleXlaCallModuleNoEntryNoQuantTrait);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  auto xla_call_module_op = FindOperationOfType<TF::XlaCallModuleOp>(main_fn);
+  ASSERT_THAT(xla_call_module_op, NotNull());
+
+  EXPECT_THAT(GetEntryFunctionName(xla_call_module_op), IsEmpty());
+}
+
+TEST_F(AttrsAndConstraintsTest, HasQuantizableTraitTrue) {
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleXlaCallModule);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  auto xla_call_module_op = FindOperationOfType<TF::XlaCallModuleOp>(main_fn);
+  ASSERT_THAT(xla_call_module_op, NotNull());
+
+  EXPECT_TRUE(HasQuantizableTrait(xla_call_module_op));
+}
+
+TEST_F(AttrsAndConstraintsTest, HasQuantizableTraitFalse) {
+  OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kModuleXlaCallModuleNoEntryNoQuantTrait);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  auto xla_call_module_op = FindOperationOfType<TF::XlaCallModuleOp>(main_fn);
+  ASSERT_THAT(xla_call_module_op, NotNull());
+
+  EXPECT_FALSE(HasQuantizableTrait(xla_call_module_op));
+}
+
+TEST_F(AttrsAndConstraintsTest, IsHybridQuantizedOpTrue) {
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleHybridQuantized);
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  Operation* dot_general = FindOperationOfType<DotGeneralOp>(main_fn);
+  EXPECT_TRUE(IsHybridQuantizedOp(dot_general));
+}
+
+TEST_F(AttrsAndConstraintsTest, IsHybridQuantizedOpFalse) {
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleXlaCallModule);
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  Operation* call_op = FindOperationOfType<TF::XlaCallModuleOp>(main_fn);
+  EXPECT_FALSE(IsHybridQuantizedOp(call_op));
+}
+
+constexpr absl::string_view kModuleDotGeneralFullyConnected = R"mlir(
+  module {
+    func.func @main(%arg0: tensor<1x1024xf32>, %arg1: tensor<1024x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
+      %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], precision = [] : (tensor<1x1024xf32>, tensor<1024x3xf32>) -> tensor<1x3xf32>
+      return %0 : tensor<1x3xf32>
+    }
+  }
+)mlir";
+
+constexpr absl::string_view kModuleDotGeneralBatchMatmul = R"mlir(
+  module {
+    func.func @main(%arg0: tensor<2x2x2xf32>, %arg1: tensor<2x2x2xf32>) -> tensor<2x2x2xf32> attributes {_from_xla_call_module} {
+      %0 = stablehlo.dot_general %arg0, %arg1,
+        batching_dims = [0] x [0],
+        contracting_dims = [2] x [1],
+        precision = [DEFAULT, DEFAULT]
+      : (tensor<2x2x2xf32>, tensor<2x2x2xf32>) -> tensor<2x2x2xf32>
+        return %0 : tensor<2x2x2xf32>
+    }
+  }
+)mlir";
+
+TEST_F(AttrsAndConstraintsTest, IsDotGeneralFullyConnectedReturnsError) {
+  DotGeneralOp dot_general_op = nullptr;
+  absl_testing::StatusIs(absl::StatusCode::kInvalidArgument,
+                         "Given dot_general op cannot be null when checking "
+                         "`IsDotGeneralBatchMatmul`");
+}
+
+TEST_F(AttrsAndConstraintsTest, IsDotGeneralFullyConnectedReturnsTrue) {
+  OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kModuleDotGeneralFullyConnected);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  auto dot_general_op = *main_fn.getOps<DotGeneralOp>().begin();
+  EXPECT_THAT(IsDotGeneralFullyConnected(dot_general_op), true);
+}
+
+TEST_F(AttrsAndConstraintsTest, IsDotGeneralFullyConnectedReturnsFalse) {
+  OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kModuleDotGeneralBatchMatmul);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  auto dot_general_op = *main_fn.getOps<DotGeneralOp>().begin();
+  EXPECT_THAT(IsDotGeneralFullyConnected(dot_general_op), false);
+}
+
+TEST_F(AttrsAndConstraintsTest, DotGeneralFullyConnectedReturnsQuantDim) {
+  OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kModuleDotGeneralFullyConnected);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  auto dot_general_op = *main_fn.getOps<DotGeneralOp>().begin();
+  EXPECT_THAT(GetDotGeneralQuantizationDim(dot_general_op), Optional(1));
+}
+
+TEST_F(AttrsAndConstraintsTest, DotGeneralBatchMatmulReturnsNullQuantDim) {
+  OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kModuleDotGeneralBatchMatmul);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  auto dot_general_op = *main_fn.getOps<DotGeneralOp>().begin();
+  EXPECT_THAT(GetDotGeneralQuantizationDim(dot_general_op), Eq(std::nullopt));
+}
+
+TEST_F(AttrsAndConstraintsTest, ContainsConvOrDotTrue) {
+  OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kModuleDotWeightOnlyPtq);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  auto call_op = *main_fn.getOps<TF::XlaCallModuleOp>().begin();
+  const StringRef function_name = GetEntryFunctionName(call_op);
+  EXPECT_TRUE(ContainsConvOrDot(function_name));
+}
+
+TEST_F(AttrsAndConstraintsTest, ContainsConvOrDotFalse) {
+  OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kModuleXlaCallModuleNoEntryNoQuantTrait);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  auto call_op = *main_fn.getOps<TF::XlaCallModuleOp>().begin();
+  const StringRef function_name = GetEntryFunctionName(call_op);
+  EXPECT_FALSE(ContainsConvOrDot(function_name));
+}
+
+}  // namespace
+}  // namespace mlir::quant

tensorflow/compiler/mlir/quantization/common/func.cc

@@ -0,0 +1,55 @@
+/* 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/common/func.h"
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/cc/saved_model/signature_constants.h"
+#include "tensorflow/compiler/mlir/tensorflow/translate/import_model.h"
+
+namespace mlir::quant {
+namespace {
+
+using ::tensorflow::kDefaultServingSignatureDefKey;
+using ::tensorflow::kImportModelDefaultGraphFuncName;
+
+// Returns true iff the function's symbol is public.
+bool IsPublicFuncOp(func::FuncOp func_op) {
+  return SymbolTable::getSymbolVisibility(&*func_op) ==
+         SymbolTable::Visibility::Public;
+}
+
+}  // namespace
+
+func::FuncOp FindMainFuncOp(ModuleOp module_op) {
+  if (const auto main_func_op = module_op.lookupSymbol<func::FuncOp>(
+          kImportModelDefaultGraphFuncName);
+      main_func_op != nullptr && IsPublicFuncOp(main_func_op)) {
+    return main_func_op;
+  }
+
+  if (const auto serving_default_func_op =
+          module_op.lookupSymbol<func::FuncOp>(kDefaultServingSignatureDefKey);
+      serving_default_func_op != nullptr &&
+      IsPublicFuncOp(serving_default_func_op)) {
+    return serving_default_func_op;
+  }
+
+  return nullptr;
+}
+
+}  // namespace mlir::quant

tensorflow/compiler/mlir/quantization/common/func.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_COMMON_FUNC_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_FUNC_H_
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+
+namespace mlir::quant {
+
+// Returns a public `func::FuncOp` in `module_op` whose name matches either
+// `main` or `serving_default`. If `func::FuncOps` with both names exist, the
+// function with name "main" takes precedence. Returns null if no such a
+// function exists.
+func::FuncOp FindMainFuncOp(ModuleOp module_op);
+
+}  // namespace mlir::quant
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_FUNC_H_

tensorflow/compiler/mlir/quantization/common/func_test.cc

@@ -0,0 +1,113 @@
+/* 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/common/func.h"
+
+#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 "tensorflow/compiler/mlir/quantization/common/test_base.h"
+
+namespace mlir::quant {
+namespace {
+
+using ::testing::IsNull;
+using ::testing::NotNull;
+
+using FindMainFuncOpTest = ::mlir::quant::QuantizationTestBase;
+
+TEST_F(FindMainFuncOpTest, ReturnsMainFuncOp) {
+  constexpr absl::string_view kModuleWithMainFunc = R"mlir(
+    module {
+      func.func @main() -> () {
+        return
+      }
+    }
+  )mlir";
+
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleWithMainFunc);
+  EXPECT_THAT(*module_op, NotNull());
+
+  func::FuncOp main_func_op = FindMainFuncOp(*module_op);
+  EXPECT_THAT(main_func_op, NotNull());
+}
+
+TEST_F(FindMainFuncOpTest, ReturnsNullWhenMainFuncOpIsPrivate) {
+  constexpr absl::string_view kModuleWithPrivateMainFunc = R"mlir(
+    module {
+      func.func private @main() -> () {
+        return
+      }
+    }
+  )mlir";
+
+  OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kModuleWithPrivateMainFunc);
+  EXPECT_THAT(*module_op, NotNull());
+
+  EXPECT_THAT(FindMainFuncOp(*module_op), IsNull());
+}
+
+TEST_F(FindMainFuncOpTest, ReturnsServingDefaultFuncOp) {
+  constexpr absl::string_view kModuleWithServingDefaultFunc = R"mlir(
+    module {
+      func.func @serving_default() -> () {
+        return
+      }
+    }
+  )mlir";
+
+  OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kModuleWithServingDefaultFunc);
+  EXPECT_THAT(*module_op, NotNull());
+
+  EXPECT_THAT(FindMainFuncOp(*module_op), NotNull());
+}
+
+TEST_F(FindMainFuncOpTest, ReturnsNullWhenServingDefaultFuncOpIsPrivate) {
+  constexpr absl::string_view kModuleWithPrivateServingDefaultFunc = R"mlir(
+    module {
+      func.func private @serving_default() -> () {
+        return
+      }
+    }
+  )mlir";
+
+  OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kModuleWithPrivateServingDefaultFunc);
+  EXPECT_THAT(*module_op, NotNull());
+
+  EXPECT_THAT(FindMainFuncOp(*module_op), IsNull());
+}
+
+TEST_F(FindMainFuncOpTest, ReturnsNullWhenMainFuncNotFound) {
+  constexpr absl::string_view kModuleWithNoMainFunc = R"mlir(
+    module {
+      func.func @foo() -> () {
+        return
+      }
+    }
+  )mlir";
+
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleWithNoMainFunc);
+  EXPECT_THAT(*module_op, NotNull());
+
+  EXPECT_THAT(FindMainFuncOp(*module_op), IsNull());
+}
+
+}  // namespace
+}  // namespace mlir::quant

tensorflow/compiler/mlir/quantization/common/ir/BUILD

@@ -0,0 +1,91 @@
+load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library")
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+td_library(
+    name = "QuantizationOpsTdFiles",
+    srcs = [
+        "QuantOps.td",
+        "QuantOpsBase.td",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        "@llvm-project//mlir:InferTypeOpInterfaceTdFiles",
+        "@llvm-project//mlir:OpBaseTdFiles",
+        "@llvm-project//mlir:QuantizationOpsTdFiles",
+        "@llvm-project//mlir:SideEffectInterfacesTdFiles",
+    ],
+)
+
+gentbl_cc_library(
+    name = "QuantOpsIncGen",
+    compatible_with = get_compatible_with_portable(),
+    tbl_outs = {
+        "QuantOps.h.inc": ["-gen-op-decls"],
+        "QuantOps.cc.inc": ["-gen-op-defs"],
+        "QuantOpsDialect.h.inc": [
+            "-gen-dialect-decls",
+            "-dialect=quantization",
+        ],
+        "QuantOpsDialect.cc.inc": [
+            "-gen-dialect-defs",
+            "-dialect=quantization",
+        ],
+    },
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "QuantOps.td",
+    deps = [":QuantizationOpsTdFiles"],
+)
+
+gentbl_cc_library(
+    name = "QuantPassIncGen",
+    compatible_with = get_compatible_with_portable(),
+    tbl_outs = {"Passes.h.inc": [
+        "-gen-pass-decls",
+        "-name=tfquant",
+    ]},
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "Passes.td",
+    deps = ["@llvm-project//mlir:PassBaseTdFiles"],
+)
+
+cc_library(
+    name = "QuantOps",
+    srcs = [
+        "ConvertConst.cc",
+        "ConvertSimQuant.cc",
+        "FakeQuantSupport.cc",
+        "QuantOps.cc",
+        "QuantizeUtils.cc",
+        "UniformSupport.cc",
+    ],
+    hdrs = [
+        "FakeQuantSupport.h",
+        "Passes.h",
+        "QuantOps.h",
+        "QuantizeUtils.h",
+        "UniformSupport.h",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        ":QuantOpsIncGen",
+        ":QuantPassIncGen",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:BytecodeOpInterface",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:InferTypeOpInterface",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:QuantOps",
+        "@llvm-project//mlir:SideEffectInterfaces",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+    ],
+)

tensorflow/compiler/mlir/quantization/common/ir/ConvertConst.cc

@@ -0,0 +1,124 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <memory>
+#include <utility>
+
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/Matchers.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/quantization/common/ir/Passes.h"
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantizeUtils.h"
+
+namespace mlir {
+namespace quant::ir {
+
+using mlir::quant::QuantizedType;
+
+namespace {
+#define GEN_PASS_DEF_QUANTCONVERTCONST
+#include "tensorflow/compiler/mlir/quantization/common/ir/Passes.h.inc"
+
+struct ConvertConstPass : public impl::QuantConvertConstBase<ConvertConstPass> {
+  void runOnOperation() override;
+};
+
+struct QuantizedConstRewrite : public OpRewritePattern<QuantizeCastOp> {
+  using OpRewritePattern<QuantizeCastOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(QuantizeCastOp qbarrier,
+                                PatternRewriter &rewriter) const override;
+};
+
+}  // namespace
+
+/// Matches a [constant] -> [qbarrier] where the qbarrier results type is
+/// quantized and the operand type is quantizable.
+
+LogicalResult QuantizedConstRewrite::matchAndRewrite(
+    QuantizeCastOp qbarrier, PatternRewriter &rewriter) const {
+  Attribute value;
+
+  // Is the operand a constant?
+  if (!matchPattern(qbarrier.getArg(), m_Constant(&value))) {
+    return failure();
+  }
+
+  // Does the qbarrier convert to a quantized type. This will not be true
+  // if a quantized type has not yet been chosen or if the cast to an equivalent
+  // storage type is not supported.
+  Type qbarrierResultType = qbarrier.getResult().getType();
+  QuantizedType quantizedElementType =
+      QuantizedType::getQuantizedElementType(qbarrierResultType);
+  if (!quantizedElementType) {
+    return failure();
+  }
+  if (!QuantizedType::castToStorageType(qbarrierResultType)) {
+    return failure();
+  }
+
+  // Is the operand type compatible with the expressed type of the quantized
+  // type? This will not be true if the qbarrier is superfluous (converts
+  // from and to a quantized type).
+  if (!quantizedElementType.isCompatibleExpressedType(
+          qbarrier.getArg().getType())) {
+    return failure();
+  }
+
+  // Is the constant value a type expressed in a way that we support?
+  if (!mlir::isa<FloatAttr, DenseElementsAttr, SparseElementsAttr>(value)) {
+    return failure();
+  }
+
+  Type newConstValueType;
+  auto newConstValue =
+      quantizeAttr(value, quantizedElementType, newConstValueType);
+  if (!newConstValue) {
+    return failure();
+  }
+
+  // When creating the new const op, use a fused location that combines the
+  // original const and the qbarrier that led to the quantization.
+  auto fusedLoc = rewriter.getFusedLoc(
+      {qbarrier.getArg().getDefiningOp()->getLoc(), qbarrier.getLoc()});
+  auto newConstOp = arith::ConstantOp::create(
+      rewriter, fusedLoc, newConstValueType, cast<TypedAttr>(newConstValue));
+  rewriter.replaceOpWithNewOp<StorageCastOp>(qbarrier, qbarrier.getType(),
+                                             newConstOp);
+  return success();
+}
+
+void ConvertConstPass::runOnOperation() {
+  RewritePatternSet patterns(&getContext());
+  auto func = getOperation();
+  auto *context = &getContext();
+  patterns.add<QuantizedConstRewrite>(context);
+  (void)applyPatternsGreedily(func, std::move(patterns));
+}
+
+std::unique_ptr<OperationPass<func::FuncOp>> createConvertConstPass() {
+  return std::make_unique<ConvertConstPass>();
+}
+
+}  // namespace quant::ir
+}  // namespace mlir

tensorflow/compiler/mlir/quantization/common/ir/ConvertSimQuant.cc

@@ -0,0 +1,158 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cassert>
+#include <memory>
+#include <utility>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/quantization/common/ir/FakeQuantSupport.h"
+#include "tensorflow/compiler/mlir/quantization/common/ir/Passes.h"
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"
+#include "tensorflow/compiler/mlir/quantization/common/ir/UniformSupport.h"
+
+namespace mlir::quant::ir {
+
+#define GEN_PASS_DEF_QUANTCONVERTSIMULATEDQUANT
+#include "tensorflow/compiler/mlir/quantization/common/ir/Passes.h.inc"
+
+struct ConvertSimulatedQuantPass
+    : public impl::QuantConvertSimulatedQuantBase<ConvertSimulatedQuantPass> {
+  void runOnOperation() override;
+};
+
+/// Base class rewrites ConstFakeQuant into a qbarrier/dbarrier pair.
+template <typename ConcreteRewriteClass, typename FakeQuantOp>
+class FakeQuantRewrite : public OpRewritePattern<FakeQuantOp> {
+ public:
+  using OpRewritePattern<FakeQuantOp>::OpRewritePattern;
+
+  FakeQuantRewrite(MLIRContext *ctx, bool *hadFailure)
+      : OpRewritePattern<FakeQuantOp>(ctx), hadFailure(hadFailure) {}
+
+  LogicalResult matchAndRewrite(FakeQuantOp op,
+                                PatternRewriter &rewriter) const override {
+    // TODO: If this pattern comes up more frequently, consider adding core
+    // support for failable rewrites.
+    if (failableRewrite(op, rewriter)) {
+      *hadFailure = true;
+      return failure();
+    }
+    return success();
+  }
+
+ private:
+  bool *hadFailure;
+
+  bool failableRewrite(FakeQuantOp op, PatternRewriter &rewriter) const {
+    auto converter =
+        mlir::quant::ir::ExpressedToQuantizedConverter::forInputType(
+            op.getType());
+    if (!converter) {
+      return (op.emitError("unsupported quantized type conversion"), true);
+    }
+
+    quant::QuantizedType elementType =
+        static_cast<const ConcreteRewriteClass *>(this)
+            ->convertFakeQuantAttrsToType(op, converter.expressed_type);
+
+    if (!elementType) {
+      // Note that the fakeQuantAttrsToType will have emitted the error.
+      return true;
+    }
+
+    Type quantizedType = converter.convert(elementType);
+    assert(quantizedType &&
+           "Converter accepted a type that it did not convert");
+
+    // TODO: Map to a qbarrier with an attribute like [Forced] to signal that
+    // this is a forced/hard-coded constraint.
+    auto qbarrier = QuantizeCastOp::create(rewriter, op.getLoc(), quantizedType,
+                                           op.getInputs());
+    rewriter.replaceOpWithNewOp<DequantizeCastOp>(op, converter.input_type,
+                                                  qbarrier.getResult());
+
+    return false;
+  }
+};
+
+class ConstFakeQuantRewrite
+    : public FakeQuantRewrite<ConstFakeQuantRewrite, ConstFakeQuant> {
+ public:
+  using BaseRewrite = FakeQuantRewrite<ConstFakeQuantRewrite, ConstFakeQuant>;
+
+  ConstFakeQuantRewrite(MLIRContext *ctx, bool *hadFailure)
+      : BaseRewrite(ctx, hadFailure) {}
+
+  quant::QuantizedType convertFakeQuantAttrsToType(ConstFakeQuant fqOp,
+                                                   Type expressedType) const {
+    return quantfork::fakeQuantAttrsToType(
+        fqOp.getLoc(), fqOp.getNumBits(), fqOp.getMin().convertToFloat(),
+        fqOp.getMax().convertToFloat(), fqOp.getNarrowRange(), expressedType,
+        fqOp.getIsSigned());
+  }
+};
+
+class ConstFakeQuantPerAxisRewrite
+    : public FakeQuantRewrite<ConstFakeQuantPerAxisRewrite,
+                              ConstFakeQuantPerAxis> {
+ public:
+  using BaseRewrite =
+      FakeQuantRewrite<ConstFakeQuantPerAxisRewrite, ConstFakeQuantPerAxis>;
+
+  ConstFakeQuantPerAxisRewrite(MLIRContext *ctx, bool *hadFailure)
+      : BaseRewrite(ctx, hadFailure) {}
+
+  quant::QuantizedType convertFakeQuantAttrsToType(ConstFakeQuantPerAxis fqOp,
+                                                   Type expressedType) const {
+    SmallVector<double, 4> min, max;
+    min.reserve(fqOp.getMin().size());
+    max.reserve(fqOp.getMax().size());
+    for (auto m : fqOp.getMin())
+      min.push_back(cast<FloatAttr>(m).getValueAsDouble());
+    for (auto m : fqOp.getMax())
+      max.push_back(cast<FloatAttr>(m).getValueAsDouble());
+
+    return quantfork::fakeQuantAttrsToType(
+        fqOp.getLoc(), fqOp.getNumBits(), fqOp.getAxis(), min, max,
+        fqOp.getNarrowRange(), expressedType, fqOp.getIsSigned());
+  }
+};
+
+void ConvertSimulatedQuantPass::runOnOperation() {
+  bool hadFailure = false;
+  auto func = getOperation();
+  RewritePatternSet patterns(func.getContext());
+  auto *ctx = func.getContext();
+  patterns.add<ConstFakeQuantRewrite, ConstFakeQuantPerAxisRewrite>(
+      ctx, &hadFailure);
+  (void)applyPatternsGreedily(func, std::move(patterns));
+  if (hadFailure) signalPassFailure();
+}
+
+std::unique_ptr<OperationPass<func::FuncOp>> createConvertSimulatedQuantPass() {
+  return std::make_unique<ConvertSimulatedQuantPass>();
+}
+
+}  // namespace mlir::quant::ir

tensorflow/compiler/mlir/quantization/common/ir/FakeQuantSupport.cc

@@ -0,0 +1,214 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/quantization/common/ir/FakeQuantSupport.h"
+
+#include <cassert>
+#include <cmath>
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Diagnostics.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace mlir::quantfork {
+
+static bool getDefaultStorageParams(unsigned numBits, bool narrowRange,
+                                    bool isSigned, MLIRContext *ctx,
+                                    Type &storageType, int64_t &qmin,
+                                    int64_t &qmax) {
+  // Hard-coded type mapping from TFLite.
+  if (numBits <= 4) {
+    storageType = IntegerType::get(ctx, 4);
+    if (isSigned) {
+      qmin = -8;
+      qmax = 7;
+    } else {
+      qmin = 0;
+      qmax = 15;
+    }
+  } else if (numBits <= 8) {
+    storageType = IntegerType::get(ctx, 8);
+    if (isSigned) {
+      qmin = -128;
+      qmax = 127;
+    } else {
+      qmin = 0;
+      qmax = 255;
+    }
+  } else if (numBits <= 16) {
+    storageType = IntegerType::get(ctx, 16);
+    if (isSigned) {
+      qmin = -32768;
+      qmax = 32767;
+    } else {
+      qmin = 0;
+      qmax = 65535;
+    }
+  } else if (numBits <= 32) {
+    storageType = IntegerType::get(ctx, 32);
+    if (isSigned) {
+      qmin = std::numeric_limits<int32_t>::min();
+      qmax = std::numeric_limits<int32_t>::max();
+    } else {
+      qmin = std::numeric_limits<uint32_t>::min();
+      qmax = std::numeric_limits<uint32_t>::max();
+    }
+  } else {
+    return true;
+  }
+
+  // Handle narrowRange.
+  if (narrowRange) {
+    qmin += 1;
+  }
+  return false;
+}
+
+// This is a specific implementation of nudging:
+// If 0.0 < rmin < rmax or rmin < rmax < 0.0, the range will be shifted
+// to include 0.0, but the range width size (rmax-rmin) isn't changed. The zero
+// point is derived from the shifted range, and the scale isn't changed. As
+// a consequence some values, which are supposed in the original [rmin, rmax]
+// range will be outside the shifted range and be clamped during quantization.
+// TODO: we should nudge the scale as well, but that requires the
+// fake quant op used in the training to use the nudged scale as well.
+static void getNudgedScaleAndZeroPoint(int64_t qmin, int64_t qmax, double rmin,
+                                       double rmax, double &scale,
+                                       int64_t &nudgedZeroPoint) {
+  // Determine the scale.
+  const double qminDouble = qmin;
+  const double qmaxDouble = qmax;
+  scale = (rmax - rmin) / (qmaxDouble - qminDouble);
+
+  // Zero point computation.
+  // In float, solve the affine equation for any known pair
+  // (real value, corresponding quantized value), of which, two such pairs
+  // are known: (rmin, qmin), (rmax, qmax).
+  // The arithmetic error on the zero point computed from either pair will be
+  // roughly machine_epsilon * (sum of absolute values of terms).
+  // Use the variant that adds the smaller error.
+  const double zeroPointFromMin = qminDouble - rmin / scale;
+  const double zeroPointFromMinError =
+      std::abs(qminDouble) + std::abs(rmin / scale);
+  const double zeroPointFromMax = qmaxDouble - rmax / scale;
+  const double zeroPointFromMaxError =
+      std::abs(qmaxDouble) + std::abs(rmax / scale);
+
+  const double zeroPointDouble = (zeroPointFromMinError < zeroPointFromMaxError)
+                                     ? zeroPointFromMin
+                                     : zeroPointFromMax;
+
+  // Now nudge the zero point to be an integer.
+  nudgedZeroPoint = 0;
+  if (zeroPointDouble < qminDouble) {
+    nudgedZeroPoint = qmin;
+  } else if (zeroPointDouble > qmaxDouble) {
+    nudgedZeroPoint = qmax;
+  } else {
+    nudgedZeroPoint = round(zeroPointDouble);
+  }
+
+  // By construction, the nudged zero point should always be in range.
+  assert(nudgedZeroPoint >= qmin);
+  assert(nudgedZeroPoint <= qmax);
+}
+
+quant::UniformQuantizedType fakeQuantAttrsToType(Location loc, unsigned numBits,
+                                                 double rmin, double rmax,
+                                                 bool narrowRange,
+                                                 Type expressedType,
+                                                 bool isSigned) {
+  MLIRContext *ctx = expressedType.getContext();
+  unsigned flags = isSigned ? quant::QuantizationFlags::Signed : 0;
+  Type storageType;
+  int64_t qmin;
+  int64_t qmax;
+  if (getDefaultStorageParams(numBits, narrowRange, isSigned, ctx, storageType,
+                              qmin, qmax)) {
+    return (emitError(loc, "unsupported FakeQuant number of bits: ") << numBits,
+            nullptr);
+  }
+
+  // Special case where min/max is close enough. The tensor contents are all
+  // 0.0s, so the scale is set to 1.0 and the tensor can be quantized to zero
+  // points and dequantized to 0.0.
+  if (std::fabs(rmax - rmin) < std::numeric_limits<double>::epsilon()) {
+    return quant::UniformQuantizedType::getChecked(
+        loc, flags, storageType, expressedType, 1.0, qmin, qmin, qmax);
+  }
+
+  double scale;
+  int64_t nudgedZeroPoint;
+  getNudgedScaleAndZeroPoint(qmin, qmax, rmin, rmax, scale, nudgedZeroPoint);
+
+  return quant::UniformQuantizedType::getChecked(loc, flags, storageType,
+                                                 expressedType, scale,
+                                                 nudgedZeroPoint, qmin, qmax);
+}
+
+quant::UniformQuantizedPerAxisType fakeQuantAttrsToType(
+    Location loc, unsigned numBits, int32_t quantizedDimension,
+    ArrayRef<double> rmins, ArrayRef<double> rmaxs, bool narrowRange,
+    Type expressedType, bool isSigned) {
+  size_t axisSize = rmins.size();
+  if (axisSize != rmaxs.size()) {
+    return (emitError(loc, "mismatched per-axis min and max size: ")
+                << axisSize << " vs. " << rmaxs.size(),
+            nullptr);
+  }
+
+  MLIRContext *ctx = expressedType.getContext();
+  Type storageType;
+  int64_t qmin;
+  int64_t qmax;
+  if (getDefaultStorageParams(numBits, narrowRange, isSigned, ctx, storageType,
+                              qmin, qmax)) {
+    return (emitError(loc, "unsupported FakeQuant number of bits: ") << numBits,
+            nullptr);
+  }
+
+  SmallVector<double, 4> scales;
+  SmallVector<int64_t, 4> zeroPoints;
+  scales.reserve(axisSize);
+  zeroPoints.reserve(axisSize);
+  for (size_t axis = 0; axis != axisSize; ++axis) {
+    double rmin = rmins[axis];
+    double rmax = rmaxs[axis];
+    if (std::fabs(rmax - rmin) < std::numeric_limits<double>::epsilon()) {
+      scales.push_back(1.0);
+      zeroPoints.push_back(qmin);
+      continue;
+    }
+
+    double scale;
+    int64_t nudgedZeroPoint;
+    getNudgedScaleAndZeroPoint(qmin, qmax, rmin, rmax, scale, nudgedZeroPoint);
+    scales.push_back(scale);
+    zeroPoints.push_back(nudgedZeroPoint);
+  }
+
+  unsigned flags = isSigned ? quant::QuantizationFlags::Signed : 0;
+  return quant::UniformQuantizedPerAxisType::getChecked(
+      loc, flags, storageType, expressedType, scales, zeroPoints,
+      quantizedDimension, qmin, qmax);
+}
+
+}  // namespace mlir::quantfork

tensorflow/compiler/mlir/quantization/common/ir/FakeQuantSupport.h

@@ -0,0 +1,79 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+//
+// This file defines support utilities for interoperating with FakeQuant* based
+// QAT (Quantized Aware Training) computations, as implemented by TFLite. Note
+// that FakeQuant* operators mix multiple concerns specific to how TFLite
+// originally implemented quantization. As such, utilities here enforce
+// opinions taken by that codebase (vs providing any amount of genericity).
+//
+// Specifically, it combines the following concerns, each of which would be
+// independent variables in a more generic setup:
+//   - numBits and isSigned imply storage data type (uint8, int8, int16)
+//   - numBits < 8 is promoted to uint8 or int8
+//   - "narrow_range" narrows the lower bound of the storage type's range by
+//     1
+//   - the specified min/max values are "nudged" so that the result has a zero
+//     that can be exactly expressed
+//   - min=max=0 implies scale=0 and zero_point=0
+//
+// With the above assumptions applied, every conforming specified FakeQuant op
+// can be represented by a UniformQuantizedType. This scheme is not expected to
+// be generalized further in the future and should be considered to be a
+// legacy set of rules.
+//
+// As canonically used in TensorFlow graphs, the presence of a FakeQuant node
+// is a hint that the specific math represented here has been simulated at
+// training time. As such, it is usually not advised to arbitrarily change
+// quantization parameters derived from FakeQuant.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_IR_FAKEQUANTSUPPORT_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_IR_FAKEQUANTSUPPORT_H_
+
+#include <cstdint>
+
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace mlir {
+namespace quantfork {
+
+/// Converts per-layer FakeQuant attributes to the corresponding type.
+/// In the event that the parameters cannot be converted, returns a nullptr
+/// convertible Type and issues an appropriate error.
+/// Note that there are multiple variants of a per-layer FakeQuant op, so
+/// this function takes the attributes discretely vs taking a reference to the
+/// originating op.
+quant::UniformQuantizedType fakeQuantAttrsToType(Location loc, unsigned numBits,
+                                                 double rmin, double rmax,
+                                                 bool narrowRange,
+                                                 Type expressedType,
+                                                 bool isSigned = false);
+
+/// Converts per-channel FakeQuant attributes to the corresponding type.
+/// In the event that the parameters cannot be converted, returns a nullptr
+/// convertible Type and issues an appropriate error.
+quant::UniformQuantizedPerAxisType fakeQuantAttrsToType(
+    Location loc, unsigned numBits, int32_t quantizedDimension,
+    ArrayRef<double> rmins, ArrayRef<double> rmax, bool narrowRange,
+    Type expressedType, bool isSigned = false);
+}  // namespace quantfork
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_IR_FAKEQUANTSUPPORT_H_

tensorflow/compiler/mlir/quantization/common/ir/Passes.h

@@ -0,0 +1,57 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+//
+// This file defines all of the passes owned by the quantization dialect. As
+// things mature, it is expected that passes specific to certain frontend or
+// backend dialects will move to those dialects directly. For now, they are
+// incubated here.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_IR_PASSES_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_IR_PASSES_H_
+
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+
+namespace mlir {
+namespace func {
+class FuncOp;
+}  // namespace func
+
+namespace quant::ir {
+
+/// Creates a pass that converts quantization simulation operations (i.e.
+/// FakeQuant and those like it) to casts into/out of supported QuantizedTypes.
+std::unique_ptr<OperationPass<func::FuncOp>> createConvertSimulatedQuantPass();
+
+/// Creates a pass that converts constants followed by a qbarrier to a
+/// constant whose value is quantized. This is typically one of the last
+/// passes done when lowering to express actual quantized arithmetic in a
+/// low level representation. Because it modifies the constant, it is
+/// destructive and cannot be undone.
+std::unique_ptr<OperationPass<func::FuncOp>> createConvertConstPass();
+
+//===----------------------------------------------------------------------===//
+// Registration
+//===----------------------------------------------------------------------===//
+
+/// Generate the code for registering passes.
+#define GEN_PASS_REGISTRATION
+#include "tensorflow/compiler/mlir/quantization/common/ir/Passes.h.inc"
+
+}  // namespace quant::ir
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_IR_PASSES_H_

tensorflow/compiler/mlir/quantization/common/ir/Passes.td

@@ -0,0 +1,34 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TF_QUANT_PASSES
+#define TF_QUANT_PASSES
+
+include "mlir/Pass/PassBase.td"
+
+def QuantConvertConst : Pass<"quant-convert-const", "func::FuncOp"> {
+  let summary = "Converts constants followed by qbarrier to actual quantized "
+                "values";
+  let constructor = "mlir::quant::ir::createConvertConstPass()";
+}
+
+def QuantConvertSimulatedQuant
+    : Pass<"quant-convert-simulated-quantization", "func::FuncOp"> {
+  let summary = "Converts training-time simulated quantization ops to "
+                "corresponding quantize/dequantize casts";
+  let constructor = "mlir::quant::ir::createConvertSimulatedQuantPass()";
+}
+
+#endif // TF_QUANT_PASSES

tensorflow/compiler/mlir/quantization/common/ir/QuantOps.cc

@@ -0,0 +1,143 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"
+
+#include <cstdint>
+#include <functional>
+#include <iterator>
+#include <numeric>
+
+#include "llvm/ADT/STLExtras.h"
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOpsDialect.cc.inc"
+
+namespace mlir::quant::ir {
+
+using mlir::quant::QuantizedType;
+
+void TFQuantDialect::initialize() {
+  addOperations<
+#define GET_OP_LIST
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.cc.inc"
+      >();
+}
+
+OpFoldResult StorageCastOp::fold(FoldAdaptor) {
+  // Matches x -> [scast -> scast] -> y, replacing the second scast with the
+  // value of x if the casts invert each other.
+  auto srcScastOp = getArg().getDefiningOp<StorageCastOp>();
+  if (!srcScastOp || srcScastOp.getArg().getType() != getType())
+    return OpFoldResult();
+  return srcScastOp.getArg();
+}
+
+/// The quantization specification should match the expressed type.
+static bool isValidQuantizationSpec(Attribute quantSpec, Type expressed) {
+  if (auto typeAttr = mlir::dyn_cast<TypeAttr>(quantSpec)) {
+    Type spec = typeAttr.getValue();
+    if (mlir::isa<TensorType, VectorType>(spec)) return false;
+
+    // The spec should be either a quantized type which is compatible to the
+    // expressed type, or a primitive type which is as same as the
+    // (element type of) the expressed type.
+    if (auto quantizedType = mlir::dyn_cast<QuantizedType>(spec))
+      return quantizedType.isCompatibleExpressedType(expressed);
+
+    if (auto tensorType = mlir::dyn_cast<TensorType>(expressed))
+      return spec == tensorType.getElementType();
+
+    if (auto vectorType = mlir::dyn_cast<VectorType>(expressed))
+      return spec == vectorType.getElementType();
+  }
+  return false;
+}
+
+LogicalResult QuantizeRegionOp::verify() {
+  // There are specifications for both inputs and outputs.
+  if (getNumOperands() != getInputSpecs().size() ||
+      getNumResults() != getOutputSpecs().size())
+    return emitOpError(
+        "has unmatched operands/results number and spec attributes number");
+
+  // Verify that quantization specifications are valid.
+  for (auto input : llvm::zip(getOperandTypes(), getInputSpecs())) {
+    Type inputType = std::get<0>(input);
+    Attribute inputSpec = std::get<1>(input);
+    if (!isValidQuantizationSpec(inputSpec, inputType)) {
+      return emitOpError() << "has incompatible specification " << inputSpec
+                           << " and input type " << inputType;
+    }
+  }
+
+  for (auto result : llvm::zip(getResultTypes(), getOutputSpecs())) {
+    Type outputType = std::get<0>(result);
+    Attribute outputSpec = std::get<1>(result);
+    if (!isValidQuantizationSpec(outputSpec, outputType)) {
+      return emitOpError() << "has incompatible specification " << outputSpec
+                           << " and output type " << outputType;
+    }
+  }
+  return success();
+}
+
+LogicalResult StatisticsOp::verify() {
+  auto tensorArg = mlir::dyn_cast<TensorType>(getArg().getType());
+  if (!tensorArg) return emitOpError("arg needs to be tensor type.");
+
+  // Verify layerStats attribute.
+  {
+    auto layerStatsType = getLayerStats().getShapedType();
+    if (!mlir::isa<FloatType>(layerStatsType.getElementType())) {
+      return emitOpError("layerStats must have a floating point element type");
+    }
+    if (layerStatsType.getRank() != 1 || layerStatsType.getDimSize(0) != 2) {
+      return emitOpError("layerStats must have shape [2]");
+    }
+  }
+  // Verify axisStats (optional) attribute.
+  if (getAxisStats()) {
+    if (!getAxis()) return emitOpError("axis must be specified for axisStats");
+
+    auto shape = tensorArg.getShape();
+    auto argSliceSize =
+        std::accumulate(std::next(shape.begin(), *getAxis()), shape.end(), 1,
+                        std::multiplies<int64_t>());
+
+    auto axisStatsType = getAxisStats()->getShapedType();
+    if (!mlir::isa<FloatType>(axisStatsType.getElementType())) {
+      return emitOpError("axisStats must have a floating point element type");
+    }
+    if (axisStatsType.getRank() != 2 || axisStatsType.getDimSize(1) != 2 ||
+        axisStatsType.getDimSize(0) != argSliceSize) {
+      return emitOpError(
+          "axisStats must have shape [N,2] "
+          "where N = the slice size defined by the axis dim");
+    }
+  }
+  return success();
+}
+
+}  // namespace mlir::quant::ir
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.cc.inc"

tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h

@@ -0,0 +1,34 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_IR_QUANTOPS_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_IR_QUANTOPS_H_
+
+#include "llvm/Support/MathExtras.h"
+#include "mlir/Bytecode/BytecodeOpInterface.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Dialect.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Interfaces/InferTypeOpInterface.h"  // from @llvm-project
+#include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOpsDialect.h.inc"
+#define GET_OP_CLASSES
+
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h.inc"
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_IR_QUANTOPS_H_

tensorflow/compiler/mlir/quantization/common/ir/QuantOps.td

@@ -0,0 +1,294 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+//
+//===----------------------------------------------------------------------===//
+//
+// This is the operation definition file for Quantization.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef QUANTIZATION_OPS
+#define QUANTIZATION_OPS
+
+include "mlir/Dialect/Quant/IR/QuantBase.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "tensorflow/compiler/mlir/quantization/common/ir/QuantOpsBase.td"
+
+class quant_TypedPrimitiveOrContainer<Type etype> :
+    Type<Or<[etype.predicate,
+                TensorOf<[etype]>.predicate,
+                VectorOfNonZeroRankOf<[etype]>.predicate]>,
+         "primitive/tensor/vector of " # etype.summary>;
+
+// A primitive type that can represent a real value. This is either a
+// floating point value or a quantized type.
+def quant_RealPrimitiveType :
+    Type<Or<[AnyFloat.predicate, quant_QuantizedType.predicate]>,
+    "real valued primitive (float or quantized type)">;
+
+// A primitive type that can represent a storage value. This is either an
+// integer or quantized type.
+def quant_StoragePrimitiveType :
+    Type<Or<[AnySignlessInteger.predicate, quant_QuantizedType.predicate]>,
+    "quantized storage primitive (integer or quantized type)">;
+
+// A primitive or container of RealPrimitiveType.
+def quant_RealValueType :
+    quant_TypedPrimitiveOrContainer<quant_RealPrimitiveType>;
+
+// A primitive or container of StoragePrimitiveType.
+def quant_StorageValueType :
+    quant_TypedPrimitiveOrContainer<quant_StoragePrimitiveType>;
+
+// Either a real valued or storage primitive or container type.
+def quant_RealOrStorageValueType :
+    Type<Or<[quant_RealValueType.predicate, quant_StorageValueType.predicate]>,
+    "real valued or storage primitive or container type">;
+
+//===----------------------------------------------------------------------===//
+// Base classes
+//===----------------------------------------------------------------------===//
+
+class Quantization_Op<string mnemonic, list<Trait> traits> :
+    Op<TF_Quant_Dialect, mnemonic, traits>;
+
+//===----------------------------------------------------------------------===//
+// Quantization casts
+//===----------------------------------------------------------------------===//
+// A QuantizeCast (qcast) represents a potential type shift from a quantizable
+// type to a quantized type.
+//
+// At runtime, a qcast will apply the transformation expressed by its
+// operand and result type. For flexibility during transformation, it is also
+// possible to have a qcast that performs no transformation (both its
+// operand and result type are quantizable).
+//
+// A qcast will typically originate from either:
+//   a) An expressed or implied constraint in the source dialect which signals
+//      that a certain level of quantization is possible or required.
+//   b) An inference made by a quantization algorithm indicating that a
+//      quantized representation may be acceptable.
+//
+// Especially early in transformation, it is common to have pairs of
+// qcast/dcast at points where a transition to a quantized type is
+// required. In addition, it is also common to have an identity qcast
+// (where the operand and result type are not quantized) at all points where
+// it is legal to use a quantized representation (but is not known to be
+// acceptable).
+def Quantization_QuantizeCastOp : Quantization_Op<"qcast", [Pure]> {
+  let arguments = (ins quant_RealValueType:$arg);
+  let results = (outs quant_RealValueType);
+}
+
+// A DequantizeCast op (dcast) represents the inverse of a qcast,
+// converting back from a quantized to quantizable (expressed) type.
+//
+// Like qcasts, a dcast is allowed to have both its operand and result
+// as non quantized types. This facilitates transformations and marks edges
+// where the computation must be carried out in the expressed type.
+//
+// Especially early in transformation, it is common to have dcasts on
+// all operands to ops that must operate with the expressed type (typically
+// math ops prior to lowering to target-specific, quantized kernels).
+def Quantization_DequantizeCastOp : Quantization_Op<"dcast", [Pure]> {
+  let arguments = (ins quant_RealValueType:$arg);
+  let results = (outs quant_RealValueType);
+}
+
+// A StorageCast (scast) represents a cast from or to a type based on the
+// storage type and a type based on a corresponding quantized type.
+//
+// This op exists to ensure type coherency for between parts of the computation
+// which are operating directly on an underlying storage type and those which
+// operate on quantized values.
+//
+// Examples from storage to quantized type:
+//   i8 -> !quant<"uniform[i8:f32]{1.0}">
+//   tensor<4xi8> -> tensor<4x!quant<"uniform[i8:f32]{1.0}">>
+//   vector<4xi8> -> vector<4x!quant<"uniform[i8:f32]{1.0}">>
+def Quantization_StorageCastOp : Quantization_Op<"scast", [Pure]> {
+  let arguments = (ins quant_RealOrStorageValueType:$arg);
+  let results = (outs quant_RealOrStorageValueType);
+  let hasFolder = 1;
+}
+
+// A QuantizeRegion (region) represents a quantization unit which wraps
+// high-precision ops with quantization specifications for all the inputs
+// and outputs. Some quantization specifications can be undetermined and
+// derived from other ports by the target specification of the kernel.
+def Quantization_QuantizeRegionOp : Quantization_Op<"region", [
+    Pure,
+    IsolatedFromAbove,
+    SingleBlockImplicitTerminator<"ReturnOp">]> {
+  let summary = [{
+    The `region` operation wraps high-precision ops as a logical low-precision
+    quantized kernel.
+  }];
+
+  let arguments = (ins Variadic<AnyType>:$inputs,
+                    TypeArrayAttr:$input_specs,
+                    TypeArrayAttr:$output_specs,
+                    StrAttr:$logical_kernel);
+  let results = (outs Variadic<AnyType>:$outputs);
+  let regions = (region SizedRegion<1>:$body);
+  let hasVerifier = 1;
+}
+
+def Quantization_ReturnOp : Quantization_Op<"return", [Terminator]> {
+  let summary = [{
+    The `return` operation terminates a quantize region and returns values.
+  }];
+
+  let arguments = (ins Variadic<AnyTensor>:$results);
+}
+
+//===----------------------------------------------------------------------===//
+// Training integration and instrumentation ops
+//===----------------------------------------------------------------------===//
+
+def Quantization_ConstFakeQuant : Quantization_Op<"const_fake_quant",
+                                    [SameOperandsAndResultType, Pure]> {
+  let summary = [{
+    Simulates the effect of uniform quantization with const range.
+  }];
+
+  let description = [{
+    Given a const min, max, num_bits and narrow_range attribute, applies the
+    same uniform quantization simulation as is done by the TensorFlow
+    fake_quant_with_min_max_args op. See the fakeQuantAttrsToType() utility
+    method and the quant-convert-simulated-quantization pass for further details.
+  }];
+
+  let arguments = (ins
+    F32Tensor:$inputs,
+    F32Attr:$min,
+    F32Attr:$max,
+    // The bitwidth of the quantization; between 2 and 16, inclusive.
+    I64Attr:$num_bits,
+    // Quantization range starts from 0 or 1; starts from 1 if true.
+    DefaultValuedOptionalAttr<BoolAttr, "false">:$narrow_range,
+    // The sign of the quantization.
+    DefaultValuedOptionalAttr<BoolAttr, "false">:$is_signed
+  );
+
+  let results = (outs
+    F32Tensor:$outputs
+  );
+}
+
+def Quantization_ConstFakeQuantPerAxis : Quantization_Op<"const_fake_quant_per_axis",
+                                    [SameOperandsAndResultType, Pure]> {
+  let summary = [{
+    Simulates the effect of per axis uniform quantization with const range.
+  }];
+
+  let description = [{
+    Given a const min, max, num_bits and narrow_range attribute, applies the
+    same per axis uniform quantization simulation as is done by the TensorFlow
+    fake_quant_with_min_max_vars_per_channel op. See the fakeQuantAttrsToType()
+    utility method and the quant-convert-simulated-quantization pass for further
+    details.
+  }];
+
+  let arguments = (ins
+    F32Tensor:$inputs,
+    F32ArrayAttr:$min,
+    F32ArrayAttr:$max,
+    // The quantized dimension of the inputs tensor.
+    I64Attr:$axis,
+    // The bitwidth of the quantization; between 2 and 16, inclusive.
+    I64Attr:$num_bits,
+    // Quantization range starts from 0 or 1; starts from 1 if true.
+    DefaultValuedOptionalAttr<BoolAttr, "false">:$narrow_range,
+    // The sign of the quantization.
+    DefaultValuedOptionalAttr<BoolAttr, "false">:$is_signed
+  );
+
+  let results = (outs
+    F32Tensor:$outputs
+  );
+}
+
+def Quantization_StatisticsRefOp : Quantization_Op<"stats_ref", [SameOperandsAndResultType]> {
+  let summary = "Indicates that statistics are resolved by reference.";
+
+  let description = [{
+    This op acts as an identity that, when encountered at runtime, should result
+    in statistics being collected about about the value of its operand/result.
+    Such statistics will be stored with the provided key, allowing this node
+    to later be converted to a 'stats' op if statistics with that key have been
+    encountered.
+  }];
+
+  let arguments = (ins
+    quant_RealValueType:$arg,
+    StrAttr:$statsKey
+  );
+  let results = (outs quant_RealValueType);
+}
+
+def Quantization_StatisticsOp : Quantization_Op<"stats", [SameOperandsAndResultType]> {
+  let summary = "Identity op which associates statistics with the value.";
+
+  let description = [{
+    Associates statistics about the runtime ranges of values observed for
+    evaluations of this node.
+
+    Statistics about the entire type are reported in the 'layerStats' attribute
+    and those for each axis, in the (optional) `axisStats` attribute. The
+    interpretation of each is determined by the last dimension of its shape.
+    Currently, only dim=2 is supported, which is interpreted as [min, max].
+
+    `layerStats` must be a rank 1 tensor: [2]
+    `axisStats` must be a rank 2 tensor: [N, 2], where N=the slice size
+      splitted by the `axis` dimension. For example:
+
+    ```
+    <?x?x3x2>, axis=3 => N=2
+    <?x?x3x2>, axis=2 => N=6
+    ```
+  }];
+
+  let arguments = (ins
+    quant_RealValueType:$arg,
+    ElementsAttr:$layerStats,
+    OptionalAttr<ElementsAttr>:$axisStats,
+    OptionalAttr<I64Attr>:$axis);
+  let results = (outs quant_RealValueType);
+  let hasVerifier = 1;
+}
+
+def Quantization_CoupledRefOp : Quantization_Op<"coupled_ref", [SameOperandsAndResultType]> {
+  let summary = [{
+    Indicates that one point of the computation is coupled to another.
+  }];
+
+  let description = [{
+    Ordinarily, relationships between ops for the purposes of determining
+    compatible quantized types is explicit based on the use-def chain. However,
+    in some situations, a use may be separated from its def by arbitrary
+    external connections. In such a case, during analysis, all coupled_ref
+    nodes in a module which share a coupledKey will be considered to be
+    directly connected as via an identity op for the purpose of type inference.
+  }];
+
+  let arguments = (ins
+    quant_RealValueType:$arg,
+    StrAttr:$coupledKey);
+  let results = (outs quant_RealValueType);
+}
+
+#endif // Quantization_OPS

tensorflow/compiler/mlir/quantization/common/ir/QuantOpsBase.td

@@ -0,0 +1,32 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+//
+//===----------------------------------------------------------------------===//
+//
+// Predicates for types in the Quantization dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef QUANTIZATION_BASE
+#define QUANTIZATION_BASE
+
+include "mlir/IR/OpBase.td"
+
+def TF_Quant_Dialect : Dialect {
+  let name = "quantization";
+  let cppNamespace = "::mlir::quant::ir";
+}
+
+#endif // QUANTIZATION_BASE

tensorflow/compiler/mlir/quantization/common/ir/QuantizeUtils.cc

@@ -0,0 +1,148 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantizeUtils.h"
+
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/quantization/common/ir/UniformSupport.h"
+
+namespace mlir {
+namespace quant::ir {
+
+/// Converts a possible primitive, real expressed value attribute to a
+/// corresponding storage attribute (typically FloatAttr -> IntegerAttr).
+/// quantizedElementType is the QuantizedType that describes the expressed
+/// origValue.
+/// Returns a converter Attribute or nullptr if conversion is not possible.
+static Attribute convertPrimitiveValueAttr(
+    Attribute origRealValue, quant::QuantizedType quantizedElementType,
+    const UniformQuantizedValueConverter &converter, Type &outConvertedType) {
+  if (mlir::isa<FloatAttr>(origRealValue)) {
+    FloatAttr floatAttr = mlir::cast<FloatAttr>(origRealValue);
+    outConvertedType = quantizedElementType.getStorageType();
+    return IntegerAttr::get(quantizedElementType.getStorageType(),
+                            converter.quantizeFloatToInt(floatAttr.getValue()));
+  }
+
+  return nullptr;
+}
+
+/// Converts a real expressed DenseFPElementsAttr to a corresponding
+/// DenseElementsAttr (typically DenseIntElementsAttr) containing quantized
+/// storage values assuming the given quantizedElementType and converter.
+static DenseElementsAttr convertDenseFPElementsAttr(
+    DenseFPElementsAttr realFPElementsAttr,
+    quant::QuantizedType quantizedElementType,
+    const UniformQuantizedValueConverter &converter) {
+  return realFPElementsAttr.mapValues(
+      quantizedElementType.getStorageType(),
+      [&converter](const APFloat &realVal) {
+        return converter.quantizeFloatToInt(realVal);
+      });
+}
+
+/// Converts a real expressed SplatElementsAttr to a corresponding
+/// SplatElementsAttr containing quantized storage values assuming the given
+/// quantizedElementType and converter.
+static SparseElementsAttr convertSparseElementsAttr(
+    SparseElementsAttr realSparseAttr,
+    quant::QuantizedType quantizedElementType,
+    const UniformQuantizedValueConverter &converter) {
+  DenseElementsAttr realDenseAttr = realSparseAttr.getValues();
+  if (!mlir::isa<DenseFPElementsAttr>(realDenseAttr)) {
+    return nullptr;
+  }
+  DenseElementsAttr quantDenseAttr =
+      convertDenseFPElementsAttr(mlir::cast<DenseFPElementsAttr>(realDenseAttr),
+                                 quantizedElementType, converter);
+  if (!quantDenseAttr) {
+    return nullptr;
+  }
+
+  // Cast from an expressed-type-based type to storage-type-based type,
+  // preserving the sparse shape (i.e. tensor<4xf32> -> tensor<4xi8>).
+  ShapedType newSparseType = mlir::dyn_cast_or_null<ShapedType>(
+      quantizedElementType.castExpressedToStorageType(
+          realSparseAttr.getType()));
+  if (!newSparseType) {
+    return nullptr;
+  }
+  return SparseElementsAttr::get(newSparseType, realSparseAttr.getIndices(),
+                                 quantDenseAttr);
+}
+
+/// Converts a real expressed Attribute to a corresponding Attribute containing
+/// quantized storage values assuming the given uniform quantizedElementType and
+/// converter.
+Attribute quantizeAttrUniform(Attribute realValue,
+                              quant::UniformQuantizedType quantizedElementType,
+                              const UniformQuantizedValueConverter &converter,
+                              Type &outConvertedType) {
+  // Fork to handle different variants of constants supported.
+  if (mlir::isa<DenseFPElementsAttr>(realValue)) {
+    // Dense tensor or vector constant.
+    auto converted =
+        convertDenseFPElementsAttr(mlir::cast<DenseFPElementsAttr>(realValue),
+                                   quantizedElementType, converter);
+    outConvertedType = converted.getType();
+    return converted;
+  }
+  if (mlir::isa<SparseElementsAttr>(realValue)) {
+    // Sparse tensor or vector constant.
+    auto converted =
+        convertSparseElementsAttr(mlir::cast<SparseElementsAttr>(realValue),
+                                  quantizedElementType, converter);
+    outConvertedType = converted.getType();
+    return converted;
+  }
+  // Nothing else matched: try to convert a primitive.
+  return convertPrimitiveValueAttr(realValue, quantizedElementType, converter,
+                                   outConvertedType);
+}
+
+/// Convert an attribute from a type based on
+/// quantizedElementType.getExpressedType() to one based on
+/// quantizedElementType.getStorageType().
+/// Returns nullptr if the conversion is not supported.
+/// On success, stores the converted type in outConvertedType.
+Attribute quantizeAttr(Attribute realValue,
+                       quant::QuantizedType quantizedElementType,
+                       Type &outConvertedType) {
+  if (auto uniformQuantized =
+          mlir::dyn_cast<quant::UniformQuantizedType>(quantizedElementType)) {
+    UniformQuantizedValueConverter converter(uniformQuantized);
+    return quantizeAttrUniform(realValue, uniformQuantized, converter,
+                               outConvertedType);
+  }
+  if (auto uniformQuantizedPerAxis =
+          mlir::dyn_cast<quant::UniformQuantizedPerAxisType>(
+              quantizedElementType)) {
+    UniformQuantizedPerAxisValueConverter converter(uniformQuantizedPerAxis);
+    auto converted = converter.convert(realValue);
+    // TODO: why we need this outConvertedType? remove it?
+    if (converted) {
+      outConvertedType = converted.getType();
+    }
+    return converted;
+  }
+  return nullptr;
+}
+
+}  // namespace quant::ir
+}  // namespace mlir

tensorflow/compiler/mlir/quantization/common/ir/QuantizeUtils.h

@@ -0,0 +1,71 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_IR_QUANTIZEUTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_IR_QUANTIZEUTILS_H_
+
+namespace mlir {
+class Attribute;
+class Type;
+
+namespace quant {
+class QuantizedType;
+
+namespace ir {
+class UniformQuantizedType;
+class UniformQuantizedValueConverter;
+
+/// Converts an attribute from a type based on
+/// quantizedElementType.getExpressedType() to one based on
+/// quantizedElementType.getStorageType(), where quantizedElementType is as from
+/// QuantizedType::getQuantizedElementType().
+/// Returns nullptr if the conversion is not supported. On success, stores the
+/// converted type in outConvertedType.
+///
+/// Examples:
+/// 1. realValue is a primitive value attribute:
+/// (realValue: FloatAttr, quantizedElementType: UniformQuantizedType[i8:f32])
+///   -> (IntegerAttr, outConvertedType: i8)
+/// 2. realValue is an elements attribute:
+/// (realValue: DenseElementsAttr[tensor<2x2xf32>],
+///  quantizedElementType: UniformQuantizedType[i8:f32])
+///   -> (DenseElementsAttr[tensor<2x2xi8>], outConvertedType: tensor<2x2xi8>)
+Attribute quantizeAttr(Attribute realValue, QuantizedType quantizedElementType,
+                       Type &outConvertedType);
+
+/// Converts an attribute from a type based on
+/// quantizedElementType.getExpressedType() to one based on
+/// quantizedElementType.getStorageType(), where quantizedElementType is as from
+/// QuantizedType::getQuantizedElementType() and casted to an
+/// UniformQuantizedType. Returns nullptr if the conversion is not supported. On
+/// success, stores the converted type in outConvertedType.
+///
+/// Examples:
+/// 1. realValue is a primitive value attribute:
+/// (realValue: FloatAttr, quantizedElementType: UniformQuantizedType[i8:f32])
+///   -> (IntegerAttr, outConvertedType: i8)
+/// 2. realValue is an elements attribute:
+/// (realValue: DenseElementsAttr[tensor<2x2xf32>],
+///  quantizedElementType: UniformQuantizedType[i8:f32])
+///   -> (DenseElementsAttr[tensor<2x2xi8>], outConvertedType: tensor<2x2xi8>)
+Attribute quantizeAttrUniform(Attribute realValue,
+                              UniformQuantizedType quantizedElementType,
+                              const UniformQuantizedValueConverter &converter,
+                              Type &outConvertedType);
+}  // namespace ir
+}  // namespace quant
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_IR_QUANTIZEUTILS_H_

tensorflow/compiler/mlir/quantization/common/ir/UniformSupport.cc

@@ -0,0 +1,112 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/quantization/common/ir/UniformSupport.h"
+
+#include <cassert>
+#include <cstddef>
+#include <cstdint>
+#include <functional>
+#include <iterator>
+#include <numeric>
+
+#include "mlir/Dialect/Quant/IR/QuantTypes.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/Types.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace mlir::quant::ir {
+
+static bool isQuantizablePrimitiveType(Type input_type) {
+  return isa<FloatType>(input_type);
+}
+
+ExpressedToQuantizedConverter ExpressedToQuantizedConverter::forInputType(
+    Type input_type) {
+  if (isa<TensorType, VectorType>(input_type)) {
+    Type element_type = cast<ShapedType>(input_type).getElementType();
+    if (!isQuantizablePrimitiveType(element_type))
+      return ExpressedToQuantizedConverter{input_type, nullptr};
+    return ExpressedToQuantizedConverter{input_type, element_type};
+  }
+  // Supported primitive type (which just is the expressed type).
+  if (isQuantizablePrimitiveType(input_type))
+    return ExpressedToQuantizedConverter{input_type, input_type};
+  // Unsupported.
+  return ExpressedToQuantizedConverter{input_type, nullptr};
+}
+
+Type ExpressedToQuantizedConverter::convert(
+    quant::QuantizedType elemental_type) const {
+  assert(expressed_type && "convert() on unsupported conversion");
+  if (auto tensor_type = dyn_cast<RankedTensorType>(input_type))
+    return RankedTensorType::get(tensor_type.getShape(), elemental_type);
+  if (auto tensor_type = dyn_cast<UnrankedTensorType>(input_type))
+    return UnrankedTensorType::get(elemental_type);
+  if (auto vector_type = dyn_cast<VectorType>(input_type))
+    return VectorType::get(vector_type.getShape(), elemental_type);
+
+  // If the expressed types match, just use the new elemental type.
+  if (elemental_type.getExpressedType() == expressed_type) {
+    return elemental_type;
+  }
+  // Unsupported.
+  return nullptr;
+}
+
+ElementsAttr UniformQuantizedPerAxisValueConverter::convert(
+    Attribute real_value) {
+  if (auto attr = dyn_cast<DenseFPElementsAttr>(real_value)) {
+    return convert(attr);
+  }
+  return nullptr;
+}
+
+DenseElementsAttr UniformQuantizedPerAxisValueConverter::convert(
+    DenseFPElementsAttr attr) {
+  // Creates the converter for each chunk. Normally the size of the
+  // quantization dim is 3, so we can cache all the converters.
+  ShapedType type = attr.getType();
+  std::size_t dim_size = type.getDimSize(quantization_dim_);
+  if (dim_size != scales_.size()) {
+    return {};
+  }
+  SmallVector<UniformQuantizedValueConverter, 4> converters;
+  converters.reserve(dim_size);
+  for (int i = 0, e = dim_size; i != e; ++i) {
+    converters.push_back(getPerChunkConverter(i));
+  }
+
+  // Scan the elements of the dense elements attributes and quantize them by
+  // using the right quantization parameters.
+  int64_t flatten_index = 0;
+  auto shape = type.getShape();
+  int64_t chunk_size =
+      std::accumulate(std::next(shape.begin(), quantization_dim_ + 1),
+                      shape.end(), 1, std::multiplies<int64_t>());
+  Type new_element_type =
+      IntegerType::get(attr.getContext(), storage_bit_width_);
+  return attr.mapValues(new_element_type, [&](const APFloat &old) {
+    int chunk_index = flatten_index / chunk_size;
+    flatten_index++;
+    return converters[chunk_index % dim_size].quantizeFloatToInt(old);
+  });
+}
+
+}  // namespace mlir::quant::ir

tensorflow/compiler/mlir/quantization/common/ir/UniformSupport.h

@@ -0,0 +1,247 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_IR_UNIFORMSUPPORT_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_IR_UNIFORMSUPPORT_H_
+
+#include <algorithm>
+#include <cassert>
+#include <cmath>
+#include <cstdint>
+#include <cstring>
+#include <utility>
+
+#include "llvm/ADT/APFloat.h"
+#include "llvm/ADT/APInt.h"
+#include "llvm/ADT/APSInt.h"
+#include "mlir/Dialect/Quant/IR/QuantTypes.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/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace mlir::quant::ir {
+
+// Performs type conversion from an arbitrary input type to a type
+// that is expressed by a QuantizedType.
+//
+// This handles cases where the inputType is a supported primitive type
+// (i.e. f32, bf16, etc) or a vector/tensor type based on a supported
+// elemental type.
+//
+// Since conversion often involves introspecting some attributes of the
+// input type in order to determine how to represent it, this is a two step
+// process.
+struct ExpressedToQuantizedConverter {
+  // Creates a converter for the given input type.
+  static ExpressedToQuantizedConverter forInputType(Type input_type);
+
+  // Converts the inputType to be based on the given elemental type,
+  // returning the new type (or nullptr and emit an error on failure).
+  Type convert(quant::QuantizedType elemental_type) const;
+
+  // Whether the conversion is legal.
+  explicit operator bool() const { return (bool)expressed_type; }
+
+  // The input type that is being converted from.
+  // This may be an elemental or composite type.
+  const Type input_type;
+
+  // Supported, elemental expressed type (i.e. f32).
+  // Will be nullptr if conversion is not supported.
+  const Type expressed_type;
+};
+
+// Reference implementation of converting between real numbers and values
+// represented by a UniformQuantizedType.
+// Note that this is not expected to be speedy and may be superseded eventually
+// by a more optimal implementation.
+// Also, the interface assumes that quantization is done per-layer and will
+// need to be wider for various per-channel schemes. As such, this is a
+// placeholder.
+class UniformQuantizedValueConverter {
+ public:
+  explicit UniformQuantizedValueConverter(
+      quant::UniformQuantizedType uniform_type)
+      : UniformQuantizedValueConverter(
+            uniform_type.getScale(),
+            static_cast<double>(uniform_type.getZeroPoint()),
+            static_cast<double>(uniform_type.getStorageTypeMin()),
+            static_cast<double>(uniform_type.getStorageTypeMax()),
+            uniform_type.getStorageTypeIntegralWidth(),
+            uniform_type.isSigned()) {
+    assert(isa<FloatType>(uniform_type.getExpressedType()));
+    assert(uniform_type.getStorageType().isSignlessInteger());
+  }
+
+  UniformQuantizedValueConverter(double scale, double zero_point,
+                                 double clamp_min, double clamp_max,
+                                 uint32_t storage_bit_width, bool is_signed)
+      : scale_(scale),
+        zero_point_(zero_point),
+        clamp_min_(clamp_min),
+        clamp_max_(clamp_max),
+        scale_double_(scale),
+        zero_point_double_(zero_point),
+        clamp_min_double_(clamp_min),
+        clamp_max_double_(clamp_max),
+        storage_bit_width_(storage_bit_width),
+        is_signed_(is_signed),
+        round_mode_(APFloat::rmNearestTiesToAway) {}
+
+  UniformQuantizedValueConverter(double scale, double zero_point,
+                                 const APFloat& clamp_min,
+                                 const APFloat& clamp_max,
+                                 uint32_t storage_bit_width, bool is_signed)
+      : scale_(scale),
+        zero_point_(zero_point),
+        clamp_min_(clamp_min),
+        clamp_max_(clamp_max),
+        scale_double_(scale),
+        zero_point_double_(zero_point),
+        clamp_min_double_(clamp_min.convertToDouble()),
+        clamp_max_double_(clamp_max.convertToDouble()),
+        storage_bit_width_(storage_bit_width),
+        is_signed_(is_signed),
+        round_mode_(APFloat::rmNearestTiesToAway) {}
+
+  virtual APInt quantizeFloatToInt(APFloat expressed_value) const {
+    // This function is a performance critical code path in quantization
+    // since it runs for each single float parameter value.
+
+    // Specialize f32->u8/i8 case to optimize performance.
+    if (&expressed_value.getSemantics() == &APFloat::IEEEsingle() &&
+        storage_bit_width_ == 8 &&
+        round_mode_ == llvm::APFloatBase::rmNearestTiesToAway) {
+      return quantizeF32ToInt8(expressed_value);
+    }
+
+    bool lossy;
+    expressed_value.convert(scale_.getSemantics(), round_mode_, &lossy);
+    // fixed_point = clamp(clamp_min, clamp_max, (
+    //   roundHalfToEven(expressed / scale) + zero_point))
+    APFloat scaled = (expressed_value / scale_);
+    scaled.roundToIntegral(round_mode_);
+    scaled.add(zero_point_, round_mode_);
+    APFloat fixed_point = llvm::minimum(scaled, clamp_max_);
+    fixed_point = llvm::maximum(fixed_point, clamp_min_);
+
+    llvm::APSInt result(storage_bit_width_, !is_signed_);
+    fixed_point.convertToInteger(result, round_mode_, &lossy);
+
+    return std::move(result);
+  }
+
+  int64_t quantizeFloatToInt64(APFloat expressed_value) const {
+    const APInt q_value = quantizeFloatToInt(std::move(expressed_value));
+    return is_signed_ ? q_value.getSExtValue() : q_value.getZExtValue();
+  }
+
+  virtual ~UniformQuantizedValueConverter() = default;
+
+ private:
+  // An optimized implementation to quantize f32 to i8/u8 with C++ native
+  // arithmetic.
+  virtual APInt quantizeF32ToInt8(const APFloat& expressed_value) const {
+    assert(&expressed_value.getSemantics() == &APFloat::IEEEsingle());
+    assert(storage_bit_width_ == 8);
+    assert(round_mode_ == llvm::APFloatBase::rmNearestTiesToAway);
+
+    const float real_value = expressed_value.convertToFloat();
+
+    const double scaled = real_value / scale_double_ + zero_point_double_;
+    // Round to nearest integer with halfway cases rounded away from zero.
+    const double scaled_rounded = std::round(scaled);
+    const double clamped = std::min(std::max(scaled_rounded, clamp_min_double_),
+                                    clamp_max_double_);
+
+    uint64_t signless_result;
+    if (is_signed_) {
+      int64_t clamped_int = static_cast<int8_t>(clamped);
+      memcpy(&signless_result, &clamped_int, sizeof(clamped_int));
+    } else {
+      signless_result = static_cast<uint8_t>(clamped);
+    }
+    return APInt(storage_bit_width_, signless_result, /*isSigned=*/is_signed_);
+  }
+
+  // Keep both APFloat and double versions of the quantization parameters
+  // around since they will be used in generic and specialized arithmetic,
+  // respectively.
+  const APFloat scale_;
+  const APFloat zero_point_;
+  const APFloat clamp_min_;
+  const APFloat clamp_max_;
+
+  const double scale_double_;
+  const double zero_point_double_;
+  const double clamp_min_double_;
+  const double clamp_max_double_;
+
+  const uint32_t storage_bit_width_;
+  const bool is_signed_;
+  const llvm::APFloat::roundingMode round_mode_;
+};
+
+// An utility class to quantize an attribute by the per-axis quantization
+// parameters. The size of the quantization dim in the converted elements
+// attribute should match the size of of scales/zero_points vectors in the
+// quantization parameters.
+class UniformQuantizedPerAxisValueConverter {
+ public:
+  explicit UniformQuantizedPerAxisValueConverter(
+      quant::UniformQuantizedPerAxisType uniform_type)
+      : scales_(uniform_type.getScales()),
+        zero_points_(uniform_type.getZeroPoints()),
+        clamp_min_(static_cast<double>(uniform_type.getStorageTypeMin())),
+        clamp_max_(static_cast<double>(uniform_type.getStorageTypeMax())),
+        storage_bit_width_(uniform_type.getStorageTypeIntegralWidth()),
+        is_signed_(uniform_type.isSigned()),
+        quantization_dim_(uniform_type.getQuantizedDimension()) {
+    assert(isa<FloatType>(uniform_type.getExpressedType()));
+    assert(uniform_type.getStorageType().isSignlessInteger());
+    assert(scales_.size() == zero_points_.size());
+  }
+
+  // Quantize an Attribute by the quantization parameters. Return nullptr if
+  // the conversion fails or the input array isn't an ElementsAttr.
+  ElementsAttr convert(Attribute real_value);
+
+ private:
+  // Quantize an DenseFPElementsAttr by the quantization parameters.
+  DenseElementsAttr convert(DenseFPElementsAttr attr);
+
+  // Get a uniform converter for the index-th chunk along the quantizationDim.
+  // All the elements in this chunk is quantized by the returned converter.
+  UniformQuantizedValueConverter getPerChunkConverter(int index) const {
+    return UniformQuantizedValueConverter(scales_[index], zero_points_[index],
+                                          clamp_min_, clamp_max_,
+                                          storage_bit_width_, is_signed_);
+  }
+
+  const ArrayRef<double> scales_;
+  const ArrayRef<int64_t> zero_points_;
+  const APFloat clamp_min_;
+  const APFloat clamp_max_;
+  const uint32_t storage_bit_width_;
+  const bool is_signed_;
+  int32_t quantization_dim_;
+};
+
+}  // namespace mlir::quant::ir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_IR_UNIFORMSUPPORT_H_

tensorflow/compiler/mlir/quantization/common/lift_as_function_call.cc

@@ -0,0 +1,551 @@
+/* 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/common/lift_as_function_call.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <optional>
+#include <queue>
+#include <stack>
+#include <string>
+
+#include "absl/algorithm/container.h"
+#include "absl/base/nullability.h"
+#include "absl/container/flat_hash_set.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/match.h"
+#include "absl/strings/str_cat.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/Sequence.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.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/Diagnostics.h"  // from @llvm-project
+#include "mlir/IR/Location.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/TypeRange.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/IR/ValueRange.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "stablehlo/dialect/Version.h"  // from @stablehlo
+#include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.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/stablehlo/utils/stablehlo_type_utils.h"
+#include "tensorflow/compiler/mlir/quantization/tensorflow/cc/quantization_unit_loc.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"
+#include "tensorflow/core/platform/mutex.h"
+#include "tsl/platform/protobuf.h"  // IWYU pragma: keep
+
+namespace mlir::quant {
+
+using ::stablehlo::quantization::Method;
+using ::tsl::protobuf::TextFormat;
+
+// Default version number for native serialization.
+constexpr int64_t kDefaultVersion = 9;
+// Default platform for XlaCallModuleOp.
+constexpr StringRef kPlatformCpu = "CPU";
+// Name of `tf.XlaCallModule`'s dictionary attribute for keeping the
+// deserialized stablehlo module's attributes.
+constexpr StringRef kStablehloModuleAttrsAttrName = "_stablehlo_module_attrs";
+// Attribute required for running shape refinement pass enabled in XlaCallModule
+// version 8 and above.
+constexpr StringRef kUsesShapePolymorphismAttr = "jax.uses_shape_polymorphism";
+
+bool IsInLiftedFunc(Operation* op) {
+  if (op == nullptr) return false;
+  return op->getParentOfType<func::FuncOp>()->hasAttr(kFusedFunctionAttr);
+}
+
+bool IsInStableHloOpRegion(Operation* op) {
+  if (op == nullptr) return false;
+  auto parent_op = op->getParentOp();
+  return parent_op != nullptr && quant::stablehlo::IsStablehloOp(parent_op);
+}
+
+// Inserts the function to the symbol table of the module thread-safely.
+StringAttr InsertToSymbolTable(Operation& module, Operation& function,
+                               const StringRef func_name) {
+  static tensorflow::mutex* mtx = new tensorflow::mutex();
+  tensorflow::mutex_lock lock(*mtx);
+
+  SymbolTable symbol_table(&module);
+  std::string unique_name = func_name.str();
+  int32_t uniquing_counter = 0;
+  while (symbol_table.lookup(unique_name) != nullptr) {
+    ++uniquing_counter;
+    unique_name = absl::StrCat(func_name.str(), "_", uniquing_counter);
+  }
+  function.setAttr("sym_name",
+                   StringAttr::get(module.getContext(), unique_name));
+  return symbol_table.insert(&function);
+}
+
+// Creates the TF::PartitionedCallOp with the given arguments and output types.
+// This function call op is for invoking the TF subgraphs.
+ValueRange CreateTFPartitionedCallOp(OpBuilder& builder,
+                                     const Location location,
+                                     const StringRef func_name,
+                                     const TypeRange output_types,
+                                     const ValueRange args) {
+  TF::PartitionedCallOp call_op = TF::PartitionedCallOp::create(
+      builder, location, output_types, args,
+      /*args_attrs=*/nullptr, /*res_attrs=*/nullptr,
+      FlatSymbolRefAttr::get(builder.getStringAttr(func_name)),
+      /*config=*/"", /*config_proto=*/"", /*executor_type=*/"");
+
+  // Set the attribute to annotate this function call op as a quantizable spot.
+  call_op->setAttr(
+      kQuantTraitAttrName,
+      builder.getStringAttr(StringRef(
+          std::string(QuantTraitValues[QuantizationTrait::FullyQuantizable]))));
+
+  return call_op.getOutput();
+}
+
+// Creates the TF::XlaCallModuleOp with the given arguments and output types.
+// This function call op is for invoking the StableHLO subgraphs.
+ValueRange CreateTFXlaCallModuleOp(OpBuilder& builder, const Location location,
+                                   const StringRef func_name,
+                                   const TypeRange output_types,
+                                   const ValueRange args) {
+  MLIRContext* ctx = builder.getContext();
+  // Collect the shapes of the output to fill up the Sout attribute.
+  SmallVector<Attribute> shape_attrs;
+  for (const Type result_type : output_types) {
+    shape_attrs.push_back(
+        tf_type::ShapeAttr::get(ctx, mlir::cast<ShapedType>(result_type)));
+  }
+  auto empty_array_attr = ArrayAttr::get(ctx, {});
+  auto platforms = ArrayAttr::get(ctx, {StringAttr::get(ctx, kPlatformCpu)});
+
+  auto call_op =
+      TF::XlaCallModuleOp::create(builder, location,
+                                  /*output=*/output_types,
+                                  /*args=*/args,
+                                  /*version=*/kDefaultVersion, /*module=*/"",
+                                  /*Sout=*/ArrayAttr::get(ctx, shape_attrs),
+                                  /*dim_args_spec=*/empty_array_attr,
+                                  /*platforms=*/platforms,
+                                  /*function_list=*/empty_array_attr,
+                                  /*has_token_input_output=*/false,
+                                  /*disabled_checks=*/empty_array_attr);
+
+  // Set the function name. This will be controlled by the
+  // XlaCallModuleSerialization related passes directly, which means that the
+  // function name can be changed by those passes.
+  call_op->setAttr(TF::kStablehloEntryFunctionAttrName,
+                   FlatSymbolRefAttr::get(builder.getStringAttr(func_name)));
+
+  // Set target version to WEEK_4 since this is an offline quantizer.
+  std::string target_version =
+      mlir::vhlo::Version::fromCompatibilityRequirement(
+          vhlo::Version::CompatibilityRequirement::WEEK_4)
+          .toString();
+  call_op->setAttr(TF::kStablehloVersionAttrName,
+                   builder.getStringAttr(target_version));
+
+  // Store the custom attribute to restore the function name when loading it
+  // back in the post calibration stage. As mentioned above, the above entry
+  // function attribute is not reliable.
+  call_op->setAttr(kOriginalStablehloEntryFunctionAttrName,
+                   builder.getStringAttr(func_name));
+
+  // Set the attribute to annotate this function call op as a quantizable spot.
+  call_op->setAttr(
+      kQuantTraitAttrName,
+      builder.getStringAttr(StringRef(
+          std::string(QuantTraitValues[QuantizationTrait::FullyQuantizable]))));
+
+  // Set jax.uses_shape_polymorphism=true to enable shape refinement at runtime.
+  // This is needed for native serialization version >= 8.
+  call_op->setAttr(kStablehloModuleAttrsAttrName,
+                   builder.getDictionaryAttr(builder.getNamedAttr(
+                       kUsesShapePolymorphismAttr, builder.getBoolAttr(true))));
+
+  return call_op.getOutput();
+}
+
+// Creates the function call op based on the given call_op_type argument.
+ValueRange CreateFunctionCallOp(OpBuilder& builder, const Location location,
+                                const FunctionCallOpType call_op_type,
+                                const StringRef func_name,
+                                const TypeRange output_types,
+                                const ValueRange args) {
+  switch (call_op_type) {
+    case FunctionCallOpType::TFXlaCallModuleOp:
+      return CreateTFXlaCallModuleOp(builder, location, func_name, output_types,
+                                     args);
+    case FunctionCallOpType::TFPartitionedCallOp:
+      return CreateTFPartitionedCallOp(builder, location, func_name,
+                                       output_types, args);
+  }
+}
+
+// Finds ops in the paths from arguments to results. The ops is listed in an
+// order that the former ops shouldn't have any dependencies on the later ones.
+SmallVector<Operation*> FindOpsFromArgumentsToResults(
+    const ArrayRef<Value> arguments, const ArrayRef<Value> results) {
+  std::queue<Value> value_queue;
+  for (Value result : results) {
+    value_queue.push(result);
+  }
+  absl::flat_hash_set<mlir::detail::ValueImpl*> argument_set;
+  for (Value argument : arguments) {
+    argument_set.insert(argument.getImpl());
+  }
+
+  // Searching for ops from results to arguments. Duplicate ops in the op stack
+  // are intentional in order to make sure the op on the top of the stack
+  // doesn't depends on any ops below it.
+  std::stack<Operation*> op_stack;
+  while (!value_queue.empty()) {
+    Value current_value = value_queue.front();
+    value_queue.pop();
+
+    Operation* defining_node = current_value.getDefiningOp();
+    if (defining_node == nullptr) continue;
+    op_stack.push(defining_node);
+    for (Value arg : defining_node->getOperands()) {
+      if (!argument_set.contains(arg.getImpl())) {
+        value_queue.push(arg);
+      }
+    }
+  }
+
+  // Remove duplicate ops from the op stack.
+  SmallVector<Operation*> sorted_ops;
+  absl::flat_hash_set<Operation*> unique_ops;
+  while (!op_stack.empty()) {
+    Operation* current_op = op_stack.top();
+    op_stack.pop();
+    if (unique_ops.contains(current_op)) continue;
+    sorted_ops.push_back(current_op);
+    unique_ops.insert(current_op);
+  }
+  return sorted_ops;
+}
+
+// Finds the name of each attribute in `attributes` and set the attr_map
+// attribute which maps an attribute identifier to its attribute name. The
+// identifier is the order of that attribute in `attributes`. This map
+// is then used to set attributes in the quantized functions in the
+// QuantizeCompositeFunctionsPass.
+// For example, for tf.MatMul with `attributes` = {{"transpose_a", false},
+// {"transpose_b", false}}, the generated attr_map is
+// "0:transpose_a,1:transpose_b", where 0 and 1 are the respective attribute
+// identifiers.
+// This function returns success if all attributes could be found.
+LogicalResult SetAttributeMap(MLIRContext& context,
+                              const ArrayRef<NamedAttribute> attributes,
+                              const ArrayRef<Operation*> ops) {
+  // A map to find which operation an attribute belongs to.
+  // The key for this map uses the entire NamedAttribute object, i.e. the
+  // {attribute_name, attribute_value} pair.
+  llvm::SmallDenseMap<NamedAttribute, Operation*> attr_to_op_map;
+  for (Operation* op : ops) {
+    for (const NamedAttribute named_attr : op->getAttrs()) {
+      attr_to_op_map.insert({named_attr, op});
+    }
+  }
+
+  for (int idx : llvm::seq<int>(0, attributes.size())) {
+    const NamedAttribute& attribute = attributes[idx];
+    // Skip the following steps if the attribute value is `NullAttribute`.
+    if (const auto string_attr =
+            mlir::dyn_cast_or_null<StringAttr>(attribute.getValue());
+        string_attr != nullptr &&
+        string_attr.getValue() == kNullAttributeValue) {
+      continue;
+    }
+
+    if (std::find_if(
+            attr_to_op_map.begin(), attr_to_op_map.end(), [&](auto attr_op) {
+              return std::get<0>(attr_op).getName() == attribute.getName();
+            }) == attr_to_op_map.end()) {
+      emitError(UnknownLoc::get(&context),
+                "Could not find attribute: " + attribute.getName().str());
+      return failure();
+    }
+
+    Operation* owner_op;
+    for (const auto& [attr, val] : attr_to_op_map) {
+      if (attr.getName() == attribute.getName()) owner_op = val;
+    }
+    if (quant::stablehlo::IsStablehloOp(owner_op)) {
+      owner_op->setAttr(StringRef(attribute.getName()), attribute.getValue());
+    } else {
+      owner_op = attr_to_op_map[attribute];
+
+      std::string new_attr_map_str{};
+      if (owner_op->hasAttr(kAttrMapAttribute)) {
+        new_attr_map_str =
+            owner_op->getAttrOfType<StringAttr>(kAttrMapAttribute).str();
+        absl::StrAppend(&new_attr_map_str, ",");
+      }
+
+      // Append "<identifier>:<attribute_name>". Ex) "0:transpose_a".
+      const std::string identifier = std::to_string(idx);
+      const StringAttr attribute_name = attribute.getName();
+      absl::StrAppend(&new_attr_map_str, identifier, ":", attribute_name.str());
+      owner_op->setAttr(kAttrMapAttribute,
+                        StringAttr::get(&context, new_attr_map_str));
+    }
+  }
+  return success();
+}
+
+// Creates a function to wrap the section between arguments and results.
+SmallVector<Value, 4> LiftAsFunctionCall(
+    OpBuilder& builder, const Location location,
+    const FunctionCallOpType call_op_type, const StringRef func_name,
+    const ArrayRef<Value> arguments, const ArrayRef<Value> results,
+    const ArrayRef<NamedAttribute> attributes) {
+  MLIRContext* context = builder.getContext();
+  if (results.empty()) {
+    emitError(UnknownLoc::get(context), "No result values specified");
+    return {};
+  }
+  Operation* result_op = results[0].getDefiningOp();
+  auto module = result_op->getParentOfType<ModuleOp>();
+
+  // Create a private function and copy all ops between arguments and results.
+  auto current_func = result_op->getParentOfType<func::FuncOp>();
+  auto guard = OpBuilder::InsertionGuard(builder);
+  builder.setInsertionPointAfter(current_func);
+  TypeRange arg_types{ValueRange{arguments}};
+  TypeRange result_types{ValueRange{results}};
+  auto func_type = FunctionType::get(context, arg_types, result_types);
+
+  SmallVector<Location> arg_locs;
+  for (Value arg : arguments) {
+    arg_locs.push_back(arg.getLoc());
+  }
+
+  auto wrap_func =
+      func::FuncOp::create(builder, location, func_name, func_type);
+  wrap_func.setVisibility(SymbolTable::Visibility::Private);
+  // The callee function for TF::XlaCallModuleOp must have this attribute.
+  if (call_op_type == FunctionCallOpType::TFXlaCallModuleOp) {
+    wrap_func->setAttr(TF::kFromXlaCallModuleAttrName, builder.getUnitAttr());
+  }
+  wrap_func->setAttr(kFusedFunctionAttr, builder.getUnitAttr());
+  builder.createBlock(&wrap_func.getBody(), wrap_func.begin(), arg_types,
+                      arg_locs);
+
+  IRMapping mapping;
+  for (int32_t i : llvm::seq<int32_t>(0, arguments.size())) {
+    mapping.map(arguments[i], wrap_func.getArgument(i));
+  }
+
+  auto cloning_ops = FindOpsFromArgumentsToResults(arguments, results);
+  // Set the location of call op to QuantizationUnitLoc if found.
+  Location call_op_loc = location;
+  for (Operation* op : cloning_ops) {
+    std::optional<quant::QuantizationUnitLoc::QuantizationUnit> unit =
+        quant::FindQuantizationUnitFromLoc(op->getLoc());
+    if (unit.has_value()) {
+      call_op_loc =
+          quant::QuantizationUnitLoc(builder.getContext(), unit.value());
+    }
+  }
+
+  if (failed(SetAttributeMap(*context, attributes, cloning_ops))) {
+    current_func.emitError() << "Some attributes couldn't be found.";
+  }
+  for (Operation* op : cloning_ops) {
+    builder.clone(*op, mapping);
+  }
+
+  SmallVector<Value> return_values;
+  for (Value result : results) {
+    return_values.push_back(mapping.lookupOrNull(result));
+  }
+  func::ReturnOp::create(builder, location, return_values);
+
+  // Create a function call to the newly created function.
+  StringAttr new_func_name =
+      InsertToSymbolTable(*module, *wrap_func, func_name);
+  builder.setInsertionPointAfter(result_op);
+  ValueRange new_results =
+      CreateFunctionCallOp(builder, call_op_loc, call_op_type,
+                           new_func_name.getValue(), result_types, arguments);
+  return SmallVector<Value, 4>(new_results.begin(), new_results.end());
+}
+
+SmallVector<Value, 4> LiftAsFunctionCall(OpBuilder& builder,
+                                         const Location location,
+                                         const FunctionCallOpType call_op_type,
+                                         const StringRef func_name,
+                                         const ArrayRef<Value> arguments,
+                                         const ArrayRef<Value> results) {
+  SmallVector<NamedAttribute> attributes;
+  return LiftAsFunctionCall(builder, location, call_op_type, func_name,
+                            arguments, results, attributes);
+}
+
+SmallVector<Value> AppendToVector(const ArrayRef<Value> arguments,
+                                  Value append) {
+  SmallVector<Value> ret(arguments);
+  ret.push_back(append);
+  return ret;
+}
+
+// Check if the given einsum equation is supported by XlaDotV2.
+// Conditions:
+// 1. Two inputs & one output.
+// 2. No ... in the equation.
+// 3. Batch dimensions should be the same, or only the left equation should have
+//    the batch dimension. This condition is from the XlaDotV2 specification. It
+//    could process the following equation by setting the attributes properly:
+//    abc,cd->abd.
+// 4. The output should be in the form: [batch dims][lhs dims][rhs dims]
+bool IsEinsumSupportedByXlaDotV2(StringAttr equation_attr) {
+  StringRef equation = equation_attr.getValue();
+
+  if (!absl::StrContains(equation, "->") || !absl::StrContains(equation, ",") ||
+      absl::StrContains(equation, ".")) {
+    return false;
+  }
+
+  // Parse equation.
+  int idx_arrow = equation.find("->");
+  StringRef calc_eq = equation.substr(0, idx_arrow);
+  StringRef out_eq = equation.substr(idx_arrow + 2);
+
+  int idx_comma = calc_eq.find(',');
+  StringRef lhs_eq = calc_eq.substr(0, idx_comma);
+  StringRef rhs_eq = calc_eq.substr(idx_comma + 1);
+
+  if (absl::StrContains(rhs_eq, ",")) return false;
+
+  int lhs_out_idx_start = out_eq.size();
+  int lhs_out_idx_end = -1;
+  int rhs_out_idx_start = out_eq.size();
+  int rhs_out_idx_end = -1;
+  int lhs_batch_dim_size = 0;
+  int rhs_batch_dim_size = 0;
+  for (const char c : lhs_eq) {
+    if (absl::StrContains(out_eq, c) && absl::StrContains(rhs_eq, c)) {
+      lhs_batch_dim_size++;
+    } else if (absl::StrContains(out_eq, c)) {
+      const int out_idx = out_eq.find(c);
+      if (out_idx < lhs_out_idx_end) {
+        // Left-hand equation is reversed in the output.
+        return false;
+      }
+      lhs_out_idx_start = std::min(lhs_out_idx_start, out_idx);
+      lhs_out_idx_end = std::max(lhs_out_idx_end, out_idx);
+    }
+  }
+
+  for (const char c : rhs_eq) {
+    if (absl::StrContains(out_eq, c) && absl::StrContains(lhs_eq, c)) {
+      rhs_batch_dim_size++;
+    } else if (absl::StrContains(out_eq, c)) {
+      int out_idx = out_eq.find(c);
+      if (out_idx < rhs_out_idx_end) {
+        return false;
+      }
+      if (out_idx < rhs_out_idx_start) rhs_out_idx_start = out_idx;
+      if (out_idx > rhs_out_idx_end) rhs_out_idx_end = out_idx;
+    }
+  }
+
+  if (lhs_batch_dim_size != rhs_batch_dim_size && lhs_batch_dim_size != 0 &&
+      rhs_batch_dim_size != 0) {
+    // Batch dimension does not match.
+    return false;
+  }
+
+  // All the lhs equations should come first.
+  if (lhs_out_idx_end > rhs_out_idx_start) return false;
+
+  // All the lhs out dim and rhs out dim should be larger than the batch dims,
+  // and they should not be mixed.
+  int batch_dim_size = std::max(rhs_batch_dim_size, lhs_batch_dim_size);
+  return lhs_out_idx_start >= batch_dim_size &&
+         rhs_out_idx_start >= batch_dim_size;
+}
+
+absl::StatusOr<Method> GetQuantizationMethod(Operation* absl_nonnull op) {
+  const auto quantization_method_attr =
+      op->getAttrOfType<StringAttr>(kQuantizationMethodAttr);
+  if (!quantization_method_attr) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Attribute ", kQuantizationMethodAttr.str(), " is not found."));
+  }
+
+  Method quantization_method;
+  const std::string method_txtpb = quantization_method_attr.getValue().str();
+  if (!TextFormat::ParseFromString(method_txtpb, &quantization_method)) {
+    return absl::InternalError(
+        absl::StrCat("Failed to parse Method from textproto: ", method_txtpb));
+  }
+
+  return quantization_method;
+}
+
+Method GetQuantizationMethodOrDefault(Operation* absl_nonnull op) {
+  absl::StatusOr<Method> method = GetQuantizationMethod(op);
+  if (method.status().code() == absl::StatusCode::kInternal) {
+    // This indicates that the `Method` protobuf string is corrupt, but this
+    // function ignores it and returns the default instance.
+    op->emitError(absl::StrCat("Failed to get quantization method: ",
+                               method.status().ToString()));
+  }
+  return method.ok() ? *method : Method::default_instance();
+}
+
+bool HasWeightOnlyPtqMethod(TF::XlaCallModuleOp xla_call_module_op) {
+  Method method = GetQuantizationMethodOrDefault(xla_call_module_op);
+  return method.has_weight_only_ptq();
+}
+
+bool IsWeightOnlyQuantizableOp(const Operation& op) {
+  if (auto call_op = dyn_cast<TF::XlaCallModuleOp>(op)) {
+    StringRef entry_function_name = GetEntryFunctionName(call_op);
+    absl::StatusOr<Method> quantization_method = GetQuantizationMethod(call_op);
+    return ContainsConvOrDot(entry_function_name) && quantization_method.ok() &&
+           quantization_method->has_weight_only_ptq();
+  }
+  return false;
+}
+
+SmallVector<func::FuncOp> GetSortedFunctions(ModuleOp module_op) {
+  auto iterator_range = module_op.getOps<func::FuncOp>();
+  SmallVector<func::FuncOp> func_ops(iterator_range.begin(),
+                                     iterator_range.end());
+  absl::c_sort(func_ops, [](func::FuncOp op1, func::FuncOp op2) {
+    return op1.getName() < op2.getName();
+  });
+  return func_ops;
+}
+
+}  // namespace mlir::quant

tensorflow/compiler/mlir/quantization/common/lift_as_function_call.h

@@ -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.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_LIFT_AS_FUNCTION_CALL_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_LIFT_AS_FUNCTION_CALL_H_
+
+#include "absl/base/nullability.h"
+#include "absl/status/statusor.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/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+
+namespace mlir::quant {
+
+// This attribute will be set for functions created by this pass.
+// Presence of this attribute will mark the function as quantization target.
+inline constexpr StringRef kFusedFunctionAttr = "tf_quant.composite_function";
+// The keyword to detect if this is a `NullAttribute`.
+inline constexpr StringRef kNullAttributeValue = "N/A";
+
+// Prefixes attached to lifted functions.
+constexpr StringRef kQuantizedFuncPrefix = "quantized_";
+constexpr StringRef kCompositeFuncPrefix = "composite_";
+
+// The attribute will be used for TF::XlaCallModuleOp to restore the original
+// function name when loading it back.
+inline constexpr StringRef kOriginalStablehloEntryFunctionAttrName =
+    "_original_entry_function";
+
+// FunctionCallOpType to be generated as the function call operator when
+// function lifting will happen.
+enum FunctionCallOpType { TFPartitionedCallOp = 0, TFXlaCallModuleOp = 1 };
+
+// Checks if an op is inside a lifted function.
+// If the given op pointer is a nullptr, returns false.
+bool IsInLiftedFunc(Operation* op);
+
+// Checks if the op is inside a StableHLO op with region.
+// If the given op pointer is a nullptr, returns false.
+bool IsInStableHloOpRegion(Operation* op);
+
+// Checks if a given einsum op is supported for XlaDotV2 quantization.
+bool IsEinsumSupportedByXlaDotV2(StringAttr equation_attr);
+
+// Gets the quantization method from `op`. It is retrieved from the
+// `kQuantizationMethodAttr` string attribute. Returns
+// `absl::InvalidArgumentError` when the attribute doesn't exist. Returns
+// `absl::InternalError` when parsing the attribute to `Method` failed.
+// `op` must be non-null.
+absl::StatusOr<::stablehlo::quantization::Method> GetQuantizationMethod(
+    Operation* absl_nonnull op);
+
+// Gets the quantization method from `op`. It is retrieved from the
+// `kQuantizationMethodAttr` string attribute. Returns a default instance of
+// `Method` iff the attribute doesn't exist or the attribute contains an invalid
+// textproto for `Method`. `op` must be non-null.
+::stablehlo::quantization::Method GetQuantizationMethodOrDefault(
+    Operation* absl_nonnull op);
+
+// Creates a function to wrap the section between arguments and results.
+// The generated function call op type will be decided by the given call_op_type
+// argument. Currently, it supports TF::XlaCallModuleOp and
+// TF::PartitionedCallOp function call op generations.
+SmallVector<Value, 4> LiftAsFunctionCall(OpBuilder& builder, Location location,
+                                         FunctionCallOpType call_op_type,
+                                         StringRef func_name,
+                                         ArrayRef<Value> arguments,
+                                         ArrayRef<Value> results,
+                                         ArrayRef<NamedAttribute> attributes);
+
+// Same as above but with empty attributes.
+SmallVector<Value, 4> LiftAsFunctionCall(OpBuilder& builder, Location location,
+                                         FunctionCallOpType call_op_type,
+                                         StringRef func_name,
+                                         ArrayRef<Value> arguments,
+                                         ArrayRef<Value> results);
+
+// Add the second argument to the first argument, which is expected to be an
+// argument list.
+// Used to attach bias to einsum argument list.
+SmallVector<Value> AppendToVector(ArrayRef<Value> arguments, Value append);
+
+// Checks if the `Method` attatched to the given `tf.XlaCallModule` op has
+// `WeightOnlyPtq`.
+bool HasWeightOnlyPtqMethod(TF::XlaCallModuleOp xla_call_module_op);
+
+// Checks if an op is a `tf.XlaCallModule` op, contains 'conv' or 'dot_general'
+// in its name and has `Method` with `WeightOnlyPtq`.
+bool IsWeightOnlyQuantizableOp(const Operation& op);
+
+// Lists the functions in a ModuleOp sorted by their names.
+SmallVector<func::FuncOp> GetSortedFunctions(ModuleOp module_op);
+
+}  // namespace mlir::quant
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_LIFT_AS_FUNCTION_CALL_H_

tensorflow/compiler/mlir/quantization/common/lift_as_function_call.td

@@ -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.
+==============================================================================*/
+
+include "mlir/IR/OpBase.td"
+include "mlir/IR/PatternBase.td"
+include "mlir/Dialect/Func/IR/FuncOps.td"
+
+// Creates a function call with TF::PartitionedCallOp and a new function to
+// wrap the section between arguments and results.
+//
+// The `returns` number indicates the number of results the function returns.
+class LiftAsTFPartitionedCall<string func_name, int returns = 1> :
+  NativeCodeCall<"LiftAsFunctionCall($_builder, $_loc, "
+                    "FunctionCallOpType::TFPartitionedCallOp, "
+                    "\""# func_name #"\", $0...)", returns>;
+
+// Creates a function call with TF::XlaCallModuleOp and a new function to wrap
+// the section between arguments and results.
+//
+// The `returns` number indicates the number of results the function returns.
+class LiftAsTFXlaCallModule<string func_name, int returns = 1> :
+  NativeCodeCall<"LiftAsFunctionCall($_builder, $_loc, "
+                    "FunctionCallOpType::TFXlaCallModuleOp, "
+                    "\""# func_name #"\", $0...)", returns>;
+
+// Add the second argument to the first argument, which is expected to be an
+// argument list.
+// bias(einsum(inputs), bias) --> einsum_with_bias(AppendToVector(inputs, bias))
+// Since inputs is a vector in case of einsum, we cannot use ArgumentList here.
+def AppendToVector : NativeCodeCall<"AppendToVector($0, $1)">;
+
+// The list of arguments of the composite function.
+def ArgumentList : NativeCodeCall<"llvm::SmallVector<Value>{$0...}">;
+
+// The list of results of the composite function.
+def ResultList : NativeCodeCall<"llvm::SmallVector<Value>{$0...}">;
+
+// Creates a list of NamedAttributes. An example usage would be:
+// (NamedAttributeList (NamedAttr<"transpose_a"> $transpose_a))
+def NamedAttributeList : NativeCodeCall<"llvm::SmallVector<NamedAttribute>{$0...}">;
+
+// Creates a NamedAttribute given its name and value. Essentially creates
+// a pair: {attribute_name, attribute_value}.
+class NamedAttr<string attr_name> :
+    NativeCodeCall<"NamedAttribute{$_builder.getStringAttr(\"" # attr_name # "\"), $0}">;
+
+// Checks if the value is not defined inside a lifted function by checking the
+// `tf_quant.composite_function` attribute.
+def IsNotInLiftedFunc :
+      Constraint<CPred<"!IsInLiftedFunc($0.getDefiningOp())">>;
+
+// Checks if the value is not inside a StableHLO op with region.
+def IsNotInStableHloOpRegion :
+      Constraint<CPred<"!IsInStableHloOpRegion($0.getDefiningOp())">>;
+
+// Checks if the given einsum op is supported for XlaDotV2 quantization.
+def IsEinsumSupportedByXlaDotV2 :
+      Constraint<CPred<"IsEinsumSupportedByXlaDotV2($0)">>;
+
+// This attribute can be used in the `AttributeList` for missing attributes. It
+// is necessary to keep other attributes in the same index as the quantized
+// composite function.
+def NullAttribute : NativeCodeCall<"$_builder.getStringAttr(\"N/A\")">;

tensorflow/compiler/mlir/quantization/common/lift_as_function_call_test.cc

@@ -0,0 +1,531 @@
+/* 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/common/lift_as_function_call.h"
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/algorithm/container.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/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/OwningOpRef.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "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/test_base.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 {
+namespace {
+
+using ::stablehlo::quantization::Method;
+using ::testing::HasSubstr;
+using ::testing::NotNull;
+using ::testing::SizeIs;
+using ::testing::StrEq;
+using ::tsl::protobuf::util::MessageDifferencer;
+
+using LiftAsFunctionCallTest = QuantizationTestBase;
+
+constexpr absl::string_view kModuleLifted = R"mlir(
+  module {
+    func.func private @composite_dot_general_fn_1(%arg0: tensor<1x1024xf32>, %arg1: tensor<1024x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module, tf_quant.composite_function} {
+      %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1024xf32>, tensor<1024x3xf32>) -> tensor<1x3xf32>
+      return %0 : tensor<1x3xf32>
+    }
+  }
+)mlir";
+
+TEST_F(LiftAsFunctionCallTest, LiftedFunctionSucceeds) {
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleLifted);
+  ASSERT_TRUE(module_op);
+
+  auto composite_dot_general_fn =
+      module_op->lookupSymbol<func::FuncOp>("composite_dot_general_fn_1");
+  ASSERT_THAT(composite_dot_general_fn, NotNull());
+
+  auto dot_general_op = FindOperationOfType<mlir::stablehlo::DotGeneralOp>(
+      composite_dot_general_fn);
+  EXPECT_TRUE(IsInLiftedFunc(dot_general_op));
+}
+
+constexpr absl::string_view kModuleStableHlo = R"mlir(
+  module {
+    func.func @main(%arg0: tensor<1x1024xf32>, %arg1: tensor<1024x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
+      %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], precision = [] : (tensor<1x1024xf32>, tensor<1024x3xf32>) -> tensor<1x3xf32>
+      return %0 : tensor<1x3xf32>
+    }
+  }
+)mlir";
+
+TEST_F(LiftAsFunctionCallTest, FunctionLiftedAsXlaCallModuleOp) {
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleStableHlo);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  auto dot_general_op =
+      FindOperationOfType<mlir::stablehlo::DotGeneralOp>(main_fn);
+
+  const SmallVector<NamedAttribute>& attributes = {
+      builder_.getNamedAttr(
+          "precision_config",
+          builder_.getArrayAttr(SmallVector<Attribute>(
+              1, mlir::stablehlo::PrecisionAttr::get(
+                     ctx_.get(), mlir::stablehlo::Precision::DEFAULT)))),
+  };
+  const SmallVector<Value> operands(dot_general_op->getOperands());
+  const SmallVector<Value> results(dot_general_op->getResults());
+  Operation* lifted_op =
+      LiftAsFunctionCall(builder_, dot_general_op->getLoc(),
+                         FunctionCallOpType::TFXlaCallModuleOp,
+                         "composite_dot_general_fn", operands, results,
+                         attributes)[0]
+          .getDefiningOp();
+  const auto entry_function_symbol_ref =
+      lifted_op->getAttrOfType<FlatSymbolRefAttr>("_entry_function");
+  SymbolTable symbol_table(*module_op);
+  auto entry_func = dyn_cast_or_null<func::FuncOp>(
+      symbol_table.lookup(entry_function_symbol_ref.getValue()));
+  auto lifted_dot_general_op =
+      FindOperationOfType<mlir::stablehlo::DotGeneralOp>(entry_func);
+
+  EXPECT_TRUE(isa<TF::XlaCallModuleOp>(lifted_op));
+  EXPECT_EQ(
+      mlir::cast<StringAttr>(lifted_op->getAttr("_original_entry_function")),
+      "composite_dot_general_fn_1");
+  EXPECT_EQ(
+      mlir::cast<ArrayAttr>(lifted_dot_general_op->getAttr("precision_config")),
+      builder_.getArrayAttr(SmallVector<Attribute>(
+          1, mlir::stablehlo::PrecisionAttr::get(
+                 ctx_.get(), mlir::stablehlo::Precision::DEFAULT))));
+}
+
+TEST_F(LiftAsFunctionCallTest, FunctionNoAttrLiftedAsXlaCallModuleOp) {
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleStableHlo);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  auto dot_general_op =
+      FindOperationOfType<mlir::stablehlo::DotGeneralOp>(main_fn);
+  const SmallVector<Value> operands(dot_general_op->getOperands());
+  const SmallVector<Value> results(dot_general_op->getResults());
+  Operation* lifted_op =
+      LiftAsFunctionCall(builder_, dot_general_op->getLoc(),
+                         FunctionCallOpType::TFXlaCallModuleOp,
+                         "composite_dot_general_fn", operands, results)[0]
+          .getDefiningOp();
+  EXPECT_TRUE(isa<TF::XlaCallModuleOp>(lifted_op));
+  EXPECT_EQ(
+      mlir::cast<StringAttr>(lifted_op->getAttr("_original_entry_function")),
+      "composite_dot_general_fn_1");
+}
+
+TEST_F(LiftAsFunctionCallTest, EinsumSupportedForXlaDotV2Succeeds) {
+  StringAttr einsum_supported_by_xla_dot_v2_attr =
+      builder_.getStringAttr("ijk,ikm->ijm");
+  StringAttr einsum_one_operand = builder_.getStringAttr("ijk->ikj");
+  StringAttr einsum_ellipsis = builder_.getStringAttr("...gse->...gs");
+  EXPECT_TRUE(IsEinsumSupportedByXlaDotV2(einsum_supported_by_xla_dot_v2_attr));
+  EXPECT_FALSE(IsEinsumSupportedByXlaDotV2(einsum_one_operand));
+  EXPECT_FALSE(IsEinsumSupportedByXlaDotV2(einsum_ellipsis));
+}
+
+TEST_F(LiftAsFunctionCallTest, GetQuantizationMethodSucceeds) {
+  // Function containing a simple `TF::XlaCallModuleOp` with a valid string
+  // attribute `_quantization_method` set to `"no_quantization {}"`.
+  constexpr absl::string_view kXlaCallModuleOpWithQuantizationMethodAttr =
+      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_dot_general_fn_1, _quantization_method = "no_quantization {}", _stablehlo_module_attrs = {jax.uses_shape_polymorphism = true}} : (tensor<1x1x3xf32>, tensor<3x4xf32>) -> tensor<1x1x4xf32>
+      return %0 : tensor<1x1x4xf32>
+    }
+  )mlir";
+
+  const OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kXlaCallModuleOpWithQuantizationMethodAttr);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  auto xla_call_module_ops = main_fn.getOps<TF::XlaCallModuleOp>();
+  ASSERT_FALSE(xla_call_module_ops.empty());
+
+  // Test that `GetQuantizationMethod` returns a valid `Method` corresponding to
+  // `"no_quantization {}"`.
+  const absl::StatusOr<Method> method =
+      GetQuantizationMethod(*xla_call_module_ops.begin());
+  ASSERT_THAT(method, absl_testing::IsOk());
+  EXPECT_TRUE(method->has_no_quantization());
+}
+
+TEST_F(LiftAsFunctionCallTest,
+       GetQuantizationMethodFailsWhenNoQuantizationMethodAttr) {
+  // Function containing a simple `TF::XlaCallModuleOp` that doesn't have the
+  // attribute "_quantization_method".
+  constexpr absl::string_view kXlaCallModuleOpWithNoQuantizationMethodAttr =
+      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_dot_general_fn_1, _stablehlo_module_attrs = {jax.uses_shape_polymorphism = true}} : (tensor<1x1x3xf32>, tensor<3x4xf32>) -> tensor<1x1x4xf32>
+      return %0 : tensor<1x1x4xf32>
+    }
+  )mlir";
+
+  const OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kXlaCallModuleOpWithNoQuantizationMethodAttr);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  auto xla_call_module_ops = main_fn.getOps<TF::XlaCallModuleOp>();
+  ASSERT_FALSE(xla_call_module_ops.empty());
+
+  // Test that `GetQuantizationMethod` returns a `absl::InvalidArgumentError`
+  // because there is no `_quantization_method` attribute.
+  const absl::StatusOr<Method> method =
+      GetQuantizationMethod(*xla_call_module_ops.begin());
+  EXPECT_THAT(method,
+              absl_testing::StatusIs(
+                  absl::StatusCode::kInvalidArgument,
+                  HasSubstr("Attribute _quantization_method is not found")));
+}
+
+TEST_F(LiftAsFunctionCallTest,
+       GetQuantizationMethodFailsWhenMalformedQuantizationMethodAttr) {
+  // Function containing a simple `TF::XlaCallModuleOp` with an invalid
+  // `_quantization_method` attribute.
+  constexpr absl::string_view kXlaCallModuleOpWithNoQuantizationMethodAttr =
+      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_dot_general_fn_1, _quantization_method = "invalid_field: 123", _stablehlo_module_attrs = {jax.uses_shape_polymorphism = true}} : (tensor<1x1x3xf32>, tensor<3x4xf32>) -> tensor<1x1x4xf32>
+      return %0 : tensor<1x1x4xf32>
+    }
+  )mlir";
+
+  const OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kXlaCallModuleOpWithNoQuantizationMethodAttr);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  auto xla_call_module_ops = main_fn.getOps<TF::XlaCallModuleOp>();
+  ASSERT_FALSE(xla_call_module_ops.empty());
+
+  const absl::StatusOr<Method> method =
+      GetQuantizationMethod(*xla_call_module_ops.begin());
+  EXPECT_THAT(method, absl_testing::StatusIs(
+                          absl::StatusCode::kInternal,
+                          HasSubstr("Failed to parse Method from textproto")));
+}
+
+constexpr absl::string_view kFunctionWithRegion =
+    R"mlir(
+  func.func @main(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
+    %if = "stablehlo.if"(%arg0) ({
+      %0 = stablehlo.add %arg1, %arg1 : tensor<f32>
+      stablehlo.return %0 : tensor<f32>
+    }, {
+      %1 = stablehlo.add %arg2, %arg2 : tensor<f32>
+      stablehlo.return %1 : tensor<f32>
+    }) : (tensor<i1>) -> (tensor<f32>)
+    %subtract = stablehlo.subtract %if, %if : tensor<f32>
+    return %subtract : tensor<f32>
+  }
+)mlir";
+
+TEST_F(LiftAsFunctionCallTest, IsInRegionSucceedsWhenOpInsideRegion) {
+  const OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kFunctionWithRegion);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  auto if_op = FindOperationOfType<mlir::stablehlo::IfOp>(main_fn);
+  Block& block = if_op->getRegion(0).front();
+  Operation& add_op = *absl::c_find_if(block, [](Operation& entry) {
+    return dyn_cast_or_null<::mlir::stablehlo::AddOp>(&entry);
+  });
+  EXPECT_TRUE(IsInStableHloOpRegion(&add_op));
+}
+
+TEST_F(LiftAsFunctionCallTest, IsInRegionFailsWhenOpNotInsideRegion) {
+  const OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kFunctionWithRegion);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  auto subtract_op = FindOperationOfType<mlir::stablehlo::SubtractOp>(main_fn);
+  EXPECT_FALSE(IsInStableHloOpRegion(subtract_op));
+}
+
+TEST_F(LiftAsFunctionCallTest,
+       GetQuantizationMethodOrDefaultReturnsCorrectMethod) {
+  // Function containing a simple `TF::XlaCallModuleOp` with a valid string
+  // attribute `_quantization_method` set to `"no_quantization { }"`.
+  constexpr absl::string_view kXlaCallModuleOpWithQuantizationMethodAttr =
+      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_dot_general_fn_1,
+            _quantization_method = "no_quantization { }",
+            _stablehlo_module_attrs = {jax.uses_shape_polymorphism = true}
+          } : (tensor<1x1x3xf32>, tensor<3x4xf32>) -> tensor<1x1x4xf32>
+      return %0 : tensor<1x1x4xf32>
+    }
+  )mlir";
+
+  const OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kXlaCallModuleOpWithQuantizationMethodAttr);
+  ASSERT_TRUE(module_op);
+
+  FailureOr<TF::XlaCallModuleOp> xla_call_module_op =
+      FindFirstOpFromMainFunc<TF::XlaCallModuleOp>(*module_op);
+  ASSERT_TRUE(succeeded(xla_call_module_op));
+
+  // Test that `GetQuantizationMethodOrDefault` returns a valid `Method`
+  // corresponding to `"no_quantization {}"`.
+  const Method method = GetQuantizationMethodOrDefault(*xla_call_module_op);
+  EXPECT_TRUE(method.has_no_quantization());
+}
+
+TEST_F(
+    LiftAsFunctionCallTest,
+    GetQuantizationMethodOrDefaultReturnsDefaultWhenNoQuantizationMethodAttr) {
+  // Function containing a simple `TF::XlaCallModuleOp` that is missing the
+  // "_quantization_method" attribute.
+  constexpr absl::string_view kXlaCallModuleOpWithoutQuantizationMethodAttr =
+      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_dot_general_fn_1,
+            _stablehlo_module_attrs = {jax.uses_shape_polymorphism = true}
+          } : (tensor<1x1x3xf32>, tensor<3x4xf32>) -> tensor<1x1x4xf32>
+      return %0 : tensor<1x1x4xf32>
+    }
+  )mlir";
+
+  const OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kXlaCallModuleOpWithoutQuantizationMethodAttr);
+  ASSERT_TRUE(module_op);
+
+  FailureOr<TF::XlaCallModuleOp> xla_call_module_op =
+      FindFirstOpFromMainFunc<TF::XlaCallModuleOp>(*module_op);
+  ASSERT_TRUE(succeeded(xla_call_module_op));
+
+  // Test that `GetQuantizationMethodOrDefault` returns the default instance.
+  const Method method = GetQuantizationMethodOrDefault(*xla_call_module_op);
+  EXPECT_TRUE(MessageDifferencer::Equals(method, Method::default_instance()));
+}
+constexpr absl::string_view kModuleDotWeightOnlyPtq = R"mlir(
+  module {
+    func.func @main(%arg0: tensor<?x2xf32> {tf_saved_model.index_path = ["input_tensor"]}) -> (tensor<?x2xf32>) {
+      %0 = stablehlo.constant dense<[-0.211145893, -0.708605706]> : tensor<2xf32>
+      %1 = stablehlo.constant dense<[[-0.630731344, 0.54962182], [0.180364341, -0.764542698]]> : tensor<2x2xf32>
+      %2 = "tf.XlaCallModule"(%arg0, %1, %0) <{Sout = [#tf_type.shape<?x2>], module = "", version = 9 : i64}> {_entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _tfl_quant_trait = "fully_quantizable", _quantization_method = "weight_only_ptq { }"} : (tensor<?x2xf32>, tensor<2x2xf32>, tensor<2xf32>) -> tensor<?x2xf32>
+      return %2 : tensor<?x2xf32>
+    }
+    func.func private @composite_dot_general_fn_1(%arg0: tensor<?x2xf32>, %arg1: tensor<2x2xf32>, %arg2: tensor<2xf32>) -> tensor<?x2xf32> attributes {_from_xla_call_module, tf_quant.composite_function} {
+      %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<?x2xf32>, tensor<2x2xf32>) -> tensor<?x2xf32>
+      return %0 : tensor<?x2xf32>
+    }
+  }
+)mlir";
+
+TEST_F(LiftAsFunctionCallTest, HasWeightOnlyPtqMethodExists) {
+  OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kModuleDotWeightOnlyPtq);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  auto call_op = *main_fn.getOps<TF::XlaCallModuleOp>().begin();
+  EXPECT_TRUE(HasWeightOnlyPtqMethod(call_op));
+}
+
+TEST_F(LiftAsFunctionCallTest, HasWeightOnlyPtqMethodDifferentMethod) {
+  const absl::string_view kModuleDotNoQuantization = R"mlir(
+    module {
+      func.func @main(%arg0: tensor<?x2xf32> {tf_saved_model.index_path = ["input_tensor"]}) -> (tensor<?x2xf32>) {
+        %0 = stablehlo.constant dense<[-0.211145893, -0.708605706]> : tensor<2xf32>
+        %1 = stablehlo.constant dense<[[-0.630731344, 0.54962182], [0.180364341, -0.764542698]]> : tensor<2x2xf32>
+        %2 = "tf.XlaCallModule"(%arg0, %1, %0) <{Sout = [#tf_type.shape<?x2>], module = "", version = 9 : i64}> {_entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1", _tfl_quant_trait = "fully_quantizable", _quantization_method = "no_quantization { }"} : (tensor<?x2xf32>, tensor<2x2xf32>, tensor<2xf32>) -> tensor<?x2xf32>
+        return %2 : tensor<?x2xf32>
+      }
+      func.func private @composite_dot_general_fn_1(%arg0: tensor<?x2xf32>, %arg1: tensor<2x2xf32>, %arg2: tensor<2xf32>) -> tensor<?x2xf32> attributes {_from_xla_call_module, tf_quant.composite_function} {
+        %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<?x2xf32>, tensor<2x2xf32>) -> tensor<?x2xf32>
+        return %0 : tensor<?x2xf32>
+      }
+    }
+  )mlir";
+  OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kModuleDotNoQuantization);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  auto call_op = *main_fn.getOps<TF::XlaCallModuleOp>().begin();
+  EXPECT_FALSE(HasWeightOnlyPtqMethod(call_op));
+}
+
+TEST_F(LiftAsFunctionCallTest, HasWeightOnlyPtqMethodNoMethod) {
+  const absl::string_view kModuleXlaCallModule = R"mlir(
+    module {
+      func.func @main(%arg0: tensor<?x2xf32> {tf_saved_model.index_path = ["input_tensor"]}) -> (tensor<?x2xf32>) {
+        %0 = stablehlo.constant dense<[-0.211145893, -0.708605706]> : tensor<2xf32>
+        %1 = stablehlo.constant dense<[[-0.630731344, 0.54962182], [0.180364341, -0.764542698]]> : tensor<2x2xf32>
+        %2 = "tf.XlaCallModule"(%arg0, %1, %0) <{Sout = [#tf_type.shape<?x2>], module = "", version = 9 : i64}> {_entry_function = @composite_fn_1, _stablehlo_version = "1.0.0", _original_entry_function = "composite_fn_1", _tfl_quant_trait = "fully_quantizable"} : (tensor<?x2xf32>, tensor<2x2xf32>, tensor<2xf32>) -> tensor<?x2xf32>
+        return %2 : tensor<?x2xf32>
+      }
+      func.func private @composite_fn_1(%arg0: tensor<?x2xf32>, %arg1: tensor<2x2xf32>, %arg2: tensor<2xf32>) -> tensor<?x2xf32> attributes {_from_xla_call_module, tf_quant.composite_function} {
+        return %arg0 : tensor<?x2xf32>
+      }
+    }
+  )mlir";
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleXlaCallModule);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  auto call_op = *main_fn.getOps<TF::XlaCallModuleOp>().begin();
+  EXPECT_FALSE(HasWeightOnlyPtqMethod(call_op));
+}
+
+TEST_F(LiftAsFunctionCallTest, IsWeightOnlyQuantizableOpDot) {
+  OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kModuleDotWeightOnlyPtq);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  auto call_op = *main_fn.getOps<TF::XlaCallModuleOp>().begin();
+  EXPECT_TRUE(IsWeightOnlyQuantizableOp(*call_op));
+}
+
+TEST_F(LiftAsFunctionCallTest, IsWeightOnlyQuantizableOpNotTfXlaCallModuleOp) {
+  const absl::string_view kModulePartitionedCallDot = R"mlir(
+    module {
+      func.func @main(%arg0: tensor<?x2xf32> {tf_saved_model.index_path = ["input_tensor"]}) -> (tensor<?x2xf32>) {
+        %0 = stablehlo.constant dense<[-0.211145893, -0.708605706]> : tensor<2xf32>
+        %1 = stablehlo.constant dense<[[-0.630731344, 0.54962182], [0.180364341, -0.764542698]]> : tensor<2x2xf32>
+        %2 = "tf.PartitionedCall"(%arg0, %1, %0)  {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_dot_general_fn_1, _quantization_method = "weight_only_ptq { }"} : (tensor<?x2xf32>, tensor<2x2xf32>, tensor<2xf32>) -> tensor<?x2xf32>
+        return %2 : tensor<?x2xf32>
+      }
+      func.func private @composite_dot_general_fn_1(%arg0: tensor<?x2xf32>, %arg1: tensor<2x2xf32>, %arg2: tensor<2xf32>) -> tensor<?x2xf32> attributes {_from_xla_call_module, tf_quant.composite_function} {
+        %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<?x2xf32>, tensor<2x2xf32>) -> tensor<?x2xf32>
+        return %0 : tensor<?x2xf32>
+      }
+    }
+  )mlir";
+  OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kModulePartitionedCallDot);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  auto call_op = *main_fn.getOps<TF::PartitionedCallOp>().begin();
+  EXPECT_FALSE(IsWeightOnlyQuantizableOp(*call_op));
+}
+
+TEST_F(LiftAsFunctionCallTest, IsWeightOnlyQuantizableOpNoConvNoDot) {
+  constexpr absl::string_view kModuleXlaCallModule = R"mlir(
+    module {
+      func.func @main(%arg0: tensor<?x2xf32> {tf_saved_model.index_path = ["input_tensor"]}) -> (tensor<?x2xf32>) {
+        %0 = stablehlo.constant dense<[-0.211145893, -0.708605706]> : tensor<2xf32>
+        %1 = stablehlo.constant dense<[[-0.630731344, 0.54962182], [0.180364341, -0.764542698]]> : tensor<2x2xf32>
+        %2 = "tf.XlaCallModule"(%arg0, %1, %0) <{Sout = [#tf_type.shape<?x2>], module = "", version = 9 : i64}> {_entry_function = @composite_fn_1, _stablehlo_version = "1.0.0", _original_entry_function = "composite_fn_1", _tfl_quant_trait = "fully_quantizable", _quantization_method = "weight_only_ptq { }"} : (tensor<?x2xf32>, tensor<2x2xf32>, tensor<2xf32>) -> tensor<?x2xf32>
+        return %2 : tensor<?x2xf32>
+      }
+      func.func private @composite_fn_1(%arg0: tensor<?x2xf32>, %arg1: tensor<2x2xf32>, %arg2: tensor<2xf32>) -> tensor<?x2xf32> attributes {_from_xla_call_module, tf_quant.composite_function} {
+        return %arg0 : tensor<?x2xf32>
+      }
+    }
+  )mlir";
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleXlaCallModule);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  auto call_op = *main_fn.getOps<TF::XlaCallModuleOp>().begin();
+  EXPECT_FALSE(IsWeightOnlyQuantizableOp(*call_op));
+}
+
+TEST_F(LiftAsFunctionCallTest, GetSortedFunctions) {
+  constexpr absl::string_view kModuleXlaCallModule = R"mlir(
+    module {
+      func.func @conv_3_fn(%arg0: tensor<1x3x3x4xf32>) -> tensor<1x3x3x4xf32> {
+        %0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x4x4xf32>
+        %1 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x3x4xf32>, tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32>
+        %2 = stablehlo.convolution(%1, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x3x4xf32>, tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32>
+        func.return %2: tensor<1x3x3x4xf32>
+      }
+
+      func.func @conv_1_fn(%arg0: tensor<1x3x3x4xf32>) -> tensor<1x3x3x4xf32> {
+        %0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x4x4xf32>
+        %1 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x3x4xf32>, tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32>
+        %2 = stablehlo.convolution(%1, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x3x4xf32>, tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32>
+        func.return %2: tensor<1x3x3x4xf32>
+      }
+
+      func.func @conv_2_fn(%arg0: tensor<1x3x3x4xf32>) -> tensor<1x3x3x4xf32> {
+        %0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x4x4xf32>
+        %1 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x3x4xf32>, tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32>
+        %2 = stablehlo.convolution(%1, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x3x4xf32>, tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32>
+        func.return %2: tensor<1x3x3x4xf32>
+      }
+    }
+  )mlir";
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleXlaCallModule);
+  ASSERT_TRUE(module_op);
+
+  SmallVector<func::FuncOp> funcs = GetSortedFunctions(*module_op);
+  ASSERT_THAT(funcs, SizeIs(3));
+  EXPECT_THAT(funcs[0].getSymName(), StrEq("conv_1_fn"));
+  EXPECT_THAT(funcs[1].getSymName(), StrEq("conv_2_fn"));
+  EXPECT_THAT(funcs[2].getSymName(), StrEq("conv_3_fn"));
+}
+
+}  // namespace
+}  // namespace mlir::quant

tensorflow/compiler/mlir/quantization/common/python/BUILD

@@ -0,0 +1,30 @@
+load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
+load(
+    "//tensorflow:tensorflow.default.bzl",
+    "tf_py_strict_test",
+)
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["@stablehlo//:license"],
+    default_visibility = [
+        "//tensorflow:__pkg__",
+        "//tensorflow/compiler/mlir/quantization/stablehlo/python:internal_visibility_allowlist_package",
+    ],
+    licenses = ["notice"],
+)
+
+pytype_strict_library(
+    name = "testing",
+    srcs = ["testing.py"],
+    tags = ["no_pip"],
+    visibility = ["//visibility:public"],
+)
+
+tf_py_strict_test(
+    name = "testing_test",
+    srcs = ["testing_test.py"],
+    deps = [
+        ":testing",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)

tensorflow/compiler/mlir/quantization/common/python/testing.py

@@ -0,0 +1,69 @@
+# 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.
+# ==============================================================================
+"""Common testing utilities for quantization libraries."""
+import itertools
+import os
+from typing import Any, Mapping, Sequence
+
+
+def parameter_combinations(
+    test_parameters: Sequence[Mapping[str, Sequence[Any]]]
+) -> Sequence[Mapping[str, Any]]:
+  """Generate all combinations of test parameters.
+
+  Args:
+    test_parameters: List of dictionaries that maps parameter keys and values.
+
+  Returns:
+    real_parameters: All possible combinations of the parameters as list of
+    dictionaries.
+  """
+  real_parameters = []
+  for parameters in test_parameters:
+    keys = parameters.keys()
+    for curr in itertools.product(*parameters.values()):
+      real_parameters.append(dict(zip(keys, curr)))
+  return real_parameters
+
+
+def get_dir_size(path: str = '.') -> int:
+  """Get the total size of files and sub-directories under the path.
+
+  Args:
+    path: Path of a directory or a file to calculate the total size.
+
+  Returns:
+    Total size of the directory or a file.
+  """
+  total = 0
+  for root, _, files in os.walk(path):
+    for filename in files:
+      total += os.path.getsize(os.path.join(root, filename))
+  return total
+
+
+def get_size_ratio(path_a: str, path_b: str) -> float:
+  """Return the size ratio of the given paths.
+
+  Args:
+    path_a: Path of a directory or a file to be the nominator of the ratio.
+    path_b: Path of a directory or a file to be the denominator of the ratio.
+
+  Returns:
+    Ratio of size of path_a / size of path_b.
+  """
+  size_a = get_dir_size(path_a)
+  size_b = get_dir_size(path_b)
+  return size_a / size_b

tensorflow/compiler/mlir/quantization/common/python/testing_test.py

@@ -0,0 +1,63 @@
+# 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.
+# ==============================================================================
+
+from tensorflow.compiler.mlir.quantization.common.python import testing
+from tensorflow.python.platform import test
+
+
+class TestingTest(test.TestCase):
+
+  def test_parameter_combinations(self):
+    """Tests that parameter_combinations returns correct combinations."""
+    test_parameters = [{
+        'shapes': [
+            [3, 3],
+            [3, None],
+        ],
+        'has_bias': [True, False],
+    }]
+    combinations = testing.parameter_combinations(test_parameters)
+
+    self.assertLen(combinations, 4)
+    self.assertIn({'shapes': [3, 3], 'has_bias': True}, combinations)
+    self.assertIn({'shapes': [3, 3], 'has_bias': False}, combinations)
+    self.assertIn({'shapes': [3, None], 'has_bias': True}, combinations)
+    self.assertIn({'shapes': [3, None], 'has_bias': False}, combinations)
+
+
+class FileSizeTestCase(test.TestCase):
+
+  def setUp(self):
+    super().setUp()
+
+    self.path_a = self.create_tempdir('dir_a').full_path
+    self.create_tempfile(file_path='dir_a/w.txt', content='abcd')
+
+    self.path_b = self.create_tempdir('dir_b').full_path
+    self.create_tempfile(file_path='dir_b/x.txt', content='1234')
+    self.create_tempfile(file_path='dir_b/y.txt', content='56')
+    self.create_tempfile(file_path='dir_b/z.txt', content='78')
+
+  def test_get_dir_size(self):
+    self.assertEqual(testing.get_dir_size(self.path_a), 4)
+    self.assertEqual(testing.get_dir_size(self.path_b), 8)
+
+  def test_get_size_ratio(self):
+    self.assertEqual(testing.get_size_ratio(self.path_a, self.path_b), 0.5)
+    self.assertEqual(testing.get_size_ratio(self.path_b, self.path_a), 2.0)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/compiler/mlir/quantization/common/quantization_lib/BUILD

@@ -0,0 +1,123 @@
+load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library")
+load("//tensorflow:tensorflow.bzl", "tf_cc_test")
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    # By default, these targets should only be used within the quantization library.
+    default_visibility = [
+        "//learning/brain/mlir/quantization:__subpackages__",
+        "//tensorflow:__subpackages__",
+    ],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "portable_tensor_utils",
+    srcs = ["portable_tensor_utils.cc"],
+    hdrs = ["portable_tensor_utils.h"],
+)
+
+cc_library(
+    name = "quantization_lib",
+    srcs = [
+        "quantization_driver.cc",
+        "quantization_interface.cc.inc",
+        "quantization_utils.cc",
+    ],
+    hdrs = [
+        "quantization_driver.h",
+        "quantization_interface.h.inc",
+        "quantization_traits.h",
+        "quantization_utils.h",
+    ],
+    deps = [
+        ":portable_tensor_utils",
+        ":quantization_config",
+        ":quantization_interfaces_inc_gen",
+        "//tensorflow/compiler/mlir/quantization/common/ir:QuantOps",
+        "//tensorflow/compiler/mlir/tools/optimize:quantization_utils",
+        "//tensorflow/core:lib_proto_parsing",
+        "//tensorflow/core:protos_all_cc",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/strings",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:QuantOps",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
+tf_cc_test(
+    name = "quantization_driver_test",
+    srcs = ["quantization_driver_test.cc"],
+    deps = [
+        ":quantization_lib",
+        "//tensorflow/compiler/mlir/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/tensorflow",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:QuantOps",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
+cc_library(
+    name = "quantization_config",
+    srcs = [
+        "quantization_config.cc",
+    ],
+    hdrs = [
+        "quantization_config.h",
+    ],
+    deps = [
+        "//tensorflow/core:protos_all_cc",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/strings",
+        "@llvm-project//llvm:Support",
+    ],
+)
+
+td_library(
+    name = "quantization_td_files",
+    srcs = [
+        "quantization.td",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        "//tensorflow/compiler/mlir/quantization/common/ir:QuantizationOpsTdFiles",
+        "@llvm-project//mlir:OpBaseTdFiles",
+    ],
+)
+
+gentbl_cc_library(
+    name = "quantization_interfaces_inc_gen",
+    compatible_with = get_compatible_with_portable(),
+    tbl_outs = {
+        "quantization_interface.h.inc": ["-gen-op-interface-decls"],
+        "quantization_interface.cc.inc": ["-gen-op-interface-defs"],
+    },
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "quantization.td",
+    deps = [
+        ":quantization_td_files",
+    ],
+)
+
+exports_files([
+    "quantization_traits.h",
+    "quantization_config.h",
+    "quantization_utils.h",
+])

tensorflow/compiler/mlir/quantization/common/quantization_lib/portable_tensor_utils.cc

@@ -0,0 +1,62 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+// This file is the MLIR copy of part of
+// third_party/tensorflow/lite/kernels/internal/reference/portable_tensor_utils.cc
+// as part of the effort to decouple TFLite from MLIR.
+
+#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/portable_tensor_utils.h"
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <cstring>
+
+namespace mlir {
+namespace quant {
+
+void PortableSymmetricQuantizeFloats(const float* values, const int size,
+                                     int8_t* quantized_values, float* min_value,
+                                     float* max_value, float* scaling_factor) {
+  auto minmax = std::minmax_element(values, values + size);
+  *min_value = *minmax.first;
+  *max_value = *minmax.second;
+
+  PortableSymmetricQuantizeFloats(values, size, quantized_values, *min_value,
+                                  *max_value, scaling_factor);
+}
+
+void PortableSymmetricQuantizeFloats(const float* values, const int size,
+                                     int8_t* quantized_values, float min_value,
+                                     float max_value, float* scaling_factor) {
+  const int32_t kScale = 127;
+  const float range = std::max(std::abs(min_value), std::abs(max_value));
+  if (range == 0) {
+    memset(quantized_values, 0, size * sizeof(int8_t));
+    *scaling_factor = 1;
+    return;
+  }
+  *scaling_factor = range / kScale;
+  const float scaling_factor_inv = kScale / range;
+  for (int i = 0; i < size; ++i) {
+    const int32_t quantized_value =
+        static_cast<int32_t>(round(values[i] * scaling_factor_inv));
+    // Clamp: just in case some odd numeric offset.
+    quantized_values[i] = static_cast<int8_t>(
+        std::min(kScale, std::max(-kScale, quantized_value)));
+  }
+}
+
+}  // namespace quant
+}  // namespace mlir

tensorflow/compiler/mlir/quantization/common/quantization_lib/portable_tensor_utils.h

@@ -0,0 +1,40 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+// This file is the MLIR copy of part of
+// third_party/tensorflow/lite/kernels/internal/reference/portable_tensor_utils.h
+// as part of the effort to decouple TFLite from MLIR.
+
+#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_QUANTIZATION_LIB_PORTABLE_TENSOR_UTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_QUANTIZATION_LIB_PORTABLE_TENSOR_UTILS_H_
+
+#include <cstdint>
+
+namespace mlir {
+namespace quant {
+
+// LINT.IfChange(portable_symmetric_quantize_floats)
+void PortableSymmetricQuantizeFloats(const float* values, int size,
+                                     int8_t* quantized_values, float* min_value,
+                                     float* max_value, float* scaling_factor);
+
+void PortableSymmetricQuantizeFloats(const float* values, int size,
+                                     int8_t* quantized_values, float min_value,
+                                     float max_value, float* scaling_factor);
+// LINT.ThenChange(//tensorflow/lite/kernels/internal/reference/portable_tensor_utils.h:portable_symmetric_quantize_floats)
+
+}  // namespace quant
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_QUANTIZATION_LIB_PORTABLE_TENSOR_UTILS_H_

tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization.td

@@ -0,0 +1,221 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This is the quantization definition file for TensorFlow.
+
+#ifdef TF_Quantization
+#else
+#define TF_Quantization
+
+include "mlir/IR/OpBase.td"
+include "mlir/Dialect/Quant/IR/QuantBase.td"
+
+//===----------------------------------------------------------------------===//
+// QuantizedType definitions.
+//===----------------------------------------------------------------------===//
+
+// The base class of a quantized type. Signed quantized types may be expressed
+// as signless integers (i.e. up to op interpretation), but we include an
+// explicit signedness check to differentiate the signed/unsigned constraints
+// predicates from one another at the TD level.
+class QuantizedType<string n, list<int> params, bit signed>
+  : Type<And<[CPred<"llvm::isa<mlir::quant::QuantizedType>($_self)">,
+              CPred<"llvm::cast<mlir::quant::QuantizedType>($_self)" #
+                    ".getStorageTypeIntegralWidth() == " # !head(params)>,
+              Or<[CPred<"llvm::cast<mlir::quant::QuantizedType>($_self)" #
+                    ".getStorageType().isSignlessInteger()">,
+                  CPred<"llvm::cast<mlir::quant::QuantizedType>($_self)" #
+                    ".getStorageType().isSignedInteger() == " # signed>]>]>,
+    "Q" # !if (signed, "I", "UI") # !head(params) # " type"> {
+  string name = n;
+  string asTraitArgsStr =
+    !interleave(params, ", ") # !if(signed, ", true", ", false");
+}
+
+// Uniform quantized types. Two integers "smantissa" and "sexp" are used to
+// express the Mantissa and Exponent components of the floating-point scale so
+// the scale of the quantized type is "smantissa * 10 ^ sexp".
+class UInt8UniformQuantizedType<int zero_pt, int smantissa, int sexp>
+    : QuantizedType<"Uniform",
+                        [8, zero_pt, smantissa, sexp, 0, 255], 0>;
+class Int8UniformQuantizedType<int zero_pt, int smantissa, int sexp>
+    : QuantizedType<"Uniform",
+                        [8, zero_pt, smantissa, sexp, -128, 127], 1>;
+
+// General uniform quantized types. The definitions can be used to specify
+// operand's tensor types.
+def QI2 : QuantizedType<"Uniform", [2], 1>;
+def QI4 : QuantizedType<"Uniform", [4], 1>;
+def QUI4 : QuantizedType<"Uniform", [4], 0>;
+def QUI8 : QuantizedType<"Uniform", [8], 0>;
+def QI8 : QuantizedType<"Uniform", [8], 1>;
+def QUI16 : QuantizedType<"Uniform", [16], 0>;
+def QI16 : QuantizedType<"Uniform", [16], 1>;
+def QUI32 : QuantizedType<"Uniform", [32], 0>;
+def QI32 : QuantizedType<"Uniform", [32], 1>;
+
+//===----------------------------------------------------------------------===//
+// TFL native op traits (for quantization).
+//
+// Ops in this link should have those traits specified:
+// https://www.tensorflow.org/lite/performance/quantization_spec
+//===----------------------------------------------------------------------===//
+
+def FixedOutputRangeInterface : OpInterface<
+  "FixedOutputRangeInterface"> {
+  let description = [{
+    Interface for defining the fixed output range.
+  }];
+
+  let methods = [
+    InterfaceMethod<
+      [{Returns the fixed output range.}],
+      "UniformQuantizedType", "GetFixedOutputRange",
+      (ins "bool":$sign, "int":$bit_width)
+    >,
+  ];
+}
+
+def AffineQuantizedOpInterface : OpInterface<
+  "AffineQuantizedOpInterface"> {
+  let description = [{
+    Interface for affine quantized ops (conv2d, fully_connected, etc.)
+  }];
+
+  let methods = [
+    InterfaceMethod<
+      [{Returns the affine operand index.}],
+      "int", "GetAffineOperandIndex",
+      (ins), [{}], [{return 1;}]>,
+    InterfaceMethod<
+      [{Returns whether narrow range is required for the affine operand.}],
+      "bool", "RequiredNarrowRangeAffineOperand",
+      (ins), [{}], [{return true;}]>,
+    InterfaceMethod<
+      [{Returns quantization dim for the affine operand.}],
+      "int", "GetQuantizationDimIndex",
+      (ins)>,
+    InterfaceMethod<
+      [{Returns the dimension index of the output channels.}],
+      "int", "GetChannelDimIndex", (ins)
+    >,
+  ];
+}
+
+def SameOperandsAndResultsScale : OpInterface<"SameScalesOpInterface"> {
+  let description = [{
+    Interface for ops potentially have same operands and results scales.
+  }];
+
+  let methods = [
+    InterfaceMethod<
+      [{Returns whether same operands and results scales are required.}],
+      "bool", "RequiredSameOperandsAndResultsScale",
+      (ins "bool":$sign, "int":$bit_width), [{}], [{return true;}]
+    >,
+    InterfaceMethod<
+      [{Returns whether operands and results must have the same quantized axis.}],
+      "bool", "RequiredSameQuantizedAxes",
+      (ins), [{}], [{return true;}]
+    >,
+  ];
+
+  let verify = [{
+    return quant::VerifySameScales($_op);
+  }];
+}
+
+def DynamicRangeQuantizedOpInterface : OpInterface<
+  "DynamicRangeQuantizedOpInterface"> {
+  let description = [{
+    Interface for ops dynamic range quantization is supported.
+
+    If the op has the kernel support for dynamic range quantization, Q/DQ op
+    pairs connected to the op are rewritten by its quantized alternatives where
+    a new op uses Q ops for its operands instead of DQ op. Otherwise, it is
+    left as is for weight-only which means the weight is dequantized at runtime.
+
+    For example, if the kernel does not support dynamic range quantization the
+    graph will be converted into the following IR:
+
+    %q_w = "tfl.pseudo_qconst"() {
+         qtype = tensor<64x3x3x3x!quant.uniform<i8<-127:127>:f32, 1.000000e+00>>
+    %w = "tfl.dequantize"(%q_w) :
+         (tensor<64x3x3x3x!quant.uniform<i8<-127:127>:f32, 1.000000e+00>>) ->
+         tensor<64x3x3x3xf32>
+    %conv = "tfl.conv_2d"(%input_act, %w, %bias)
+
+    but if it is supported, it will be rewritten as:
+
+    %q_w = "tfl.pseudo_qconst"() {
+         qtype = tensor<64x3x3x3x!quant.uniform<i8<-127:127>:f32, 1.000000e+00>>
+    %conv = "tfl.conv_2d"(%input_act, %q_w, %bias)
+
+    Note that this is part of reaching feature parity with the old quantizer for
+    dynamic range quantization except:
+    - Only use_updated_hybrid_scheme=True is supported which means the ops with
+    the asymmetrically quantizing input support is enabled to use this feature
+    during MLIR graph rewriting passes while it is configurable in the old
+    quantizer. So when those ops are matched during graph rewriting passes,
+    MLIR quantizer will always ignore the pre-set value of the attribute, if
+    there's any, and set it to True. The reason behind this decision is that
+    generally activations of these ops show better accuracy with asymmetric
+    input quantization so we want to deprecate symmetric activation quantization
+    for those ops eventually.
+    - Unlike to the old quantizer, per-channel quantization is supported for
+    weight-only TransposeConvOp.
+  }];
+
+  let methods = [
+    InterfaceMethod<
+      [{Returns the quantizable operand indices of the op.}],
+      "std::vector<int>", "GetQuantizableOperandIndices",
+      (ins), [{}], [{return {};}]>,
+    InterfaceMethod<
+      [{Returns whether the op has the kernel support for dynamic range
+      quantization.}],
+      "bool", "GetDynamicRangeQuantKernelSupport",
+      (ins), [{}], [{return false;}]>,
+    InterfaceMethod<
+      [{Returns whether the op requires asymmetric quantize input attribute
+      setting.}],
+      "bool", "RequireAsymmetricQuantizeInputsAttr",
+      (ins), [{}], [{return false;}]>,
+  ];
+}
+
+// Specify this trait if the op has a fixed output value range.
+class FixedResultScale<QuantizedType qt> : NativeOpTrait<!strconcat(
+  "quant::FixedResult", qt.name, "Scale<", qt.asTraitArgsStr, ">::Impl")>;
+
+// Specify this trait if the bias-th input of the op is a bias input, which
+// needs a scale based on the scales of op1 and op2.
+class AccumulatorUniformScale<int bias, int op1, int op2> : NativeOpTrait<
+  !strconcat("quant::AccumulatorUniformScale<",
+             !interleave([bias, op1, op2], ", "),
+             ">::Impl")>;
+
+// Specify the operand index of the coefficient operand for an affine op
+// and also the quantization dimension if per-axis quantization is support.
+// If the quantization dimension is -1, per-axis quantization isn't supported.
+class AffineOpCoefficient<int dim, int index> : NativeOpTrait<
+  !strconcat("quant::AffineOpCoefficient<",
+             !interleave([dim, index], ", "),
+             ">::Impl")>;
+
+// Specify this trait if the op does have quantizable output. Quantizers will
+// apply quantization on this op.
+def QuantizableResult : NativeOpTrait<"quant::QuantizableResult">;
+#endif // TF_Quantization

tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_config.cc

@@ -0,0 +1,183 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_config.h"
+
+#include <ios>
+#include <optional>
+#include <sstream>
+#include <string>
+#include <vector>
+
+#include "absl/strings/numbers.h"
+#include "absl/strings/str_split.h"
+#include "absl/strings/string_view.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/raw_ostream.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+// Returns whether the given dtype is a quantization type in TensorFlow.
+static bool IsQuantizationType(tensorflow::DataType dtype) {
+  switch (dtype) {
+    case tensorflow::DT_QINT8:
+    case tensorflow::DT_QUINT8:
+    case tensorflow::DT_QINT16:
+    case tensorflow::DT_QUINT16:
+    case tensorflow::DT_QINT32:
+      return true;
+    default:
+      return false;
+  }
+}
+
+namespace mlir {
+namespace quant {
+namespace {
+bool GetBooleanSpecs(const std::string& bool_val) {
+  bool result;
+  std::stringstream iss(bool_val);
+  iss >> std::boolalpha >> result;
+  return result;
+}
+}  // namespace
+
+void ParseCustomOpSpecs(const absl::string_view node_names,
+                        const CustomOpUpdateOptions& update_option,
+                        CustomOpMap& custom_op_map) {
+  if (node_names.empty()) return;
+
+  const std::vector<std::string> custom_nodes = absl::StrSplit(node_names, ',');
+
+  for (const std::string& cur_node : custom_nodes) {
+    const std::vector<std::string> node_infos = absl::StrSplit(cur_node, '=');
+    const std::string& node_name = node_infos[0];
+    const std::string& node_specification = node_infos[1];
+    CustomOpInfo new_node_info;
+    switch (update_option) {
+      case CustomOpUpdateOptions::kInputIndices: {
+        const std::vector<std::string> indices =
+            absl::StrSplit(node_specification, '-');
+        for (const std::string& cur_index : indices) {
+          custom_op_map[node_name].quantizable_input_indices.push_back(
+              std::stoi(cur_index));
+        }
+        break;
+      }
+      case CustomOpUpdateOptions::kWeightOnly:
+        custom_op_map[node_name].is_weight_only =
+            GetBooleanSpecs(node_specification);
+        break;
+      case CustomOpUpdateOptions::kNoSideEffect:
+        custom_op_map[node_name].no_side_effect =
+            GetBooleanSpecs(node_specification);
+        break;
+    }
+  }
+}
+
+bool ParseInputNodeQuantSpecs(const absl::string_view node_names,
+                              const absl::string_view min_values,
+                              const absl::string_view max_values,
+                              const absl::string_view inference_type,
+                              QuantizationSpecs* quant_specs) {
+  const std::vector<std::string> input_nodes = absl::StrSplit(node_names, ',');
+  std::vector<std::optional<double>> node_mins;
+  if (!min_values.empty()) {
+    std::vector<std::string> node_mins_str = absl::StrSplit(min_values, ',');
+    for (const std::string& node_mins_str : node_mins_str) {
+      double value;
+      if (!absl::SimpleAtod(node_mins_str, &value)) {
+        llvm::errs() << "Unexpected mins: " << node_mins_str << "\n";
+        return true;
+      }
+      node_mins.push_back(value);
+    }
+  }
+
+  std::vector<std::optional<double>> node_maxs;
+  if (!max_values.empty()) {
+    const std::vector<std::string> node_maxs_str =
+        absl::StrSplit(max_values, ',');
+    for (const std::string& node_maxs_str : node_maxs_str) {
+      double value;
+      if (!absl::SimpleAtod(node_maxs_str, &value)) {
+        llvm::errs() << "Unexpected mins: " << node_maxs_str << "\n";
+        return true;
+      }
+      node_maxs.push_back(value);
+    }
+  }
+
+  tensorflow::DataType final_type = tensorflow::DT_FLOAT;
+  if (!inference_type.empty() && !DataType_Parse(inference_type, &final_type)) {
+    return true;
+  }
+  return GetInputNodeQuantSpecs(input_nodes, node_mins, node_maxs, final_type,
+                                quant_specs);
+}
+
+bool GetInputNodeQuantSpecs(const std::vector<std::string>& node_names,
+                            const std::vector<std::optional<double>>& node_mins,
+                            const std::vector<std::optional<double>>& node_maxs,
+                            const tensorflow::DataType inference_type,
+                            QuantizationSpecs* quant_specs) {
+  quant_specs->inference_type = inference_type;
+
+  // If min/max are not specified, just return;
+  if (node_mins.empty() || node_maxs.empty()) return false;
+
+  // Otherwise make sure min/max has the same size as inputs.
+  if (IsQuantizationType(inference_type)) {
+    // min/max should have same size as inputs, or shouldn't be specified.
+    if (node_names.size() != node_mins.size() ||
+        node_names.size() != node_maxs.size()) {
+      return true;
+    }
+    for (int i = 0; i < node_names.size(); ++i) {
+      quant_specs->input_ranges.push_back({node_mins[i], node_maxs[i]});
+    }
+    return false;
+  }
+  if (!node_mins.empty()) {
+    llvm::dbgs() << "Ignored input_min_values.";
+  }
+  if (!node_maxs.empty()) {
+    llvm::dbgs() << "Ignored input_max_values.";
+  }
+  return false;
+}
+
+std::string GetQDQQuantModeString(const QDQConversionMode mode) {
+  switch (mode) {
+    case QDQConversionMode::kQDQStatic:
+      return "Static";
+    case QDQConversionMode::kQDQDynamic:
+      return "Dynamic";
+    case QDQConversionMode::kQDQStrict:
+      return "Strict";
+    default:
+      return "NoQDQ";
+  }
+}
+
+QDQConversionMode GetQDQQuantModeFromString(const std::string& mode_str) {
+  if (mode_str == "Static") return QDQConversionMode::kQDQStatic;
+  if (mode_str == "Dynamic") return QDQConversionMode::kQDQDynamic;
+  if (mode_str == "Strict") return QDQConversionMode::kQDQStrict;
+  return QDQConversionMode::kQDQNone;
+}
+
+}  // namespace quant
+}  // namespace mlir

tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_config.h

@@ -0,0 +1,250 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This header file defines node specs for quantization and the methods to parse
+// command line flags to these specs.
+
+#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_QUANTIZATION_LIB_QUANTIZATION_CONFIG_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_QUANTIZATION_LIB_QUANTIZATION_CONFIG_H_
+
+#include <cstdint>
+#include <optional>
+#include <string>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "absl/container/flat_hash_set.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace mlir {
+namespace quant {
+
+// Stores information about how to quantize a user-specified custom operation.
+struct CustomOpInfo {
+  std::vector<std::int32_t> quantizable_input_indices;
+  bool is_weight_only = false;
+  bool no_side_effect = true;
+};
+
+using CustomOpMap = std::unordered_map<std::string, CustomOpInfo>;
+enum CustomOpUpdateOptions { kInputIndices, kWeightOnly, kNoSideEffect };
+enum class QDQConversionMode { kQDQNone, kQDQStatic, kQDQDynamic, kQDQStrict };
+
+struct QuantizationSpecs {
+  // Which function this node quant specifications belong to.
+  std::string target_func = "main";
+
+  // Whether to trigger quantization passses for post-training quantization.
+  // If true, the model input doesn't require user specified input ranges.
+  bool post_training_quantization = false;
+
+  // Whether to allow dynamic range quantization. This is the easiest
+  // quantization mode which doesn't require QAT or sample inputs.
+  // This option only targets `DT_HALF` and `DT_QINT8` inference type.
+  bool weight_quantization = false;
+
+  // Whether to use the MLIR dynamic range quantizer instead of TOCO.
+  bool enable_mlir_dynamic_range_quantizer = false;
+
+  // Whether to allow weight-only quantization. This scheme quantizes
+  // weights but will dequantize them back at runtime which is useful for
+  // memory bound case without kernel support available in lower precisions.
+  // Used in MLIR dynamic range quantizer.
+  bool weight_only_quantization = false;
+
+  // The minimum number of elements in a weights array required to apply
+  // quantization. This is especially useful not to quantize small tensors as
+  // it is hard to get performance benefits from them with quantization. Used
+  // in MLIR dynamic range quantizer with int8 weight data type.
+  int64_t minimum_elements_for_weights = 1024;
+
+  // Whether to calculate scales in float to keep quantized values the same with
+  // old TOCO quantizer.
+  bool legacy_float_scale = false;
+
+  // Whether to perform per-tensor quantization. Currently, this option is only
+  // valid when the quantization parameters need to be created by scanning the
+  // constant content (post-training quantization or QAT without weight
+  // FakeQuant).
+  bool disable_per_channel = false;
+
+  // Whether to disable per-channel weight quantization and enable legacy per
+  // tensor quantization. The legacy quantization for Dense layers is
+  // inconsistent with Conv 1x1 which always performs per channel quantization.
+  bool disable_per_channel_for_dense_layers = false;
+
+  // Whether to use fixed output ranges of the activation ops (tanh, sigmoid,
+  // etc.) and not infer weight constants.
+  // If this option is set, quantization emulation ops should be placed after
+  // the ops in the input graph. This flag should be set to false for
+  // post-training quantization.
+  bool disable_infer_tensor_range = false;
+
+  // Whether to use the unfrozen variable quantization in MLIR. Typically,
+  // variables are frozen for passing passes, but some variables aren't frozen.
+  // If it is true, QuantizeVariables pass will be added after the
+  // PrepareQuantizePass.
+  bool enable_mlir_variable_quantization = false;
+
+  // The node type when the model is exported. Currently this is limited to
+  // DT_FLOAT, DT_HALF, DT_QINT8, and DT_QUINT8. When DT_HALF is used, the
+  // `weight_quantization` flag needs to set to true. When DT_QUINT8 is used,
+  // the `weight_quantization` flag needs to set to false.
+  tensorflow::DataType inference_type = tensorflow::DT_FLOAT;
+
+  // The input and output data type during inference. This flag is only used
+  // when `inference_type` is different from DT_FLOAT. This flag can only be set
+  // to DT_FLOAT or as same as `inference_type`. If this flag is different
+  // from `inference_type`, adaptor ops are inserted as heading and tailing ops
+  // in the result model.
+  tensorflow::DataType inference_input_type = tensorflow::DT_FLOAT;
+
+  // Input node ranges. These ranges are stored as the same order of function
+  // arguments. They are only used when `weight_quantization` is set to false,
+  // and the model is required to have quantization parameters, either from
+  // quantization aware training or calibration, for the remaining tensors.
+  std::vector<std::pair<std::optional<double>, std::optional<double>>>
+      input_ranges;
+
+  // Whether to disable setting the quantization parameters of the input nodes
+  // using input ranges.
+  bool disable_set_input_nodes_quantization_params = false;
+
+  // The default ranges can be used when a tensor doesn't have quantization
+  // parameters and couldn't be quantized. Used only for latency tests.
+  std::pair<std::optional<double>, std::optional<double>> default_ranges;
+
+  // A serialized "QuantizationInfo" object to specify value ranges for some of
+  // the tensors with known names.
+  std::string serialized_quant_stats = "";
+
+  // Whether to run the passes to propagate the quantization parameters and
+  // graph rewrites. Returns false if the inference_type is DT_FLOAT or
+  // `weight_quantization` flag is set.
+  bool RunPropagationAndRewriteQuantizationPasses() const {
+    return inference_type != tensorflow::DT_FLOAT && !weight_quantization;
+  }
+
+  // TODO: b/202075505 - make implicit weight type clearer
+  // Whether run the passes and graph rewrites for dynamic range quantization.
+  bool RunAndRewriteDynamicRangeQuantizationPasses() const {
+    bool dynamic_range_quantize =
+        (inference_type != tensorflow::DT_FLOAT) && weight_quantization &&
+        !post_training_quantization && !disable_infer_tensor_range &&
+        enable_mlir_dynamic_range_quantizer;
+    return dynamic_range_quantize;
+  }
+
+  // Returns whether this inference type represents a signed storage type.
+  bool IsSignedInferenceType() const {
+    switch (inference_type) {
+      case tensorflow::DT_QUINT8:
+      case tensorflow::DT_QUINT16:
+        return false;
+      default:
+        return true;
+    }
+  }
+
+  // Gets the width of this quantization type. Returns 0 if it isn't a
+  // quantization type.
+  int64_t GetQuantizationTypeWidth() const {
+    switch (inference_type) {
+      case tensorflow::DT_INT8:
+      case tensorflow::DT_UINT8:
+      case tensorflow::DT_QINT8:
+      case tensorflow::DT_QUINT8:
+        return 8;
+      case tensorflow::DT_INT16:
+      case tensorflow::DT_UINT16:
+      case tensorflow::DT_QINT16:
+      case tensorflow::DT_QUINT16:
+        return 16;
+      case tensorflow::DT_INT32:
+      case tensorflow::DT_QINT32:
+        return 32;
+      default:
+        return 0;
+    }
+  }
+
+  // Whether to add the NumericVerify ops to verify numbers before and after
+  // quantization.
+  bool verify_numeric = false;
+  // Whether to add verification for layer by layer, or on whole model. When
+  // disabled (per-layer) float and quantized ops will be run from same input
+  // (output of previous quantized layer). When enabled, float and quantized ops
+  // will run with respective float and quantized output of previous ops.
+  bool whole_model_verify = false;
+
+  // Whether to use fake quant attributes to calculate quantization parameters.
+  bool use_fake_quant_num_bits = false;
+
+  // Names of ops to block from quantization. Used in QuantizePass.
+  // For dynamic range quantization, ops in blocklist are quantized in weight-
+  // only manner.
+  absl::flat_hash_set<std::string> ops_blocklist;
+
+  // Names of locations to block from quantization. Used in QuantizePass.
+  absl::flat_hash_set<std::string> nodes_blocklist;
+
+  // Map from custom op code to custom op quantization information.
+  // For dynamic range quantization, among the custom ops in the graph those
+  // specified in this map are subject to quantization.
+  CustomOpMap custom_map;
+
+  // If other than kQDQNone, the model is a floating point graph with QDQ ops
+  // to be eliminated and fused into quantized kernels.
+  QDQConversionMode qdq_conversion_mode = QDQConversionMode::kQDQNone;
+
+  // When set, adheres to the QDQ annotations added by the framework when
+  // possible rather than quantizing any op that is possible to quantize.
+  bool strict_qdq_mode = false;
+};
+
+// Parses the command line flag strings to the CustomOpMap specification.
+void ParseCustomOpSpecs(absl::string_view node_names,
+                        const CustomOpUpdateOptions& update_option,
+                        CustomOpMap& custom_op_map);
+
+// Parses the command line flag strings to the quantization specification for
+// input arrays of a graph. The array names are not stored in the spec, and will
+// be matched by position. Returns true if failed.
+bool ParseInputNodeQuantSpecs(absl::string_view node_names,
+                              absl::string_view min_values,
+                              absl::string_view max_values,
+                              absl::string_view inference_type,
+                              QuantizationSpecs* quant_specs);
+
+// Gets the quantization specification for input arrays. The array names are not
+// stored in the spec, and will be matched by position. The min/max will be
+// ignored if the inference_type isn't a quantized type. Returns true if failed.
+bool GetInputNodeQuantSpecs(const std::vector<std::string>& node_names,
+                            const std::vector<std::optional<double>>& node_mins,
+                            const std::vector<std::optional<double>>& node_maxs,
+                            tensorflow::DataType inference_type,
+                            QuantizationSpecs* quant_specs);
+
+// Returns a human-readable string of the QDQQuantMode enum class
+std::string GetQDQQuantModeString(QDQConversionMode mode);
+
+// Returns the QDQQuantMode enum class from a human-readable string
+QDQConversionMode GetQDQQuantModeFromString(const std::string& mode_str);
+}  // namespace quant
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_QUANTIZATION_LIB_QUANTIZATION_CONFIG_H_

tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_driver.cc

@@ -0,0 +1,959 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_driver.h"
+
+#include <cmath>
+#include <cstdint>
+#include <iterator>
+#include <limits>
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Matchers.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"
+#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_traits.h"
+#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"
+
+namespace mlir {
+namespace quant {
+namespace {
+
+constexpr int32_t kBiasMax = std::numeric_limits<int32_t>::max() / 2;
+
+// Uses the type of `value` to set the initial state of the index-th result if
+// `as_result` is true or index-th operand if `as_result` is false. The state
+// is immutable if the type is a quantized type. Returns the index of this
+// new state in the state vector.
+void InitializeStateForValue(
+    Operation* op, const int index, const Value value, const bool as_result,
+    std::vector<QuantState>& states,
+    DenseMap<Value, QuantizationDriver::QuantStateIndex>& value_to_state,
+    DenseMap<QuantizationDriver::OpWithOperandIndex,
+             QuantizationDriver::QuantStateIndex>& operand_states,
+    DenseMap<QuantizationDriver::OpWithResultIndex,
+             QuantizationDriver::QuantStateIndex>& result_states) {
+  const auto [cached, inserted] = value_to_state.try_emplace(value, 0);
+  if (!inserted) {
+    if (as_result) {
+      result_states[{op, index}] = cached->second;
+    } else {
+      operand_states[{op, index}] = cached->second;
+    }
+    return;
+  }
+
+  const QuantizedType quantized_type =
+      QuantizedType::getQuantizedElementType(value.getType());
+
+  const bool immutable = quantized_type != nullptr;
+  const QuantizationDriver::QuantStateIndex next_state_index = states.size();
+  states.push_back({quantized_type, immutable});
+  if (as_result) {
+    result_states[{op, index}] = next_state_index;
+  } else {
+    operand_states[{op, index}] = next_state_index;
+  }
+
+  cached->second = next_state_index;
+}
+
+bool HasPerAxisQuantizedOperand(Operation* op) {
+  for (int i = 0; i < op->getNumOperands(); ++i) {
+    if (auto dq_op = dyn_cast_or_null<mlir::quant::ir::DequantizeCastOp>(
+            op->getOperand(i).getDefiningOp())) {
+      auto type =
+          mlir::cast<TensorType>(dq_op.getArg().getType()).getElementType();
+      if (auto per_axis_qtype =
+              mlir::dyn_cast_or_null<quant::UniformQuantizedPerAxisType>(
+                  QuantizedType::getQuantizedElementType(type))) {
+        return true;
+      }
+    }
+  }
+  return false;
+}
+
+}  // namespace
+
+void QuantizationDriver::InitializeArgState(const BlockArgument arg,
+                                            const Value arg_value) {
+  const auto [cached, inserted] = value_to_state_.try_emplace(arg_value, 0);
+  if (!inserted) {
+    arg_states_[arg] = cached->second;
+    return;
+  }
+
+  const QuantizedType quantized_type =
+      QuantizedType::getQuantizedElementType(arg_value.getType());
+  const bool immutable = quantized_type != nullptr;
+  const QuantizationDriver::QuantStateIndex next_state_index = states_.size();
+  states_.push_back({quantized_type, immutable});
+  arg_states_[arg] = next_state_index;
+  cached->second = next_state_index;
+}
+
+void QuantizationDriver::InitializeOperandState(Operation* op, const int index,
+                                                const Value value) {
+  InitializeStateForValue(op, index, value, /*as_result=*/false, states_,
+                          value_to_state_, operand_states_, result_states_);
+}
+
+void QuantizationDriver::InitializeResultState(Operation* op, const int index,
+                                               const Value value) {
+  InitializeStateForValue(op, index, value, /*as_result=*/true, states_,
+                          value_to_state_, operand_states_, result_states_);
+}
+
+std::unique_ptr<OpQuantSpec> QuantizationDriver::GetQuantSpec(Operation* op) {
+  return op_quant_spec_getter_(op);
+}
+
+std::unique_ptr<OpQuantScaleSpec> QuantizationDriver::GetQuantScaleSpec(
+    Operation* op) {
+  return op_quant_scale_spec_getter_(op);
+}
+
+bool QuantizationDriver::IsQuantized(Operation* op) {
+  for (int i = 0; i < op->getNumResults(); ++i) {
+    if (GetResultQuantState(op, i).IsEmpty()) return false;
+  }
+  return true;
+}
+
+bool QuantizationDriver::SetConstantResultParams(Operation* op) {
+  DenseFPElementsAttr attr;
+  const Value result = op->getResult(0);
+  if (!matchPattern(result, m_Constant(&attr))) {
+    return false;
+  }
+  // TODO: b/323478683 - Make storage_type_width and narrow_range configurable.
+  Type final_type;
+  const auto it = optimized_weights_.find(op);
+  const bool is_weight = it != optimized_weights_.end();
+  const bool is_weight_with_per_channel_support =
+      is_weight && it->second != -1 && is_signed_;
+
+  if (is_weight_with_per_channel_support && !disable_per_channel_) {
+    // When `disable_per_channel_` is false, per-channel symmetric quantization
+    // parameters are created from the weights when the ops support per-channel
+    // quantization. Otherwise, uses per-tensor asymmetric quantization with
+    // narrow range.
+
+    // per-axis quantization weight, with symmetric min/max enforced.
+    final_type = GetUniformQuantizedPerAxisTypeForWeight(
+        attr, it->second, /*symmetric=*/true, /*num_bits=*/8, is_signed_,
+        /*narrow_range=*/true, legacy_float_scale_);
+  } else {
+    // per-tensor quantization weight
+    final_type = GetUniformQuantizedTypeForWeight(
+        attr, /*symmetric=*/is_weight && is_signed_,
+        /*num_bits=*/8, is_signed_,
+        /*narrow_range=*/is_weight, legacy_float_scale_);
+  }
+  if (const auto quant_type = mlir::dyn_cast_or_null<QuantizedType>(final_type);
+      quant_type != nullptr) {
+    return SetResultParams(op, /*result_index=*/0, quant_type);
+  }
+  return false;
+}
+
+bool QuantizationDriver::SetResultParams(Operation* op, const int result_index,
+                                         const QuantizedType quantized_type) {
+  QuantState& state = GetResultQuantState(op, result_index);
+  if (state.params == quantized_type) {
+    return false;
+  }
+  if (!state.IsEmpty()) {
+    RequantizeStates& rescales = GetResultRequantizeStates(op, result_index);
+    RequantizeState& rescale = rescales.emplace_back();
+    rescale.pos = RequantizeState::ON_INPUT;
+    rescale.params = quantized_type;
+    return true;
+  }
+  state.params = quantized_type;
+  AddUserToList(op, result_index);
+  return true;
+}
+
+QuantizedType QuantizationDriver::GetBiasParams(
+    Operation* op, const int bias_index,
+    const ArrayRef<int> non_bias_operand_indices,
+    const AccumulatorScaleFunc func) {
+  QuantState& bias_state = GetOperandQuantState(op, bias_index);
+  if (!bias_state.IsEmpty()) {
+    return bias_state.params;
+  }
+  std::vector<QuantizedType> op_types{};
+  op_types.reserve(non_bias_operand_indices.size());
+
+  int adjusted_quant_dim = -1;
+  if (op->getNumOperands() > bias_index) {
+    // Some kernels allow 1D bias, broadcasting it inside the kernel. In this
+    // case, the `quantizedDimension=0` when quantizing per-channel.
+    // However, for some kernels which require bias to be already broadcasted
+    // to match the accumulation shape, the very last index should be used.
+    Operation* bias_op = op->getOperand(bias_index).getDefiningOp();
+    if (bias_op != nullptr) {
+      Type bias_type = bias_op->getResult(0).getType();
+      if (bias_type != builder_.getNoneType()) {
+        const int bias_rank = mlir::dyn_cast<ShapedType>(bias_type).getRank();
+        adjusted_quant_dim = bias_rank > 1 ? bias_rank - 1 : 0;
+      }
+    }
+  }
+
+  for (const int non_bias_operand_index : non_bias_operand_indices) {
+    const QuantState& non_bias_state =
+        GetOperandQuantState(op, non_bias_operand_index);
+    op_types.push_back(non_bias_state.params);
+  }
+  return func(op_types, adjusted_quant_dim, legacy_float_scale_);
+}
+
+bool QuantizationDriver::SetOperandParams(Operation* op,
+                                          const int operand_index,
+                                          const QuantizedType quantized_type,
+                                          const bool override) {
+  QuantState& state = GetOperandQuantState(op, operand_index);
+  if (state.params == quantized_type) {
+    return false;
+  }
+
+  if (!state.IsEmpty() && !override) {
+    RequantizeStates& rescales = GetOperandRequantizeStates(op, operand_index);
+    for (RequantizeState& rescale : rescales) {
+      if (rescale.params == quantized_type) {
+        rescale.users.emplace_back(op, operand_index);
+        return true;
+      }
+    }
+    RequantizeState& rescale = rescales.emplace_back();
+    rescale.pos = RequantizeState::ON_OUTPUT;
+    rescale.params = quantized_type;
+    rescale.users.emplace_back(op, operand_index);
+    return true;
+  }
+
+  state.params = quantized_type;
+  AddOperandToList(op, operand_index);
+  return true;
+}
+
+void QuantizationDriver::QuantizeOpResult(Operation* op, const int result_index,
+                                          const QuantizedType quantized_type) {
+  builder_.setInsertionPointAfter(op);
+  const Value original_result = op->getResult(result_index);
+  QuantizeValue(original_result, quantized_type, op->getLoc());
+}
+
+void QuantizationDriver::QuantizeArg(BlockArgument arg,
+                                     const QuantizedType quantized_type) {
+  builder_.setInsertionPointToStart(arg.getOwner());
+  QuantizeValue(arg, quantized_type, builder_.getUnknownLoc());
+}
+
+void QuantizationDriver::QuantizeValue(Value value,
+                                       QuantizedType quantized_type,
+                                       const Location loc) {
+  const Type expressed_type = value.getType();
+  const Type new_value_type =
+      quantized_type.castFromExpressedType(expressed_type);
+  // Skip if `value` or `value`'s element type doesn't match the expressed type
+  // of `quantized_type`.
+  if (new_value_type == nullptr) return;
+
+  auto quantize = mlir::quant::ir::QuantizeCastOp::create(
+      builder_, loc, new_value_type, value);
+  auto dequantize = mlir::quant::ir::DequantizeCastOp::create(
+      builder_, loc, expressed_type, quantize.getResult());
+
+  // This attribute is set to distinguish the quantize ops being added by the
+  // quantization pass. These ops can be removed without losing original
+  // program accuracy.
+  // TODO: b/323478683 - Make the attribute being part of op definition.
+  quantize->setAttr(kVolatileOpAttrName, builder_.getUnitAttr());
+
+  // `original_result` has a use to `quantize`, so this will replace that use
+  // by the result of `dequantize`. Remember to reset that use afterwards
+  value.replaceAllUsesWith(dequantize);
+  quantize.getOperation()->replaceUsesOfWith(dequantize, value);
+}
+
+void QuantizationDriver::RequantizeOpResult(Operation* op,
+                                            const int result_index,
+                                            RequantizeStates& states) {
+  if (states.empty()) return;
+
+  builder_.setInsertionPointAfter(op);
+  Value value = op->getResult(result_index);
+  RequantizeState::RequantizePosition pos = states.front().pos;
+  if (pos == RequantizeState::NO_REQUANTIZE) {
+    return;
+  }
+  for (const RequantizeState& state : states) {
+    // Check that all requantization positions are the same for each state.
+    // Unsure if this check is required.
+    if (state.pos != pos) {
+      return;
+    }
+  }
+  if (pos == RequantizeState::ON_OUTPUT) {
+    Operation* user = value.getUses().begin().getUser();
+    if (isa<mlir::quant::ir::QuantizeCastOp>(user)) {
+      // The requantize op is inserted between `quantize` and `dequantize` ops.
+      value = user->getResult(0);
+      builder_.setInsertionPointAfter(user);
+    }
+  }
+  RequantizeValue(value, states, op->getLoc());
+}
+
+void QuantizationDriver::RequantizeArg(const BlockArgument arg,
+                                       RequantizeStates& states) {
+  Value value = arg;
+  builder_.setInsertionPointToStart(arg.getOwner());
+  if (value.hasOneUse()) {
+    Operation* user = value.use_begin().getUser();
+    if (auto q = dyn_cast<mlir::quant::ir::QuantizeCastOp>(user)) {
+      value = q.getResult();
+      builder_.setInsertionPoint(arg.getOwner(), ++Block::iterator(user));
+    }
+  }
+  RequantizeValue(value, states, builder_.getUnknownLoc());
+}
+
+void QuantizationDriver::RequantizeValue(Value value, RequantizeStates& states,
+                                         const Location loc) {
+  if (states.empty() || states.front().pos == RequantizeState::NO_REQUANTIZE) {
+    return;
+  }
+  if (states.front().pos == RequantizeState::ON_INPUT) {
+    RequantizeState& state = states.front();
+    const Type expressed_type = value.getType();
+    // The value needs to be requantized. A Quantize op will be created to use
+    // it as the operand and replace its uses.
+    const Type new_type = state.params.castFromExpressedType(expressed_type);
+    if (!new_type) return;
+    auto requantize_op =
+        mlir::quant::ir::QuantizeCastOp::create(builder_, loc, new_type, value);
+    value.replaceAllUsesWith(requantize_op);
+    requantize_op.getOperation()->replaceUsesOfWith(requantize_op, value);
+    // This requantization was defined as required for the result value, so
+    // there should be only one requant state.
+    return;
+  }
+
+  // If this is an operand that requires requantization, then the value should
+  // only have one `DequantizeCastOp` user which produces the operand value.
+  if (!value.hasOneUse()) {
+    return;
+  }
+  auto dequant_op = dyn_cast_or_null<mlir::quant::ir::DequantizeCastOp>(
+      value.use_begin().getUser());
+  if (!dequant_op) {
+    return;
+  }
+  // It is possible that the dequant value is used by a op that doesn't require
+  // requant, so only overwrite the first if that is not the case.
+  const int num_uses = std::distance(dequant_op.getResult().use_begin(),
+                                     dequant_op.getResult().use_end());
+
+  // Whether to replace quantization params of the first dequantize op
+  // after the quantized value is produced.
+  // If there is a use other than the requantize states, then we can't clobber.
+  bool clobber_first = num_uses <= states.size();
+  for (RequantizeState& state : states) {
+    Type expressed_type = QuantizedType::castToExpressedType(value.getType());
+    if (!expressed_type) continue;
+    // The value needs to be requantized. A Quantize op will be created to use
+    // it as the operand and replace its uses.
+    const Type new_type = state.params.castFromExpressedType(expressed_type);
+    // This value isn't an expressed type (float), skip.
+    if (!new_type) continue;
+
+    auto requantize_op =
+        mlir::quant::ir::QuantizeCastOp::create(builder_, loc, new_type, value);
+
+    if (clobber_first) {
+      dequant_op.setOperand(requantize_op.getResult());
+      // All ops requiring this value already use the result of dequant.
+      clobber_first = false;
+    } else {
+      auto new_dequant_op = mlir::quant::ir::DequantizeCastOp::create(
+          builder_, loc, dequant_op.getResult().getType(),
+          requantize_op.getResult());
+      for (auto [op, operand_idx] : state.users) {
+        op->setOperand(operand_idx, new_dequant_op.getResult());
+      }
+    }
+  }
+}
+
+// A heuristic to get quantization parameters satisfies the same scale
+// constraints:
+// - If there are immutable states,
+//   - use the single input, or,
+//   - use the single output, or,
+//   - use the first one in the collection,
+// - use the single input if it is ready, or,
+// - use the single output if it is ready, or,
+// - use the first ready one in the collection.
+QuantizedType QuantizationDriver::GetQuantParamsForSameScaleConstraint(
+    Operation* op) {
+  // Two vector to collect Non-empty operands and results states.
+  std::vector<QuantState*> mutable_states, immutable_states;
+  for (int i = 0; i < op->getNumOperands(); ++i) {
+    QuantState& state = GetOperandQuantState(op, i);
+    if (state.immutable) {
+      immutable_states.push_back(&state);
+    } else if (!state.IsEmpty()) {
+      mutable_states.push_back(&state);
+    }
+  }
+
+  const int immutable_operands_num = immutable_states.size();
+  const int mutable_operands_num = mutable_states.size();
+  // Use the operand's state if it is immutable and it is the only one
+  // operand.
+  if (op->getNumOperands() == 1 && immutable_operands_num == 1) {
+    return immutable_states.front()->params;
+  }
+
+  for (int i = 0; i < op->getNumResults(); ++i) {
+    QuantState& state = GetResultQuantState(op, i);
+    if (state.immutable) {
+      immutable_states.push_back(&state);
+    } else if (!state.IsEmpty()) {
+      mutable_states.push_back(&state);
+    }
+  }
+
+  const int immutable_results_num =
+      immutable_states.size() - immutable_operands_num;
+  const int mutable_results_num = mutable_states.size() - mutable_operands_num;
+  // Use the result's state if it is immutable and it is the only one result.
+  if (op->getNumResults() == 1 && immutable_results_num == 1) {
+    return immutable_states.back()->params;
+  }
+
+  // Use the first immutable state to quantize the rest operands and results.
+  if (!immutable_states.empty()) return immutable_states.front()->params;
+
+  // If there are no immutable states, use the operand's state if it is the
+  // only one operand and has parameters propagated.
+  if (op->getNumOperands() == 1 && mutable_operands_num == 1) {
+    return mutable_states.front()->params;
+  }
+
+  // If there are no immutable states, use the result's state if it is the
+  // only one result and has parameters propagated.
+  if (op->getNumResults() == 1 && mutable_results_num == 1) {
+    return mutable_states.back()->params;
+  }
+
+  // Use the first propagated state to quantize the rest operands and results.
+  if (!mutable_states.empty()) return mutable_states.front()->params;
+
+  // None operands/results have parameters propagated, skip this node for now.
+  return {};
+}
+
+void QuantizationDriver::PreprocessConstantOps() {
+  fn_.walk([&](arith::ConstantOp cst) {
+    // Non-float tensors are neither weights nor require quantization.
+    const auto type = mlir::dyn_cast<ShapedType>(cst.getType());
+    if (!type || !mlir::isa<FloatType>(type.getElementType())) return;
+
+    // Skip if the value is NaN or INF.
+    // Otherwise the illegal scale/zp will be calculated.
+    auto float_attr = mlir::dyn_cast<DenseFPElementsAttr>(cst.getValueAttr());
+    if (float_attr && (float_attr.getValues<APFloat>().empty() ||
+                       !float_attr.getValues<APFloat>()[0].isFinite())) {
+      return;
+    }
+
+    const Value value = cst.getResult();
+    builder_.setInsertionPoint(cst);
+
+    // The following loop will change the value uses, thus we cache all the uses
+    // needs to be changed.
+    SmallVector<std::pair<Operation*, int>> uses;
+    for (OpOperand& use : value.getUses()) {
+      uses.push_back({use.getOwner(), use.getOperandNumber()});
+    }
+    for (const auto [user, operand_num] : uses) {
+      const std::unique_ptr<OpQuantSpec> spec = GetQuantSpec(user);
+      const std::unique_ptr<OpQuantScaleSpec> scale_spec =
+          GetQuantScaleSpec(user);
+      const BiasParamsMap biases = spec->biases_params;
+
+      // The quantization parameters of a `weight` shouldn't be determined by
+      // other values. So any constants which are not bias, an operand of an
+      // op with same scale requirements, and haven't been quantized are
+      // weights.
+      if (!biases.contains(operand_num) &&
+          !scale_spec->has_same_scale_requirement &&
+          !dyn_cast<mlir::quant::ir::QuantizeCastOp>(user)) {
+        // Needs to scan the content of weights to get the quantization
+        // parameters if there are no quantization parameters (FakeQuant ops).
+        // For this case, the weight will not be duplicated.
+        weights_.insert(cst);
+        if (spec->coeff_op_quant_dim.find(operand_num) !=
+            spec->coeff_op_quant_dim.end()) {
+          optimized_weights_.insert(
+              {cst, spec->coeff_op_quant_dim[operand_num]});
+        }
+      } else {
+        // This is a bias or an operand of an op with same scale requirements,
+        // so the quantization parameter are propagated from or determined by
+        // other values. Duplicate this constant in case it is shared by
+        // different users.
+        if (uses.size() > 1) {
+          auto new_constant_op =
+              arith::ConstantOp::create(builder_, cst.getLoc(), cst.getValue());
+          user->setOperand(operand_num, new_constant_op);
+        }
+      }
+    }
+  });
+}
+
+void QuantizationDriver::SetupAllStates() {
+  for (BlockArgument arg : fn_.getArguments()) {
+    args_.push_back(arg);
+    Value value = arg;
+    // If the argument is quantized, it should only has one user.
+    if (arg.hasOneUse()) {
+      Operation* user = value.use_begin().getUser();
+      if (auto q = dyn_cast<mlir::quant::ir::QuantizeCastOp>(user)) {
+        value = q.getResult();
+      }
+    }
+    InitializeArgState(arg, value);
+  }
+
+  fn_.walk([&](Operation* op) {
+    std::unique_ptr<OpQuantScaleSpec> scale_spec = GetQuantScaleSpec(op);
+    if (!IsOpQuantizable(op) && !scale_spec->has_same_scale_requirement) {
+      return;
+    }
+    work_list_.push_back(op);
+
+    for (int i = 0; i < op->getNumOperands(); ++i) {
+      Value operand = op->getOperand(i);
+      if (Operation* inst = operand.getDefiningOp()) {
+        // If the operand comes from a `mlir::quant::ir::DequantizeCastOp`, we
+        // use the quantized input of this `mlir::quant::ir::DequantizeCastOp`
+        // to set the state.
+        if (auto dq = dyn_cast<mlir::quant::ir::DequantizeCastOp>(inst)) {
+          operand = dq.getArg();
+        }
+      }
+      InitializeOperandState(op, i, operand);
+    }
+
+    for (int i = 0; i < op->getNumResults(); ++i) {
+      Value result = op->getResult(i);
+      // If the result has been quantized, it should only be used by a
+      // `mlir::quant::ir::QuantizeCastOp`. For this case, we uses the quantized
+      // result to create the state and mark it immutable.
+      if (result.hasOneUse()) {
+        Operation* user = result.use_begin().getUser();
+        if (auto q = dyn_cast<mlir::quant::ir::QuantizeCastOp>(user)) {
+          result = q.getResult();
+        }
+      }
+      InitializeResultState(op, i, result);
+    }
+  });
+}
+
+arith::ConstantOp QuantizationDriver::DuplicateConstantOpIfNeeded(
+    arith::ConstantOp op, Operation* target_op, const int operand_index) {
+  if (op.getResult().hasOneUse()) {
+    return op;
+  }
+  OpBuilder builder(op->getContext());
+  builder.setInsertionPointAfter(op);
+  arith::ConstantOp new_op = cast<arith::ConstantOp>(builder.clone(*op));
+  target_op->getOpOperand(operand_index).set(new_op.getResult());
+  InitializeOperandState(target_op, operand_index, new_op.getResult());
+  InitializeResultState(new_op, 0, new_op.getResult());
+  return new_op;
+}
+
+bool QuantizationDriver::ShouldCheckBiasScale(
+    Operation* op, const int bias_index, ArrayRef<int> input_indices,
+    const QuantizedType quantized_type, int& input_index, int& filter_index) {
+  // For now, restrict scale adjustment to ops with affine quantized weights,
+  // and having weights and biases as constants. This currently only applies to
+  // FC and Conv* ops. Restriction for the weight can be relaxed if there are
+  // needs for adjusting scale of variable weights.
+  auto affine_op = dyn_cast<AffineQuantizedOpInterface>(op);
+  auto bias_op = op->getOperand(bias_index).getDefiningOp<arith::ConstantOp>();
+  if (!affine_op || !bias_op || input_indices.size() != 2) return false;
+  if (!mlir::isa<DenseFPElementsAttr>(bias_op.getValue())) return false;
+  filter_index = affine_op.GetAffineOperandIndex();
+  if (!op->getOperand(filter_index).getDefiningOp<arith::ConstantOp>()) {
+    return false;
+  }
+  if (filter_index == input_indices[0]) {
+    input_index = input_indices[1];
+  } else if (filter_index == input_indices[1]) {
+    input_index = input_indices[0];
+  } else {
+    return false;
+  }
+
+  const QuantState& input_state = GetOperandQuantState(op, input_index);
+  const QuantState& filter_state = GetOperandQuantState(op, filter_index);
+  // If quantization parameter for the filter is fixed, should return it as-is.
+  // Only checks ops with 8-bit input and weights, and 32-bit biases.
+  return input_state.params.getStorageTypeIntegralWidth() == 8 &&
+         filter_state.params.getStorageTypeIntegralWidth() == 8 &&
+         quantized_type.getStorageTypeIntegralWidth() == 32;
+}
+
+bool QuantizationDriver::SetBiasParamsWithAdjustments(
+    Operation* op, const int bias_index, ArrayRef<int> input_indices,
+    const QuantizedType params) {
+  bool changed = false;
+
+  int input_index;
+  int filter_index;
+  if (!ShouldCheckBiasScale(op, bias_index, input_indices, params, input_index,
+                            filter_index)) {
+    return SetOperandParams(op, bias_index, params);
+  }
+
+  QuantState input_state = GetOperandQuantState(op, input_index);
+  QuantState filter_state = GetOperandQuantState(op, filter_index);
+  auto bias_op = op->getOperand(bias_index).getDefiningOp<arith::ConstantOp>();
+  const double input_scale =
+      mlir::cast<UniformQuantizedType>(input_state.params).getScale();
+
+  auto bias_values = mlir::cast<DenseFPElementsAttr>(bias_op.getValue());
+  // Restrict maximum absolute value of bias within INT_MAX / 2, to make some
+  // room for accumulator.
+  if (auto bias_quantized_type = mlir::dyn_cast<UniformQuantizedType>(params);
+      bias_quantized_type != nullptr) {
+    double bias_half_range = 0.0f;
+    for (auto bias : bias_values.getValues<APFloat>()) {
+      if (bias_half_range < std::abs(bias.convertToFloat())) {
+        bias_half_range = std::abs(bias.convertToFloat());
+      }
+    }
+    if (bias_half_range / bias_quantized_type.getScale() < kBiasMax) {
+      return SetOperandParams(op, bias_index, params);
+    }
+    const double new_bias_scale =
+        static_cast<double>(bias_half_range) / kBiasMax;
+
+    changed |= SetOperandParams(
+        op, bias_index,
+        UniformQuantizedType::getChecked(
+            bias_op->getLoc(), params.getFlags(), params.getStorageType(),
+            params.getExpressedType(), new_bias_scale, 0,
+            params.getStorageTypeMin(), params.getStorageTypeMax()));
+    arith::ConstantOp filter_op = DuplicateConstantOpIfNeeded(
+        op->getOperand(filter_index).getDefiningOp<arith::ConstantOp>(), op,
+        filter_index);
+    if (!filter_op) {
+      return SetOperandParams(op, bias_index, params);
+    }
+
+    const auto filter_quantized_type =
+        mlir::cast<UniformQuantizedType>(filter_state.params);
+    changed |= SetOperandParams(
+        op, filter_index,
+        UniformQuantizedType::getChecked(
+            filter_op->getLoc(), filter_quantized_type.getFlags(),
+            filter_quantized_type.getStorageType(),
+            filter_quantized_type.getExpressedType(),
+            new_bias_scale / input_scale, 0,
+            filter_quantized_type.getStorageTypeMin(),
+            filter_quantized_type.getStorageTypeMax()),
+        /*override=*/true);
+  } else if (auto bias_quantized_type =
+                 mlir::dyn_cast<quant::UniformQuantizedPerAxisType>(params);
+             bias_quantized_type != nullptr) {
+    const auto filter_quantized_type =
+        mlir::cast<quant::UniformQuantizedPerAxisType>(filter_state.params);
+    std::vector<double> new_bias_scales = bias_quantized_type.getScales().vec();
+    std::vector<double> new_filter_scales =
+        filter_quantized_type.getScales().vec();
+
+    bool needs_adjustment = false;
+    for (int i = 0; i < bias_quantized_type.getScales().size(); ++i) {
+      const float abs_bias = std::abs(bias_values.getValues<float>()[i]);
+      if (abs_bias / new_bias_scales[i] > kBiasMax) {
+        new_bias_scales[i] = static_cast<double>(abs_bias) / kBiasMax;
+        new_filter_scales[i] = new_bias_scales[i] / input_scale;
+        needs_adjustment = true;
+      }
+    }
+    if (!needs_adjustment) {
+      return SetOperandParams(op, bias_index, params);
+    }
+    changed |= SetOperandParams(
+        op, bias_index,
+        quant::UniformQuantizedPerAxisType::getChecked(
+            bias_op->getLoc(), params.getFlags(), params.getStorageType(),
+            params.getExpressedType(), new_bias_scales,
+            bias_quantized_type.getZeroPoints(),
+            bias_quantized_type.getQuantizedDimension(),
+            params.getStorageTypeMin(), params.getStorageTypeMax()));
+
+    arith::ConstantOp filter_op = DuplicateConstantOpIfNeeded(
+        op->getOperand(filter_index).getDefiningOp<arith::ConstantOp>(), op,
+        filter_index);
+    changed |= SetOperandParams(
+        op, filter_index,
+        quant::UniformQuantizedPerAxisType::getChecked(
+            filter_op->getLoc(), filter_quantized_type.getFlags(),
+            filter_quantized_type.getStorageType(),
+            filter_quantized_type.getExpressedType(), new_filter_scales,
+            filter_quantized_type.getZeroPoints(),
+            filter_quantized_type.getQuantizedDimension(),
+            filter_quantized_type.getStorageTypeMin(),
+            filter_quantized_type.getStorageTypeMax()),
+        /*override=*/true);
+  }
+  return changed;
+}
+
+// This method scans the operations in the function to setup the initial
+// states for quantization parameter propagation.
+// TODO: b/323478683 - This algorithm assumes there are only one pair of
+// `mlir::quant::ir::QuantizeCastOp` and `mlir::quant::ir::DequantizeCastOp` ops
+// between two quantizable ops. A sanity check should be applied.
+void QuantizationDriver::Initialize() {
+  // Duplicate the bias constant, so the states can be setup correctly.
+  // TODO: b/323478683 - Function definition should also be duplicated if there
+  // are multiple call sites.
+  PreprocessConstantOps();
+
+  // Setup all the internal states.
+  SetupAllStates();
+}
+
+// Propagates the quantization parameters to the operands, results, and biases.
+// TODO: b/323478683 - Do not use while loop to handle this logic.
+bool QuantizationDriver::PropagateParamsAndReturnIfChanged() {
+  // TODO: b/323478683 - Use a typed indicator instead of a bool value.
+  bool changed = false;
+  while (!work_list_.empty()) {
+    Operation* op = work_list_.back();
+    work_list_.pop_back();
+
+    // This op has been quantized, so we should not consider it again.
+    if (quantized_.contains(op)) continue;
+    quantized_.insert(op);
+
+    if (auto constant_op = dyn_cast<arith::ConstantOp>(op); constant_op) {
+      // If the workflow requires inferring ranges from the content
+      // (post-training quantization) and it is weight (filter) and hasn't
+      // been quantized, we infer the quantization parameters from the content.
+      if (infer_tensor_range_ && IsWeight(constant_op) && !IsQuantized(op)) {
+        // The quantization parameters are determined by the content of the
+        // constant.
+        changed |= SetConstantResultParams(op);
+      }
+      continue;
+    }
+
+    std::unique_ptr<OpQuantScaleSpec> scale_spec = GetQuantScaleSpec(op);
+
+    if (scale_spec->has_same_scale_requirement) {
+      const QuantizedType params = GetQuantParamsForSameScaleConstraint(op);
+      // The quantization parameters haven't been propagated to any operands
+      // or results. Skip this node for now.
+      if (!params) {
+        quantized_.erase(op);
+        continue;
+      }
+
+      // If this is a QDQ conversion only, the op could have a same-scale
+      // requirement for the floating point kernel but allow per-axis
+      // quantization for the quantized kernel. If the quantized dimension
+      // changes, the following logic no longer works as the same `params`
+      // shouldn't be used for both input and output quantization params.
+      // E.g. During TransposeOp's quantization propagation in
+      // PrepareQuantize, if the quantization is per-axis and the
+      // QuantizedDimension is transposed, then the output q-dq params must
+      // reflect the new QuantizedDimension. So, check and skip the
+      // propagation if any of the operands has a per-axis quantized type param
+      // and `RequiredSameQuantizedAxes` set to false.
+      // Currently, these lines of code are only applicable to TFL_TransposeOp
+      // and TFL_ReshapeOp. And the output q-dq propagation for this Op is
+      // performed in `PropagateTransposedPerAxisQuantDim` and
+      // `PropagateReshapedPerAxisQuantDim` respectively.
+      if (is_qdq_conversion_ &&
+          !scale_spec->required_same_quantized_axes_func()) {
+        if (HasPerAxisQuantizedOperand(op)) continue;
+      }
+
+      // Use the final state to set all the operands' parameters.
+      for (int i = 0; i < op->getNumOperands(); ++i) {
+        if (auto type =
+                mlir::dyn_cast<ShapedType>(op->getOperand(i).getType())) {
+          // Without this check, it will accidentally propagate the quantization
+          // information by the shared non-float tensors.
+          if (mlir::isa<FloatType>(type.getElementType()))
+            changed |= SetOperandParams(op, i, params);
+        }
+      }
+
+      // Use the final state to set all the results' parameters.
+      for (int i = 0; i < op->getNumResults(); ++i)
+        if (auto type = mlir::dyn_cast<ShapedType>(op->getResult(i).getType());
+            type != nullptr) {
+          // Without this check, it will accidentally propagate the quantization
+          // information by the shared non-float-tensors.
+          if (mlir::isa<FloatType>(type.getElementType()))
+            changed |= SetResultParams(op, i, params);
+        }
+    }
+
+    // If the model already contains immutable QDQs, require upstream to
+    // explicitly fix output range instead.
+    if (scale_spec->has_fixed_output_range && infer_tensor_range_ &&
+        !is_qdq_conversion_) {
+      // Infer ranges from the activation ops. This is usually required for
+      // the post-training quantization workflow.
+      // TODO: b/323478683 - Different result can have different fixed range.
+      const QuantizedType params =
+          scale_spec->fixed_output_range_func(is_signed_, bit_width_);
+      for (auto i = 0; i < op->getNumResults(); ++i) {
+        // The range is null if the result has been quantized.
+        if (params) {
+          changed |= SetResultParams(op, i, params);
+        }
+      }
+    }
+
+    const std::unique_ptr<OpQuantSpec> spec = GetQuantSpec(op);
+    for (const auto& [bias_operand_idx, non_bias_params] :
+         spec->biases_params) {
+      const auto& [non_bias_operand_indices, accumulator_scale_func] =
+          non_bias_params;
+      const QuantizedType params =
+          GetBiasParams(op, bias_operand_idx, non_bias_operand_indices,
+                        accumulator_scale_func);
+      if (!params) {
+        quantized_.erase(op);
+        continue;
+      }
+      changed |= SetBiasParamsWithAdjustments(op, bias_operand_idx,
+                                              non_bias_operand_indices, params);
+    }
+  }
+
+  return changed;
+}
+
+// Finalizes the arguments and result states in the function.
+void QuantizationDriver::Finalize() {
+  for (BlockArgument arg : args_) {
+    const QuantState& state = GetArgQuantState(arg);
+    RequantizeStates& requantizes = GetArgRequantizeStates(arg);
+    if (state.IsEmpty() || (state.immutable && requantizes.empty())) {
+      continue;
+    }
+
+    if (!state.immutable) {
+      QuantizeArg(arg, state.params);
+    }
+
+    if (!requantizes.empty()) {
+      RequantizeArg(arg, requantizes);
+    }
+  }
+
+  for (const auto& [op_with_result_idx, quant_state_idx] : result_states_) {
+    const auto [op, result_idx] = op_with_result_idx;
+    const QuantState& state = GetResultQuantState(op, result_idx);
+    RequantizeStates& requantizes = GetResultRequantizeStates(op, result_idx);
+    if (state.IsEmpty() || (state.immutable && requantizes.empty())) {
+      continue;
+    }
+
+    if (!state.immutable) {
+      QuantizeOpResult(op, result_idx, state.params);
+    }
+
+    if (!requantizes.empty()) {
+      RequantizeOpResult(op, result_idx, requantizes);
+    }
+  }
+}
+
+// Runs quantization in following steps:
+//   1. Scans the operations in the function to setup the initial
+//      states for quantization parameter propagation.
+//   2. Propagates the quantization parameters to the operands, results, and
+//      biases.
+//   3. Finalizes the arguments and result states in the function.
+void QuantizationDriver::Run() {
+  Initialize();
+  if (PropagateParamsAndReturnIfChanged()) {
+    Finalize();
+  }
+}
+
+void ApplyQuantizationParamsPropagation(
+    const func::FuncOp func, const bool is_signed, const int bit_width,
+    const bool disable_per_channel,
+    const OpQuantSpecGetter op_quant_spec_getter,
+    const bool infer_tensor_ranges, const bool legacy_float_scale,
+    const bool is_qdq_conversion) {
+  ApplyQuantizationParamsPropagation(
+      func, is_signed, bit_width, disable_per_channel, op_quant_spec_getter,
+      GetDefaultQuantScaleSpec, infer_tensor_ranges, legacy_float_scale,
+      is_qdq_conversion);
+}
+
+void ApplyQuantizationParamsPropagation(
+    const func::FuncOp func, const bool is_signed, const int bit_width,
+    const bool disable_per_channel,
+    const OpQuantSpecGetter op_quant_spec_getter,
+    const OpQuantScaleSpecGetter op_quant_scale_spec_getter,
+    const bool infer_tensor_ranges, const bool legacy_float_scale,
+    const bool is_qdq_conversion) {
+  QuantizationDriver(func, is_signed, bit_width, disable_per_channel,
+                     op_quant_spec_getter, op_quant_scale_spec_getter,
+                     infer_tensor_ranges, legacy_float_scale, is_qdq_conversion)
+      .Run();
+}
+
+}  // namespace quant
+}  // namespace mlir

tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_driver.h

@@ -0,0 +1,387 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_QUANTIZATION_LIB_QUANTIZATION_DRIVER_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_QUANTIZATION_LIB_QUANTIZATION_DRIVER_H_
+
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/container/flat_hash_set.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"
+
+namespace mlir {
+namespace quant {
+
+// The state for each op result during the quantization parameters propagation.
+struct QuantState {
+  // Quantization parameters propagated to an op result.
+  QuantizedType params;
+  // A flag indicates this state (the params) shouldn't be changed after it is
+  // initialized. This flag will be set to true if the quantization parameters
+  // are from the quantization-aware training.
+  const bool immutable;
+
+  bool IsEmpty() const { return params == nullptr; }
+};
+
+// The state for rescaling the propagated quantization parameters. This can be
+// on the input side to satisfy the constraint of previous operation, or on the
+// output side to satisfy the constraint of the next operation.
+struct RequantizeState {
+  // Sometimes, we have to "requantize" the quantization result to satisfy all
+  // the constraints. The "requantize" can happen either on the input or output
+  // of the quantization result.
+  enum RequantizePosition {
+    NO_REQUANTIZE,
+    ON_INPUT,
+    ON_OUTPUT
+  } pos = NO_REQUANTIZE;
+
+  // Quantization parameters will be used to add the requantize ops.
+  QuantizedType params;
+
+  // Avoid clobbering all uses of the value, limit to just these ops.
+  SmallVector<std::pair<Operation*, int>> users;
+};
+
+using RequantizeStates = SmallVector<RequantizeState>;
+
+// This is a worklist-driven driver for propagating quantization parameters
+// across operations.
+//
+// The initial quantization parameters are extracted from the quantized type
+// between adjacent `mlir::quant::ir::QuantizeCastOp` and
+// `mlir::quant::ir::DequantizeCastOp`s. All these initial parameters are marked
+// as immutable because they are from quantization-aware training.
+//
+// The algorithm traverses each op and sets the quantization parameters of its
+// operands and results, according to its quantization specification, and then
+// adds the operands and results to the worklist. If there are any conflicts
+// (for example, there are quantization parameters propagated from the previous
+// iteration), this process stops if the existing parameters are the immutable,
+// or adding `requantize` op to resolve the conflicts.
+//
+// After the algorithm is converged, pairs of `mlir::quant::ir::QuantizeCastOp`
+// and `mlir::quant::ir::DequantizeCastOp` are inserted to the right position to
+// materialize the propagation and requantize results.
+//
+class QuantizationDriver {
+ public:
+  // Type alias of int used to access `states_`.
+  using QuantStateIndex = int;
+
+  // (op, operand index) pair.
+  using OpWithOperandIndex = std::pair<Operation*, int>;
+
+  // (op, result index) pair.
+  using OpWithResultIndex = std::pair<Operation*, int>;
+
+  explicit QuantizationDriver(func::FuncOp func_op, const bool is_signed,
+                              const int bit_width,
+                              const bool disable_per_channel,
+                              OpQuantSpecGetter op_quant_spec_getter,
+                              OpQuantScaleSpecGetter op_quant_scale_spec_getter,
+                              const bool infer_tensor_range,
+                              const bool legacy_float_scale = false,
+                              const bool is_qdq_conversion = false)
+      : fn_(func_op),
+        builder_(func_op.getBody()),
+        is_signed_(is_signed),
+        bit_width_(bit_width),
+        disable_per_channel_(disable_per_channel),
+        op_quant_spec_getter_(op_quant_spec_getter),
+        op_quant_scale_spec_getter_(op_quant_scale_spec_getter),
+        infer_tensor_range_(infer_tensor_range),
+        legacy_float_scale_(legacy_float_scale),
+        is_qdq_conversion_(is_qdq_conversion) {}
+
+  // The entry point of the quantization parameters propagation.
+  void Run();
+
+  // Sets up the states for all the op results in the function.
+  void Initialize();
+
+  // Propagates the quantization parameters across all the ops.
+  bool PropagateParamsAndReturnIfChanged();
+
+  // Inserts the Quantize and Dequantize ops according to the propagation
+  // result.
+  void Finalize();
+
+  SmallVector<BlockArgument, 4> GetArgs() { return args_; }
+
+  llvm::DenseMap<std::pair<mlir::Operation*, int>, int> GetResultStates() {
+    return result_states_;
+  }
+
+  DenseMap<OpWithResultIndex, QuantStateIndex> result_states_;
+
+  // Returns the state of the block argument.
+  QuantState& GetArgQuantState(BlockArgument arg) {
+    return states_[arg_states_[arg]];
+  }
+
+  // Returns the state of the index-th result of the op.
+  QuantState& GetResultQuantState(Operation* op, const int index) {
+    return states_[result_states_[{op, index}]];
+  }
+
+ private:
+  // Duplicates the constant op if it has multiple uses, and replaces
+  // target_op->operand[operand_index] with the newly created op. This also
+  // replaces corresponsing quantization states.
+  arith::ConstantOp DuplicateConstantOpIfNeeded(arith::ConstantOp op,
+                                                Operation* target_op,
+                                                int operand_index);
+
+  // Adjusts bias scale that is derived from other scales (fc, conv ops) to
+  // prevent overflow of quantized bias values. This also changes quantization
+  // state of other inputs when needed.
+  bool SetBiasParamsWithAdjustments(Operation* op, int bias_index,
+                                    ArrayRef<int> input_indices,
+                                    QuantizedType params);
+
+  // Checks preconditions to adjust bias scale.
+  bool ShouldCheckBiasScale(Operation* op, int bias_index,
+                            ArrayRef<int> input_indices,
+                            QuantizedType quantized_type, int& input_index,
+                            int& filter_index);
+
+  // Preprocesses the constants by doing the following:
+  //   - Duplicates constants if it is used by multiple ops. For example, if a
+  //     constant is used by multiple ops as a bias, duplicate constants and
+  //     let each op assign its own quantization parameter for bias.
+  //   - Adds all the non-bias constants (weights) to a set for looking up
+  //     later.
+  //   - Adds all per-channel weights to a set for looking up later.
+  void PreprocessConstantOps();
+
+  // Sets up all the data structures for quantization propagation.
+  void SetupAllStates();
+
+  // Returns Whether the constant is a weight, which shouldn't be shared by
+  // different ops.
+  bool IsWeight(Operation* cst) { return llvm::is_contained(weights_, cst); }
+
+  // Returns all the related quantization constraints of the op.
+  std::unique_ptr<OpQuantSpec> GetQuantSpec(Operation* op);
+  std::unique_ptr<OpQuantScaleSpec> GetQuantScaleSpec(Operation* op);
+
+  // Returns whether quantization parameters have been propagated to the results
+  // of this op.
+  bool IsQuantized(Operation* op);
+
+  // Adds all the users of index-th result of op to the work list.
+  void AddUserToList(Operation* op, const int index) {
+    for (Operation* user : op->getResult(index).getUsers()) {
+      work_list_.push_back(user);
+    }
+  }
+
+  // Adds the defining op of index-th operand of op to the work list.
+  void AddOperandToList(Operation* op, const int index) {
+    if (Operation* operand_op = op->getOperand(index).getDefiningOp();
+        operand_op != nullptr) {
+      work_list_.push_back(operand_op);
+    }
+  }
+
+  // Returns the quantization params for the bias input from the non-bias
+  // operands which have their indexes in the `non_biases` vector. The returned
+  // parameters are calculated by `func`.
+  QuantizedType GetBiasParams(Operation* op, int bias_index,
+                              ArrayRef<int> non_bias_operand_indices,
+                              AccumulatorScaleFunc func);
+
+  // Sets the quantization parameters of the result to `quantized_type`. If
+  // any quantization parameters have been propagated, a requantize will
+  // happen on the input of propagated quantization. Returns `true` if internal
+  // state has been modified.
+  bool SetResultParams(Operation* op, int result_index,
+                       QuantizedType quantized_type);
+
+  // Sets the quantization parameters of the operand to `quantized_type`. If any
+  // quantization parameters have been propagated, a `requantize` will happen on
+  // the output of propagated quantization. When `override` is set, quantization
+  // state of the value is replaced instead of adding requantization. Returns
+  // `true` if internal state has been modified.
+  bool SetOperandParams(Operation* op, int operand_index,
+                        QuantizedType quantized_type, bool override = false);
+
+  // Sets the quantization parameters of the constant result according to its
+  // content.
+  bool SetConstantResultParams(Operation* op);
+
+  // Inserts the Quantize and Dequantize ops after `op`'s `index`-th result. The
+  // quantized element type for the result is `quantized_type`.
+  void QuantizeOpResult(Operation* op, int result_index,
+                        QuantizedType quantized_type);
+
+  // Inserts the Quantize and Dequantize ops after `arg`. The quantized element
+  // type for `arg` is `quantized_type`.
+  void QuantizeArg(BlockArgument arg, QuantizedType quantized_type);
+
+  // Inserts the Quantize and Dequantize ops (i.e. QDQ) after `value`. The
+  // quantized element type for `value` is `quantized_type`.
+  void QuantizeValue(Value value, QuantizedType quantized_type, Location loc);
+
+  // Inserts the Quantize ops for requantizing the index-th result of the op.
+  void RequantizeOpResult(Operation* op, int result_index,
+                          RequantizeStates& states);
+
+  // Inserts the Quantize ops for requantizing a block argument.
+  void RequantizeArg(BlockArgument arg, RequantizeStates& states);
+
+  // Inserts the Quantize and Dequantize ops to quantize the value and returns
+  // the Quantize op.
+  void RequantizeValue(Value value, RequantizeStates& states, Location loc);
+
+  // Returns the quantization parameter satisfies the same scale
+  // constraints for the op. Returns an empty option if this quantization
+  // parameter doesn't exist.
+  QuantizedType GetQuantParamsForSameScaleConstraint(Operation* op);
+
+  // Returns the state of the index-th operand of the op.
+  QuantState& GetOperandQuantState(Operation* op, const int index) {
+    return states_[operand_states_[{op, index}]];
+  }
+
+  // Returns the states of the index-th operand of the op.
+  RequantizeStates& GetOperandRequantizeStates(Operation* op, const int index) {
+    return rescale_states_[operand_states_[{op, index}]];
+  }
+
+  // Returns the states of the index-th result of the op.
+  RequantizeStates& GetResultRequantizeStates(Operation* op, const int index) {
+    return rescale_states_[result_states_[{op, index}]];
+  }
+
+  // Returns the states of the arg.
+  RequantizeStates& GetArgRequantizeStates(BlockArgument arg) {
+    return rescale_states_[arg_states_[arg]];
+  }
+
+  // Sets the state of an argument. If this value is cached, uses the cached
+  // result without creating new entry in the state vector. Otherwise, allocate
+  // a new entry in the state vector.
+  void InitializeArgState(BlockArgument arg, Value arg_value);
+
+  // Sets the state of the index-th operand of the op. If this operand is
+  // cached, uses the cached result without creating new entry in the state
+  // vector. Otherwise, allocate a new entry in the state vector.
+  void InitializeOperandState(Operation* op, int index, Value value);
+
+  // Sets the state of the index-th result of the op. If this result is cached,
+  // uses the cached result without creating new entry in the state vector.
+  // Otherwise, allocate a new entry in the state vector.
+  void InitializeResultState(Operation* op, int index, Value value);
+
+  func::FuncOp fn_;
+  OpBuilder builder_;
+  const bool is_signed_;
+  const int bit_width_;
+  const bool disable_per_channel_;
+
+  // We should distinguish weights and bias constants. Biases are specified by
+  // the quantization spec or are the operands of ops with same scale spec. The
+  // rest are weights.
+  DenseSet<Operation*> weights_;
+
+  // The weights require narrow_range quantization. This map collects all the
+  // weight operands defined by the op quant spec. The value of each entry is
+  // the quantization dimension. If it is positive, per-channel quantization is
+  // required.
+  DenseMap<Operation*, int> optimized_weights_;
+
+  // All the ops needs to propagate the quantization parameters to.
+  std::vector<Operation*> work_list_;
+  absl::flat_hash_set<Operation*> quantized_;
+
+  // The vector contains all the quantization parameters propagated from the
+  // defining operations of the value, or from the quantization aware training.
+  std::vector<QuantState> states_;
+
+  // The map contains all the quantization parameters which are required to
+  // satisfy the same operands and results constraint. The keys of this map are
+  // the values from `operand_states_` and `result_state_`.
+  absl::flat_hash_map<QuantStateIndex, RequantizeStates> rescale_states_;
+
+  // Maps of indexes to the propagation state vector from the ops operands,
+  // results and arguments.
+  DenseMap<OpWithOperandIndex, QuantStateIndex> operand_states_;
+  DenseMap<BlockArgument, QuantStateIndex> arg_states_;
+  DenseMap<Value, QuantStateIndex> value_to_state_;
+
+  // This vector is to preserve the arguments order, so the newly inserted
+  // quantized ops for the arguments are deterministically ordered.
+  SmallVector<BlockArgument, 4> args_;
+
+  OpQuantSpecGetter op_quant_spec_getter_;
+  OpQuantScaleSpecGetter op_quant_scale_spec_getter_;
+
+  // Infer output ranges for activation ops and constants. This is usually
+  // required for post-training quantization.
+  const bool infer_tensor_range_;
+
+  // Calculate scales in float instead of double, so that the scales and
+  // quantized values are exactly the same with the TOCO quantizer.
+  const bool legacy_float_scale_;
+
+  // If true, the model is a floating point graph with QDQ ops to be eliminated
+  // and fused into quantized kernels.
+  const bool is_qdq_conversion_;
+};
+
+// Propagates quantization parameters across ops in this function and satisfies
+// the quantization specification of the ops. This methods assumes the initial
+// quantization parameters are stored as adjacent quantize and dequantize ops
+// and the propagation results are materialized by inserting pairs of quantize
+// and dequantize ops to this function. Set `disable_per_channel` to true to not
+// use per channel quantization even the op supports it.
+// Setting `infer_tensor_range` to true, to infer quantization parameters from
+// the activation ops and weight constants. This is only used for post-training
+// quantization.
+void ApplyQuantizationParamsPropagation(func::FuncOp func, bool is_signed,
+                                        int bit_width, bool disable_per_channel,
+                                        OpQuantSpecGetter op_quant_spec_getter,
+                                        bool infer_tensor_ranges,
+                                        bool legacy_float_scale,
+                                        bool is_qdq_conversion);
+
+void ApplyQuantizationParamsPropagation(
+    func::FuncOp func, bool is_signed, int bit_width, bool disable_per_channel,
+    OpQuantSpecGetter op_quant_spec_getter,
+    OpQuantScaleSpecGetter op_quant_scale_spec_getter, bool infer_tensor_ranges,
+    bool legacy_float_scale, bool is_qdq_conversion);
+
+}  // namespace quant
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_QUANTIZATION_LIB_QUANTIZATION_DRIVER_H_

tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_driver_test.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 "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_driver.h"
+
+#include <cstdint>
+#include <functional>
+#include <memory>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/strings/string_view.h"
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.h"  // IWYU pragma: keep
+#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"
+
+namespace mlir::quant {
+namespace {
+
+using ApplyQuantizationParamsPropagationTest =
+    mlir::quant::QuantizationTestBase;
+using ::testing::IsEmpty;
+using ::testing::Not;
+
+constexpr absl::string_view kModuleTFLite = R"mlir(
+  module {
+    func.func @main(%arg0: tensor<1x4x4x3xf32>) -> tensor<1x4x4x3xf32> attributes {_from_xla_call_module} {
+      %cst_0 = arith.constant dense<1.0> : tensor<3x1x1x3xf32>
+      %cst_1 = arith.constant dense<2.0> : tensor<3xf32>
+      %0 = "tf.XlaCallModule"(%arg0, %cst_0, %cst_1) <{Sout = [#tf_type.shape<1x4x4x3>], module = "", version = 9 : i64}> {_entry_function = @composite_fn_1, _stablehlo_version = "1.0.0", _original_entry_function = "composite_fn_1", _tfl_quant_trait = "fully_quantizable"} : (tensor<1x4x4x3xf32>, tensor<3x1x1x3xf32>, tensor<3xf32>) -> tensor<1x4x4x3xf32>
+      %1 = "tf.XlaCallModule"(%0, %cst_0, %cst_1) <{Sout = [#tf_type.shape<1x4x4x3>], module = "", version = 9 : i64}> {_entry_function = @composite_fn_2, _stablehlo_version = "1.0.0", _original_entry_function = "composite_fn_2", _tfl_quant_trait = "fully_quantizable"} : (tensor<1x4x4x3xf32>, tensor<3x1x1x3xf32>, tensor<3xf32>) -> tensor<1x4x4x3xf32>
+      return %1 : tensor<1x4x4x3xf32>
+    }
+    func.func private @composite_fn_1(%arg0: tensor<1x4x4x3xf32>, %arg1: tensor<3x1x1x3xf32>, %arg2: tensor<3xf32>) -> tensor<1x4x4x3xf32> attributes {tf_quant.composite_function} {
+      %perm = "tf.Const"() { value = dense<[1, 2, 3, 0]> : tensor<4xi32> } : () -> tensor<4xi32>
+      %filter_hwio = "tf.Transpose"(%arg1, %perm) : (tensor<3x1x1x3xf32>, tensor<4xi32>) -> tensor<1x1x3x3xf32>
+      %0 = "tf.Conv2D"(%arg0, %filter_hwio) {
+        dilations = [1, 1, 1, 1],
+        strides = [1, 1, 1, 1],
+        padding = "VALID",
+        data_format = "NHWC"
+      } : (tensor<1x4x4x3xf32>, tensor<1x1x3x3xf32>) -> tensor<1x4x4x3xf32>
+      %1 = "tf.BiasAdd"(%0, %arg2) {
+        data_format = "NHWC"
+      } : (tensor<1x4x4x3xf32>, tensor<3xf32>) -> tensor<1x4x4x3xf32>
+      %2 = "tf.Relu"(%1) : (tensor<1x4x4x3xf32>) -> tensor<1x4x4x3xf32>
+      return %2 : tensor<1x4x4x3xf32>
+    }
+    func.func private @composite_fn_2(%arg0: tensor<1x4x4x3xf32>, %arg1: tensor<3x1x1x3xf32>, %arg2: tensor<3xf32>) -> tensor<1x4x4x3xf32> attributes {tf_quant.composite_function} {
+      %perm = "tf.Const"() { value = dense<[1, 2, 3, 0]> : tensor<4xi32> } : () -> tensor<4xi32>
+      %filter_hwio = "tf.Transpose"(%arg1, %perm) : (tensor<3x1x1x3xf32>, tensor<4xi32>) -> tensor<1x1x3x3xf32>
+      %0 = "tf.Conv2D"(%arg0, %filter_hwio) {
+        dilations = [1, 1, 1, 1],
+        strides = [1, 1, 1, 1],
+        padding = "VALID",
+        data_format = "NHWC"
+      } : (tensor<1x4x4x3xf32>, tensor<1x1x3x3xf32>) -> tensor<1x4x4x3xf32>
+      %1 = "tf.BiasAdd"(%0, %arg2) {
+        data_format = "NHWC"
+      } : (tensor<1x4x4x3xf32>, tensor<3xf32>) -> tensor<1x4x4x3xf32>
+      %2 = "tf.Relu"(%1) : (tensor<1x4x4x3xf32>) -> tensor<1x4x4x3xf32>
+      return %2 : tensor<1x4x4x3xf32>
+    }
+  }
+)mlir";
+
+// TOOD: b/323478683 - Directly use types rather than creating a `unique_ptr`.
+std::unique_ptr<OpQuantSpec> GetOpQuantSpec(
+    const mlir::Operation* op,
+    bool disable_per_channel_for_dense_layers = false) {
+  auto spec = std::make_unique<OpQuantSpec>();
+  spec->coeff_op_quant_dim[1] = 3;
+  spec->biases_params[2] = {{0, 1}, GetUniformQuantizedTypeForBias};
+  for (const auto& [key, value] : spec->coeff_op_quant_dim) {
+    spec->quantizable_operands.insert(key);
+  }
+  return spec;
+}
+
+TEST_F(ApplyQuantizationParamsPropagationTest,
+       ConstsUsedMultipleTimesAreDuplicated) {
+  const OwningOpRef<ModuleOp> module_op_ref =
+      mlir::quant::QuantizationTestBase::ParseModuleOpString(kModuleTFLite);
+  func::FuncOp main_fn = mlir::quant::FindMainFuncOp(*module_op_ref);
+
+  auto op_quant_spec_getter = [&](mlir::Operation* op) {
+    return GetOpQuantSpec(op, /*disable_per_channel_for_dense_layers=*/false);
+  };
+  QuantizationDriver quantization_driver(
+      main_fn, /*is_signed=*/true, /*bit_width=*/8,
+      /*disable_per_channel=*/false, op_quant_spec_getter,
+      GetDefaultQuantScaleSpec,
+      /*infer_tensor_range=*/true, /*legacy_float_scale=*/false,
+      /*is_qdq_conversion=*/false);
+
+  quantization_driver.Initialize();
+
+  int64_t num_constant_op = 0;
+  main_fn.walk([&](arith::ConstantOp cst) { ++num_constant_op; });
+  EXPECT_EQ(num_constant_op, 4);
+}
+
+TEST_F(ApplyQuantizationParamsPropagationTest,
+       PropagateParamsCreatesQuantState) {
+  const OwningOpRef<ModuleOp> module_op_ref =
+      ParseModuleOpString(kModuleTFLite);
+  func::FuncOp main_fn = mlir::quant::FindMainFuncOp(*module_op_ref);
+
+  auto op_quant_spec_getter = [&](mlir::Operation* op) {
+    return GetOpQuantSpec(op, /*disable_per_channel_for_dense_layers=*/false);
+  };
+  QuantizationDriver quantization_driver(
+      main_fn, /*is_signed=*/true, /*bit_width=*/8,
+      /*disable_per_channel=*/false, op_quant_spec_getter,
+      GetDefaultQuantScaleSpec,
+      /*infer_tensor_range=*/true, /*legacy_float_scale=*/false,
+      /*is_qdq_conversion=*/false);
+
+  quantization_driver.Initialize();
+  ASSERT_TRUE(quantization_driver.PropagateParamsAndReturnIfChanged());
+  EXPECT_THAT(quantization_driver.GetArgs(), Not(IsEmpty()));
+
+  for (const auto& arg : quantization_driver.GetArgs()) {
+    const QuantState& state = quantization_driver.GetArgQuantState(arg);
+    EXPECT_TRUE(isa<QuantizedType>(state.params));
+  }
+  for (const auto& result : quantization_driver.GetResultStates()) {
+    Operation* op = result.first.first;
+    const int res_index = result.first.second;
+    const QuantState state =
+        quantization_driver.GetResultQuantState(op, res_index);
+    EXPECT_TRUE(isa<QuantizedType>(state.params));
+  }
+}
+
+TEST_F(ApplyQuantizationParamsPropagationTest, FinalizeInsertsQDQOps) {
+  const OwningOpRef<ModuleOp> module_op_ref =
+      ParseModuleOpString(kModuleTFLite);
+  func::FuncOp main_fn = mlir::quant::FindMainFuncOp(*module_op_ref);
+
+  auto op_quant_spec_getter = [&](mlir::Operation* op) {
+    return GetOpQuantSpec(op, /*disable_per_channel_for_dense_layers=*/false);
+  };
+  ApplyQuantizationParamsPropagation(
+      main_fn, /*is_signed=*/true, /*bit_width=*/8,
+      /*disable_per_channel=*/false, op_quant_spec_getter,
+      /*infer_tensor_ranges=*/true, /*legacy_float_scale=*/false,
+      /*is_qdq_conversion=*/false);
+  Operation* xla_call_module_op =
+      mlir::quant::FindOperationOfType<TF::XlaCallModuleOp>(main_fn);
+  Operation* filter_dcast_op =
+      xla_call_module_op->getOperand(1).getDefiningOp();
+  Operation* filter_qcast_op = filter_dcast_op->getOperand(0).getDefiningOp();
+  ASSERT_NE(filter_qcast_op, nullptr);
+  EXPECT_TRUE(isa<mlir::quant::ir::QuantizeCastOp>(filter_qcast_op));
+  EXPECT_TRUE(isa<mlir::quant::ir::DequantizeCastOp>(filter_dcast_op));
+  EXPECT_TRUE(isa<UniformQuantizedPerAxisType>(
+      mlir::cast<TensorType>(filter_qcast_op->getResult(0).getType())
+          .getElementType()));
+}
+
+}  // namespace
+}  // namespace mlir::quant

tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_traits.h

@@ -0,0 +1,152 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This file defines the op traits used in the MLIR TensorFlow Lite dialect.
+
+#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_QUANTIZATION_LIB_QUANTIZATION_TRAITS_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_QUANTIZATION_LIB_QUANTIZATION_TRAITS_H_
+
+#include <cmath>
+#include <cstdint>
+#include <vector>
+
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+
+using QuantizedType = mlir::quant::QuantizedType;
+using UniformQuantizedType = mlir::quant::UniformQuantizedType;
+
+namespace mlir {
+namespace quant {
+// Verifies that the op satisfies the same operands and results scales
+// constraints. Note that this constraint can only be applied on some
+// storage types of the op.
+LogicalResult VerifySameScales(Operation* op);
+}  // namespace quant
+
+// This includes the interface class definition. It couldn't be in a namespace
+// because the table gen doesn't emit the namespace when it is used.
+#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_interface.h.inc"
+
+namespace OpTrait {
+namespace quant {
+
+// The base class that all the quantization related OpTrait implements.
+template <typename ConcreteType, template <typename> class TraitType>
+struct QuantizationSpecTraitBase : public TraitBase<ConcreteType, TraitType> {
+  static bool IsBias(int index) { return false; }
+  static bool IsQuantizable() { return true; }
+};
+
+// This class provides the API for ops that has a fixed output value range.
+// This is used as a trait like this:
+//
+//   class SoftmaxOp
+//       : public Op<SoftmaxOp,
+//           OpTrait::quant::FixedResultUniformScale<
+//               8, -128, 390625, -8, 0, 255, false>::Impl> {
+//
+// TODO(fengliuai): create a better way to express floating point scale in the
+// template argument list.
+template <unsigned BitWidth, int ZeroPoint, int ScaleMantissa, int ScaleExp,
+          int64_t StorageTypeMin, int64_t StorageTypeMax, bool Sign>
+class FixedResultUniformScale {
+ public:
+  template <typename ConcreteType>
+  class Impl
+      : public QuantizationSpecTraitBase<
+            ConcreteType, FixedResultUniformScale<
+                              BitWidth, ZeroPoint, ScaleMantissa, ScaleExp,
+                              StorageTypeMin, StorageTypeMax, Sign>::Impl> {
+   public:
+    QuantizedType GetResultQuantizedType(int index) {
+      auto op = this->getOperation();
+      const auto result_type =
+          op->getResult(index).getType().template cast<ShapedType>();
+      if (!result_type.getElementType().template isa<FloatType>()) return {};
+      Builder builder(op->getContext());
+      const IntegerType storage_type = builder.getIntegerType(BitWidth);
+      const double scale = static_cast<double>(ScaleMantissa) *
+                           std::pow(10.0, static_cast<double>(ScaleExp));
+      return UniformQuantizedType::getChecked(
+          Sign, storage_type, result_type.getElementType(), scale, ZeroPoint,
+          StorageTypeMin, StorageTypeMax, builder.getUnknownLoc());
+    }
+  };
+};
+
+// This class provides the API for ops that has input as bias. This is used
+// as a trait like this:
+//
+//   class Conv2DOp
+//       : public Op<Conv2DOp, OpTrait::quant::AccumulatorScale<2, 0, 1>::Impl>
+//
+// TODO(fengliuai): supports a configurable accumulator bit width.
+template <int Bias, int... Operands>
+class AccumulatorUniformScale {
+ public:
+  template <typename ConcreteType>
+  class Impl
+      : public QuantizationSpecTraitBase<
+            ConcreteType, AccumulatorUniformScale<Bias, Operands...>::Impl> {
+   public:
+    // Whether the index-th operand is a bias.
+    static bool IsBias(int index) { return index == Bias; }
+
+    // Returns the indexes of all the non-bias operands.
+    static std::vector<int> GetAllNonBiasOperands() {
+      return std::vector<int>({Operands...});
+    }
+  };
+};
+
+// The trait to specify the operand index of the coefficient for an affine op
+// and also the quantization dimension if per-axis quantization is support.
+// If the quantization dimension is -1, per-axis quantization isn't supported.
+//
+//   class Conv2DOp
+//       : public Op<Conv2DOp, OpTrait::quant::AffineOpCoefficient<0>::Impl>
+//
+template <int QuantDim, int OperandIndex = 1>
+class AffineOpCoefficient {
+ public:
+  template <typename ConcreteType>
+  class Impl
+      : public TraitBase<ConcreteType,
+                         AffineOpCoefficient<QuantDim, OperandIndex>::Impl> {
+   public:
+    static int GetCoefficientOperandIndex() { return OperandIndex; }
+    static int GetQuantizationDim() { return QuantDim; }
+  };
+};
+
+// This class provides the API for ops that can be quantized.
+// This is as a trait like this:
+//
+//   class LessOp : public Op<LessOp, OpTrait::quant::QuantizableResult> {
+//
+template <typename ConcreteType>
+class QuantizableResult
+    : public QuantizationSpecTraitBase<ConcreteType, QuantizableResult> {};
+
+}  // namespace quant
+}  // namespace OpTrait
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_QUANTIZATION_LIB_QUANTIZATION_TRAITS_H_

tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h

@@ -0,0 +1,972 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This header file defines common utils used by TFLite transformation
+// passes to work with op attributes.
+
+#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_QUANTIZATION_LIB_QUANTIZATION_UTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_QUANTIZATION_LIB_QUANTIZATION_UTILS_H_
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <functional>
+#include <memory>
+#include <string>
+#include <type_traits>
+#include <utility>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/container/flat_hash_set.h"
+#include "absl/strings/string_view.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/IRMapping.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Matchers.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/quantization/common/ir/FakeQuantSupport.h"
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"
+#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_config.h"
+#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_traits.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace mlir {
+namespace quant {
+
+// A unit attribute can be attached to the quantize/dequantize ops which are
+// added by the quantization passes. These ops can be removed erased without
+// losing accuracy.
+inline constexpr char kVolatileOpAttrName[] = "volatile";
+
+// Following attributes are used to mark ops that are not quantizable during
+// debug model generation process for whole-model verify mode. If these
+// attributes are attached, the upstream float/quantized ops know which ops to
+// connect to, and it also prevents these ops from being copied again.
+inline constexpr char kDebugModeOpFloatAttrName[] = "debug_float";
+inline constexpr char kDebugModeOpQuantAttrName[] = "debug_quant";
+
+// Used to annotate custom ops if they are quantizable.
+inline constexpr char kQuantTraitAttrName[] = "_tfl_quant_trait";
+enum QuantizationTrait { FullyQuantizable = 0, NotQuantizable = 1 };
+inline constexpr absl::string_view QuantTraitValues[] = {"fully_quantizable",
+                                                         "not_quantizable"};
+inline constexpr char kOutputQuantized[] = "_output_quantized";
+
+inline constexpr double kNearZeroTolerance = 1.0e-6;
+
+using QuantParams = QuantizedType;
+using QuantSpec = QuantizationSpecs;
+using SignedInteger = std::pair<unsigned, unsigned>;  // bitwidth and sign
+using QuantParamsForResults = llvm::SmallVector<QuantizedType, 4>;
+using AccumulatorScaleFunc =
+    std::function<QuantizedType(const std::vector<QuantizedType>&, int, bool)>;
+using BiasParamsMap =
+    absl::flat_hash_map<int, std::pair<std::vector<int>, AccumulatorScaleFunc>>;
+// UniformQuantizedType GetFixedOutputRange(bool sign, int bit_width)
+using GetFixedOutputRangeFunc = std::function<UniformQuantizedType(bool, int)>;
+// bool RequiredSameOperandsAndResultsScale(bool sign, int $bit_width)
+using RequiredSameOperandsAndResultsScaleFunc = std::function<bool(bool, int)>;
+// bool RequiredSameQuantizedAxes()
+using RequiredSameQuantizedAxesFunc = std::function<bool()>;
+
+using CustomMap = CustomOpMap;
+
+// Quantization spec of an op, driving the quantization algorithm.
+struct OpQuantSpec {
+  // Maps the operand index of a bias input to its quantization specifications,
+  // including the non-bias operand indexes and the method retrieving
+  // quantization parameters from list of parameters of the non-bias operands.
+  // This map is empty if the op doesn't have a bias operand.
+  BiasParamsMap biases_params;
+
+  // Quantization parameters for value restricted outputs. This is the
+  // "hard-coded" parameters and should be used unconditionally for the
+  // quantized op. This vector is empty if the op doesn't have value restricted
+  // outputs.
+  llvm::DenseMap<SignedInteger, QuantParamsForResults> restricted_output_params;
+
+  // Coefficient operand index and whether supporting per-channel quantization.
+  // For QAT, this information is carried by the FakeQuant*/Quantize/Dequantize
+  // ops, but post-training quantization, the quantization parameters need to be
+  // inferred from the tensor content and op property. A "-1" value indicates
+  // the operand doesn't support per-channel quantization.
+  llvm::DenseMap<int, int> coeff_op_quant_dim;
+
+  // Indices of quantizable operands. Biases are not included in this field,
+  // the indices of biases can be found in the `biases_params`.
+  absl::flat_hash_set<int> quantizable_operands;
+};
+
+// A function signature for getting the particular OpQuantSpec for the provided
+// op.
+using OpQuantSpecGetter =
+    std::function<std::unique_ptr<OpQuantSpec>(Operation*)>;
+
+// Quantization scale spec of an op. The information defined in the MLIR
+// interfaces FixedOutputRangeInterface and SameOperandsAndResultsScale should
+// be checked first if present.
+// TODO: b/323478683: Consider deprecating this.
+struct OpQuantScaleSpec {
+  // Whether this op has a fixed range requirement (e.g. sigmoid)
+  bool has_fixed_output_range = false;
+  // Whether this op should have same operand and result scales (e.g. concat)
+  bool has_same_scale_requirement = false;
+  // Whether this op should have same operand and result type (e.g. gather)
+  bool has_same_operand_and_result_type_requirement = false;
+  // Returns the fixed output range, when has_fixed_output_range is set.
+  GetFixedOutputRangeFunc fixed_output_range_func;
+  // Returns whether same operands and results scales are required.
+  RequiredSameOperandsAndResultsScaleFunc required_same_scale_func =
+      [](bool sign, int bit_width) { return true; };
+  // Returns whether operands and results must have the same quantized axis.
+  RequiredSameQuantizedAxesFunc required_same_quantized_axes_func = []() {
+    return true;
+  };
+};
+
+// A function signature for getting the particular OpQuantScaleSpec for the
+// provided op.
+using OpQuantScaleSpecGetter =
+    std::function<std::unique_ptr<OpQuantScaleSpec>(Operation*)>;
+
+// Used in TFL Numeric Verify
+struct NumericVerifySpec {
+  // Whether to enable numeric verification
+  bool verify_numeric = false;
+
+  // Tolerance level from the quantized value for verification. If the tolerance
+  // is very small(<0.1), only the stats of the diff is displayed.
+  float error_tolerance = 5.0f;
+
+  // Whether to verify numerical correctness layer by layer or by whole model
+  bool whole_model_verify = false;
+
+  // Whether to enable log for failures
+  bool log_if_failed_flag = false;
+};
+
+// Used in TFL Quantize Pass
+struct QuantPassSpec {
+  // Variables to control TFL Numeric Verify
+  NumericVerifySpec numeric_verify_spec;
+
+  // Variables related to quantization
+  QuantSpec quant_spec;
+};
+
+// Re-calculates scales again in float instead of simply downcasting existing
+// scales.
+quant::QuantizedType DownCastScale(quant::QuantizedType type,
+                                   const SmallVectorImpl<double>& mins,
+                                   const SmallVectorImpl<double>& maxs,
+                                   Location loc);
+
+quant::QuantizedType DownCastScale(quant::QuantizedType type, double min,
+                                   double max, Location loc);
+
+bool IsOpQuantizable(Operation* op);
+bool QuantizableOpSupportsFloatOutputType(Operation* op);
+
+// Specialized version of location to string for flatbuffer exported locations.
+inline std::string GetTensorNameFromLoc(Location loc) {
+  if (auto name_loc = llvm::dyn_cast<NameLoc>(loc)) {
+    return name_loc.getName().str();
+  }
+  return "";
+}
+
+template <typename QuantizeOpT, typename DequantizeOpT>
+struct ConvertStatsToQDQs
+    : public OpRewritePattern<mlir::quant::ir::StatisticsOp> {
+  ConvertStatsToQDQs(int num_bits, bool narrow_range, bool is_signed,
+                     bool legacy_float_scale, MLIRContext* context)
+      : OpRewritePattern<mlir::quant::ir::StatisticsOp>(context),
+        num_bits(num_bits),
+        narrow_range(narrow_range),
+        is_signed(is_signed),
+        legacy_float_scale(legacy_float_scale) {}
+
+  LogicalResult matchAndRewrite(mlir::quant::ir::StatisticsOp op,
+                                PatternRewriter& rewriter) const override {
+    Type expressed = llvm::cast<ShapedType>(op.getType()).getElementType();
+    quant::QuantizedType quant_type;
+    SmallVector<double, 4> mins, maxs;
+
+    if (op.getAxisStats().has_value()) {
+      // Per axis quantization (or per channel quantization)
+      int stats_num = op.getAxisStats()->getNumElements();
+      if (stats_num == 0 || stats_num % 2 != 0) return failure();
+      auto stats =
+          llvm::dyn_cast<DenseFPElementsAttr>(op.getAxisStats().value());
+      if (!stats) return failure();
+
+      for (auto it = stats.begin(), e = stats.end(); it != e; ++it) {
+        double rmin = FloatAttr::getValueAsDouble(*it++);
+        double rmax = FloatAttr::getValueAsDouble(*it);
+        // The default nudging implementation of mlir quant library might cause
+        // clamping during inference if the calibration range isn't wide enough.
+        // So here we adjust the range to include 0.0.
+        rmin = std::min(rmin, 0.0);
+        rmax = std::max(rmax, 0.0);
+        if (num_bits == 16) {
+          // TODO: b/266536261 - Since the kernel implementation assumes that
+          // 16x8 integer quantization is symmetric, this MLIR quantizer
+          // supports only symmetric quantization.
+          rmax = std::max(std::abs(rmin), std::abs(rmax));
+          rmin = -rmax;
+        }
+        TensorRangeSanityCheck(op, rmin, rmax);
+        mins.push_back(rmin);
+        maxs.push_back(rmax);
+      }
+      quant_type = quantfork::fakeQuantAttrsToType(
+          op.getLoc(), num_bits, *op.getAxis(), mins, maxs, narrow_range,
+          expressed, is_signed);
+      if (legacy_float_scale) {
+        quant_type = DownCastScale(quant_type, mins, maxs, op->getLoc());
+      }
+    } else if (auto stats =
+                   llvm::dyn_cast<DenseFPElementsAttr>(op.getLayerStats())) {
+      // Per tensor quantization
+      auto statValues = stats.getValues<APFloat>();
+      double rmin = FloatAttr::getValueAsDouble(statValues[0]);
+      double rmax = FloatAttr::getValueAsDouble(statValues[1]);
+      // The default nudging implementation of mlir quant library might cause
+      // clamping during inference if the calibration range isn't wide enough.
+      // So here we adjust the range to include 0.0.
+      rmin = std::min(rmin, 0.0);
+      rmax = std::max(rmax, 0.0);
+      if (num_bits == 16) {
+        // TODO: b/266536261 - Since the kernel implementation assumes that
+        // 16x8 integer quantization is symmetric, this MLIR quantizer supports
+        // only symmetric quantization.
+        rmax = std::max(std::abs(rmin), std::abs(rmax));
+        rmin = -rmax;
+      }
+      TensorRangeSanityCheck(op, rmin, rmax);
+      quant_type =
+          quantfork::fakeQuantAttrsToType(op.getLoc(), num_bits, rmin, rmax,
+                                          narrow_range, expressed, is_signed);
+      if (legacy_float_scale) {
+        quant_type = DownCastScale(quant_type, rmin, rmax, op->getLoc());
+      }
+    } else {
+      return failure();
+    }
+
+    rewriter.setInsertionPointAfter(op.getOperation());
+    Type result_type = quant_type.castFromExpressedType(op.getType());
+    auto q =
+        QuantizeOpT::create(rewriter, op.getLoc(), result_type, op.getArg());
+    q->setAttr(kVolatileOpAttrName, rewriter.getUnitAttr());
+
+    auto dq = DequantizeOpT::create(rewriter, op.getLoc(), op.getType(), q);
+    op.getResult().replaceAllUsesWith(dq);
+    q.getOperation()->replaceUsesOfWith(dq, op.getArg());
+    op.erase();
+
+    return success();
+  }
+
+ private:
+  int num_bits;
+  bool narrow_range;
+  bool is_signed;
+  bool legacy_float_scale;
+
+  // Emits an op warning message if the calibrated range is larger than 10.0 and
+  // the storage type is less than or equal to 8 bits.
+  void TensorRangeSanityCheck(mlir::quant::ir::StatisticsOp op, double& min,
+                              double& max) const {
+    double range = std::fabs(max - min);
+    if (num_bits <= 8 && range >= 10.0) {
+      op.emitWarning()
+          << "Tensor range is too wide to be quantized. Use tf.clip_by_value "
+             "or tf.relu6 to narrow the tensor range. Range: "
+          << range << ", bit width: " << num_bits;
+    }
+    if (std::abs(max - min) < kNearZeroTolerance) {
+      op.emitWarning() << "Tensor range (" << min << ", " << max
+                       << ") is too narrow and it might cause overflow. "
+                          "Expanding range symmetrically by "
+                       << kNearZeroTolerance;
+      min -= kNearZeroTolerance;
+      max += kNearZeroTolerance;
+    }
+  }
+};
+
+template <typename VerifierT>
+bool UsedBy(Operation* op) {
+  for (Operation* user : op->getUsers()) {
+    if (llvm::isa_and_nonnull<VerifierT>(user)) return true;
+  }
+  return false;
+}
+
+template <typename VerifierT>
+void CreateVerifier(Operation* quantizing_op, Operation* quantized_op,
+                    PatternRewriter& rewriter, int result_idx,
+                    const QuantPassSpec& quant_params) {
+  rewriter.setInsertionPointAfter(quantized_op);
+  FloatAttr tolerance = rewriter.getF32FloatAttr(
+      quant_params.numeric_verify_spec.error_tolerance);
+  BoolAttr log =
+      rewriter.getBoolAttr(quant_params.numeric_verify_spec.log_if_failed_flag);
+  // Verify the quantized value by sending the result to the verifier.
+  VerifierT::create(rewriter, quantizing_op->getLoc(),
+                    quantized_op->getResult(result_idx).getType(),
+                    quantized_op->getResult(result_idx),
+                    quantizing_op->getResult(result_idx), tolerance, log);
+}
+
+template <>
+inline bool UsedBy<void>(Operation* op) {
+  return false;
+}
+
+// This specialization is not going to be called, but needed for compilation.
+template <>
+inline void CreateVerifier<void>(Operation* quantizing_op,
+                                 Operation* quantized_op,
+                                 PatternRewriter& rewriter, int result_idx,
+                                 const QuantPassSpec& quant_params) {}
+
+// A base rewrite pattern which matches any N-in-M-out operations with
+// quantization parameters propagated to at least one of its operands. The
+// quantization parameters are annotated by the QuantizeOp/DequantizeOp pairs.
+// Each matched pattern are rewritten by its quantized alternatives.
+//
+// The concrete pattern, extends from this base pattern, can specify whether it
+// allows dynamic range quantized operands and results for the operations in the
+// current context. These "DynamicRangeQuantized" operands and results don't
+// have quantization parameters propagated to, so will be in float in the
+// quantized results. The concrete pattern should define the following two
+// functions:
+//
+//   bool AllowDynamicRangeQuantizedOperand(Operation *) const
+//   bool AllowDynamicRangeQuantizedResult(Operation *) const
+//
+// Full integer quantization disallows "DynamicRangeQuantized" operands or
+// results. Dynamic range quantization allows "DynamicRangeQuantized" operands
+// and results.
+template <typename ConcreteT, typename QuantizeOpT, typename DequantizeOpT,
+          typename VerifierT, typename RootOpT = DequantizeOpT>
+class QuantizationPattern : public RewritePattern {
+ public:
+  using BaseType = QuantizationPattern<ConcreteT, QuantizeOpT, DequantizeOpT,
+                                       VerifierT, RootOpT>;
+
+  explicit QuantizationPattern(MLIRContext* context,
+                               const QuantPassSpec& quant_params)
+      // Set the score to a large number so it is always preferred.
+      : RewritePattern(RootOpT::getOperationName(), 300, context),
+        quant_params_(quant_params) {}
+
+  LogicalResult matchAndRewrite(Operation* op,
+                                PatternRewriter& rewriter) const override {
+    llvm::SmallVector<Operation*, 4> quantizing_ops;
+
+    // Collect all the ops to quantize, as the user / producer of the root op.
+    if constexpr (std::is_same_v<RootOpT, DequantizeOpT>) {
+      if (op->getNumResults() != 1) {
+        return failure();
+      }
+      auto users = op->getResult(0).getUsers();
+      quantizing_ops.append(users.begin(), users.end());
+    } else if constexpr (std::is_same_v<RootOpT, QuantizeOpT>) {
+      if (op->getNumOperands() != 1) {
+        return failure();
+      }
+      Value quantize_operand = op->getOperand(0);
+      if (QuantizedType::getQuantizedElementType(quantize_operand.getType())) {
+        // The input of this QuantizeOp has already been quantized, i.e.
+        // rescale.
+        return failure();
+      }
+      DenseFPElementsAttr attr;
+      if (matchPattern(quantize_operand, m_Constant(&attr))) {
+        // Const-> QuantizeOp pattern will be handled separately.
+        return failure();
+      }
+      if (Operation* quantizing_op = quantize_operand.getDefiningOp()) {
+        quantizing_ops.push_back(quantizing_op);
+      }
+    }
+
+    tensorflow::DataType inference_type =
+        quant_params_.quant_spec.inference_type;
+    bool weight_only_quantization =
+        quant_params_.quant_spec.weight_only_quantization;
+    bool enable_verify = quant_params_.numeric_verify_spec.verify_numeric;
+    bool enable_whole_model_verify =
+        quant_params_.numeric_verify_spec.whole_model_verify;
+    absl::flat_hash_set<std::string> ops_blocklist =
+        quant_params_.quant_spec.ops_blocklist;
+    absl::flat_hash_set<std::string> nodes_blocklist =
+        quant_params_.quant_spec.nodes_blocklist;
+    CustomMap custom_map = quant_params_.quant_spec.custom_map;
+
+    // Rewrite the floating-point ops to the quantized version, by fusing
+    // preceding dequantize ops and succeding quantize ops.
+    for (Operation* quantizing_op : quantizing_ops) {
+      // If it is requantize op, we shouldn't rewrite this op.
+      if (llvm::isa<QuantizeOpT, DequantizeOpT>(quantizing_op)) {
+        return failure();
+      }
+
+      // If the op is terminator, not quantizable or any ops from the mlir quant
+      // ops dialect, we shouldn't rewrite. In case of whole-model verify debug
+      // mode, not-quantizable ops should be duplicated to keep parallel
+      // float/quant model execution.
+      if (quantizing_op->hasTrait<OpTrait::IsTerminator>()) {
+        return failure();
+      }
+
+      if (!IsOpQuantizable(quantizing_op) &&
+          !static_cast<const ConcreteT*>(this)->IsQuantizableCustomOp(
+              quantizing_op, custom_map)) {
+        if (!(enable_verify && enable_whole_model_verify)) {
+          return failure();
+        }
+        if (quantizing_op->hasAttr(kDebugModeOpQuantAttrName) ||
+            quantizing_op->hasAttr(kDebugModeOpFloatAttrName)) {
+          return failure();
+        }
+
+        rewriter.setInsertionPoint(quantizing_op);
+        Operation* float_op = rewriter.clone(*quantizing_op);
+        quantizing_op->setAttr(kDebugModeOpQuantAttrName,
+                               rewriter.getUnitAttr());
+        float_op->setAttr(kDebugModeOpFloatAttrName, rewriter.getUnitAttr());
+        RewireFloatModelBackbone(quantizing_op, float_op);
+        return success();
+      }
+
+      // Blocklist op is checked in advance for non-dynamic range quantization
+      // case.
+      if (!quant_params_.quant_spec.weight_quantization &&
+          (ops_blocklist.find(quantizing_op->getName().getStringRef().str()) !=
+           ops_blocklist.end())) {
+        return failure();
+      }
+
+      if (!nodes_blocklist.empty()) {
+        if (auto name_loc = llvm::dyn_cast<NameLoc>(quantizing_op->getLoc())) {
+          std::string sloc = name_loc.getName().str();
+          if (!sloc.empty() &&
+              (nodes_blocklist.find(sloc) != nodes_blocklist.end())) {
+            return failure();
+          }
+        }
+      }
+
+      // An op with float inputs and outputs are expected when it's used by a
+      // NumericVerify op. Skip this op.
+      if (enable_verify && UsedBy<VerifierT>(quantizing_op)) {
+        continue;
+      }
+
+      bool is_operand_or_result_modified = false;
+      // Collect all the quantized inputs and "clone" the matched op by these
+      // inputs.
+      SmallVector<Value, 4> inputs;
+      inputs.reserve(quantizing_op->getNumOperands());
+      for (auto operand : quantizing_op->getOperands()) {
+        Type operand_type = operand.getType();
+        if (isa<NoneType>(operand_type)) {
+          inputs.push_back(operand);
+          continue;
+        }
+
+        auto ele_type =
+            llvm::cast<TensorType>(operand.getType()).getElementType();
+        if (static_cast<const ConcreteT*>(this)
+                ->AllowDynamicRangeQuantizedOperand(quantizing_op,
+                                                    custom_map)) {
+          auto dq_op = dyn_cast_or_null<DequantizeOpT>(operand.getDefiningOp());
+
+          if (dq_op && inference_type == tensorflow::DT_QINT8 &&
+              !static_cast<const ConcreteT*>(this)->IsWeightOnlyOp(
+                  quantizing_op, ops_blocklist, weight_only_quantization,
+                  custom_map)) {
+            // Dynamic range quantization is applied by having QuantizeOp as an
+            // input. Only int8 weight is supported for now.
+            inputs.push_back(dq_op.getOperand());
+            is_operand_or_result_modified = true;
+          } else {
+            // Otherwise, it's the case where the operand is activations or the
+            // quantizing_op is non-supported/weight-only.
+            inputs.push_back(operand);
+          }
+        } else {
+          if (auto dq_op =
+                  dyn_cast_or_null<DequantizeOpT>(operand.getDefiningOp())) {
+            is_operand_or_result_modified = true;
+            inputs.push_back(dq_op.getOperand());
+          } else if (!ele_type.isF32()) {
+            // If the operand is an integer tensor, then it doesn't require the
+            // DequantizeOp in the pattern.
+            inputs.push_back(operand);
+          } else {
+            return failure();
+          }
+        }
+      }
+
+      Operation* quantized_op;
+      if (QuantizableOpSupportsFloatOutputType(quantizing_op)) {
+        rewriter.setInsertionPointAfter(quantizing_op);
+        OperationState new_state(
+            quantizing_op->getLoc(), quantizing_op->getName().getStringRef(),
+            inputs, quantizing_op->getResultTypes(), quantizing_op->getAttrs());
+        for (const auto& indexed_regions :
+             llvm::enumerate(quantizing_op->getRegions())) {
+          Region* target_region = new_state.addRegion();
+          IRMapping mapping;
+          indexed_regions.value().cloneInto(target_region, mapping);
+        }
+        quantized_op = rewriter.create(new_state);
+        rewriter.replaceOp(quantizing_op, quantized_op);
+      } else {
+        // Collect all the quantized outputs and replace them by the results of
+        // the new quantized op.
+        llvm::SmallDenseMap<Value, int> outputs_replaced;
+        SmallVector<Type, 4> output_types;
+        output_types.reserve(quantizing_op->getNumResults());
+        for (const auto& enumerated_result :
+             llvm::enumerate(quantizing_op->getResults())) {
+          Value result = enumerated_result.value();
+          Type result_type = result.getType();
+          // Add this to the test coverage once we create test ops with none
+          // type results.
+          if (isa<NoneType>(result_type)) {
+            outputs_replaced.insert({result, enumerated_result.index()});
+            output_types.push_back(result_type);
+            continue;
+          }
+          Type result_ele_type =
+              llvm::cast<TensorType>(result.getType()).getElementType();
+          // If the user is the QuantizeOp, it must be the only user.
+          if (result.hasOneUse() &&
+              llvm::isa<QuantizeOpT>(*result.user_begin())) {
+            auto user = llvm::cast<QuantizeOpT>(*result.user_begin());
+            outputs_replaced.insert(
+                {user.getResult(), enumerated_result.index()});
+            output_types.push_back(user.getType());
+            is_operand_or_result_modified = true;
+          } else if (!result_ele_type.isF32()) {
+            // If the result is an integer tensor, then it doesn't require the
+            // D op in the pattern.
+            outputs_replaced.insert({result, enumerated_result.index()});
+            output_types.push_back(result.getType());
+          } else if (static_cast<const ConcreteT*>(this)
+                         ->AllowDynamicRangeQuantizedResult(quantizing_op,
+                                                            custom_map)) {
+            outputs_replaced.insert({result, enumerated_result.index()});
+            output_types.push_back(result.getType());
+          } else {
+            return failure();
+          }
+        }
+
+        // For float16 quantization if none of the operand or result is
+        // modified, replacing the op. See b/335025403.
+        if (inference_type == tensorflow::DT_HALF &&
+            !is_operand_or_result_modified) {
+          return failure();
+        }
+
+        rewriter.setInsertionPointAfter(quantizing_op);
+        OperationState new_state(
+            quantizing_op->getLoc(), quantizing_op->getName().getStringRef(),
+            inputs, output_types, quantizing_op->getAttrs());
+        for (int i = 0; i < quantizing_op->getNumRegions(); ++i) {
+          new_state.addRegion();
+        }
+        quantized_op = rewriter.create(new_state);
+        if (quantizing_op->getNumRegions() != 0) {
+          for (const auto& indexed_regions :
+               llvm::enumerate(quantizing_op->getRegions())) {
+            Region& target_region =
+                quantized_op->getRegion(indexed_regions.index());
+            IRMapping mapping;
+            indexed_regions.value().cloneInto(&target_region, mapping);
+          }
+        }
+        for (auto output : outputs_replaced) {
+          output.getFirst().replaceAllUsesWith(
+              quantized_op->getResult(output.getSecond()));
+        }
+      }
+
+      // To verify the numericals, the original floating-point ops are
+      // preserved in the graph. The result of these floating-point ops are sent
+      // to a numeric verifier op as the reference.
+      if (enable_verify && !std::is_same_v<VerifierT, void>) {
+        // For constant operands, the floating-point constant is duplicated in
+        // case it is quantized.
+        for (int i = 0, e = quantized_op->getNumOperands(); i < e; ++i) {
+          auto def = quantized_op->getOperand(i).getDefiningOp();
+          if (auto q = llvm::dyn_cast_or_null<QuantizeOpT>(def)) {
+            DenseFPElementsAttr attr;
+            if (!matchPattern(q.getOperand(), m_Constant(&attr))) {
+              continue;
+            }
+            auto cst = arith::ConstantOp::create(rewriter,
+                                                 quantized_op->getLoc(), attr);
+            quantizing_op->setOperand(i, cst.getResult());
+          }
+        }
+
+        for (int i = 0, e = quantized_op->getNumResults(); i < e; ++i) {
+          if (!isa<FloatType>(
+                  cast<ShapedType>(quantizing_op->getResult(i).getType())
+                      .getElementType())) {
+            continue;
+          }
+          CreateVerifier<VerifierT>(quantizing_op, quantized_op, rewriter, i,
+                                    quant_params_);
+
+          if (enable_whole_model_verify) {
+            RewireFloatModelBackbone(quantized_op, quantizing_op);
+          }
+        }
+      }
+    }
+    return success();
+  }
+
+ private:
+  // Reconnects float ops in the whole-model verify mode. Works for both
+  // Quantizable ops and Unquantizable ops
+  void RewireFloatModelBackbone(Operation* quantized_op,
+                                Operation* float_op) const {
+    for (int i = 0, e = quantized_op->getNumResults(); i < e; ++i) {
+      if (!llvm::cast<ShapedType>(float_op->getResult(i).getType())
+               .getElementType()
+               .isF32()) {
+        continue;
+      }
+      // Find the Quantize/Dequantize users of the new op results, and replace
+      // the usage. Then all the floating-point ops are connected, forming a
+      // separate float "backbone" model that the quantized model can be
+      // compared against in parallel.
+      // N.B. the return op will use this floating-point result.
+      Value result;
+      if (!IsOpQuantizable(float_op)) {
+        // For not quantizable ops, search for dequantize attached to the
+        // quantized op of the output.
+        if (Operation* quantize_op = dyn_cast_or_null<QuantizeOpT>(
+                *quantized_op->getResult(i).getUsers().begin())) {
+          result = quantize_op->getResult(0);
+        } else {
+          quantized_op->emitError()
+              << "Output[" << i
+              << "] is expected to have only one user [QUANTIZE]";
+          return;
+        }
+      } else {
+        result = quantized_op->getResult(i);
+      }
+      for (auto user : result.getUsers()) {
+        // Skip the Requantize op and set the user to the following dequantize
+        // op. This happens when the quantizer tries to match the scale conflict
+        // with QuantizeOp - QuantizeOp(requant) - DequantizeOp triples. The
+        // correct float op should be the user of the last DequantizeOp.
+        if (llvm::isa<QuantizeOpT>(user)) {
+          user = *user->getResult(0).getUsers().begin();
+        }
+        if (auto dequantize = llvm::dyn_cast<DequantizeOpT>(user)) {
+          // Replace all uses, except not quantizable ops that are being used in
+          // the float backbone.
+          dequantize.getResult().replaceUsesWithIf(
+              float_op->getResult(i), [&](OpOperand& use) {
+                return !use.getOwner()->hasAttr(kDebugModeOpQuantAttrName);
+              });
+        }
+      }
+    }
+  }
+
+  QuantPassSpec quant_params_;
+};
+
+// A pattern that removes debug attributes that are annotated to ops during
+// the debug model creation.
+class RemoveDebugAttrPattern : public RewritePattern {
+ public:
+  explicit RemoveDebugAttrPattern(MLIRContext* context)
+      : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}
+  LogicalResult matchAndRewrite(Operation* op,
+                                PatternRewriter& rewriter) const override;
+};
+
+// Converts quantized tensor type with signed integer type to quantized tensor
+// type with unsigned integer type.
+Type ConvertSignedQuantizedToUnsigned(Type signed_tensor_type, Location loc);
+
+// Converts quantize ops with unsigned quantized types to these with signed
+// quantized types and preserves the scales.
+template <typename QuantizeOpT>
+struct ConvertUnsignedToSigned : public OpRewritePattern<QuantizeOpT> {
+  using BaseType = ConvertUnsignedToSigned<QuantizeOpT>;
+  using QType = quant::QuantizedType;
+
+  explicit ConvertUnsignedToSigned(MLIRContext* context)
+      : OpRewritePattern<QuantizeOpT>(context, 1) {}
+
+  LogicalResult matchAndRewrite(QuantizeOpT op,
+                                PatternRewriter& rewriter) const override {
+    Type output_type = op.getResult().getType();
+    auto qtype = QType::getQuantizedElementType(output_type);
+    if (!qtype || qtype.isSigned()) return failure();
+
+    int num_bits = qtype.getStorageTypeIntegralWidth();
+    if (num_bits == 8) {
+      // If storage is 8-bit, trained num bits may be less than 8 so check here.
+      num_bits =
+          static_cast<int>(std::ceil(std::log2(qtype.getStorageTypeMax())));
+    }
+    // This is a positive value, and will be applied on zero points and fixed
+    // point ranges.
+    int64_t offset =
+        QType::getDefaultMinimumForInteger(/*isSigned=*/false, num_bits) -
+        QType::getDefaultMinimumForInteger(/*isSigned=*/true, num_bits);
+
+    auto flags = quant::QuantizationFlags::Signed;
+    QType new_qtype;
+    if (auto uqtype = llvm::dyn_cast<quant::UniformQuantizedType>(qtype)) {
+      new_qtype = quant::UniformQuantizedType::getChecked(
+          op.getLoc(), flags, qtype.getStorageType(), qtype.getExpressedType(),
+          uqtype.getScale(), uqtype.getZeroPoint() - offset,
+          uqtype.getStorageTypeMin() - offset,
+          uqtype.getStorageTypeMax() - offset);
+    } else if (auto aqtype =
+                   llvm::dyn_cast<quant::UniformQuantizedPerAxisType>(qtype)) {
+      auto zero_points = aqtype.getZeroPoints();
+      llvm::SmallVector<int64_t, 4> new_zero_points(zero_points.begin(),
+                                                    zero_points.end());
+      for (int i = 0, e = new_zero_points.size(); i < e; ++i) {
+        new_zero_points[i] -= offset;
+      }
+      new_qtype = quant::UniformQuantizedPerAxisType::getChecked(
+          op.getLoc(), flags, qtype.getStorageType(), qtype.getExpressedType(),
+          aqtype.getScales(), new_zero_points, aqtype.getQuantizedDimension(),
+          aqtype.getStorageTypeMin() - offset,
+          aqtype.getStorageTypeMax() - offset);
+    } else {
+      return failure();
+    }
+
+    if (!new_qtype) return failure();
+    Type new_output_type = new_qtype.castFromExpressedType(
+        QType::castToExpressedType(output_type));
+    rewriter.replaceOpWithNewOp<QuantizeOpT>(op, new_output_type, op.getArg());
+    return success();
+  }
+};
+
+// Fold Extra Requantize ops if the preceding ops has free scale requirement.
+template <typename RequantizeOpT>
+struct FoldTrivalRequantizeOp : public OpRewritePattern<RequantizeOpT> {
+  explicit FoldTrivalRequantizeOp(MLIRContext* context)
+      : OpRewritePattern<RequantizeOpT>(context, 1) {}
+
+  LogicalResult matchAndRewrite(RequantizeOpT op,
+                                PatternRewriter& rewriter) const override {
+    Value pre_quantized = op->getOperand(0);
+    auto pre_quantized_type =
+        quant::QuantizedType::getQuantizedElementType(pre_quantized.getType());
+    if (!pre_quantized_type) return failure();
+
+    Operation* def = pre_quantized.getDefiningOp();
+    if (!def) return failure();
+    if (llvm::isa<FixedOutputRangeInterface, SameScalesOpInterface>(def) ||
+        !def->hasTrait<OpTrait::quant::QuantizableResult>()) {
+      return failure();
+    }
+
+    // This op should not clobber def, if more than one requant of this value.
+    if (!pre_quantized.hasOneUse()) {
+      return failure();
+    }
+
+    op.emitWarning("Remove trivial `rescale` op. Please fix the source graph.");
+
+    llvm::SmallVector<Type, 4> new_output_types;
+    for (auto result : def->getResults()) {
+      if (result.hasOneUse() && *result.getUsers().begin() == op) {
+        new_output_types.push_back(op.getResult().getType());
+      } else {
+        new_output_types.push_back(result.getType());
+      }
+    }
+
+    // Remove this rescale op.
+    rewriter.replaceOp(op, {pre_quantized});
+
+    // Replace the output scale of the preceding op.
+    rewriter.setInsertionPointAfter(def);
+    OperationState new_state(def->getLoc(), def->getName().getStringRef(),
+                             def->getOperands(), new_output_types,
+                             def->getAttrs());
+    Operation* new_op = rewriter.create(new_state);
+
+    rewriter.replaceOp(def, new_op->getResults());
+    return success();
+  }
+};
+
+// Given a quantized type `input`, magnifying its scales by the factor stored in
+// `factor`. If `input` isn't a quantized type or the `factor` doesn't match the
+// dimension size of `input` or isn't floating-point, nullptr will be returned.
+TypeAttr RescaleQuantizedType(Type input, Attribute factor);
+
+// Converts the min/max/num_bits/narrow_range information to a
+// QuantizedType, and then returns the attribute containing the QuantizedType.
+// The `min` and `max` arguments can be FloatAttr or DenseFPElementsAttr and
+// returns UniformQuantizedType or UniformQuantizedPerAxisType respectively.
+// `narrow_range` is set to true for weights and `is_signed` is set to true
+// if it is using signed int symmetric quantization.
+//
+// Note that this method may broadcast min and max to match the dimension length
+// of `input_type`, if the `quant_dim` is valid. On the other hand, the
+// symmetry of min and max is not adjusted by this method. The QAT workflow
+// should set min/max correctly (and use `narrow_range`=true, `is_signed`=true)
+// if symmetric quantization is required.
+TypeAttr GetQuantizedTypeAttr(Builder builder, Type input_type, Attribute min,
+                              Attribute max, int quant_dim,
+                              IntegerAttr num_bits, BoolAttr narrow_range,
+                              bool is_signed, bool legacy_float_scale = false,
+                              bool use_fake_quant_num_bits = false);
+
+// Casts the `target` type to a quantized type by using the quantization
+// parameters from the type in the `source` type attribute.
+// Examples:
+//   f32 -> !quant.uniform<i8:f32, 1.0>
+//   tensor<4xf32> -> tensor<4x!quant.uniform<i8:f32, 1.0>>
+// The result is wrapped by a type attribute. Returns nullptr if the cast
+// isn't valid.
+//
+// `axis` is to specify the quantization dimension in the `target` and only
+// used if the element type of `source` is a per-channel quantized type. During
+// the casting, the quantization dimension of the result type needs to be set
+// this new `axis` value.
+TypeAttr CastQuantizedTypeAttrFromExpressedType(Builder builder,
+                                                TypeAttr source, Type target,
+                                                int axis);
+
+// Quantizes the elements in the attribute `real_value` by the quantization
+// parameters in `tensor_type`. Returns empty Attribute if the
+// `tensor_type` is not a QuantizedType or the quantization fails.
+ElementsAttr Quantize(Attribute real_value, Type tensor_type);
+
+// Quantizes the elements in "legacy mode", where it calls TOCO's methods to
+// to quantize values with float scale.
+ElementsAttr QuantizeLegacy(Attribute real_value, Type tensor_type);
+
+// Returns the quantized type for an element attribute. The quantization
+// parameters in this type is based on the min and max element of the
+// attribute. When the elements in the `attr` are not in floating-point, or
+// the value range isn't straddling zero, an empty type is returned. The min/max
+// are adjusted to be symmetric if `symmetric` flag is set to True. And
+// `symmetric` can only be set to true when it is signed and narrow_range.
+Type GetUniformQuantizedTypeForWeight(ElementsAttr attr, bool symmetric,
+                                      unsigned num_bits, bool is_signed,
+                                      bool narrow_range,
+                                      bool legacy_float_scale = false,
+                                      bool use_fake_quant_num_bits = false);
+
+// Returns the per channel quantized type for an element attribute.
+// `quant_dim` defines the quantization axis. The channel min/max are adjusted
+// to be symmetric if `symmetric` flag is set to True. And `symmetric` can only
+// be set to true when it is signed and narrow_range.
+Type GetUniformQuantizedPerAxisTypeForWeight(
+    ElementsAttr attr, int quant_dim, bool symmetric, unsigned num_bits,
+    bool is_signed, bool narrow_range, bool legacy_float_scale = false,
+    bool use_fake_quant_num_bits = false);
+
+// Returns the quantized type of a bias input, given the quantized types of
+// other operands which are multiply-accumulated (the bias is added to the
+// accumulated value).
+quant::QuantizedType GetUniformQuantizedTypeForBias(
+    const std::vector<quant::QuantizedType>& op_types, int adjusted_quant_dim,
+    bool legacy_float_scale = false);
+
+// Gets quantization scale specs (e.g. fixed output range, same result and
+// operand scales) from the default quantization interfaces. The op should
+// outlive returned spec for its interface methods to be properly referenced.
+std::unique_ptr<OpQuantScaleSpec> GetDefaultQuantScaleSpec(Operation* op);
+
+// The function might contain more stats ops than required, and it will
+// introduce requantize if the calibration stats have conflicts. This method
+// tries to remove all the redundant stats ops.
+bool RemoveRedundantStatsOps(mlir::func::FuncOp func,
+                             OpQuantSpecGetter op_quant_spec_getter,
+                             OpQuantScaleSpecGetter op_quant_scale_spec_getter =
+                                 GetDefaultQuantScaleSpec);
+
+// Given quantization parameters for int8, compute the quantization parameters
+// for uint if it is required, and wrap the result in an UniformQuantizedType.
+quant::UniformQuantizedType GetFixedOutputRange(bool is_signed, int bit_width,
+                                                Type tensor_type, double scale,
+                                                int64_t zero_point,
+                                                int64_t storage_min,
+                                                int64_t storage_max);
+
+quant::UniformQuantizedType GetFixedOutputRange(bool is_signed, int bit_width,
+                                                Type tensor_type, double scale,
+                                                int64_t zero_point);
+
+// Extracts min and max values from the DenseFPElementsAttr, and stores them
+// into `mins` and `maxs`. When mins and maxs are extracted per-channel,
+// `dim_size` is number of channels and `slice_size` is the size of slice per
+// each channel. When `symmetric` is true, the range is expanded to [-M, M].
+void ExtractMinMaxFromAttr(DenseFPElementsAttr values, int dim_size,
+                           int slice_size, bool symmetric,
+                           SmallVectorImpl<double>& mins,
+                           SmallVectorImpl<double>& maxs);
+
+// Returns the quantized type for the
+// input_type/min/max/storage_type_width/narrow_range.
+Type GetQuantizedType(Builder builder, Type input_type, ArrayRef<double> min,
+                      ArrayRef<double> max, int quant_dim,
+                      int storage_type_width, bool narrow_range, bool is_signed,
+                      bool legacy_float_scale = false,
+                      bool use_fake_quant_num_bits = false);
+}  // namespace quant
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_QUANTIZATION_LIB_QUANTIZATION_UTILS_H_

tensorflow/compiler/mlir/quantization/common/test_base.h

@@ -0,0 +1,85 @@
+/* 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_COMMON_TEST_BASE_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_TEST_BASE_H_
+
+#include <memory>
+
+#include <gtest/gtest.h>
+#include "absl/strings/string_view.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/Quant.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/Parser/Parser.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo
+#include "tensorflow/compiler/mlir/quantization/common/func.h"
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/context.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace mlir::quant {
+
+using ::testing::Test;
+
+class QuantizationTestBase : public Test {
+ protected:
+  QuantizationTestBase()
+      : ctx_(stablehlo::CreateMlirContextForQuantization()),
+        builder_(ctx_.get()) {
+    ctx_->loadDialect<arith::ArithDialect, mlir::stablehlo::StablehloDialect,
+                      func::FuncDialect, TF::TensorFlowDialect,
+                      tf_saved_model::TensorFlowSavedModelDialect,
+                      tf_executor::TensorFlowExecutorDialect,
+                      quant::QuantDialect, ir::TFQuantDialect>();
+  }
+
+  // Parses `module_op_str` to create a `ModuleOp`.
+  OwningOpRef<ModuleOp> ParseModuleOpString(
+      const absl::string_view module_op_str) {
+    return parseSourceString<ModuleOp>(module_op_str, ctx_.get());
+  }
+
+  // Convenience function that returns the first operation of type `OpT` from
+  // the `@main` function in `module_op`. Useful when testing with a text
+  // representation of a `ModuleOp` containing a single function `@main`.
+  // Returns `failure` iff there is no `@main` or no such operation is found in
+  // `@main`.
+  template <typename OpT>
+  FailureOr<OpT> FindFirstOpFromMainFunc(ModuleOp module_op) {
+    func::FuncOp main_func_op = FindMainFuncOp(module_op);
+    if (main_func_op == nullptr) return failure();
+
+    auto ops = main_func_op.getOps<OpT>();
+    if (ops.empty()) return failure();
+
+    return *ops.begin();
+  }
+
+  std::unique_ptr<MLIRContext> ctx_;
+  OpBuilder builder_;
+};
+
+}  // namespace mlir::quant
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_TEST_BASE_H_

tensorflow/compiler/mlir/quantization/common/tf_uniform_quantized_types.cc

@@ -0,0 +1,232 @@
+/* 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/common/tf_uniform_quantized_types.h"
+
+#include <cstdint>
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/MathExtras.h"
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+#define DEBUG_TYPE "uniform-quantized-types"
+
+namespace mlir {
+namespace tf_quant {
+
+using quant::QuantizedType;
+using quant::UniformQuantizedPerAxisType;
+using quant::UniformQuantizedType;
+
+UniformQuantizedType CreateI8F32UniformQuantizedType(const Location loc,
+                                                     MLIRContext& context,
+                                                     const double scale,
+                                                     const int64_t zero_point,
+                                                     const bool narrow_range) {
+  return UniformQuantizedType::getChecked(
+      loc, /*flags=*/quant::QuantizationFlags::Signed,
+      /*storageType=*/IntegerType::get(&context, /*width=*/8),
+      /*expressedType=*/Float32Type::get(&context), scale, zero_point,
+      /*storageTypeMin=*/llvm::minIntN(8) + (narrow_range ? 1 : 0),
+      /*storageTypeMax=*/llvm::maxIntN(8));
+}
+
+UniformQuantizedType CreateI32F32UniformQuantizedType(
+    const Location loc, MLIRContext& context, const double scale,
+    const int64_t zero_point) {
+  return UniformQuantizedType::getChecked(
+      loc, /*flags=*/quant::QuantizationFlags::Signed,
+      /*storageType=*/IntegerType::get(&context, /*width=*/32),
+      /*expressedType=*/Float32Type::get(&context), scale, zero_point,
+      /*storageTypeMin=*/llvm::minIntN(32),
+      /*storageTypeMax=*/llvm::maxIntN(32));
+}
+
+UniformQuantizedPerAxisType CreateI8F32UniformQuantizedPerAxisType(
+    const Location loc, MLIRContext& context, const ArrayRef<double> scales,
+    const ArrayRef<int64_t> zero_points, const int quantization_dimension,
+    const bool narrow_range) {
+  return UniformQuantizedPerAxisType::getChecked(
+      loc, /*flags=*/quant::QuantizationFlags::Signed,
+      /*storageType=*/IntegerType::get(&context, /*width=*/8),
+      /*expressedType=*/Float32Type::get(&context), SmallVector<double>(scales),
+      SmallVector<int64_t>(zero_points), quantization_dimension,
+      /*storageTypeMin=*/llvm::minIntN(8) + (narrow_range ? 1 : 0),
+      /*storageTypeMax=*/llvm::maxIntN(8));
+}
+
+UniformQuantizedPerAxisType CreateI32F32UniformQuantizedPerAxisType(
+    const Location loc, MLIRContext& context, const ArrayRef<double> scales,
+    const ArrayRef<int64_t> zero_points, const int quantization_dimension) {
+  return UniformQuantizedPerAxisType::getChecked(
+      loc, /*flags=*/quant::QuantizationFlags::Signed,
+      /*storageType=*/IntegerType::get(&context, /*width=*/32),
+      /*expressedType=*/Float32Type::get(&context), SmallVector<double>(scales),
+      SmallVector<int64_t>(zero_points), quantization_dimension,
+      /*storageTypeMin=*/llvm::minIntN(32),
+      /*storageTypeMax=*/llvm::maxIntN(32));
+}
+
+bool IsStorageTypeI8(const QuantizedType quantized_type) {
+  const Type storage_type = quantized_type.getStorageType();
+  return storage_type.isInteger(/*width=*/8);
+}
+
+bool IsStorageTypeI32(const QuantizedType quantized_type) {
+  const Type storage_type = quantized_type.getStorageType();
+  return storage_type.isInteger(/*width=*/32);
+}
+
+bool IsExpressedTypeF32(const QuantizedType quantized_type) {
+  const Type expressed_type = quantized_type.getExpressedType();
+  return mlir::isa<Float32Type>(expressed_type);
+}
+
+bool IsI8F32UniformQuantizedType(const Type type) {
+  const UniformQuantizedType quantized_type =
+      mlir::dyn_cast_or_null<UniformQuantizedType>(type);
+  if (!quantized_type) {
+    LLVM_DEBUG(llvm::dbgs()
+               << "Expected a uniform quantized type. Got: " << type << ".\n");
+    return false;
+  }
+
+  if (!IsStorageTypeI8(quantized_type)) {
+    LLVM_DEBUG(llvm::dbgs() << "Expected an i8 storage type. Got: "
+                            << quantized_type << ".\n");
+    return false;
+  }
+
+  if (!IsExpressedTypeF32(quantized_type)) {
+    LLVM_DEBUG(llvm::dbgs() << "Expected an f32 expressed type. Got: "
+                            << quantized_type << ".\n");
+    return false;
+  }
+
+  return true;
+}
+
+bool IsI8F32UniformQuantizedPerAxisType(const Type type) {
+  const UniformQuantizedPerAxisType quantized_per_axis_type =
+      mlir::dyn_cast_or_null<UniformQuantizedPerAxisType>(type);
+  if (!quantized_per_axis_type) {
+    LLVM_DEBUG(llvm::dbgs()
+               << "Expected a uniform quantized type. Got: " << type << ".\n");
+    return false;
+  }
+
+  if (!IsStorageTypeI8(quantized_per_axis_type)) {
+    LLVM_DEBUG(llvm::dbgs() << "Expected an i8 storage type. Got: "
+                            << quantized_per_axis_type << ".\n");
+    return false;
+  }
+
+  if (!IsExpressedTypeF32(quantized_per_axis_type)) {
+    LLVM_DEBUG(llvm::dbgs() << "Expected an f32 expressed type. Got: "
+                            << quantized_per_axis_type << ".\n");
+    return false;
+  }
+
+  return true;
+}
+
+bool IsI32F32UniformQuantizedType(const Type type) {
+  const UniformQuantizedType quantized_type =
+      mlir::dyn_cast_or_null<UniformQuantizedType>(type);
+  if (!quantized_type) {
+    LLVM_DEBUG(llvm::dbgs()
+               << "Expected a uniform quantized type. Got: " << type << ".\n");
+    return false;
+  }
+
+  if (!IsStorageTypeI32(quantized_type)) {
+    LLVM_DEBUG(llvm::dbgs() << "Expected an i32 storage type. Got: "
+                            << quantized_type << ".\n");
+    return false;
+  }
+
+  if (!IsExpressedTypeF32(quantized_type)) {
+    LLVM_DEBUG(llvm::dbgs() << "Expected an f32 expressed type. Got: "
+                            << quantized_type << ".\n");
+    return false;
+  }
+
+  return true;
+}
+
+bool IsI32F32UniformQuantizedPerAxisType(const Type type) {
+  const UniformQuantizedPerAxisType quantized_per_axis_type =
+      mlir::dyn_cast_or_null<UniformQuantizedPerAxisType>(type);
+  if (!quantized_per_axis_type) {
+    LLVM_DEBUG(llvm::dbgs()
+               << "Expected a uniform quantized type. Got: " << type << ".\n");
+    return false;
+  }
+
+  if (!IsStorageTypeI32(quantized_per_axis_type)) {
+    LLVM_DEBUG(llvm::dbgs() << "Expected an i32 storage type. Got: "
+                            << quantized_per_axis_type << ".\n");
+    return false;
+  }
+
+  if (!IsExpressedTypeF32(quantized_per_axis_type)) {
+    LLVM_DEBUG(llvm::dbgs() << "Expected an f32 expressed type. Got: "
+                            << quantized_per_axis_type << ".\n");
+    return false;
+  }
+
+  return true;
+}
+
+// Determines whether the storage type of a quantized type is supported by
+// `tfl.quantize` or `tfl.dequantize` ops. ui8, i8 and i16 are supported.
+bool IsSupportedByTfliteQuantizeOrDequantizeOps(IntegerType storage_type) {
+  if (storage_type.getWidth() == 8 ||
+      (storage_type.isSigned() && storage_type.getWidth() == 16)) {
+    return true;
+  }
+  LLVM_DEBUG(llvm::dbgs()
+             << "Uniform quantize / dequantize op only supports ui8, i8 or "
+                "i16 for the storage type of uniform quantized type. Got: "
+             << storage_type << ".\n");
+  return false;
+}
+
+bool IsQuantizedTensorType(Type type) {
+  if (!mlir::isa<TensorType>(type)) {
+    return false;
+  }
+  Type element_type = mlir::cast<TensorType>(type).getElementType();
+  return mlir::isa<QuantizedType>(element_type);
+}
+
+bool IsOpFullyQuantized(Operation* op) {
+  return llvm::all_of(op->getOperandTypes(), IsQuantizedTensorType) &&
+         llvm::all_of(op->getResultTypes(), IsQuantizedTensorType);
+}
+
+bool IsOpNotQuantized(Operation* op) {
+  return !llvm::any_of(op->getOperandTypes(), IsQuantizedTensorType) &&
+         !llvm::any_of(op->getResultTypes(), IsQuantizedTensorType);
+}
+
+}  // namespace tf_quant
+}  // namespace mlir

tensorflow/compiler/mlir/quantization/common/tf_uniform_quantized_types.h

@@ -0,0 +1,116 @@
+/* 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_COMMON_TF_UNIFORM_QUANTIZED_TYPES_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_TF_UNIFORM_QUANTIZED_TYPES_H_
+
+#include <cstdint>
+
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace mlir {
+namespace tf_quant {
+
+// Creates a `UniformQuantizedType` with the given `scale` and `zero_point`
+// values. The produced type has f32 as its expressed type and i8 as its
+// storage type. The available values use the full range of the storage value,
+// i.e. [-128, 127]. Assumes asymmetric quantization, meaning the zero point
+// value can be a non-zero value.
+// If `narrow_range` is set true (ex: for weights), a restricted range of
+// integers will be used for symmetric mapping, i.e. [-127, 127].
+quant::UniformQuantizedType CreateI8F32UniformQuantizedType(
+    Location loc, MLIRContext& context, double scale, int64_t zero_point,
+    bool narrow_range = false);
+
+// Creates a `UniformQuantizedType` with the given `scale` and `zero_point`
+// values. The produced type has f32 as its expressed type and i32 as its
+// storage type. The available values use the full range of the storage value.
+// Assumes asymmetric quantization, meaning the zero point value can be
+// a non-zero value.
+quant::UniformQuantizedType CreateI32F32UniformQuantizedType(
+    Location loc, MLIRContext& context, double scale, int64_t zero_point);
+
+// Creates a `UniformQuantizedPerAxisType` with the given `scales` and
+// `zero_points` values. The produced type has f32 as its expressed type and
+// i8 as its storage type. The available values use the full range of the
+// storage value, i.e. [-128, 127]. Assumes asymmetric quantization, meaning the
+// zero point values can be non-zero values.
+// If `narrow_range` is set true (ex: for weights), a restricted range of
+// integers will be used for symmetric mapping, i.e. [-127, 127].
+quant::UniformQuantizedPerAxisType CreateI8F32UniformQuantizedPerAxisType(
+    Location loc, MLIRContext& context, ArrayRef<double> scales,
+    ArrayRef<int64_t> zero_points, int quantization_dimension,
+    bool narrow_range = false);
+
+// Creates a `UniformQuantizedPerAxisType` with the given `scales` and
+// `zero_points` values. The produced type has f32 as its expressed type and
+// i32 as its storage type. The available values use the full range of the
+// storage value. Assumes asymmetric quantization, meaning the
+// zero point values can be non-zero values.
+quant::UniformQuantizedPerAxisType CreateI32F32UniformQuantizedPerAxisType(
+    Location loc, MLIRContext& context, ArrayRef<double> scales,
+    ArrayRef<int64_t> zero_points, int quantization_dimension);
+
+bool IsStorageTypeI8(quant::QuantizedType quantized_type);
+
+bool IsStorageTypeI32(quant::QuantizedType quantized_type);
+
+bool IsExpressedTypeF32(quant::QuantizedType quantized_type);
+
+// Given a value, extract the `ElementType`.
+// `value` should be a non-null `TensorType`.
+inline Type GetElementType(const Value value) {
+  return mlir::cast<TensorType>(value.getType()).getElementType();
+}
+
+// Returns true iff `type` is a uniform quantized type whose storage type is
+// 8-bit integer and expressed type is f32.
+bool IsI8F32UniformQuantizedType(Type type);
+
+// Returns true iff `type` is a uniform quantized per-axis (per-channel) type
+// whose storage type is 8-bit integer and expressed type is f32.
+bool IsI8F32UniformQuantizedPerAxisType(Type type);
+
+// Returns true iff `type` is a uniform quantized type whose storage type is
+// 32-bit integer and expressed type is f32.
+bool IsI32F32UniformQuantizedType(Type type);
+
+// Returns true iff `type` is a uniform quantized per-axis (per-channel) type
+// whose storage type is 32-bit integer and expressed type is f32.
+bool IsI32F32UniformQuantizedPerAxisType(Type type);
+
+// Determines whether the storage type of a quantized type is supported by
+// `tfl.quantize` or `tfl.dequantize` ops. ui8, i8 and i16 are supported.
+bool IsSupportedByTfliteQuantizeOrDequantizeOps(IntegerType storage_type);
+
+// Returns true if a type is quantized tensor type.
+bool IsQuantizedTensorType(Type type);
+
+// Returns true if all operands and results are quantized.
+bool IsOpFullyQuantized(Operation* op);
+
+// Returns true iff none among operand and result tensors are quantized.
+bool IsOpNotQuantized(Operation* op);
+
+}  // namespace tf_quant
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_TF_UNIFORM_QUANTIZED_TYPES_H_

tensorflow/compiler/mlir/quantization/common/uniform_quantized_types.cc

@@ -0,0 +1,229 @@
+/* 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/common/uniform_quantized_types.h"
+
+#include <cstdint>
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/MathExtras.h"
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+#define DEBUG_TYPE "uniform-quantized-types"
+
+namespace mlir {
+namespace quant {
+
+UniformQuantizedType CreateI8F32UniformQuantizedType(const Location loc,
+                                                     MLIRContext& context,
+                                                     const double scale,
+                                                     const int64_t zero_point,
+                                                     const bool narrow_range) {
+  return UniformQuantizedType::getChecked(
+      loc, /*flags=*/QuantizationFlags::Signed,
+      /*storageType=*/IntegerType::get(&context, /*width=*/8),
+      /*expressedType=*/Float32Type::get(&context), scale, zero_point,
+      /*storageTypeMin=*/llvm::minIntN(8) + (narrow_range ? 1 : 0),
+      /*storageTypeMax=*/llvm::maxIntN(8));
+}
+
+UniformQuantizedType CreateI32F32UniformQuantizedType(
+    const Location loc, MLIRContext& context, const double scale,
+    const int64_t zero_point) {
+  return UniformQuantizedType::getChecked(
+      loc, /*flags=*/QuantizationFlags::Signed,
+      /*storageType=*/IntegerType::get(&context, /*width=*/32),
+      /*expressedType=*/Float32Type::get(&context), scale, zero_point,
+      /*storageTypeMin=*/llvm::minIntN(32),
+      /*storageTypeMax=*/llvm::maxIntN(32));
+}
+
+UniformQuantizedPerAxisType CreateI8F32UniformQuantizedPerAxisType(
+    const Location loc, MLIRContext& context, const ArrayRef<double> scales,
+    const ArrayRef<int64_t> zero_points, const int quantization_dimension,
+    const bool narrow_range) {
+  return UniformQuantizedPerAxisType::getChecked(
+      loc, /*flags=*/QuantizationFlags::Signed,
+      /*storageType=*/IntegerType::get(&context, /*width=*/8),
+      /*expressedType=*/Float32Type::get(&context),
+      SmallVector<double>(scales), SmallVector<int64_t>(zero_points),
+      quantization_dimension,
+      /*storageTypeMin=*/llvm::minIntN(8) + (narrow_range ? 1 : 0),
+      /*storageTypeMax=*/llvm::maxIntN(8));
+}
+
+UniformQuantizedPerAxisType CreateI32F32UniformQuantizedPerAxisType(
+    const Location loc, MLIRContext& context, const ArrayRef<double> scales,
+    const ArrayRef<int64_t> zero_points, const int quantization_dimension) {
+  return UniformQuantizedPerAxisType::getChecked(
+      loc, /*flags=*/QuantizationFlags::Signed,
+      /*storageType=*/IntegerType::get(&context, /*width=*/32),
+      /*expressedType=*/Float32Type::get(&context),
+      SmallVector<double>(scales), SmallVector<int64_t>(zero_points),
+      quantization_dimension, /*storageTypeMin=*/llvm::minIntN(32),
+      /*storageTypeMax=*/llvm::maxIntN(32));
+}
+
+bool IsStorageTypeI8(const QuantizedType quantized_type) {
+  const Type storage_type = quantized_type.getStorageType();
+  return storage_type.isInteger(/*width=*/8);
+}
+
+bool IsStorageTypeI32(const QuantizedType quantized_type) {
+  const Type storage_type = quantized_type.getStorageType();
+  return storage_type.isInteger(/*width=*/32);
+}
+
+bool IsExpressedTypeF32(const QuantizedType quantized_type) {
+  const Type expressed_type = quantized_type.getExpressedType();
+  return mlir::isa<Float32Type>(expressed_type);
+}
+
+bool IsI8F32UniformQuantizedType(const Type type) {
+  const UniformQuantizedType quantized_type =
+      mlir::dyn_cast_or_null<UniformQuantizedType>(type);
+  if (!quantized_type) {
+    LLVM_DEBUG(llvm::dbgs()
+               << "Expected a uniform quantized type. Got: " << type << ".\n");
+    return false;
+  }
+
+  if (!IsStorageTypeI8(quantized_type)) {
+    LLVM_DEBUG(llvm::dbgs() << "Expected an i8 storage type. Got: "
+                            << quantized_type << ".\n");
+    return false;
+  }
+
+  if (!IsExpressedTypeF32(quantized_type)) {
+    LLVM_DEBUG(llvm::dbgs() << "Expected an f32 expressed type. Got: "
+                            << quantized_type << ".\n");
+    return false;
+  }
+
+  return true;
+}
+
+bool IsI8F32UniformQuantizedPerAxisType(const Type type) {
+  const UniformQuantizedPerAxisType quantized_per_axis_type =
+      mlir::dyn_cast_or_null<UniformQuantizedPerAxisType>(type);
+  if (!quantized_per_axis_type) {
+    LLVM_DEBUG(llvm::dbgs()
+               << "Expected a uniform quantized type. Got: " << type << ".\n");
+    return false;
+  }
+
+  if (!IsStorageTypeI8(quantized_per_axis_type)) {
+    LLVM_DEBUG(llvm::dbgs() << "Expected an i8 storage type. Got: "
+                            << quantized_per_axis_type << ".\n");
+    return false;
+  }
+
+  if (!IsExpressedTypeF32(quantized_per_axis_type)) {
+    LLVM_DEBUG(llvm::dbgs() << "Expected an f32 expressed type. Got: "
+                            << quantized_per_axis_type << ".\n");
+    return false;
+  }
+
+  return true;
+}
+
+bool IsI32F32UniformQuantizedType(const Type type) {
+  const UniformQuantizedType quantized_type =
+      mlir::dyn_cast_or_null<UniformQuantizedType>(type);
+  if (!quantized_type) {
+    LLVM_DEBUG(llvm::dbgs()
+               << "Expected a uniform quantized type. Got: " << type << ".\n");
+    return false;
+  }
+
+  if (!IsStorageTypeI32(quantized_type)) {
+    LLVM_DEBUG(llvm::dbgs() << "Expected an i32 storage type. Got: "
+                            << quantized_type << ".\n");
+    return false;
+  }
+
+  if (!IsExpressedTypeF32(quantized_type)) {
+    LLVM_DEBUG(llvm::dbgs() << "Expected an f32 expressed type. Got: "
+                            << quantized_type << ".\n");
+    return false;
+  }
+
+  return true;
+}
+
+bool IsI32F32UniformQuantizedPerAxisType(const Type type) {
+  const UniformQuantizedPerAxisType quantized_per_axis_type =
+      mlir::dyn_cast_or_null<UniformQuantizedPerAxisType>(type);
+  if (!quantized_per_axis_type) {
+    LLVM_DEBUG(llvm::dbgs()
+               << "Expected a uniform quantized type. Got: " << type << ".\n");
+    return false;
+  }
+
+  if (!IsStorageTypeI32(quantized_per_axis_type)) {
+    LLVM_DEBUG(llvm::dbgs() << "Expected an i32 storage type. Got: "
+                            << quantized_per_axis_type << ".\n");
+    return false;
+  }
+
+  if (!IsExpressedTypeF32(quantized_per_axis_type)) {
+    LLVM_DEBUG(llvm::dbgs() << "Expected an f32 expressed type. Got: "
+                            << quantized_per_axis_type << ".\n");
+    return false;
+  }
+
+  return true;
+}
+
+// Determines whether the storage type of a quantized type is supported by
+// `tfl.quantize` or `tfl.dequantize` ops. ui8, i8 and i16 are supported.
+bool IsSupportedByTfliteQuantizeOrDequantizeOps(IntegerType storage_type) {
+  if (storage_type.getWidth() == 8 ||
+      (storage_type.isSigned() && storage_type.getWidth() == 16)) {
+    return true;
+  }
+  LLVM_DEBUG(llvm::dbgs()
+             << "Uniform quantize / dequantize op only supports ui8, i8 or "
+                "i16 for the storage type of uniform quantized type. Got: "
+             << storage_type << ".\n");
+  return false;
+}
+
+bool IsQuantizedTensorType(Type type) {
+  if (!mlir::isa<TensorType>(type)) {
+    return false;
+  }
+  Type element_type = mlir::cast<TensorType>(type).getElementType();
+  return mlir::isa<QuantizedType>(element_type);
+}
+
+bool IsOpFullyQuantized(Operation* op) {
+  return llvm::all_of(op->getOperandTypes(), IsQuantizedTensorType) &&
+         llvm::all_of(op->getResultTypes(), IsQuantizedTensorType);
+}
+
+bool IsOpNotQuantized(Operation* op) {
+  return !llvm::any_of(op->getOperandTypes(), IsQuantizedTensorType) &&
+         !llvm::any_of(op->getResultTypes(), IsQuantizedTensorType);
+}
+
+}  // namespace quant
+}  // namespace mlir

tensorflow/compiler/mlir/quantization/common/uniform_quantized_types.h

@@ -0,0 +1,120 @@
+/* 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_COMMON_UNIFORM_QUANTIZED_TYPES_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_UNIFORM_QUANTIZED_TYPES_H_
+
+#include <cstdint>
+
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace mlir {
+namespace quant {
+
+// Creates a `UniformQuantizedType` with the given `scale` and `zero_point`
+// values. The produced type has f32 as its expressed type and i8 as its
+// storage type. The available values use the full range of the storage value,
+// i.e. [-128, 127]. Assumes asymmetric quantization, meaning the zero point
+// value can be a non-zero value.
+// If `narrow_range` is set true (ex: for weights), a restricted range of
+// integers will be used for symmetric mapping, i.e. [-127, 127].
+UniformQuantizedType CreateI8F32UniformQuantizedType(Location loc,
+                                                     MLIRContext& context,
+                                                     double scale,
+                                                     int64_t zero_point,
+                                                     bool narrow_range = false);
+
+// Creates a `UniformQuantizedType` with the given `scale` and `zero_point`
+// values. The produced type has f32 as its expressed type and i32 as its
+// storage type. The available values use the full range of the storage value.
+// Assumes asymmetric quantization, meaning the zero point value can be
+// a non-zero value.
+UniformQuantizedType CreateI32F32UniformQuantizedType(Location loc,
+                                                      MLIRContext& context,
+                                                      double scale,
+                                                      int64_t zero_point);
+
+// Creates a `UniformQuantizedPerAxisType` with the given `scales` and
+// `zero_points` values. The produced type has f32 as its expressed type and
+// i8 as its storage type. The available values use the full range of the
+// storage value, i.e. [-128, 127]. Assumes asymmetric quantization, meaning the
+// zero point values can be non-zero values.
+// If `narrow_range` is set true (ex: for weights), a restricted range of
+// integers will be used for symmetric mapping, i.e. [-127, 127].
+UniformQuantizedPerAxisType CreateI8F32UniformQuantizedPerAxisType(
+    Location loc, MLIRContext& context, ArrayRef<double> scales,
+    ArrayRef<int64_t> zero_points, int quantization_dimension,
+    bool narrow_range = false);
+
+// Creates a `UniformQuantizedPerAxisType` with the given `scales` and
+// `zero_points` values. The produced type has f32 as its expressed type and
+// i32 as its storage type. The available values use the full range of the
+// storage value. Assumes asymmetric quantization, meaning the
+// zero point values can be non-zero values.
+UniformQuantizedPerAxisType CreateI32F32UniformQuantizedPerAxisType(
+    Location loc, MLIRContext& context, ArrayRef<double> scales,
+    ArrayRef<int64_t> zero_points, int quantization_dimension);
+
+bool IsStorageTypeI8(QuantizedType quantized_type);
+
+bool IsStorageTypeI32(QuantizedType quantized_type);
+
+bool IsExpressedTypeF32(QuantizedType quantized_type);
+
+// Given a value, extract the `ElementType`.
+// `value` should be a non-null `TensorType`.
+inline Type GetElementType(const Value value) {
+  return mlir::cast<TensorType>(value.getType()).getElementType();
+}
+
+// Returns true iff `type` is a uniform quantized type whose storage type is
+// 8-bit integer and expressed type is f32.
+bool IsI8F32UniformQuantizedType(Type type);
+
+// Returns true iff `type` is a uniform quantized per-axis (per-channel) type
+// whose storage type is 8-bit integer and expressed type is f32.
+bool IsI8F32UniformQuantizedPerAxisType(Type type);
+
+// Returns true iff `type` is a uniform quantized type whose storage type is
+// 32-bit integer and expressed type is f32.
+bool IsI32F32UniformQuantizedType(Type type);
+
+// Returns true iff `type` is a uniform quantized per-axis (per-channel) type
+// whose storage type is 32-bit integer and expressed type is f32.
+bool IsI32F32UniformQuantizedPerAxisType(Type type);
+
+// Determines whether the storage type of a quantized type is supported by
+// `tfl.quantize` or `tfl.dequantize` ops. ui8, i8 and i16 are supported.
+bool IsSupportedByTfliteQuantizeOrDequantizeOps(IntegerType storage_type);
+
+// Returns true if a type is quantized tensor type.
+bool IsQuantizedTensorType(Type type);
+
+// Returns true if all operands and results are quantized.
+bool IsOpFullyQuantized(Operation* op);
+
+// Returns true iff none among operand and result tensors are quantized.
+bool IsOpNotQuantized(Operation* op);
+
+}  // namespace quant
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_UNIFORM_QUANTIZED_TYPES_H_

tensorflow/compiler/mlir/quantization/common/uniform_quantized_types_test.cc

@@ -0,0 +1,756 @@
+/* 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/common/uniform_quantized_types.h"
+
+#include <cstdint>
+#include <limits>
+
+#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/Dialect/Quant/IR/Quant.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/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/test_base.h"
+
+namespace mlir {
+namespace quant {
+namespace {
+
+using ::testing::ElementsAreArray;
+using ::testing::IsNull;
+using ::testing::Ne;
+using ::testing::NotNull;
+using ::testing::Test;
+
+class CreateI8F32UniformQuantizedTypeTest : public Test {
+ protected:
+  CreateI8F32UniformQuantizedTypeTest() : ctx_() {
+    ctx_.loadDialect<quant::QuantDialect>();
+  }
+
+  MLIRContext ctx_;
+};
+
+TEST_F(CreateI8F32UniformQuantizedTypeTest, I8StorageTypeSucceeds) {
+  const UniformQuantizedType quantized_type =
+      CreateI8F32UniformQuantizedType(UnknownLoc::get(&ctx_), ctx_,
+                                      /*scale=*/1.0, /*zero_point=*/0);
+  // Storage type of `i8` is currently verifiable as `unsigned` in `Types.cpp`.
+  EXPECT_TRUE(quantized_type.getStorageType().isSignlessInteger(8));
+}
+
+TEST_F(CreateI8F32UniformQuantizedTypeTest, F32ExpressedTypeSucceeds) {
+  const UniformQuantizedType quantized_type =
+      CreateI8F32UniformQuantizedType(UnknownLoc::get(&ctx_), ctx_,
+                                      /*scale=*/1.0, /*zero_point=*/0);
+
+  EXPECT_TRUE(quantized_type.getExpressedType().isF32());
+}
+
+TEST_F(CreateI8F32UniformQuantizedTypeTest, SignedQuantizedTypeSucceeds) {
+  const UniformQuantizedType quantized_type =
+      CreateI8F32UniformQuantizedType(UnknownLoc::get(&ctx_), ctx_,
+                                      /*scale=*/1.0, /*zero_point=*/0);
+
+  EXPECT_TRUE(quantized_type.isSigned());
+}
+
+TEST_F(CreateI8F32UniformQuantizedTypeTest, StorageTypeMinMaxEqualToI8MinMax) {
+  const UniformQuantizedType quantized_type =
+      CreateI8F32UniformQuantizedType(UnknownLoc::get(&ctx_), ctx_,
+                                      /*scale=*/1.0, /*zero_point=*/0);
+
+  EXPECT_EQ(quantized_type.getStorageTypeMin(), -128);
+  EXPECT_EQ(quantized_type.getStorageTypeMax(), 127);
+}
+
+TEST_F(CreateI8F32UniformQuantizedTypeTest, StorageTypeMinMaxNarrowRange) {
+  const UniformQuantizedType quantized_type = CreateI8F32UniformQuantizedType(
+      UnknownLoc::get(&ctx_), ctx_,
+      /*scale=*/1.0, /*zero_point=*/0, /*narrow_range=*/true);
+
+  EXPECT_EQ(quantized_type.getStorageTypeMin(), -127);
+  EXPECT_EQ(quantized_type.getStorageTypeMax(), 127);
+}
+
+TEST_F(CreateI8F32UniformQuantizedTypeTest, HasScaleAndZeroPointProperlySet) {
+  const UniformQuantizedType quantized_type =
+      CreateI8F32UniformQuantizedType(UnknownLoc::get(&ctx_), ctx_,
+                                      /*scale=*/8.0, /*zero_point=*/99);
+
+  EXPECT_EQ(quantized_type.getScale(), 8.0);
+  EXPECT_EQ(quantized_type.getZeroPoint(), 99);
+}
+
+class CreateI32F32UniformQuantizedTypeTest : public Test {
+ protected:
+  CreateI32F32UniformQuantizedTypeTest() : ctx_() {
+    ctx_.loadDialect<quant::QuantDialect>();
+  }
+
+  MLIRContext ctx_;
+};
+
+TEST_F(CreateI32F32UniformQuantizedTypeTest, I32StorageTypeSucceeds) {
+  const UniformQuantizedType quantized_type =
+      CreateI32F32UniformQuantizedType(UnknownLoc::get(&ctx_), ctx_,
+                                       /*scale=*/1.0, /*zero_point=*/0);
+
+  // Storage type of `i32` is currently verifiable as `unsigned` in `Types.cpp`.
+  EXPECT_TRUE(quantized_type.getStorageType().isSignlessInteger(32));
+}
+
+TEST_F(CreateI32F32UniformQuantizedTypeTest, F32ExpressedTypeSucceeds) {
+  const UniformQuantizedType quantized_type =
+      CreateI32F32UniformQuantizedType(UnknownLoc::get(&ctx_), ctx_,
+                                       /*scale=*/1.0, /*zero_point=*/0);
+
+  EXPECT_TRUE(quantized_type.getExpressedType().isF32());
+}
+
+TEST_F(CreateI32F32UniformQuantizedTypeTest, SignedQuantizedTypeSucceeds) {
+  const UniformQuantizedType quantized_type =
+      CreateI32F32UniformQuantizedType(UnknownLoc::get(&ctx_), ctx_,
+                                       /*scale=*/1.0, /*zero_point=*/0);
+
+  EXPECT_TRUE(quantized_type.isSigned());
+}
+
+TEST_F(CreateI32F32UniformQuantizedTypeTest,
+       StorageTypeMinMaxEqualToI32MinMax) {
+  const UniformQuantizedType quantized_type =
+      CreateI32F32UniformQuantizedType(UnknownLoc::get(&ctx_), ctx_,
+                                       /*scale=*/1.0, /*zero_point=*/0);
+
+  EXPECT_EQ(quantized_type.getStorageTypeMin(),
+            std::numeric_limits<int32_t>::min());
+  EXPECT_EQ(quantized_type.getStorageTypeMax(),
+            std::numeric_limits<int32_t>::max());
+}
+
+TEST_F(CreateI32F32UniformQuantizedTypeTest, HasScaleAndZeroPointProperlySet) {
+  const UniformQuantizedType quantized_type =
+      CreateI32F32UniformQuantizedType(UnknownLoc::get(&ctx_), ctx_,
+                                       /*scale=*/8.0, /*zero_point=*/1111);
+
+  EXPECT_EQ(quantized_type.getScale(), 8.0);
+  EXPECT_EQ(quantized_type.getZeroPoint(), 1111);
+}
+
+class CreateI8F32UniformQuantizedPerAxisTypeTest : public Test {
+ protected:
+  CreateI8F32UniformQuantizedPerAxisTypeTest() : ctx_() {
+    ctx_.loadDialect<quant::QuantDialect>();
+  }
+
+  MLIRContext ctx_;
+};
+
+TEST_F(CreateI8F32UniformQuantizedPerAxisTypeTest, I8StorageTypeSucceeds) {
+  const UniformQuantizedPerAxisType quantized_type =
+      CreateI8F32UniformQuantizedPerAxisType(
+          UnknownLoc::get(&ctx_), ctx_,
+          /*scales=*/SmallVector<double, 2>{1.0, 1.0},
+          /*zero_points=*/SmallVector<int64_t, 2>{0, 0},
+          /*quantization_dimension=*/0);
+
+  // Storage type of `i8` is currently verifiable as `unsigned` in `Types.cpp`.
+  EXPECT_TRUE(quantized_type.getStorageType().isSignlessInteger(8));
+}
+
+TEST_F(CreateI8F32UniformQuantizedPerAxisTypeTest, F32ExpressedTypeSucceeds) {
+  const UniformQuantizedPerAxisType quantized_type =
+      CreateI8F32UniformQuantizedPerAxisType(
+          UnknownLoc::get(&ctx_), ctx_,
+          /*scales=*/SmallVector<double, 2>{1.0, 1.0},
+          /*zero_points=*/SmallVector<int64_t, 2>{0, 0},
+          /*quantization_dimension=*/0);
+
+  EXPECT_TRUE(quantized_type.getExpressedType().isF32());
+}
+
+TEST_F(CreateI8F32UniformQuantizedPerAxisTypeTest,
+       SignedQuantizedTypeSucceeds) {
+  const UniformQuantizedPerAxisType quantized_type =
+      CreateI8F32UniformQuantizedPerAxisType(
+          UnknownLoc::get(&ctx_), ctx_,
+          /*scales=*/SmallVector<double, 2>{1.0, 1.0},
+          /*zero_points=*/SmallVector<int64_t, 2>{0, 0},
+          /*quantization_dimension=*/0);
+
+  EXPECT_TRUE(quantized_type.isSigned());
+}
+
+TEST_F(CreateI8F32UniformQuantizedPerAxisTypeTest,
+       StorageTypeMinMaxEqualToI8MinMax) {
+  const UniformQuantizedPerAxisType quantized_type =
+      CreateI8F32UniformQuantizedPerAxisType(
+          UnknownLoc::get(&ctx_), ctx_,
+          /*scales=*/SmallVector<double, 2>{1.0, 1.0},
+          /*zero_points=*/SmallVector<int64_t, 2>{0, 0},
+          /*quantization_dimension=*/0);
+
+  EXPECT_EQ(quantized_type.getStorageTypeMin(), -128);
+  EXPECT_EQ(quantized_type.getStorageTypeMax(), 127);
+}
+
+TEST_F(CreateI8F32UniformQuantizedPerAxisTypeTest,
+       StorageTypeMinMaxNarrowRange) {
+  const UniformQuantizedPerAxisType quantized_type =
+      CreateI8F32UniformQuantizedPerAxisType(
+          UnknownLoc::get(&ctx_), ctx_,
+          /*scales=*/SmallVector<double, 2>{1.0, 1.0},
+          /*zero_points=*/SmallVector<int64_t, 2>{0, 0},
+          /*quantization_dimension=*/0, /*narrow_range=*/true);
+
+  EXPECT_EQ(quantized_type.getStorageTypeMin(), -127);
+  EXPECT_EQ(quantized_type.getStorageTypeMax(), 127);
+}
+
+TEST_F(CreateI8F32UniformQuantizedPerAxisTypeTest,
+       HasQuantizationDimensionProperlySet) {
+  const UniformQuantizedPerAxisType quantized_type =
+      CreateI8F32UniformQuantizedPerAxisType(
+          UnknownLoc::get(&ctx_), ctx_,
+          /*scales=*/SmallVector<double, 2>{1.0, 1.0},
+          /*zero_points=*/SmallVector<int64_t, 2>{0, 0},
+          /*quantization_dimension=*/3);
+
+  EXPECT_EQ(quantized_type.getQuantizedDimension(), 3);
+}
+
+TEST_F(CreateI8F32UniformQuantizedPerAxisTypeTest,
+       HasScaleAndZeroPointProperlySet) {
+  const UniformQuantizedPerAxisType quantized_type =
+      CreateI8F32UniformQuantizedPerAxisType(
+          UnknownLoc::get(&ctx_), ctx_,
+          /*scales=*/SmallVector<double, 2>{8.0, 9.0},
+          /*zero_points=*/SmallVector<int64_t, 2>{98, 99},
+          /*quantization_dimension=*/0);
+
+  EXPECT_THAT(quantized_type.getScales(), ElementsAreArray({8.0, 9.0}));
+  EXPECT_THAT(quantized_type.getZeroPoints(), ElementsAreArray({98, 99}));
+}
+
+class CreateI32F32UniformQuantizedPerAxisTypeTest : public Test {
+ protected:
+  CreateI32F32UniformQuantizedPerAxisTypeTest() : ctx_() {
+    ctx_.loadDialect<quant::QuantDialect>();
+  }
+
+  MLIRContext ctx_;
+};
+
+TEST_F(CreateI32F32UniformQuantizedPerAxisTypeTest, I32StorageTypeSucceeds) {
+  const UniformQuantizedPerAxisType quantized_type =
+      CreateI32F32UniformQuantizedPerAxisType(
+          UnknownLoc::get(&ctx_), ctx_,
+          /*scales=*/SmallVector<double, 2>{1.0, 1.0},
+          /*zero_points=*/SmallVector<int64_t, 2>{0, 0},
+          /*quantization_dimension=*/0);
+
+  // Storage type of `i32` is currently verifiable as `unsigned` in `Types.cpp`.
+  EXPECT_TRUE(quantized_type.getStorageType().isSignlessInteger(32));
+}
+
+TEST_F(CreateI32F32UniformQuantizedPerAxisTypeTest, F32ExpressedTypeSucceeds) {
+  const UniformQuantizedPerAxisType quantized_type =
+      CreateI32F32UniformQuantizedPerAxisType(
+          UnknownLoc::get(&ctx_), ctx_,
+          /*scales=*/SmallVector<double, 2>{1.0, 1.0},
+          /*zero_points=*/SmallVector<int64_t, 2>{0, 0},
+          /*quantization_dimension=*/0);
+
+  EXPECT_TRUE(quantized_type.getExpressedType().isF32());
+}
+
+TEST_F(CreateI32F32UniformQuantizedPerAxisTypeTest,
+       StorageTypeMinMaxEqualToI32MinMax) {
+  const UniformQuantizedPerAxisType quantized_type =
+      CreateI32F32UniformQuantizedPerAxisType(
+          UnknownLoc::get(&ctx_), ctx_,
+          /*scales=*/SmallVector<double, 2>{1.0, 1.0},
+          /*zero_points=*/SmallVector<int64_t, 2>{0, 0},
+          /*quantization_dimension=*/0);
+
+  EXPECT_EQ(quantized_type.getStorageTypeMin(),
+            std::numeric_limits<int32_t>::min());
+  EXPECT_EQ(quantized_type.getStorageTypeMax(),
+            std::numeric_limits<int32_t>::max());
+}
+
+TEST_F(CreateI32F32UniformQuantizedPerAxisTypeTest,
+       HasQuantizationDimensionProperlySet) {
+  const UniformQuantizedPerAxisType quantized_type =
+      CreateI32F32UniformQuantizedPerAxisType(
+          UnknownLoc::get(&ctx_), ctx_,
+          /*scales=*/SmallVector<double, 2>{1.0, 1.0},
+          /*zero_points=*/SmallVector<int64_t, 2>{0, 0},
+          /*quantization_dimension=*/3);
+
+  EXPECT_EQ(quantized_type.getQuantizedDimension(), 3);
+}
+
+TEST_F(CreateI32F32UniformQuantizedPerAxisTypeTest,
+       HasScaleAndZeroPointProperlySet) {
+  const UniformQuantizedPerAxisType quantized_type =
+      CreateI32F32UniformQuantizedPerAxisType(
+          UnknownLoc::get(&ctx_), ctx_,
+          /*scales=*/SmallVector<double, 2>{8.0, 9.0},
+          /*zero_points=*/SmallVector<int64_t, 2>{98, 99},
+          /*quantization_dimension=*/0);
+
+  EXPECT_THAT(quantized_type.getScales(), ElementsAreArray({8.0, 9.0}));
+  EXPECT_THAT(quantized_type.getZeroPoints(), ElementsAreArray({98, 99}));
+}
+
+class IsI8F32UniformQuantizedTypeTest : public Test {
+ protected:
+  IsI8F32UniformQuantizedTypeTest() : builder_(&ctx_) {
+    ctx_.loadDialect<quant::QuantDialect>();
+  }
+
+  MLIRContext ctx_;
+  OpBuilder builder_;
+};
+
+TEST_F(IsI8F32UniformQuantizedTypeTest, I8F32UniformQuantizedTypeSucceeds) {
+  const UniformQuantizedType qi8_type = quant::UniformQuantizedType::get(
+      /*flags=*/QuantizationFlags::Signed, builder_.getI8Type(),
+      builder_.getF32Type(), /*scale=*/1.0,
+      /*zeroPoint=*/0, /*storageTypeMin=*/-128, /*storageTypeMax=*/127);
+  EXPECT_TRUE(IsI8F32UniformQuantizedType(qi8_type));
+}
+
+TEST_F(IsI8F32UniformQuantizedTypeTest, UniformQuantizedTypeSucceeds) {
+  const UniformQuantizedType qi8_type = quant::UniformQuantizedType::get(
+      /*flags=*/QuantizationFlags::Signed, builder_.getI8Type(),
+      builder_.getF32Type(), /*scale=*/1.0,
+      /*zeroPoint=*/0, /*storageTypeMin=*/-128, /*storageTypeMax=*/127);
+  EXPECT_THAT(mlir::dyn_cast_or_null<UniformQuantizedType>(qi8_type),
+              NotNull());
+}
+
+TEST_F(IsI8F32UniformQuantizedTypeTest, StorageTypeI8Succeeds) {
+  const UniformQuantizedType qi8_type = quant::UniformQuantizedType::get(
+      /*flags=*/QuantizationFlags::Signed, builder_.getI8Type(),
+      builder_.getF32Type(), /*scale=*/1.0,
+      /*zeroPoint=*/0, /*storageTypeMin=*/-128, /*storageTypeMax=*/127);
+  EXPECT_TRUE(IsStorageTypeI8(qi8_type));
+}
+
+TEST_F(IsI8F32UniformQuantizedTypeTest, ExpressedTypeF32Succeeds) {
+  const UniformQuantizedType qi8_type = quant::UniformQuantizedType::get(
+      /*flags=*/QuantizationFlags::Signed, builder_.getI8Type(),
+      builder_.getF32Type(), /*scale=*/1.0,
+      /*zeroPoint=*/0, /*storageTypeMin=*/-128, /*storageTypeMax=*/127);
+  EXPECT_TRUE(IsExpressedTypeF32(qi8_type));
+}
+
+class IsI8F32UniformQuantizedPerAxisTypeTest : public Test {
+ protected:
+  IsI8F32UniformQuantizedPerAxisTypeTest() : builder_(&ctx_) {
+    ctx_.loadDialect<quant::QuantDialect>();
+  }
+
+  MLIRContext ctx_;
+  OpBuilder builder_;
+};
+
+TEST_F(IsI8F32UniformQuantizedPerAxisTypeTest,
+       I8F32UniformQuantizedPerAxisTypeSucceeds) {
+  const UniformQuantizedPerAxisType qi8_per_axis_type =
+      quant::UniformQuantizedPerAxisType::get(
+          /*flags=*/QuantizationFlags::Signed, builder_.getI8Type(),
+          builder_.getF32Type(),
+          /*scales=*/{1.0},
+          /*zeroPoints=*/{0}, /*quantizedDimension=*/0, /*storageTypeMin=*/-128,
+          /*storageTypeMax=*/127);
+  EXPECT_TRUE(IsI8F32UniformQuantizedPerAxisType(qi8_per_axis_type));
+  EXPECT_FALSE(IsI8F32UniformQuantizedType(qi8_per_axis_type));
+}
+
+TEST_F(IsI8F32UniformQuantizedTypeTest, UniformQuantizedPerAxisTypeSucceeds) {
+  const UniformQuantizedPerAxisType qi8_per_axis_type =
+      quant::UniformQuantizedPerAxisType::get(
+          /*flags=*/QuantizationFlags::Signed, builder_.getI8Type(),
+          builder_.getF32Type(),
+          /*scales=*/{1.0},
+          /*zeroPoints=*/{0}, /*quantizedDimension=*/0, /*storageTypeMin=*/-128,
+          /*storageTypeMax=*/127);
+  EXPECT_THAT(
+      mlir::dyn_cast_or_null<UniformQuantizedPerAxisType>(qi8_per_axis_type),
+      NotNull());
+}
+
+TEST_F(IsI8F32UniformQuantizedPerAxisTypeTest, StorageTypeI8Succeeds) {
+  const UniformQuantizedPerAxisType qi8_per_axis_type =
+      quant::UniformQuantizedPerAxisType::get(
+          /*flags=*/QuantizationFlags::Signed, builder_.getI8Type(),
+          builder_.getF32Type(),
+          /*scales=*/{1.0},
+          /*zeroPoints=*/{0}, /*quantizedDimension=*/0, /*storageTypeMin=*/-128,
+          /*storageTypeMax=*/127);
+  EXPECT_TRUE(IsStorageTypeI8(qi8_per_axis_type));
+}
+
+TEST_F(IsI8F32UniformQuantizedPerAxisTypeTest, ExpressedTypeF32Succeeds) {
+  const UniformQuantizedPerAxisType qi8_per_axis_type =
+      quant::UniformQuantizedPerAxisType::get(
+          /*flags=*/QuantizationFlags::Signed, builder_.getI8Type(),
+          builder_.getF32Type(),
+          /*scales=*/{1.0},
+          /*zeroPoints=*/{0}, /*quantizedDimension=*/0, /*storageTypeMin=*/-128,
+          /*storageTypeMax=*/127);
+  EXPECT_TRUE(IsExpressedTypeF32(qi8_per_axis_type));
+}
+
+class IsI32F32UniformQuantizedTypeTest : public Test {
+ protected:
+  IsI32F32UniformQuantizedTypeTest() : builder_(&ctx_) {
+    ctx_.loadDialect<quant::QuantDialect>();
+  }
+
+  MLIRContext ctx_;
+  OpBuilder builder_;
+};
+
+TEST_F(IsI32F32UniformQuantizedTypeTest, I32F32UniformQuantizedTypeSucceeds) {
+  const UniformQuantizedType qi32_type = quant::UniformQuantizedType::get(
+      /*flags=*/QuantizationFlags::Signed, builder_.getI32Type(),
+      builder_.getF32Type(),
+      /*scale=*/1.0,
+      /*zeroPoint=*/0, /*storageTypeMin=*/-2147483647,
+      /*storageTypeMax=*/2147483646);
+  EXPECT_TRUE(IsI32F32UniformQuantizedType(qi32_type));
+}
+
+TEST_F(IsI32F32UniformQuantizedTypeTest, UniformQuantizedTypeSucceeds) {
+  const UniformQuantizedType qi32_type = quant::UniformQuantizedType::get(
+      /*flags=*/QuantizationFlags::Signed, builder_.getI32Type(),
+      builder_.getF32Type(),
+      /*scale=*/1.0,
+      /*zeroPoint=*/0, /*storageTypeMin=*/-2147483647,
+      /*storageTypeMax=*/2147483646);
+  EXPECT_TRUE(IsI32F32UniformQuantizedType(qi32_type));
+  EXPECT_THAT(mlir::dyn_cast_or_null<UniformQuantizedType>(qi32_type),
+              NotNull());
+}
+
+TEST_F(IsI32F32UniformQuantizedTypeTest, StorageTypeI32Succeeds) {
+  const UniformQuantizedType qi32_type = quant::UniformQuantizedType::get(
+      /*flags=*/QuantizationFlags::Signed, builder_.getI32Type(),
+      builder_.getF32Type(),
+      /*scale=*/1.0,
+      /*zeroPoint=*/0, /*storageTypeMin=*/-2147483647,
+      /*storageTypeMax=*/2147483646);
+  EXPECT_TRUE(IsI32F32UniformQuantizedType(qi32_type));
+  EXPECT_TRUE(IsStorageTypeI32(qi32_type));
+}
+
+TEST_F(IsI32F32UniformQuantizedTypeTest, ExpressedTypeF32Succeeds) {
+  const UniformQuantizedType qi32_per_axis_type =
+      quant::UniformQuantizedType::get(
+          /*flags=*/QuantizationFlags::Signed, builder_.getI32Type(),
+          builder_.getF32Type(),
+          /*scale=*/1.0,
+          /*zeroPoint=*/0, /*storageTypeMin=*/-2147483647,
+          /*storageTypeMax=*/2147483646);
+  EXPECT_TRUE(IsExpressedTypeF32(qi32_per_axis_type));
+}
+
+class IsI32F32UniformQuantizedPerAxisTypeTest : public Test {
+ protected:
+  IsI32F32UniformQuantizedPerAxisTypeTest() : builder_(&ctx_) {
+    ctx_.loadDialect<quant::QuantDialect>();
+  }
+
+  MLIRContext ctx_;
+  OpBuilder builder_;
+};
+
+TEST_F(IsI32F32UniformQuantizedPerAxisTypeTest,
+       I32F32UniformQuantizedPerAxisTypeSucceeds) {
+  const UniformQuantizedPerAxisType qi32_per_axis_type =
+      quant::UniformQuantizedPerAxisType::get(
+          /*flags=*/QuantizationFlags::Signed, builder_.getI32Type(),
+          builder_.getF32Type(),
+          /*scales=*/{1.0},
+          /*zeroPoints=*/{0}, /*quantizedDimension=*/0,
+          /*storageTypeMin=*/-2147483647, /*storageTypeMax=*/2147483646);
+  EXPECT_TRUE(IsI32F32UniformQuantizedPerAxisType(qi32_per_axis_type));
+  EXPECT_FALSE(IsI32F32UniformQuantizedType(qi32_per_axis_type));
+}
+
+TEST_F(IsI32F32UniformQuantizedPerAxisTypeTest,
+       I8F32UniformQuantizedTypeFails) {
+  const UniformQuantizedType qi8_type = quant::UniformQuantizedType::get(
+      /*flags=*/QuantizationFlags::Signed, builder_.getI8Type(),
+      builder_.getF32Type(),
+      /*scale=*/1.0, /*zeroPoint=*/0, /*storageTypeMin=*/-128,
+      /*storageTypeMax=*/127);
+  EXPECT_FALSE(IsI32F32UniformQuantizedPerAxisType(qi8_type));
+  EXPECT_FALSE(IsStorageTypeI32(qi8_type));
+  EXPECT_THAT(mlir::dyn_cast_or_null<UniformQuantizedPerAxisType>(qi8_type),
+              IsNull());
+}
+
+TEST_F(IsI32F32UniformQuantizedTypeTest, UniformQuantizedPerAxisTypeSucceeds) {
+  const UniformQuantizedPerAxisType qi32_per_axis_type =
+      quant::UniformQuantizedPerAxisType::get(
+          /*flags=*/QuantizationFlags::Signed, builder_.getI32Type(),
+          builder_.getF32Type(),
+          /*scales=*/{1.0},
+          /*zeroPoints=*/{0}, /*quantizedDimension=*/0,
+          /*storageTypeMin=*/-2147483647, /*storageTypeMax=*/2147483646);
+
+  EXPECT_THAT(
+      mlir::dyn_cast_or_null<UniformQuantizedPerAxisType>(qi32_per_axis_type),
+      NotNull());
+}
+
+TEST_F(IsI32F32UniformQuantizedPerAxisTypeTest, StorageTypeI8Succeeds) {
+  const UniformQuantizedPerAxisType qi32_per_axis_type =
+      quant::UniformQuantizedPerAxisType::get(
+          /*flags=*/QuantizationFlags::Signed, builder_.getI32Type(),
+          builder_.getF32Type(),
+          /*scales=*/{1.0},
+          /*zeroPoints=*/{0}, /*quantizedDimension=*/0,
+          /*storageTypeMin=*/-2147483647, /*storageTypeMax=*/2147483646);
+
+  EXPECT_TRUE(IsStorageTypeI32(qi32_per_axis_type));
+}
+
+TEST_F(IsI32F32UniformQuantizedPerAxisTypeTest, ExpressedTypeF32Succeeds) {
+  const UniformQuantizedPerAxisType qi32_per_axis_type =
+      quant::UniformQuantizedPerAxisType::get(
+          /*flags=*/QuantizationFlags::Signed, builder_.getI32Type(),
+          builder_.getF32Type(),
+          /*scales=*/{1.0},
+          /*zeroPoints=*/{0}, /*quantizedDimension=*/0,
+          /*storageTypeMin=*/-2147483647, /*storageTypeMax=*/2147483646);
+  EXPECT_TRUE(IsExpressedTypeF32(qi32_per_axis_type));
+}
+
+class IsSupportedByTfliteQuantizeOrDequantizeOpsTest : public Test {
+ protected:
+  IsSupportedByTfliteQuantizeOrDequantizeOpsTest() : builder_(&ctx_) {
+    ctx_.loadDialect<quant::QuantDialect>();
+  }
+
+  MLIRContext ctx_;
+  OpBuilder builder_;
+};
+
+TEST_F(IsSupportedByTfliteQuantizeOrDequantizeOpsTest, StorageTypeI8Succeeds) {
+  auto qi8_type = quant::UniformQuantizedType::get(
+      /*flags=*/QuantizationFlags::Signed, builder_.getI8Type(),
+      builder_.getF32Type(),
+      /*scale=*/1.0,
+      /*zeroPoint=*/0, /*storageTypeMin=*/-128, /*storageTypeMax=*/127);
+  EXPECT_TRUE(IsSupportedByTfliteQuantizeOrDequantizeOps(
+      dyn_cast_or_null<IntegerType>(qi8_type.getStorageType())));
+}
+
+TEST_F(IsSupportedByTfliteQuantizeOrDequantizeOpsTest, StorageTypeI16Succeeds) {
+  auto qi16_type = quant::UniformQuantizedType::get(
+      /*flags=*/QuantizationFlags::Signed, builder_.getI8Type(),
+      builder_.getF32Type(),
+      /*scale=*/1.0,
+      /*zeroPoint=*/0, /*storageTypeMin=*/-128, /*storageTypeMax=*/127);
+  EXPECT_TRUE(IsSupportedByTfliteQuantizeOrDequantizeOps(
+      dyn_cast_or_null<IntegerType>(qi16_type.getStorageType())));
+}
+
+TEST_F(IsSupportedByTfliteQuantizeOrDequantizeOpsTest, StorageTypeUI8Succeeds) {
+  auto qi8_type = quant::UniformQuantizedType::get(
+      /*flags=*/QuantizationFlags::Signed, builder_.getI8Type(),
+      builder_.getF32Type(),
+      /*scale=*/1.0,
+      /*zeroPoint=*/0, /*storageTypeMin=*/-128, /*storageTypeMax=*/127);
+  EXPECT_TRUE(IsSupportedByTfliteQuantizeOrDequantizeOps(
+      dyn_cast_or_null<IntegerType>(qi8_type.getStorageType())));
+}
+
+using IsOpFullyQuantizedTest = QuantizationTestBase;
+
+TEST_F(IsOpFullyQuantizedTest, TrueIfOpFullyQuantized) {
+  constexpr absl::string_view kFullyQuantizedAdd = R"mlir(
+    func.func @fully_quantized_add(%arg0: tensor<2x!quant.uniform<i8:f32, 1.000000e+00:0>>) -> tensor<2x!quant.uniform<i8:f32, 1.000000e+00:0>> {
+      %0 = stablehlo.add %arg0, %arg0 : tensor<2x!quant.uniform<i8:f32, 1.000000e+00:0>>
+      return %0 : tensor<2x!quant.uniform<i8:f32, 1.000000e+00:0>>
+    }
+  )mlir";
+
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kFullyQuantizedAdd);
+  ASSERT_TRUE(module_op);
+
+  auto func_op = module_op->lookupSymbol<func::FuncOp>("fully_quantized_add");
+  ASSERT_THAT(func_op, NotNull());
+
+  auto add_op_itr = func_op.getBody().op_begin<mlir::stablehlo::AddOp>();
+  ASSERT_THAT(add_op_itr,
+              Ne(func_op.getBody().op_end<mlir::stablehlo::AddOp>()));
+
+  EXPECT_TRUE(IsOpFullyQuantized(*add_op_itr));
+}
+
+TEST_F(IsOpFullyQuantizedTest, FalseIfOpNotQuantized) {
+  constexpr absl::string_view kNotQuantizedAdd = R"mlir(
+    func.func @not_quantized_add(%arg0: tensor<2xf32>) -> tensor<2xf32> {
+      %0 = stablehlo.add %arg0, %arg0 : tensor<2xf32>
+      return %0 : tensor<2xf32>
+    }
+  )mlir";
+
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kNotQuantizedAdd);
+  ASSERT_TRUE(module_op);
+
+  auto func_op = module_op->lookupSymbol<func::FuncOp>("not_quantized_add");
+  ASSERT_THAT(func_op, NotNull());
+
+  auto add_op_itr = func_op.getBody().op_begin<mlir::stablehlo::AddOp>();
+  ASSERT_THAT(add_op_itr,
+              Ne(func_op.getBody().op_end<mlir::stablehlo::AddOp>()));
+
+  EXPECT_FALSE(IsOpFullyQuantized(*add_op_itr));
+}
+
+TEST_F(IsOpFullyQuantizedTest, FalseIfOpPartiallyQuantized) {
+  constexpr absl::string_view kQuantizeOp = R"mlir(
+    func.func @quantize(%arg0: tensor<2xf32>) -> tensor<2x!quant.uniform<i8:f32, 1.000000e+00:0>> {
+      %0 = stablehlo.uniform_quantize %arg0 : (tensor<2xf32>) -> tensor<2x!quant.uniform<i8:f32, 1.000000e+00:0>>
+      return %0 : tensor<2x!quant.uniform<i8:f32, 1.000000e+00:0>>
+    }
+  )mlir";
+
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kQuantizeOp);
+  ASSERT_TRUE(module_op);
+
+  auto func_op = module_op->lookupSymbol<func::FuncOp>("quantize");
+  ASSERT_THAT(func_op, NotNull());
+
+  auto uniform_quantize_op_itr =
+      func_op.getBody().op_begin<mlir::stablehlo::UniformQuantizeOp>();
+  ASSERT_THAT(
+      uniform_quantize_op_itr,
+      Ne(func_op.getBody().op_end<mlir::stablehlo::UniformQuantizeOp>()));
+
+  EXPECT_FALSE(IsOpFullyQuantized(*uniform_quantize_op_itr));
+}
+
+using IsOpNotQuantizedTest = QuantizationTestBase;
+
+TEST_F(IsOpNotQuantizedTest, TrueIfOpNotQuantized) {
+  constexpr absl::string_view kNotQuantizedAdd = R"mlir(
+    func.func @not_quantized_add(%arg0: tensor<2xf32>) -> tensor<2xf32> {
+      %0 = stablehlo.add %arg0, %arg0 : tensor<2xf32>
+      return %0 : tensor<2xf32>
+    }
+  )mlir";
+
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kNotQuantizedAdd);
+  ASSERT_TRUE(module_op);
+
+  auto func_op = module_op->lookupSymbol<func::FuncOp>("not_quantized_add");
+  ASSERT_THAT(func_op, NotNull());
+
+  auto add_op_itr = func_op.getBody().op_begin<mlir::stablehlo::AddOp>();
+  ASSERT_THAT(add_op_itr,
+              Ne(func_op.getBody().op_end<mlir::stablehlo::AddOp>()));
+
+  EXPECT_TRUE(IsOpNotQuantized(*add_op_itr));
+}
+
+TEST_F(IsOpNotQuantizedTest, FalseIfOpQuantized) {
+  constexpr absl::string_view kQuantizedAdd = R"mlir(
+    func.func @quantized_add(%arg0: tensor<2x!quant.uniform<i8:f32, 1.000000e+00:0>>) -> tensor<2x!quant.uniform<i8:f32, 1.000000e+00:0>> {
+      %0 = stablehlo.add %arg0, %arg0 : tensor<2x!quant.uniform<i8:f32, 1.000000e+00:0>>
+      return %0 : tensor<2x!quant.uniform<i8:f32, 1.000000e+00:0>>
+    }
+  )mlir";
+
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kQuantizedAdd);
+  ASSERT_TRUE(module_op);
+
+  auto func_op = module_op->lookupSymbol<func::FuncOp>("quantized_add");
+  ASSERT_THAT(func_op, NotNull());
+
+  auto add_op_itr = func_op.getBody().op_begin<mlir::stablehlo::AddOp>();
+  ASSERT_THAT(add_op_itr,
+              Ne(func_op.getBody().op_end<mlir::stablehlo::AddOp>()));
+
+  EXPECT_FALSE(IsOpNotQuantized(*add_op_itr));
+}
+
+TEST_F(IsOpNotQuantizedTest, FalseIfOpPartiallyQuantized) {
+  constexpr absl::string_view kQuantizeOp = R"mlir(
+    func.func @quantize(%arg0: tensor<2xf32>) -> tensor<2x!quant.uniform<i8:f32, 1.000000e+00:0>> {
+      %0 = stablehlo.uniform_quantize %arg0 : (tensor<2xf32>) -> tensor<2x!quant.uniform<i8:f32, 1.000000e+00:0>>
+      return %0 : tensor<2x!quant.uniform<i8:f32, 1.000000e+00:0>>
+    }
+  )mlir";
+
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kQuantizeOp);
+  ASSERT_TRUE(module_op);
+
+  auto func_op = module_op->lookupSymbol<func::FuncOp>("quantize");
+  ASSERT_THAT(func_op, NotNull());
+
+  auto uniform_quantize_op_itr =
+      func_op.getBody().op_begin<mlir::stablehlo::UniformQuantizeOp>();
+  ASSERT_THAT(
+      uniform_quantize_op_itr,
+      Ne(func_op.getBody().op_end<mlir::stablehlo::UniformQuantizeOp>()));
+
+  // `uniform_quantize` is considered partially quantized because its output is
+  // a quantized tensor whereas its input is not quantized.
+  EXPECT_FALSE(IsOpNotQuantized(*uniform_quantize_op_itr));
+}
+
+using UniformQuantizedTypeTest = QuantizationTestBase;
+
+TEST_F(UniformQuantizedTypeTest, GetElementTypeSucceeds) {
+  constexpr absl::string_view kQuantizeOp = R"mlir(
+    func.func @quantize(%arg0: tensor<2xf32>) -> tensor<2x!quant.uniform<i8:f32, 1.000000e+00:0>> {
+      %0 = stablehlo.uniform_quantize %arg0 : (tensor<2xf32>) -> tensor<2x!quant.uniform<i8:f32, 1.000000e+00:0>>
+      return %0 : tensor<2x!quant.uniform<i8:f32, 1.000000e+00:0>>
+    }
+  )mlir";
+
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kQuantizeOp);
+  ASSERT_TRUE(module_op);
+
+  auto func_op = module_op->lookupSymbol<func::FuncOp>("quantize");
+  ASSERT_THAT(func_op, NotNull());
+
+  auto uniform_quantize_op =
+      *func_op.getOps<::mlir::stablehlo::UniformQuantizeOp>().begin();
+  Value result = uniform_quantize_op.getResult();
+  EXPECT_THAT(GetElementType(result), NotNull());
+}
+
+}  // namespace
+}  // namespace quant
+}  // namespace mlir