chore: import upstream snapshot with attribution

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

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

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

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

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

@@ -0,0 +1,41 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_OPS_STABLEHLO_OP_QUANT_SPEC_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_OPS_STABLEHLO_OP_QUANT_SPEC_H_
+
+#include <memory>
+
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"
+#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
+
+namespace mlir::quant::stablehlo {
+
+// Returns StableHLO quantization specs for an op.
+std::unique_ptr<OpQuantSpec> GetStableHloOpQuantSpec(Operation* op);
+
+// Returns quantization constraints (ex: fixed output, same scale) given
+// a StableHLO op.
+std::unique_ptr<OpQuantScaleSpec> GetStableHloQuantConstraints(Operation* op);
+
+// Checks if an op is quantizable in StableHLO quantizer. Argument op is not
+// necessarily a StableHLO op.
+bool IsOpQuantizableStableHlo(Operation* op);
+
+}  // namespace mlir::quant::stablehlo
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_OPS_STABLEHLO_OP_QUANT_SPEC_H_

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

@@ -0,0 +1,332 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/quantization/stablehlo/ops/stablehlo_op_quant_spec.h"
+
+#include <memory>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/strings/string_view.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo
+#include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.h"
+#include "tensorflow/compiler/mlir/quantization/common/func.h"
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"
+#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"
+#include "tensorflow/compiler/mlir/quantization/common/test_base.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace mlir::quant::stablehlo {
+namespace {
+
+using ::mlir::stablehlo::GatherOp;
+using ::testing::IsEmpty;
+using ::testing::IsTrue;
+using ::testing::NotNull;
+using ::testing::Pair;
+using ::testing::UnorderedElementsAre;
+
+using IsOpQuantizableStableHloTest = ::mlir::quant::QuantizationTestBase;
+
+// Quantizable ops: constants
+// Non-quantizable ops: normal StableHLO ops and terminators
+constexpr absl::string_view kModuleConstantAdd = R"mlir(
+  module {
+    func.func @constant_add() -> (tensor<3x2xf32>) {
+      %cst1 = stablehlo.constant dense<2.4> : tensor<3x2xf32>
+      %cst2 = stablehlo.constant dense<5.7> : tensor<3x2xf32>
+      %add = stablehlo.add %cst1, %cst2 : (tensor<3x2xf32>, tensor<3x2xf32>) -> tensor<3x2xf32>
+      func.return %add : tensor<3x2xf32>
+    }
+  }
+)mlir";
+
+// Quantizable ops: XlaCallModule op with "fully_quantizable" attribute and
+// same-scale StableHLO ops
+// Non-quantizable ops: quantize/dequantize ops
+constexpr absl::string_view kModuleCompositeSameScale = R"mlir(
+  module {
+    func.func @same_scale_after_composite() -> tensor<3x1xf32> {
+      %0 = "tf.XlaCallModule"() {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn_1, _stablehlo_version = "1.0.0", _original_entry_function = "composite_dot_general_fn_1", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable",   device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : () -> tensor<1x3xf32>
+      %1 = "quantization.qcast"(%0) {volatile} : (tensor<1x3xf32>) -> tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
+      %2 = "quantization.dcast"(%1) : (tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<1x3xf32>
+      %3 = stablehlo.reshape %2 : (tensor<1x3xf32>) -> tensor<3x1xf32>
+      %4 = "quantization.qcast"(%3) {volatile} : (tensor<3x1xf32>) -> tensor<3x1x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
+      %5 = "quantization.dcast"(%4) : (tensor<3x1x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<3x1xf32>
+      return %5 : tensor<3x1xf32>
+    }
+  }
+)mlir";
+
+// Non-quantizable ops: XlaCallModule op without "fully_quantizable" attribute
+constexpr absl::string_view kModuleCompositeNoAttr = R"mlir(
+  module {
+    func.func @composite_without_attr() -> tensor<1x3xf32> {
+      %0 = "tf.XlaCallModule"() {Sout = [#tf_type.shape<1x3>], _entry_function = @non_quantizable_composite, _stablehlo_version = "1.0.0", _original_entry_function = "non_quantizable_composite", _stablehlo_module_attrs = {}, device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : () -> tensor<1x3xf32>
+      return %0 : tensor<1x3xf32>
+    }
+  }
+)mlir";
+
+TEST_F(IsOpQuantizableStableHloTest, ConstantOpQuantizable) {
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleConstantAdd);
+  ASSERT_TRUE(module_op);
+
+  auto test_func = module_op->lookupSymbol<func::FuncOp>("constant_add");
+  ASSERT_THAT(test_func, NotNull());
+
+  auto constant_op =
+      FindOperationOfType<mlir::stablehlo::ConstantOp>(test_func);
+  EXPECT_TRUE(IsOpQuantizableStableHlo(constant_op));
+}
+
+TEST_F(IsOpQuantizableStableHloTest, TerminatorOpNotQuantizable) {
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleConstantAdd);
+  ASSERT_TRUE(module_op);
+
+  auto test_func = module_op->lookupSymbol<func::FuncOp>("constant_add");
+  ASSERT_THAT(test_func, NotNull());
+
+  auto return_op = FindOperationOfType<func::ReturnOp>(test_func);
+  EXPECT_FALSE(IsOpQuantizableStableHlo(return_op));
+}
+
+TEST_F(IsOpQuantizableStableHloTest, SameScaleOpQuantizable) {
+  OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kModuleCompositeSameScale);
+  ASSERT_TRUE(module_op);
+
+  auto test_func =
+      module_op->lookupSymbol<func::FuncOp>("same_scale_after_composite");
+  ASSERT_THAT(test_func, NotNull());
+
+  auto reshape_op = FindOperationOfType<mlir::stablehlo::ReshapeOp>(test_func);
+  EXPECT_TRUE(IsOpQuantizableStableHlo(reshape_op));
+}
+
+TEST_F(IsOpQuantizableStableHloTest, NonSameScaleOpNotQuantizable) {
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleConstantAdd);
+  ASSERT_TRUE(module_op);
+
+  auto test_func = module_op->lookupSymbol<func::FuncOp>("constant_add");
+  ASSERT_THAT(test_func, NotNull());
+
+  auto add_op = FindOperationOfType<mlir::stablehlo::AddOp>(test_func);
+  EXPECT_FALSE(IsOpQuantizableStableHlo(add_op));
+}
+
+TEST_F(IsOpQuantizableStableHloTest, ValidXlaCallModuleOpQuantizable) {
+  OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kModuleCompositeSameScale);
+  ASSERT_TRUE(module_op);
+
+  auto test_func =
+      module_op->lookupSymbol<func::FuncOp>("same_scale_after_composite");
+  ASSERT_THAT(test_func, NotNull());
+
+  auto xla_call_module_op = FindOperationOfType<TF::XlaCallModuleOp>(test_func);
+  EXPECT_TRUE(IsOpQuantizableStableHlo(xla_call_module_op));
+}
+
+TEST_F(IsOpQuantizableStableHloTest, InvalidXlaCallModuleOpNotQuantizable) {
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleCompositeNoAttr);
+  ASSERT_TRUE(module_op);
+
+  auto test_func =
+      module_op->lookupSymbol<func::FuncOp>("composite_without_attr");
+  ASSERT_THAT(test_func, NotNull());
+
+  auto xla_call_module_op = FindOperationOfType<TF::XlaCallModuleOp>(test_func);
+  EXPECT_FALSE(IsOpQuantizableStableHlo(xla_call_module_op));
+}
+
+TEST_F(IsOpQuantizableStableHloTest, QuantizeDequantizeOpNotQuantizable) {
+  OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kModuleCompositeSameScale);
+  ASSERT_TRUE(module_op);
+
+  auto test_func =
+      module_op->lookupSymbol<func::FuncOp>("same_scale_after_composite");
+  ASSERT_THAT(test_func, NotNull());
+
+  auto quantize_op =
+      FindOperationOfType<mlir::quant::ir::QuantizeCastOp>(test_func);
+  EXPECT_FALSE(IsOpQuantizableStableHlo(quantize_op));
+
+  auto dequantize_op =
+      FindOperationOfType<mlir::quant::ir::DequantizeCastOp>(test_func);
+  EXPECT_FALSE(IsOpQuantizableStableHlo(dequantize_op));
+}
+
+TEST_F(IsOpQuantizableStableHloTest,
+       XlaCallModuleOpQuantizableWhenNotDenylisted) {
+  // A `TF::XlaCallModuleOp` with `_quantization_method = ""`.
+  constexpr absl::string_view
+      kModuleXlaCallModuleOpWithDefaultQuantizationMethod = R"mlir(
+    func.func @xla_call_module_default_quantization_method(%arg0: tensor<1x1x3xf32>, %arg1: tensor<3x4xf32>) -> tensor<1x1x4xf32> {
+      %0 = "tf.XlaCallModule"(%arg0, %arg1) <{Sout = [#tf_type.shape<1x1x4>], dim_args_spec = [], disabled_checks = [], function_list = [], has_token_input_output = false, module = "", platforms = ["CPU"], version = 9 : i64}> {_entry_function = @composite_dot_general_fn_1, _stablehlo_version = "1.0.0", _original_entry_function = "composite_dot_general_fn_1", _quantization_method = "", _stablehlo_module_attrs = {jax.uses_shape_polymorphism = true}, _tfl_quant_trait = "fully_quantizable"} : (tensor<1x1x3xf32>, tensor<3x4xf32>) -> tensor<1x1x4xf32>
+      return %0 : tensor<1x1x4xf32>
+    }
+  )mlir";
+
+  OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kModuleXlaCallModuleOpWithDefaultQuantizationMethod);
+  ASSERT_TRUE(module_op);
+
+  auto test_func = module_op->lookupSymbol<func::FuncOp>(
+      "xla_call_module_default_quantization_method");
+  ASSERT_THAT(test_func, NotNull());
+
+  auto xla_call_module_op = FindOperationOfType<TF::XlaCallModuleOp>(test_func);
+  EXPECT_TRUE(IsOpQuantizableStableHlo(xla_call_module_op));
+}
+
+TEST_F(IsOpQuantizableStableHloTest, DenylistedXlaCallModuleOpNotQuantizable) {
+  // A `TF::XlaCallModuleOp` with `_quantization_method = "no_quantization {}"`,
+  // indicating it has been explicitly denylisted by the user.
+  constexpr absl::string_view kModuleDenylistedXlaCallModuleOp = R"mlir(
+    func.func @xla_call_module_denylisted(%arg0: tensor<1x1x3xf32>, %arg1: tensor<3x4xf32>) -> tensor<1x1x4xf32> {
+      %0 = "tf.XlaCallModule"(%arg0, %arg1) <{Sout = [#tf_type.shape<1x1x4>], dim_args_spec = [], disabled_checks = [], function_list = [], has_token_input_output = false, module = "", platforms = ["CPU"], version = 9 : i64}> {_entry_function = @composite_dot_general_fn_1, _stablehlo_version = "1.0.0", _original_entry_function = "composite_dot_general_fn_1", _quantization_method = "no_quantization {}", _stablehlo_module_attrs = {jax.uses_shape_polymorphism = true}, _tfl_quant_trait = "fully_quantizable"} : (tensor<1x1x3xf32>, tensor<3x4xf32>) -> tensor<1x1x4xf32>
+      return %0 : tensor<1x1x4xf32>
+    }
+  )mlir";
+
+  OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kModuleDenylistedXlaCallModuleOp);
+  ASSERT_TRUE(module_op);
+
+  auto test_func =
+      module_op->lookupSymbol<func::FuncOp>("xla_call_module_denylisted");
+  ASSERT_THAT(test_func, NotNull());
+
+  auto xla_call_module_op = FindOperationOfType<TF::XlaCallModuleOp>(test_func);
+  EXPECT_FALSE(IsOpQuantizableStableHlo(xla_call_module_op));
+}
+
+using GetStableHloOpQuantSpecTest = ::mlir::quant::QuantizationTestBase;
+
+TEST_F(GetStableHloOpQuantSpecTest,
+       EmptyCoeffOpQuantDimForPerTensorQuantizedConvolution) {
+  // A `TF::XlaCallModuleOp` with `_quantization_method = "static_range_ptq
+  // {}"`, representing a per-tensor static-range PTQ quantization.
+  constexpr absl::string_view
+      kXlaCallModuleOpWithPerTensorQuantizedConvolution = R"mlir(
+    func.func @main(%arg0: tensor<1x1x3xf32>, %arg1: tensor<3x4xf32>) -> tensor<1x1x4xf32> {
+      %0 = "tf.XlaCallModule"(%arg0, %arg1) <{Sout = [#tf_type.shape<1x1x4>], dim_args_spec = [], disabled_checks = [], function_list = [], has_token_input_output = false, module = "", platforms = ["CPU"], version = 9 : i64}>
+          {
+            _entry_function = @composite_conv_fn_1,
+            _stablehlo_version = "1.0.0",
+            _original_entry_function = "composite_conv_fn_1",
+            _quantization_method = "static_range_ptq {}",
+            _stablehlo_module_attrs = {jax.uses_shape_polymorphism = true},
+            _tfl_quant_trait = "fully_quantizable"
+          } : (tensor<1x1x3xf32>, tensor<3x4xf32>) -> tensor<1x1x4xf32>
+      return %0 : tensor<1x1x4xf32>
+    }
+  )mlir";
+
+  const OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kXlaCallModuleOpWithPerTensorQuantizedConvolution);
+  ASSERT_TRUE(module_op);
+
+  const FailureOr<TF::XlaCallModuleOp> xla_call_module_op =
+      FindFirstOpFromMainFunc<TF::XlaCallModuleOp>(*module_op);
+  ASSERT_TRUE(succeeded(xla_call_module_op));
+
+  const std::unique_ptr<OpQuantSpec> op_quant_spec =
+      GetStableHloOpQuantSpec(*xla_call_module_op);
+  ASSERT_THAT(op_quant_spec, NotNull());
+
+  EXPECT_THAT(op_quant_spec->coeff_op_quant_dim, IsEmpty());
+}
+
+TEST_F(GetStableHloOpQuantSpecTest,
+       EmptyCoeffOpQuantDimForPerChannelQuantizedConvolution) {
+  constexpr absl::string_view
+      kXlaCallModuleOpWithPerChannelQuantizedConvolution = R"mlir(
+    func.func @main(%arg0: tensor<1x1x3xf32>, %arg1: tensor<3x4xf32>) -> tensor<1x1x4xf32> {
+      %0 = "tf.XlaCallModule"(%arg0, %arg1) <{Sout = [#tf_type.shape<1x1x4>], dim_args_spec = [], disabled_checks = [], function_list = [], has_token_input_output = false, module = "", platforms = ["CPU"], version = 9 : i64}>
+          {
+            _entry_function = @composite_conv_fn_1,
+            _stablehlo_version = "1.0.0",
+            _original_entry_function = "composite_conv_fn_1",
+            _quantization_method = "static_range_ptq {input_quantized_types {key: 1, value {dimension_specs {dimension: 3}}}}",
+            _stablehlo_module_attrs = {jax.uses_shape_polymorphism = true},
+            _tfl_quant_trait = "fully_quantizable"
+          } : (tensor<1x1x3xf32>, tensor<3x4xf32>) -> tensor<1x1x4xf32>
+      return %0 : tensor<1x1x4xf32>
+    }
+  )mlir";
+
+  const OwningOpRef<ModuleOp> module_op =
+      ParseModuleOpString(kXlaCallModuleOpWithPerChannelQuantizedConvolution);
+  ASSERT_TRUE(module_op);
+
+  const FailureOr<TF::XlaCallModuleOp> xla_call_module_op =
+      FindFirstOpFromMainFunc<TF::XlaCallModuleOp>(*module_op);
+  ASSERT_TRUE(succeeded(xla_call_module_op));
+
+  const std::unique_ptr<OpQuantSpec> op_quant_spec =
+      GetStableHloOpQuantSpec(*xla_call_module_op);
+  ASSERT_THAT(op_quant_spec, NotNull());
+
+  EXPECT_THAT(op_quant_spec->coeff_op_quant_dim,
+              UnorderedElementsAre(Pair(1, 3)));
+}
+
+using GetStableHloQuantConstraintsTest = ::mlir::quant::QuantizationTestBase;
+
+TEST_F(GetStableHloQuantConstraintsTest,
+       HasSameOperandAndResultTypeRequirementSucceeds) {
+  // Quantizable ops: constants
+  // Non-quantizable ops: normal StableHLO ops and terminators
+  constexpr absl::string_view kModuleGather = R"mlir(
+    module {
+      func.func @main() -> (tensor<2x3x2x2xf32>) {
+        %0 = stablehlo.constant dense<1.0> : tensor<3x4x2xf32>
+        %1 = stablehlo.constant dense<2> : tensor<2x3x2xi64>
+        %2 = "stablehlo.gather"(%0, %1) {
+          dimension_numbers = #stablehlo.gather<
+            offset_dims = [2, 3],
+            collapsed_slice_dims = [0],
+            start_index_map = [1, 0],
+            index_vector_dim = 2>,
+          slice_sizes = array<i64: 1, 2, 2>,
+          indices_are_sorted = false
+        } : (tensor<3x4x2xf32>, tensor<2x3x2xi64>) -> tensor<2x3x2x2xf32>
+        func.return %2 : tensor<2x3x2x2xf32>
+      }
+    }
+  )mlir";
+  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleGather);
+  ASSERT_TRUE(module_op);
+
+  func::FuncOp main_fn = FindMainFuncOp(*module_op);
+  ASSERT_THAT(main_fn, NotNull());
+
+  Operation* gather_op = FindOperationOfType<GatherOp>(main_fn);
+  const auto spec = GetStableHloQuantConstraints(gather_op);
+
+  EXPECT_THAT(spec, NotNull());
+  EXPECT_THAT(spec->has_same_operand_and_result_type_requirement, IsTrue());
+}
+
+}  // namespace
+}  // namespace mlir::quant::stablehlo