chore: import upstream snapshot with attribution

tensorflow/compiler/mlir/quantization/stablehlo/utils/bfloat16_type.cc

@@ -0,0 +1,43 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/quantization/stablehlo/utils/bfloat16_type.h"
+
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace mlir::quant::stablehlo {
+
+bool IsLargeFloatType(Type type) {
+  type = getElementTypeOrSelf(type);
+  return isa<FloatType>(type) && type.getIntOrFloatBitWidth() > 16;
+}
+
+Type ToBfloat16Type(Type type) {
+  if (auto shaped = mlir::dyn_cast<ShapedType>(type)) {
+    const Type elem = shaped.getElementType();
+    if (IsLargeFloatType(elem)) {
+      return shaped.clone(BFloat16Type::get(type.getContext()));
+    }
+  } else if (IsLargeFloatType(type)) {
+    return BFloat16Type::get(type.getContext());
+  }
+  return type;
+}
+
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/utils/bfloat16_type.h

@@ -0,0 +1,32 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_UTILS_BFLOAT16_TYPE_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_UTILS_BFLOAT16_TYPE_H_
+
+#include "mlir/IR/Types.h"  // from @llvm-project
+
+namespace mlir::quant::stablehlo {
+
+// Returns true if the type or its element type is a float type with bit_width
+// > 16.
+bool IsLargeFloatType(Type type);
+
+// Converts large float type to bfloat16. Otherwise returns original type.
+Type ToBfloat16Type(Type type);
+
+}  // namespace mlir::quant::stablehlo
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_UTILS_BFLOAT16_TYPE_H_

tensorflow/compiler/mlir/quantization/stablehlo/utils/bfloat16_type_test.cc

@@ -0,0 +1,168 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/compiler/mlir/quantization/stablehlo/utils/bfloat16_type.h"
+
+#include <memory>
+
+#include <gtest/gtest.h>
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/register_common_dialects.h"
+
+namespace mlir::quant::stablehlo {
+namespace {
+
+std::unique_ptr<MLIRContext> CreateContext() {
+  auto context = std::make_unique<MLIRContext>();
+  DialectRegistry mlir_registry;
+  RegisterCommonToolingDialects(mlir_registry);
+  context->appendDialectRegistry(mlir_registry);
+  return context;
+}
+
+TEST(IsLargeFloatTypeTest, scalars) {
+  auto context = CreateContext();
+
+  EXPECT_FALSE(IsLargeFloatType(Float8E4M3FNType::get(context.get())));
+  EXPECT_FALSE(IsLargeFloatType(Float8E4M3FNUZType::get(context.get())));
+  EXPECT_FALSE(IsLargeFloatType(Float8E4M3B11FNUZType::get(context.get())));
+  EXPECT_FALSE(IsLargeFloatType(Float8E5M2FNUZType::get(context.get())));
+  EXPECT_FALSE(IsLargeFloatType(Float8E5M2Type::get(context.get())));
+  EXPECT_FALSE(IsLargeFloatType(Float16Type::get(context.get())));
+  EXPECT_FALSE(IsLargeFloatType(BFloat16Type::get(context.get())));
+  EXPECT_TRUE(IsLargeFloatType(Float32Type::get(context.get())));
+  EXPECT_TRUE(IsLargeFloatType(Float64Type::get(context.get())));
+  EXPECT_TRUE(IsLargeFloatType(Float80Type::get(context.get())));
+
+  EXPECT_FALSE(IsLargeFloatType(IntegerType::get(context.get(), 8)));
+  EXPECT_FALSE(IsLargeFloatType(IntegerType::get(context.get(), 16)));
+  EXPECT_FALSE(IsLargeFloatType(IntegerType::get(context.get(), 32)));
+}
+
+TEST(IsLargeFloatTypeTest, tensors) {
+  auto context = CreateContext();
+
+  EXPECT_FALSE(IsLargeFloatType(
+      RankedTensorType::get({2, 2}, Float8E4M3FNType::get(context.get()))));
+  EXPECT_FALSE(IsLargeFloatType(
+      RankedTensorType::get({2, 2}, Float16Type::get(context.get()))));
+  EXPECT_FALSE(IsLargeFloatType(
+      RankedTensorType::get({2, 2}, Float8E4M3FNUZType::get(context.get()))));
+  EXPECT_FALSE(IsLargeFloatType(RankedTensorType::get(
+      {2, 2}, Float8E4M3B11FNUZType::get(context.get()))));
+  EXPECT_FALSE(IsLargeFloatType(
+      RankedTensorType::get({2, 2}, Float8E5M2FNUZType::get(context.get()))));
+  EXPECT_FALSE(IsLargeFloatType(
+      RankedTensorType::get({2, 2}, Float8E5M2Type::get(context.get()))));
+  EXPECT_FALSE(IsLargeFloatType(
+      RankedTensorType::get({2, 2}, BFloat16Type::get(context.get()))));
+  EXPECT_TRUE(IsLargeFloatType(
+      RankedTensorType::get({2, 2}, Float32Type::get(context.get()))));
+  EXPECT_TRUE(IsLargeFloatType(
+      RankedTensorType::get({2, 2}, Float64Type::get(context.get()))));
+  EXPECT_TRUE(IsLargeFloatType(
+      RankedTensorType::get({2, 2}, Float80Type::get(context.get()))));
+
+  EXPECT_FALSE(IsLargeFloatType(
+      RankedTensorType::get({2, 2}, IntegerType::get(context.get(), 8))));
+  EXPECT_FALSE(IsLargeFloatType(
+      RankedTensorType::get({2, 2}, IntegerType::get(context.get(), 16))));
+  EXPECT_FALSE(IsLargeFloatType(
+      RankedTensorType::get({2, 2}, IntegerType::get(context.get(), 32))));
+}
+
+TEST(ToBfloat16TypeTest, scalars) {
+  auto context = CreateContext();
+
+  EXPECT_EQ(ToBfloat16Type(Float8E4M3FNType::get(context.get())),
+            Float8E4M3FNType::get(context.get()));
+  EXPECT_EQ(ToBfloat16Type(Float8E4M3FNUZType::get(context.get())),
+            Float8E4M3FNUZType::get(context.get()));
+  EXPECT_EQ(ToBfloat16Type(Float8E4M3B11FNUZType::get(context.get())),
+            Float8E4M3B11FNUZType::get(context.get()));
+  EXPECT_EQ(ToBfloat16Type(Float8E5M2FNUZType::get(context.get())),
+            Float8E5M2FNUZType::get(context.get()));
+  EXPECT_EQ(ToBfloat16Type(Float8E5M2Type::get(context.get())),
+            Float8E5M2Type::get(context.get()));
+  EXPECT_EQ(ToBfloat16Type(Float16Type::get(context.get())),
+            Float16Type::get(context.get()));
+  EXPECT_EQ(ToBfloat16Type(BFloat16Type::get(context.get())),
+            BFloat16Type::get(context.get()));
+  EXPECT_EQ(ToBfloat16Type(Float32Type::get(context.get())),
+            BFloat16Type::get(context.get()));
+  EXPECT_EQ(ToBfloat16Type(Float64Type::get(context.get())),
+            BFloat16Type::get(context.get()));
+  EXPECT_EQ(ToBfloat16Type(Float80Type::get(context.get())),
+            BFloat16Type::get(context.get()));
+
+  EXPECT_EQ(ToBfloat16Type(IntegerType::get(context.get(), 8)),
+            IntegerType::get(context.get(), 8));
+  EXPECT_EQ(ToBfloat16Type(IntegerType::get(context.get(), 16)),
+            IntegerType::get(context.get(), 16));
+  EXPECT_EQ(ToBfloat16Type(IntegerType::get(context.get(), 32)),
+            IntegerType::get(context.get(), 32));
+}
+
+TEST(ToBfloat16TypeTest, tensors) {
+  auto context = CreateContext();
+
+  EXPECT_EQ(
+      ToBfloat16Type(
+          RankedTensorType::get({2, 2}, Float8E4M3FNType::get(context.get()))),
+      RankedTensorType::get({2, 2}, Float8E4M3FNType::get(context.get())));
+  EXPECT_EQ(
+      ToBfloat16Type(RankedTensorType::get(
+          {2, 2}, Float8E4M3FNUZType::get(context.get()))),
+      RankedTensorType::get({2, 2}, Float8E4M3FNUZType::get(context.get())));
+  EXPECT_EQ(
+      ToBfloat16Type(RankedTensorType::get(
+          {2, 2}, Float8E4M3B11FNUZType::get(context.get()))),
+      RankedTensorType::get({2, 2}, Float8E4M3B11FNUZType::get(context.get())));
+  EXPECT_EQ(
+      ToBfloat16Type(RankedTensorType::get(
+          {2, 2}, Float8E5M2FNUZType::get(context.get()))),
+      RankedTensorType::get({2, 2}, Float8E5M2FNUZType::get(context.get())));
+  EXPECT_EQ(ToBfloat16Type(RankedTensorType::get(
+                {2, 2}, Float8E5M2Type::get(context.get()))),
+            RankedTensorType::get({2, 2}, Float8E5M2Type::get(context.get())));
+  EXPECT_EQ(ToBfloat16Type(
+                RankedTensorType::get({2, 2}, Float16Type::get(context.get()))),
+            RankedTensorType::get({2, 2}, Float16Type::get(context.get())));
+  EXPECT_EQ(ToBfloat16Type(RankedTensorType::get(
+                {2, 2}, BFloat16Type::get(context.get()))),
+            RankedTensorType::get({2, 2}, BFloat16Type::get(context.get())));
+  EXPECT_EQ(ToBfloat16Type(
+                RankedTensorType::get({2, 2}, Float32Type::get(context.get()))),
+            RankedTensorType::get({2, 2}, BFloat16Type::get(context.get())));
+  EXPECT_EQ(ToBfloat16Type(
+                RankedTensorType::get({2, 2}, Float64Type::get(context.get()))),
+            RankedTensorType::get({2, 2}, BFloat16Type::get(context.get())));
+  EXPECT_EQ(ToBfloat16Type(
+                RankedTensorType::get({2, 2}, Float80Type::get(context.get()))),
+            RankedTensorType::get({2, 2}, BFloat16Type::get(context.get())));
+
+  EXPECT_EQ(ToBfloat16Type(RankedTensorType::get(
+                {2, 2}, IntegerType::get(context.get(), 8))),
+            RankedTensorType::get({2, 2}, IntegerType::get(context.get(), 8)));
+  EXPECT_EQ(ToBfloat16Type(RankedTensorType::get(
+                {2, 2}, IntegerType::get(context.get(), 16))),
+            RankedTensorType::get({2, 2}, IntegerType::get(context.get(), 16)));
+  EXPECT_EQ(ToBfloat16Type(RankedTensorType::get(
+                {2, 2}, IntegerType::get(context.get(), 32))),
+            RankedTensorType::get({2, 2}, IntegerType::get(context.get(), 32)));
+}
+
+}  // namespace
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/utils/fill_quantization_options.cc

@@ -0,0 +1,134 @@
+/* Copyright 2023 The StableHLO Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_options.pb.h"
+
+namespace mlir::quant::stablehlo {
+
+using ::stablehlo::quantization::CustomQuantizationMethod;
+using ::stablehlo::quantization::PresetQuantizationMethod;
+using ::stablehlo::quantization::QuantizationComponentSpec;
+using ::stablehlo::quantization::QuantizationOptions;
+using QuantizationComponent =
+    ::stablehlo::quantization::QuantizationComponentSpec_QuantizationComponent;
+using BitType = ::stablehlo::quantization::QuantizationComponentSpec_BitType;
+using BitWidth = ::stablehlo::quantization::QuantizationComponentSpec_BitWidth;
+
+// Sets component, bit type and bit width information to the given spec
+// instance.
+void SetQuantizationComponentSpec(QuantizationComponentSpec* spec,
+                                  const QuantizationComponent& component,
+                                  const BitType bit_type,
+                                  const BitWidth bit_width) {
+  spec->set_quantization_component(component);
+  spec->set_bit_type(bit_type);
+  spec->set_bit_width(bit_width);
+}
+
+::stablehlo::quantization::QuantizationOptions FillPresetQuantizationOptions(
+    ::stablehlo::quantization::QuantizationOptions quantization_options_) {
+  CustomQuantizationMethod custom_method =
+      quantization_options_.quantization_method().custom_quantization_method();
+  QuantizationComponentSpec *activation_component, *weight_component,
+      *bias_component;
+  const auto preset_method = quantization_options_.quantization_method()
+                                 .preset_quantization_method()
+                                 .preset_method();
+  if (!preset_method) return quantization_options_;
+  switch (preset_method) {
+    case PresetQuantizationMethod::FLOAT16:
+      weight_component = custom_method.add_quantization_component_spec();
+      SetQuantizationComponentSpec(weight_component,
+                                   QuantizationComponentSpec::COMPONENT_WEIGHT,
+                                   QuantizationComponentSpec::BIT_TYPE_FLOAT,
+                                   QuantizationComponentSpec::BIT_WIDTH_16);
+      bias_component = custom_method.add_quantization_component_spec();
+      SetQuantizationComponentSpec(bias_component,
+                                   QuantizationComponentSpec::COMPONENT_BIAS,
+                                   QuantizationComponentSpec::BIT_TYPE_FLOAT,
+                                   QuantizationComponentSpec::BIT_WIDTH_16);
+      break;
+    // Note: This is weight-only quantization by default, but with the legacy
+    // flag "--force_dynamic_range_in_kernel", a DRQ behavior will be forced
+    // in the kernel.
+    case PresetQuantizationMethod::WEIGHT_ONLY:
+      weight_component = custom_method.add_quantization_component_spec();
+      SetQuantizationComponentSpec(weight_component,
+                                   QuantizationComponentSpec::COMPONENT_WEIGHT,
+                                   QuantizationComponentSpec::BIT_TYPE_INT,
+                                   QuantizationComponentSpec::BIT_WIDTH_8);
+      break;
+    case PresetQuantizationMethod::POST_TRAINING_QUANTIZATION_STATIC_RANGE:
+      activation_component = custom_method.add_quantization_component_spec();
+      SetQuantizationComponentSpec(
+          activation_component, QuantizationComponentSpec::COMPONENT_ACTIVATION,
+          QuantizationComponentSpec::BIT_TYPE_INT,
+          QuantizationComponentSpec::BIT_WIDTH_8);
+      weight_component = custom_method.add_quantization_component_spec();
+      SetQuantizationComponentSpec(weight_component,
+                                   QuantizationComponentSpec::COMPONENT_WEIGHT,
+                                   QuantizationComponentSpec::BIT_TYPE_INT,
+                                   QuantizationComponentSpec::BIT_WIDTH_8);
+      bias_component = custom_method.add_quantization_component_spec();
+      SetQuantizationComponentSpec(bias_component,
+                                   QuantizationComponentSpec::COMPONENT_BIAS,
+                                   QuantizationComponentSpec::BIT_TYPE_INT,
+                                   QuantizationComponentSpec::BIT_WIDTH_32);
+      break;
+    default:
+      break;
+  }
+  *quantization_options_.mutable_quantization_method()
+       ->mutable_custom_quantization_method() = custom_method;
+  return quantization_options_;
+}
+
+LogicalResult GetActivationBitWidth(QuantizationOptions quantization_options,
+                                    int* bit_width) {
+  CustomQuantizationMethod custom_method =
+      quantization_options.quantization_method().custom_quantization_method();
+
+  // TODO: b/288046643 - Look up bit width for each op/op instance instead of
+  // global configuration per component.
+  for (const auto& component : custom_method.quantization_component_spec()) {
+    if (component.quantization_component() ==
+        QuantizationComponentSpec::COMPONENT_ACTIVATION) {
+      switch (component.bit_width()) {
+        case QuantizationComponentSpec::BIT_WIDTH_4:
+          *bit_width = 4;
+          return success();
+          break;
+        case QuantizationComponentSpec::BIT_WIDTH_8:
+          *bit_width = 8;
+          return success();
+          break;
+        case QuantizationComponentSpec::BIT_WIDTH_16:
+          *bit_width = 16;
+          return success();
+          break;
+        case QuantizationComponentSpec::BIT_WIDTH_32:
+          *bit_width = 32;
+          return success();
+          break;
+        default:
+          break;
+      }
+    }
+  }
+  return failure();
+}
+
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/utils/fill_quantization_options.h

@@ -0,0 +1,41 @@
+/* Copyright 2023 The StableHLO Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_UTILS_FILL_QUANTIZATION_OPTIONS_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_UTILS_FILL_QUANTIZATION_OPTIONS_H_
+
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_options.pb.h"
+
+namespace mlir::quant::stablehlo {
+
+using ::stablehlo::quantization::QuantizationOptions;
+
+// Returns QuantizationOptions filled with detailed specs when user specifies
+// an optional preset method name. The preset methods are defined in
+// quantization_options.proto. This function will only be executed if a user
+// gives a preset method, not a custom method.
+QuantizationOptions FillPresetQuantizationOptions(
+    QuantizationOptions quantization_options);
+
+// Returns LogicalResult depending on the look up of activation bit width in the
+// custom quantization method. If such information exists, returns success,
+// otherwise, returns false.
+LogicalResult GetActivationBitWidth(QuantizationOptions quantization_options,
+                                    int* bit_width);
+
+}  // namespace mlir::quant::stablehlo
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_UTILS_FILL_QUANTIZATION_OPTIONS_H_

tensorflow/compiler/mlir/quantization/stablehlo/utils/math_utils.cc

@@ -0,0 +1,54 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/quantization/stablehlo/utils/math_utils.h"
+
+#include <cmath>
+#include <cstdint>
+
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+
+namespace mlir::quant::stablehlo {
+
+// Borrowed from:
+// https://github.com/tensorflow/tensorflow/blob/57946ceb4b6119d6d0f49abbb2e3d1636a3b83a0/tensorflow/lite/kernels/internal/quantization_util.cc#L53
+// https://github.com/tensorflow/tensorflow/blob/f5c9cbb1c462912030bd845244118f952cbbbd5e/tensorflow/core/kernels/uniform_quant_ops/math_utils.cc#L29
+// And then modified so that it doesn't overflow int32 values.
+LogicalResult QuantizeMultiplier(double double_multiplier,
+                                 int32_t& quantized_fraction, int32_t& shift) {
+  if (!std::isfinite(double_multiplier) || double_multiplier <= 0) {
+    return failure();
+  }
+  const double fraction = std::frexp(double_multiplier, &shift);
+  quantized_fraction = static_cast<int32_t>(std::round(fraction * (1L << 15)));
+  // Clip extreme values.  These are more than enough to overflow int8, the
+  // storage type for quantized values, and the final values will be clamped
+  // no matter what.
+  if (quantized_fraction == (1L << 15)) {
+    quantized_fraction /= 2;
+    ++shift;
+  }
+  if (shift < -15) {
+    shift = 0;
+    quantized_fraction = 0;
+  }
+  if (shift > 14) {
+    shift = 14;
+    quantized_fraction = (1LL << 15) - 1;
+  }
+  return success();
+}
+
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/utils/math_utils.h

@@ -0,0 +1,32 @@
+#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_UTILS_MATH_UTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_UTILS_MATH_UTILS_H_
+
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cstdint>
+
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+
+namespace mlir::quant::stablehlo {
+
+// Decomposes a given floating point value num into a normalized and quantized
+// fraction and an integral power of two.
+LogicalResult QuantizeMultiplier(double double_multiplier,
+                                 int32_t& quantized_fraction, int32_t& shift);
+
+}  // namespace mlir::quant::stablehlo
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_UTILS_MATH_UTILS_H_

tensorflow/compiler/mlir/quantization/stablehlo/utils/math_utils_test.cc

@@ -0,0 +1,65 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/quantization/stablehlo/utils/math_utils.h"
+
+#include <cstdint>
+
+#include <gtest/gtest.h>
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+
+namespace mlir::quant::stablehlo {
+namespace {
+
+TEST(UtilsTest, QuantizeMultiplierNormalMultipliers) {
+  int32_t quantized_fraction;
+  int32_t shift;
+
+  EXPECT_TRUE(succeeded(QuantizeMultiplier(1.2, quantized_fraction, shift)));
+  EXPECT_EQ(quantized_fraction, 19661);
+  EXPECT_EQ(shift, 1);
+
+  EXPECT_TRUE(succeeded(QuantizeMultiplier(15.5, quantized_fraction, shift)));
+  EXPECT_EQ(quantized_fraction, 31744);
+  EXPECT_EQ(shift, 4);
+
+  EXPECT_TRUE(succeeded(QuantizeMultiplier(1, quantized_fraction, shift)));
+  EXPECT_EQ(quantized_fraction, 16384);
+  EXPECT_EQ(shift, 1);
+}
+
+TEST(UtilsTest, QuantizeMultiplierExtremeMultipliers) {
+  int32_t quantized_fraction;
+  int32_t shift;
+
+  EXPECT_TRUE(
+      succeeded(QuantizeMultiplier(0.00001f, quantized_fraction, shift)));
+  EXPECT_EQ(quantized_fraction, 0);
+  EXPECT_EQ(shift, 0);
+
+  EXPECT_TRUE(succeeded(QuantizeMultiplier(40000, quantized_fraction, shift)));
+  EXPECT_EQ(quantized_fraction, 32767);
+  EXPECT_EQ(shift, 14);
+}
+
+TEST(UtilsTest, QuantizeMultiplierInvalidArgument) {
+  int32_t quantized_fraction;
+  int32_t shift;
+
+  EXPECT_FALSE(succeeded(QuantizeMultiplier(0, quantized_fraction, shift)));
+}
+
+}  // namespace
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/utils/stablehlo_type_utils.h

@@ -0,0 +1,34 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_UTILS_STABLEHLO_TYPE_UTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_UTILS_STABLEHLO_TYPE_UTILS_H_
+
+#include "llvm/ADT/StringRef.h"
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo
+
+namespace mlir::quant::stablehlo {
+
+// Checks if an op is from StableHLO dialect.
+inline bool IsStablehloOp(Operation* op) {
+  return op->getDialect()->getNamespace() ==
+         mlir::stablehlo::StablehloDialect::getDialectNamespace();
+}
+
+}  // namespace mlir::quant::stablehlo
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_UTILS_STABLEHLO_TYPE_UTILS_H_

tensorflow/compiler/mlir/quantization/stablehlo/utils/stablehlo_type_utils_test.cc

@@ -0,0 +1,57 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/quantization/stablehlo/utils/stablehlo_type_utils.h"
+
+#include <gtest/gtest.h>
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo
+
+namespace mlir::quant::stablehlo {
+namespace {
+
+using ::testing::Test;
+
+class StablehloTypeUtilsTest : public Test {
+ protected:
+  StablehloTypeUtilsTest() {
+    ctx_.loadDialect<mlir::stablehlo::StablehloDialect,
+                     mlir::arith::ArithDialect, mlir::func::FuncDialect>();
+  }
+
+  MLIRContext ctx_;
+  OpBuilder builder_{&ctx_};
+};
+
+TEST_F(StablehloTypeUtilsTest, IsStablehloOpSucceedsWithStablehloOp) {
+  const OwningOpRef<mlir::stablehlo::ConstantOp> constant_op =
+      mlir::stablehlo::ConstantOp::create(builder_, builder_.getUnknownLoc(),
+                                          builder_.getI32IntegerAttr(0));
+  EXPECT_TRUE(IsStablehloOp(*constant_op));
+}
+
+TEST_F(StablehloTypeUtilsTest, IsStablehloOpFailsWithArithOp) {
+  const OwningOpRef<mlir::arith::ConstantOp> constant_op =
+      mlir::arith::ConstantOp::create(builder_, builder_.getUnknownLoc(),
+                                      builder_.getI32IntegerAttr(0));
+  EXPECT_FALSE(IsStablehloOp(*constant_op));
+}
+
+}  // namespace
+}  // namespace mlir::quant::stablehlo

tensorflow/compiler/mlir/quantization/stablehlo/utils/tf_type_utils.cc

@@ -0,0 +1,110 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/quantization/stablehlo/utils/tf_type_utils.h"
+
+#include "absl/status/status.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/mangling_util.h"
+#include "tensorflow/core/framework/numeric_types.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace mlir::quant::tensorflow {
+
+bool IsTFQintType(const Type type) {
+  return mlir::isa<TF::Qint8Type, TF::Qint16Type, TF::Qint32Type,
+                   TF::Quint8Type, TF::Quint16Type>(type);
+}
+
+Type GetIntTypeFromTFQint(const Type type) {
+  return TypeSwitch<Type, Type>(type)
+      .Case<TF::Qint8Type>(
+          [&type](Type) { return IntegerType::get(type.getContext(), 8); })
+      .Case<TF::Qint16Type>(
+          [&type](Type) { return IntegerType::get(type.getContext(), 16); })
+      .Case<TF::Qint32Type>(
+          [&type](Type) { return IntegerType::get(type.getContext(), 32); })
+      .Case<TF::Quint8Type>([&type](Type) {
+        return IntegerType::get(type.getContext(), 8,
+                                IntegerType::SignednessSemantics::Unsigned);
+      })
+      .Case<TF::Quint16Type>([&type](Type) {
+        return IntegerType::get(type.getContext(), 16,
+                                IntegerType::SignednessSemantics::Unsigned);
+      })
+      .Default([&type](Type) { return type; });
+}
+
+FailureOr<mlir::DenseElementsAttr> GetDenseAttrFromTensorProtoAttr(
+    const llvm::StringRef mangled_tensor_proto, TensorType tensor_type) {
+  ::tensorflow::TensorProto tensor_proto;
+  absl::Status status = ::tensorflow::mangling_util::DemangleTensor(
+      mangled_tensor_proto, &tensor_proto);
+  if (!status.ok()) {
+    return failure();
+  }
+
+  ::tensorflow::Tensor t;
+  if (!t.FromProto(tensor_proto)) {
+    return failure();
+  }
+
+  if (t.dtype() == ::tensorflow::DT_QINT8) {
+    const auto arr = t.flat<::tensorflow::qint8>();
+    return mlir::DenseElementsAttr::get(
+        tensor_type.clone(IntegerType::get(tensor_type.getContext(), 8)),
+        llvm::ArrayRef(arr.data(), arr.size()));
+  } else if (t.dtype() == ::tensorflow::DT_QINT32) {
+    const auto arr = t.flat<::tensorflow::qint32>();
+    return mlir::DenseElementsAttr::get(
+        tensor_type.clone(IntegerType::get(tensor_type.getContext(), 32)),
+        llvm::ArrayRef(arr.data(), arr.size()));
+  } else {
+    return failure();
+  }
+}
+
+bool IsTFUniformQuantizedOp(Operation *op) {
+  return llvm::isa<
+      // clang-format off
+      // go/keep-sorted start
+      TF::UniformDequantizeOp,
+      TF::UniformQuantizeOp,
+      TF::UniformQuantizedAddOp,
+      TF::UniformQuantizedClipByValueOp,
+      TF::UniformQuantizedConvolutionHybridOp,
+      TF::UniformQuantizedConvolutionOp,
+      TF::UniformQuantizedDotHybridOp,
+      TF::UniformQuantizedDotOp,
+      TF::UniformRequantizeOp
+      // go/keep-sorted end
+      // clang-format on
+      >(op);
+}
+
+}  // namespace mlir::quant::tensorflow

tensorflow/compiler/mlir/quantization/stablehlo/utils/tf_type_utils.h

@@ -0,0 +1,43 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_UTILS_TF_TYPE_UTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_UTILS_TF_TYPE_UTILS_H_
+
+#include "llvm/ADT/StringRef.h"
+#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/LogicalResult.h"  // from @llvm-project
+
+namespace mlir::quant::tensorflow {
+
+// GetDenseAttrFromTensorProtoAttr returns DenseElementsAttr from tensor proto.
+FailureOr<mlir::DenseElementsAttr> GetDenseAttrFromTensorProtoAttr(
+    llvm::StringRef mangled_tensor_proto, TensorType result_tensor_type);
+
+// Check if a type is TF qint type.
+bool IsTFQintType(Type type);
+
+// Convert qint type to the corresponding int type. Return original type if it
+// is not qint type.
+Type GetIntTypeFromTFQint(Type type);
+
+// Check if an op is TF UniformQuantized op.
+bool IsTFUniformQuantizedOp(Operation* op);
+
+}  // namespace mlir::quant::tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_UTILS_TF_TYPE_UTILS_H_

tensorflow/compiler/mlir/quantization/stablehlo/utils/tf_type_utils_test.cc

@@ -0,0 +1,224 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/quantization/stablehlo/utils/tf_type_utils.h"
+
+#include <cstdint>
+#include <memory>
+#include <string>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Quant/IR/Quant.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/register_common_dialects.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/mangling_util.h"
+#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
+#include "xla/tsl/framework/numeric_types.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/ir/types/dialect.h"
+
+namespace mlir::quant::tensorflow {
+namespace {
+
+std::string GetQint8Tensor() {
+  ::tensorflow::Tensor tensor(::tensorflow::DT_QINT8, {2, 2});
+  tensor.matrix<tsl::qint8>()(0, 0) = tsl::qint8(1);
+  tensor.matrix<tsl::qint8>()(0, 1) = tsl::qint8(2);
+  tensor.matrix<tsl::qint8>()(1, 0) = tsl::qint8(3);
+  tensor.matrix<tsl::qint8>()(1, 1) = tsl::qint8(4);
+
+  ::tensorflow::TensorProto tensor_proto;
+  tensor.AsProtoTensorContent(&tensor_proto);
+  return ::tensorflow::mangling_util::MangleTensor(tensor_proto);
+}
+
+std::string GetQint16Tensor() {
+  ::tensorflow::Tensor tensor(::tensorflow::DT_QINT16, {2, 2});
+  tensor.matrix<tsl::qint16>()(0, 0) = tsl::qint16(1);
+  tensor.matrix<tsl::qint16>()(0, 1) = tsl::qint16(2);
+  tensor.matrix<tsl::qint16>()(1, 0) = tsl::qint16(3);
+  tensor.matrix<tsl::qint16>()(1, 1) = tsl::qint16(4);
+
+  ::tensorflow::TensorProto tensor_proto;
+  tensor.AsProtoTensorContent(&tensor_proto);
+  return ::tensorflow::mangling_util::MangleTensor(tensor_proto);
+}
+
+std::string GetQint32Tensor() {
+  ::tensorflow::Tensor tensor(::tensorflow::DT_QINT32, {2, 2});
+  tensor.matrix<tsl::qint32>()(0, 0) = tsl::qint32(1);
+  tensor.matrix<tsl::qint32>()(0, 1) = tsl::qint32(2);
+  tensor.matrix<tsl::qint32>()(1, 0) = tsl::qint32(3);
+  tensor.matrix<tsl::qint32>()(1, 1) = tsl::qint32(4);
+
+  ::tensorflow::TensorProto tensor_proto;
+  tensor.AsProtoTensorContent(&tensor_proto);
+  return ::tensorflow::mangling_util::MangleTensor(tensor_proto);
+}
+
+std::unique_ptr<MLIRContext> CreateContext() {
+  auto context = std::make_unique<MLIRContext>();
+  DialectRegistry mlir_registry;
+  RegisterCommonToolingDialects(mlir_registry);
+  context->appendDialectRegistry(mlir_registry);
+  context->getOrLoadDialect<tf_type::TFTypeDialect>();
+  context->getOrLoadDialect<quant::QuantDialect>();
+  context->getOrLoadDialect<mlir::mhlo::MhloDialect>();
+  context->getOrLoadDialect<sparse_tensor::SparseTensorDialect>();
+  return context;
+}
+
+TEST(GetDenseAttrFromTensorProtoAttrTest, Qint8ToUQ8Succeeds) {
+  auto context = CreateContext();
+  TensorType result_tensor_type = RankedTensorType::get(
+      {2, 2}, quant::UniformQuantizedType::get(
+                  quant::QuantizationFlags::FlagValue::Signed,
+                  IntegerType::get(context.get(), 8),
+                  Float32Type::get(context.get()), 3.0, 2, -128, 127));
+
+  auto dense_attr =
+      GetDenseAttrFromTensorProtoAttr(GetQint8Tensor(), result_tensor_type);
+
+  ASSERT_TRUE(succeeded(dense_attr));
+  EXPECT_THAT(dense_attr->getValues<int8_t>(), testing::SizeIs(4));
+  EXPECT_EQ(dense_attr->getValues<int8_t>()[0], 1);
+  EXPECT_EQ(dense_attr->getValues<int8_t>()[1], 2);
+  EXPECT_EQ(dense_attr->getValues<int8_t>()[2], 3);
+  EXPECT_EQ(dense_attr->getValues<int8_t>()[3], 4);
+}
+
+TEST(GetDenseAttrFromTensorProtoAttrTest, Qint8ToInt8Succeeds) {
+  auto context = CreateContext();
+  TensorType result_tensor_type =
+      RankedTensorType::get({2, 2}, IntegerType::get(context.get(), 8));
+
+  auto dense_attr =
+      GetDenseAttrFromTensorProtoAttr(GetQint8Tensor(), result_tensor_type);
+
+  ASSERT_TRUE(succeeded(dense_attr));
+  EXPECT_THAT(dense_attr->getValues<int8_t>(), testing::SizeIs(4));
+  EXPECT_EQ(dense_attr->getValues<int8_t>()[0], 1);
+  EXPECT_EQ(dense_attr->getValues<int8_t>()[1], 2);
+  EXPECT_EQ(dense_attr->getValues<int8_t>()[2], 3);
+  EXPECT_EQ(dense_attr->getValues<int8_t>()[3], 4);
+}
+
+TEST(GetDenseAttrFromTensorProtoAttrTest, Qint32ToUQ32Succeeds) {
+  auto context = CreateContext();
+  TensorType result_tensor_type = RankedTensorType::get(
+      {2, 2},
+      quant::UniformQuantizedType::get(
+          quant::QuantizationFlags::FlagValue::Signed,
+          IntegerType::get(context.get(), 32), Float32Type::get(context.get()),
+          3.0, 2, -2147483648, 2147483647));
+
+  auto dense_attr =
+      GetDenseAttrFromTensorProtoAttr(GetQint32Tensor(), result_tensor_type);
+
+  ASSERT_TRUE(succeeded(dense_attr));
+  EXPECT_THAT(dense_attr->getValues<int32_t>(), testing::SizeIs(4));
+  EXPECT_EQ(dense_attr->getValues<int32_t>()[0], 1);
+  EXPECT_EQ(dense_attr->getValues<int32_t>()[1], 2);
+  EXPECT_EQ(dense_attr->getValues<int32_t>()[2], 3);
+  EXPECT_EQ(dense_attr->getValues<int32_t>()[3], 4);
+}
+
+TEST(GetDenseAttrFromTensorProtoAttrTest, Qint32ToInt32Succeeds) {
+  auto context = CreateContext();
+  TensorType result_tensor_type =
+      RankedTensorType::get({2, 2}, IntegerType::get(context.get(), 32));
+
+  auto dense_attr =
+      GetDenseAttrFromTensorProtoAttr(GetQint32Tensor(), result_tensor_type);
+
+  ASSERT_TRUE(succeeded(dense_attr));
+  EXPECT_THAT(dense_attr->getValues<int32_t>(), testing::SizeIs(4));
+  EXPECT_EQ(dense_attr->getValues<int32_t>()[0], 1);
+  EXPECT_EQ(dense_attr->getValues<int32_t>()[1], 2);
+  EXPECT_EQ(dense_attr->getValues<int32_t>()[2], 3);
+  EXPECT_EQ(dense_attr->getValues<int32_t>()[3], 4);
+}
+
+TEST(GetDenseAttrFromTensorProtoAttrTest, UnsupportedQint16Fails) {
+  auto context = CreateContext();
+  TensorType result_tensor_type =
+      RankedTensorType::get({2, 2}, IntegerType::get(context.get(), 16));
+
+  EXPECT_TRUE(failed(
+      GetDenseAttrFromTensorProtoAttr(GetQint16Tensor(), result_tensor_type)));
+}
+
+TEST(IsTFQintTypeTest, ValidTFQintTypeSucceeds) {
+  auto context = CreateContext();
+
+  EXPECT_TRUE(IsTFQintType(TF::Qint8Type::get(context.get())));
+  EXPECT_TRUE(IsTFQintType(TF::Qint16Type::get(context.get())));
+  EXPECT_TRUE(IsTFQintType(TF::Qint32Type::get(context.get())));
+  EXPECT_TRUE(IsTFQintType(TF::Quint8Type::get(context.get())));
+  EXPECT_TRUE(IsTFQintType(TF::Quint16Type::get(context.get())));
+
+  EXPECT_FALSE(IsTFQintType(TF::Int8RefType::get(context.get())));
+  EXPECT_FALSE(IsTFQintType(TF::Float8E5M2RefType::get(context.get())));
+}
+
+TEST(GetIntTypeFromTFQintTest, ChecksIntTypesFromTFQint) {
+  auto context = CreateContext();
+
+  auto type = GetIntTypeFromTFQint(TF::Qint8Type::get(context.get()));
+  EXPECT_TRUE(llvm::isa<IntegerType>(type));
+  EXPECT_EQ(mlir::dyn_cast<IntegerType>(type).getWidth(), 8);
+  EXPECT_FALSE(mlir::dyn_cast<IntegerType>(type).isSigned());
+  EXPECT_FALSE(mlir::dyn_cast<IntegerType>(type).isUnsigned());
+
+  type = GetIntTypeFromTFQint(TF::Qint16Type::get(context.get()));
+  EXPECT_TRUE(llvm::isa<IntegerType>(type));
+  EXPECT_EQ(mlir::dyn_cast<IntegerType>(type).getWidth(), 16);
+  EXPECT_FALSE(mlir::dyn_cast<IntegerType>(type).isSigned());
+  EXPECT_FALSE(mlir::dyn_cast<IntegerType>(type).isUnsigned());
+
+  type = GetIntTypeFromTFQint(TF::Qint32Type::get(context.get()));
+  EXPECT_TRUE(llvm::isa<IntegerType>(type));
+  EXPECT_EQ(mlir::dyn_cast<IntegerType>(type).getWidth(), 32);
+  EXPECT_FALSE(mlir::dyn_cast<IntegerType>(type).isSigned());
+  EXPECT_FALSE(mlir::dyn_cast<IntegerType>(type).isUnsigned());
+
+  type = GetIntTypeFromTFQint(TF::Quint8Type::get(context.get()));
+  EXPECT_TRUE(llvm::isa<IntegerType>(type));
+  EXPECT_EQ(mlir::dyn_cast<IntegerType>(type).getWidth(), 8);
+  EXPECT_TRUE(mlir::dyn_cast<IntegerType>(type).isUnsigned());
+
+  type = GetIntTypeFromTFQint(TF::Quint16Type::get(context.get()));
+  EXPECT_TRUE(llvm::isa<IntegerType>(type));
+  EXPECT_EQ(mlir::dyn_cast<IntegerType>(type).getWidth(), 16);
+  EXPECT_TRUE(mlir::dyn_cast<IntegerType>(type).isUnsigned());
+
+  // Non qint types are returned as is.
+  EXPECT_EQ(GetIntTypeFromTFQint(IntegerType::get(type.getContext(), 32)),
+            IntegerType::get(type.getContext(), 32));
+}
+
+}  // namespace
+}  // namespace mlir::quant::tensorflow