chore: import upstream snapshot with attribution

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

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

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

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

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

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

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

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

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

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

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

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

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

@@ -0,0 +1,50 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+include "mlir/IR/OpBase.td"
+include "mlir/IR/PatternBase.td"
+include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"
+include "mhlo/IR/hlo_ops.td"
+include "stablehlo/dialect/ChloOps.td"
+
+def IsDefinedByConvOrDotGeneralOp : Constraint<Or<[
+  CPred<"$0.getDefiningOp<mhlo::ConvolutionOp>()">,
+  CPred<"$0.getDefiningOp<mhlo::DotGeneralOp>()">]>>;
+
+def IsNull : Constraint<CPred<"!$0">>;
+
+// This pattern optimizes:
+//   conv/dot_general + a + b -> conv/dot_general + (a + b)
+//   conv/dot_general - a - b -> conv/dot_general - (a + b)
+//   conv/dot_general + a - b -> conv/dot_general + (a - b)
+//   conv/dot_general - a + b -> conv/dot_general - (a - b)
+foreach OpsTuple = [
+    [CHLO_BroadcastAddOp, CHLO_BroadcastAddOp, CHLO_BroadcastAddOp],
+    [CHLO_BroadcastSubOp, CHLO_BroadcastSubOp, CHLO_BroadcastAddOp],
+    [CHLO_BroadcastAddOp, CHLO_BroadcastSubOp, CHLO_BroadcastSubOp],
+    [CHLO_BroadcastSubOp, CHLO_BroadcastAddOp, CHLO_BroadcastSubOp]] in {
+  def optimizeConsecutiveConv#OpsTuple[0]#OpsTuple[1] : Pat<
+    (OpsTuple[1]
+      (OpsTuple[0] $input, $zp_offset, $broadcast_dims_1),
+      $bias, $broadcast_dims_2),
+    (OpsTuple[0]
+      $input,
+      (OpsTuple[2] $zp_offset, $bias, $broadcast_dims_2), $broadcast_dims_1),
+    [
+      (IsNull $broadcast_dims_1),
+      (IsNull $broadcast_dims_2),
+      (TensorOf<[AnyInteger]> $input),
+      (IsDefinedByConvOrDotGeneralOp $input)]>;
+}

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

@@ -0,0 +1,40 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/passes.h"
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Pass/PassOptions.h"  // from @llvm-project
+#include "mlir/Transforms/Passes.h"  // from @llvm-project
+#include "stablehlo/transforms/Passes.h"  // from @stablehlo
+#include "xla/mlir_hlo/mhlo/transforms/passes.h"
+
+namespace mlir::quant::stablehlo {
+
+void AddQuantizationLoweringPasses(mlir::OpPassManager& pm) {
+  // These passes are grouped together and must run in this specific order.
+  pm.addNestedPass<mlir::func::FuncOp>(CreateConvertTFQuantOpsToMHLOPass());
+  pm.addNestedPass<mlir::func::FuncOp>(mhlo::createChloLegalizeToHloPass());
+  pm.addNestedPass<mlir::func::FuncOp>(mlir::createCanonicalizerPass());
+  pm.addPass(mhlo::createHloLegalizeToStablehloPass());
+  pm.addNestedPass<mlir::func::FuncOp>(
+      mlir::stablehlo::createStablehloLegalizeQuantToMathPass());
+  pm.addPass(mhlo::createStablehloLegalizeToHloPass());
+  pm.addNestedPass<mlir::func::FuncOp>(mlir::createCanonicalizerPass());
+  pm.addNestedPass<mlir::func::FuncOp>(CreateVerifyQuantLegalizationPass());
+}
+
+}  // namespace mlir::quant::stablehlo

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

@@ -0,0 +1,62 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_PASSES_BRIDGE_PASSES_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_PASSES_BRIDGE_PASSES_H_
+
+#include <memory>
+
+#define GEN_PASS_DECL
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+
+namespace mlir::quant::stablehlo {
+
+// Creates an instance of the ConvertTFQuantOpsToMHLOPass pass, which will
+// convert TF uniform quantized ops to the corresponding quantized MHLO ops.
+std::unique_ptr<OperationPass<func::FuncOp>>
+CreateConvertTFQuantOpsToMHLOPass();
+
+// TODO(b/288094093): Migrate uniform quantization legalization in a separate
+// pass.
+void PopulateLegalizeTfQuantizationPatterns(MLIRContext *context,
+                                            RewritePatternSet *patterns);
+
+// Creates an instance of the ConvertTFQuantTypes pass, which will convert TF
+// qint types to int types and surround TF UniformQuantized ops with qint <->
+// int casts.
+std::unique_ptr<OperationPass<func::FuncOp>> CreateConvertTFQuantTypesPass();
+
+// Creates an instance of the VerifyQuantLegalization pass, which verifies all
+// quant ops and types are lowered.
+std::unique_ptr<OperationPass<func::FuncOp>>
+CreateVerifyQuantLegalizationPass();
+
+// Add all passes for lowering TF quant ops and types to MHLO int.
+void AddQuantizationLoweringPasses(mlir::OpPassManager &pm);
+
+// Creates an instance of OptimizeIntGraphPass, which optimizes the int graph
+// lowered from the quantized graph.
+std::unique_ptr<OperationPass<func::FuncOp>> CreateOptimizeIntGraphPass();
+
+#define GEN_PASS_REGISTRATION
+#define GEN_PASS_DECL_CONVERTTFQUANTOPSTOMHLO
+#define GEN_PASS_DECL_CONVERTTFQUANTTYPES
+#define GEN_PASS_DECL_VERIFYQUANTLEGALIZATION
+#define GEN_PASS_DECL_OPTIMIZEINTGRAPH
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/passes.h.inc"
+}  // namespace mlir::quant::stablehlo
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_PASSES_BRIDGE_PASSES_H_

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

@@ -0,0 +1,69 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// Declare bridge passes that perform legalization and/or lowering.
+
+include "mlir/Pass/PassBase.td"
+
+def ConvertTFQuantOpsToMHLO : Pass<"quant-convert-tf-quant-ops-to-mhlo", "mlir::func::FuncOp"> {
+  let summary = "Convert TF Quant ops to MHLO quantizated ops.";
+
+  let description = [{
+    Convert TF Quant ops to MHLO quant ops.
+  }];
+
+  let constructor = "mlir::quant::stablehlo::CreateConvertTFQuantOpsToMHLOPass()";
+  let dependentDialects = ["TF::TensorFlowDialect", "chlo::ChloDialect",
+                           "mhlo::MhloDialect", "tf_type::TFTypeDialect",
+                           "quant::QuantDialect"];
+}
+
+def ConvertTFQuantTypes : Pass<"convert-tf-quant-types", "mlir::func::FuncOp"> {
+  let summary = "Replace TensorFlow qint types with int types.";
+
+  let description = [{
+    Converts TF ops with qint types to int types. Some UniformQuantized ops
+    argument/result allow qint type only. For such cases, add qint <-> int
+    tf.Cast around the ops so that they are still valid.
+  }];
+
+  let constructor = "mlir::quant::stablehlo::CreateConvertTFQuantTypesPass()";
+  let dependentDialects = ["TF::TensorFlowDialect", "tf_type::TFTypeDialect"];
+}
+
+def VerifyQuantLegalization : Pass<"verify-quant-legalization", "mlir::func::FuncOp"> {
+  let summary = "Verifies that all TF quant ops and types have been legalized.";
+
+  let description = [{
+    Ensures that all TF quant ops and types have been legalized to HLO
+    and reports an error about which op failed to legalize. This pass
+    does not transform any ops and is checking.}];
+
+  let constructor = "mlir::quant::stablehlo::CreateVerifyQuantLegalizationPass()";
+  let dependentDialects =  ["tf_type::TFTypeDialect",
+                            "quant::QuantDialect"];
+
+}
+
+def OptimizeIntGraph : Pass<"optimize-int-graph", "mlir::func::FuncOp"> {
+  let summary = "Optimization patterns for quantized integer graph";
+
+  let description = [{
+    This includes patterns for merging addition of zp offset and bias.
+  }];
+
+  let constructor = "mlir::quant::stablehlo::CreateOptimizeIntGraphPass()";
+  let dependentDialects = ["mhlo::MhloDialect"];
+}

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

@@ -0,0 +1,93 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// The TF dialect uses some TF types that are illegal in the MHLO dialect and
+// some generic types that are legal in MHLO. This pass legalizes TF types into
+// types that are legal in MHLO. For example, TF::Qint8Type is converted to i8.
+// Rewrites here should run before TF to MHLO op legalizations are run.
+
+#include <memory>
+
+#include "absl/log/log.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/Quant.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/IR/Visitors.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/quantization/stablehlo/utils/tf_type_utils.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
+
+namespace mlir::quant::stablehlo {
+namespace {
+
+using quant::tensorflow::IsTFQintType;
+using quant::tensorflow::IsTFUniformQuantizedOp;
+
+#define GEN_PASS_DEF_VERIFYQUANTLEGALIZATION
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/passes.h.inc"
+
+class VerifyQuantLegalization
+    : public impl::VerifyQuantLegalizationBase<VerifyQuantLegalization> {
+ public:
+  void runOnOperation() override;
+};
+
+bool IsQuantType(Type type) {
+  auto element_type = getElementTypeOrSelf(type);
+  return mlir::isa<quant::UniformQuantizedType>(element_type) ||
+         IsTFQintType(element_type);
+}
+
+bool IsMhloUniformQuantizedOp(Operation& op) {
+  return llvm::isa<mhlo::UniformQuantizeOp, mhlo::UniformDequantizeOp>(op);
+}
+
+void VerifyQuantLegalization::runOnOperation() {
+  Operation* func_op = getOperation();
+
+  auto walk_result = func_op->walk([&](Operation* op) {
+    // Verify all uq and qint types are lowered.
+    if (llvm::any_of(op->getOperandTypes(), IsQuantType) ||
+        llvm::any_of(op->getResultTypes(), IsQuantType) ||
+        IsTFUniformQuantizedOp(op) || IsMhloUniformQuantizedOp(*op)) {
+      op->emitOpError("is illegal as it is a UQ op or contains uq/qint types");
+      LOG(ERROR) << "Found illegal op containing uq/qint type: "
+                 << op->getName().getStringRef().str();
+      return WalkResult::interrupt();
+    }
+    return WalkResult::advance();
+  });
+
+  if (walk_result.wasInterrupted()) {
+    signalPassFailure();
+  }
+}
+
+}  // namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>>
+CreateVerifyQuantLegalizationPass() {
+  return std::make_unique<VerifyQuantLegalization>();
+}
+
+}  // namespace mlir::quant::stablehlo

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

@@ -0,0 +1,234 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <memory>
+#include <utility>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/TypeRange.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h"  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/quantization/stablehlo/utils/bfloat16_type.h"
+#include "tensorflow/core/platform/bfloat16.h"
+
+namespace mlir::quant::stablehlo {
+namespace {
+
+class BFloat16TypeConverter : public TypeConverter {
+ public:
+  BFloat16TypeConverter() {
+    addConversion([](const Type type) -> Type {
+      return quant::stablehlo::IsLargeFloatType(type)
+                 ? quant::stablehlo::ToBfloat16Type(type)
+                 : type;
+    });
+  }
+};
+
+// This helper function makes legality check easier. Both convert ops in the
+// patterns below are considered legal:
+//  - `BitcastConvertOp` (i32 -> f32) + `ConvertOp` (f32 -> bf16)
+//  - `ConvertOp` (bf16 -> f32) -> `BitcastConvertOp` (f32 -> i32)
+template <typename ConvertOp, typename OtherConvertOp>
+bool IsConvertOpLegal(ConvertOp convert_op, BFloat16TypeConverter& converter) {
+  if (!converter.isLegal(convert_op.getOperand().getType())) {
+    auto other_convert_op = dyn_cast_or_null<OtherConvertOp>(
+        convert_op.getOperand().getDefiningOp());
+    return other_convert_op &&
+           converter.isLegal(other_convert_op.getOperand().getType());
+  } else if (!converter.isLegal(convert_op.getResult().getType())) {
+    if (!convert_op.getResult().hasOneUse()) {
+      return false;
+    }
+    auto other_convert_op = dyn_cast_or_null<OtherConvertOp>(
+        *convert_op.getResult().getUsers().begin());
+    return other_convert_op &&
+           converter.isLegal(other_convert_op.getResult().getType());
+  }
+  return true;
+}
+
+class BFloat16TypeConversionTarget : public ConversionTarget {
+ public:
+  explicit BFloat16TypeConversionTarget(MLIRContext& ctx,
+                                        BFloat16TypeConverter& converter)
+      : ConversionTarget(ctx), converter_(converter) {
+    markUnknownOpDynamicallyLegal([this](Operation* op) {
+      // The FuncOp type can contain types that the op's operand and result
+      // types do not contain.
+      if (auto func = dyn_cast<func::FuncOp>(op)) {
+        if (!converter_.isSignatureLegal(func.getFunctionType())) return false;
+      } else if (auto bitcast_convert_op =
+                     dyn_cast<mlir::stablehlo::BitcastConvertOp>(op)) {
+        return IsConvertOpLegal<mlir::stablehlo::BitcastConvertOp,
+                                mlir::stablehlo::ConvertOp>(bitcast_convert_op,
+                                                            converter_);
+      } else if (auto convert_op = dyn_cast<mlir::stablehlo::ConvertOp>(op)) {
+        return IsConvertOpLegal<mlir::stablehlo::ConvertOp,
+                                mlir::stablehlo::BitcastConvertOp>(convert_op,
+                                                                   converter_);
+      }
+      return converter_.isLegal(op);
+    });
+  }
+
+ private:
+  BFloat16TypeConverter& converter_;
+};
+
+class BFloat16TypePattern : public ConversionPattern {
+ public:
+  BFloat16TypePattern(TypeConverter& converter, MLIRContext* ctx)
+      : ConversionPattern(converter, MatchAnyOpTypeTag(), /*benefit=*/1, ctx) {}
+
+  LogicalResult matchAndRewrite(
+      Operation* op, const ArrayRef<Value> operands,
+      ConversionPatternRewriter& rewriter) const override {
+    if (getTypeConverter()->isLegal(op)) {
+      return failure();
+    }
+    if (isa<mlir::stablehlo::BitcastConvertOp>(op)) {
+      // Skip `BitcastConvertOp`, which is handled by the other pattern.
+      return failure();
+    }
+
+    // Update the results.
+    SmallVector<Type, 4> new_results;
+    if (failed(getTypeConverter()->convertTypes(op->getResultTypes(),
+                                                new_results)))
+      return failure();
+
+    // Update the regions. The dialect conversion framework wants new regions to
+    // be created and updated, rather than updating the old op. Thus we use an
+    // OperationState so we can add regions to the new op.
+    OperationState state(op->getLoc(), op->getName().getStringRef(), operands,
+                         new_results, op->getAttrs(), op->getSuccessors());
+    for (Region& region : op->getRegions()) {
+      auto new_region = std::make_unique<Region>(op);
+      rewriter.inlineRegionBefore(region, *new_region, new_region->begin());
+      if (failed(rewriter.convertRegionTypes(new_region.get(),
+                                             *getTypeConverter()))) {
+        return failure();
+      }
+      state.addRegion(std::move(new_region));
+    }
+
+    // Convert value of ConstantOp to bfloat16.
+    if (auto const_op = dyn_cast<mlir::stablehlo::ConstantOp>(op)) {
+      const auto values = const_op.getValue().tryGetValues<float>();
+      if (!values.has_value()) {
+        return failure();
+      }
+      const SmallVector<tensorflow::bfloat16> bfloat16_values(values->begin(),
+                                                              values->end());
+      state.attributes.set(
+          const_op.getValueAttrName(),
+          DenseFPElementsAttr::get(
+              mlir::dyn_cast<ShapedType>(const_op.getValue().getType())
+                  .clone(rewriter.getBF16Type()),
+              bfloat16_values));
+    }
+
+    rewriter.replaceOp(op, rewriter.create(state)->getResults());
+
+    return success();
+  }
+};
+
+class BitcastConvertOpPattern
+    : public OpConversionPattern<mlir::stablehlo::BitcastConvertOp> {
+ public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      mlir::stablehlo::BitcastConvertOp op,
+      mlir::stablehlo::BitcastConvertOpAdaptor adaptor,
+      ConversionPatternRewriter& rewriter) const override {
+    const bool is_input_legal =
+        getTypeConverter()->isLegal(op.getOperand().getType());
+    const bool is_output_legal =
+        getTypeConverter()->isLegal(op.getResult().getType());
+    if (is_input_legal && is_output_legal) {
+      return failure();
+    } else if (is_input_legal) {
+      // output is f32, we bitcast_convert to f32 and then convert to bf16.
+      const Value output = mlir::stablehlo::BitcastConvertOp::create(
+          rewriter, op->getLoc(), op.getResult().getType(),
+          adaptor.getOperand());
+      rewriter.replaceOpWithNewOp<mlir::stablehlo::ConvertOp>(
+          op, getTypeConverter()->convertType(op.getResult().getType()),
+          output);
+    } else if (is_output_legal) {
+      // input is f32, we convert from bf16 and then bitcast_convert.
+      const Value output = mlir::stablehlo::ConvertOp::create(
+          rewriter, op->getLoc(), op.getOperand().getType(),
+          adaptor.getOperand());
+      rewriter.replaceOpWithNewOp<mlir::stablehlo::BitcastConvertOp>(
+          op, op.getResult().getType(), output);
+    } else {
+      // Both input/output are f32. Convert to no-op.
+      rewriter.replaceOp(op, adaptor.getOperand());
+    }
+    return success();
+  }
+};
+}  // namespace
+
+#define GEN_PASS_DEF_CONVERTFUNCTOBFLOAT16PASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"
+namespace {
+class ConvertFuncToBfloat16Pass
+    : public impl::ConvertFuncToBfloat16PassBase<ConvertFuncToBfloat16Pass> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ConvertFuncToBfloat16Pass)
+
+  explicit ConvertFuncToBfloat16Pass() = default;
+
+ private:
+  void runOnOperation() override;
+};
+
+void ConvertFuncToBfloat16Pass::runOnOperation() {
+  func::FuncOp func_op = getOperation();
+  MLIRContext* context = func_op.getContext();
+  RewritePatternSet patterns(context);
+
+  BFloat16TypeConverter converter;
+  patterns.add<BFloat16TypePattern, BitcastConvertOpPattern>(converter,
+                                                             context);
+  populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(patterns,
+                                                                 converter);
+  BFloat16TypeConversionTarget target(*context, converter);
+  if (failed(applyPartialConversion(func_op.getOperation(), target,
+                                    std::move(patterns)))) {
+    return signalPassFailure();
+  }
+}
+
+}  // namespace
+}  // namespace mlir::quant::stablehlo

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

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

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

@@ -0,0 +1,149 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <string>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "llvm/Support/raw_ostream.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/Quant.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Dialect/Shape/IR/Shape.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/IR/Visitors.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "shardy/dialect/sdy/ir/register.h"  // from @shardy
+#include "stablehlo/dialect/Serialization.h"  // from @stablehlo
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h"  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/quantization/stablehlo/utils/bfloat16_type.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+
+namespace mlir::quant::stablehlo {
+
+absl::StatusOr<std::string> ConvertSerializedStableHloModuleToBfloat16(
+    const StringRef serialized_stablehlo_module) {
+  // StableHLO module is empty often because the XlaCallModuleOp is already
+  // deserialized, e.g. after invoking XlaCallModuleDeserializationPass. We
+  // don't handle this situation.
+  if (serialized_stablehlo_module.empty()) {
+    return absl::InvalidArgumentError("StableHLO module is empty.");
+  }
+
+  MLIRContext context;
+  mlir::sdy::loadAllRequiredDialects(&context);
+  OwningOpRef<ModuleOp> stablehlo_module_op =
+      mlir::stablehlo::deserializePortableArtifact(serialized_stablehlo_module,
+                                                   &context);
+  auto version =
+      mlir::stablehlo::getPortableArtifactVersion(serialized_stablehlo_module);
+  if (failed(version)) {
+    return absl::InternalError(
+        "Failed to get the deserialized StableHLO version, XlaCallModuleOp "
+        "must have a valid StableHLO module serialized using "
+        "stablehlo::serializePortableArtifact APIs.");
+  }
+
+  // Convert the StableHLO module to bfloat16.
+  PassManager pm(&context);
+  pm.addNestedPass<func::FuncOp>(createConvertFuncToBfloat16Pass());
+  if (failed(pm.run(stablehlo_module_op.get()))) {
+    return absl::InternalError(
+        "Failed to convert StableHLO module to bfloat16.");
+  }
+
+  std::string bytecode;
+  llvm::raw_string_ostream os(bytecode);
+  if (failed(mlir::stablehlo::serializePortableArtifact(
+          stablehlo_module_op.get(), version.value().toString(), os,
+          /*allowOtherDialects=*/true))) {
+    return absl::InternalError("Failed to serialize StableHLO module.");
+  }
+  return bytecode;
+}
+
+#define GEN_PASS_DEF_CONVERTXLACALLMODULEOPTOBFLOAT16PASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"
+
+namespace {
+class ConvertXlaCallModuleOpToBfloat16Pass
+    : public impl::ConvertXlaCallModuleOpToBfloat16PassBase<
+          ConvertXlaCallModuleOpToBfloat16Pass> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
+      ConvertXlaCallModuleOpToBfloat16Pass)
+
+  explicit ConvertXlaCallModuleOpToBfloat16Pass() = default;
+
+ private:
+  void runOnOperation() override;
+};
+
+void ConvertXlaCallModuleOpToBfloat16Pass::runOnOperation() {
+  Operation* func_op = getOperation();
+  SymbolTableCollection symbol_table;
+  OpBuilder builder(&getContext());
+
+  auto result = func_op->walk([&](TF::XlaCallModuleOp op) {
+    // Converts the serialized StableHLO module to bfloat16.
+    auto result =
+        ConvertSerializedStableHloModuleToBfloat16(op.getModuleAttr());
+    if (!result.ok()) {
+      llvm::errs() << "Failed to convert StableHLO module to bfloat16: "
+                   << result.status().message();
+      return WalkResult::interrupt();
+    }
+    op.setModuleAttr(StringAttr::get(&getContext(), *result));
+
+    // Convert the `tf.XlaCallModuleOp` to bfloat16 and add casts around it.
+    builder.setInsertionPoint(op);
+    for (auto& op_operand : op->getOpOperands()) {
+      if (quant::stablehlo::IsLargeFloatType(op_operand.get().getType())) {
+        op_operand.set(TF::CastOp::create(
+            builder, op->getLoc(),
+            quant::stablehlo::ToBfloat16Type(op_operand.get().getType()),
+            op_operand.get()));
+      }
+    }
+    builder.setInsertionPointAfter(op);
+    for (auto op_result : op->getOpResults()) {
+      if (quant::stablehlo::IsLargeFloatType(op_result.getType())) {
+        const Type original_type = op_result.getType();
+        op_result.setType(quant::stablehlo::ToBfloat16Type(original_type));
+        const Value cast =
+            TF::CastOp::create(builder, op->getLoc(), original_type, op_result);
+        op_result.replaceAllUsesExcept(cast, cast.getDefiningOp());
+      }
+    }
+    return WalkResult::advance();
+  });
+
+  if (result.wasInterrupted()) return signalPassFailure();
+}
+
+}  // namespace
+}  // namespace mlir::quant::stablehlo

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

@@ -0,0 +1,151 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <utility>
+
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "stablehlo/dialect/ChloOps.h"  // from @stablehlo
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo
+
+namespace mlir::quant::stablehlo {
+
+#define GEN_PASS_DEF_MERGEFUSIONWITHDEQUANTIZEPASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"
+
+namespace {
+
+class MergeFusionWithDequantizePass
+    : public impl::MergeFusionWithDequantizePassBase<
+          MergeFusionWithDequantizePass> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MergeFusionWithDequantizePass)
+
+  explicit MergeFusionWithDequantizePass() = default;
+
+ private:
+  void runOnOperation() override;
+};
+
+class MergeFusionWithUniformDequantizePattern
+    : public OpRewritePattern<func::CallOp> {
+ public:
+  explicit MergeFusionWithUniformDequantizePattern(MLIRContext* context)
+      : OpRewritePattern<func::CallOp>(context) {}
+  LogicalResult matchAndRewrite(func::CallOp call_op,
+                                PatternRewriter& rewriter) const override {
+    if (call_op.getNumResults() != 1) return failure();
+    auto users = call_op->getUsers();
+    for (auto user : users) {
+      if (!llvm::isa<mlir::stablehlo::UniformDequantizeOp>(user)) {
+        return failure();
+      }
+    }
+    auto func_name = call_op.getCallee();
+    if (!func_name.starts_with("quantized_")) return failure();
+    if (call_op->getNumResults() != 1) return failure();
+    if (!mlir::isa<quant::UniformQuantizedType>(
+            getElementTypeOrSelf(call_op->getResult(0).getType())))
+      return failure();
+
+    // Fetch the callee function.
+    SymbolTable symbol_table(call_op->getParentOfType<ModuleOp>());
+    auto func_op =
+        dyn_cast_or_null<func::FuncOp>(symbol_table.lookup(func_name));
+    if (!func_op) return failure();
+    // The quantized fusion should have requantize and return ops at the end.
+    auto return_op = dyn_cast_or_null<func::ReturnOp>(
+        func_op.getRegion().getBlocks().front().getTerminator());
+    if (!return_op) return failure();
+    auto req_op = llvm::dyn_cast_or_null<mlir::stablehlo::UniformQuantizeOp>(
+        return_op.getOperands()[0].getDefiningOp());
+    if (!req_op) return failure();
+
+    // Create a new func.call op with f32 output.
+    auto new_call_op = call_op.clone();
+    new_call_op->getResult(0).setType(
+        mlir::cast<ShapedType>(call_op.getResult(0).getType())
+            .clone(rewriter.getF32Type()));
+    rewriter.setInsertionPoint(call_op);
+    rewriter.insert(new_call_op);
+
+    // Remove the dequantize ops and replace uses by the new func.call op.
+    SmallVector<Operation*> users_to_erase;
+    for (auto user : users) {
+      llvm::dyn_cast<mlir::stablehlo::UniformDequantizeOp>(user)
+          .replaceAllUsesWith(new_call_op.getResult(0));
+      users_to_erase.push_back(user);
+    }
+    for (auto user : users_to_erase) rewriter.eraseOp(user);
+    rewriter.eraseOp(call_op);
+    if (failed(func_op.eraseResult(0))) {
+      return failure();
+    }
+    if (failed(func_op.insertResult(0, new_call_op.getResult(0).getType(),
+                                    /*resultAttrs=*/nullptr))) {
+      return failure();
+    }
+
+    // Modify the quantized fused function to do dequantize+relu(6).
+    rewriter.setInsertionPoint(req_op);
+    Value new_result = mlir::stablehlo::UniformDequantizeOp::create(
+        rewriter, req_op.getLoc(), func_op.getResultTypes()[0],
+        req_op.getOperand());
+    if (func_name.contains("_relu6_")) {
+      auto min = mlir::stablehlo::ConstantOp::create(
+          rewriter, req_op.getLoc(), rewriter.getF32FloatAttr(0));
+      auto max = mlir::stablehlo::ConstantOp::create(
+          rewriter, req_op.getLoc(), rewriter.getF32FloatAttr(6));
+      new_result = mlir::stablehlo::ClampOp::create(rewriter, req_op.getLoc(),
+                                                    min, new_result, max);
+    } else if (func_name.contains("_relu_")) {
+      auto min = mlir::stablehlo::ConstantOp::create(
+          rewriter, req_op.getLoc(), rewriter.getF32FloatAttr(0));
+      new_result = mlir::chlo::BroadcastMaxOp::create(rewriter, req_op.getLoc(),
+                                                      min, new_result, nullptr);
+    }
+    return_op->setOperand(0, new_result);
+    rewriter.eraseOp(req_op);
+
+    return success();
+  }
+};
+
+void MergeFusionWithDequantizePass::runOnOperation() {
+  ModuleOp module_op = getOperation();
+  MLIRContext* ctx = module_op.getContext();
+  RewritePatternSet patterns(ctx);
+  patterns.add<MergeFusionWithUniformDequantizePattern>(ctx);
+  if (failed(applyPatternsGreedily(module_op, std::move(patterns)))) {
+    signalPassFailure();
+  }
+}
+
+}  // namespace
+}  // namespace mlir::quant::stablehlo

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

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

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

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

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

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

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

@@ -0,0 +1,61 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_PASSES_PASSES_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_PASSES_PASSES_H_
+
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "absl/status/statusor.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_options.pb.h"
+
+namespace mlir::quant::stablehlo {
+
+// Creates a pass that quantizes weight component of StableHLO graph.
+std::unique_ptr<OperationPass<func::FuncOp>> CreateQuantizeWeightPass(
+    const ::stablehlo::quantization::QuantizationComponentSpec&
+        quantization_component_spec = {});
+
+// Converts a serialized StableHLO module to bfloat16 and output serialized
+// module.
+absl::StatusOr<std::string> ConvertSerializedStableHloModuleToBfloat16(
+    StringRef serialized_stablehlo_module);
+
+std::unique_ptr<OperationPass<ModuleOp>>
+CreateLiftQuantizableSpotsAsFunctionsPass(
+    const ::stablehlo::quantization::QuantizationSpecs& quantization_specs);
+
+// Creates a pass that inserts CalibrationStatisticsSaverOp.
+std::unique_ptr<OperationPass<ModuleOp>>
+CreateInsertCalibrationStatisticsSaverPass(
+    StringRef calibration_data_dir,
+    const std::vector<std::string>& aggregator_ops_to_ignore);
+
+// Adds generated pass default constructors or options definitions.
+#define GEN_PASS_DECL
+// Adds generated pass registration functions.
+#define GEN_PASS_REGISTRATION
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"
+
+}  // namespace mlir::quant::stablehlo
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_PASSES_PASSES_H_

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

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

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

@@ -0,0 +1,160 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <utility>
+
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Matchers.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"
+#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h"  // IWYU pragma: keep
+
+namespace mlir::quant::stablehlo {
+
+#define GEN_PASS_DEF_POSTQUANTIZEPASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"
+
+namespace {
+
+// Applies clean-up patterns after quantization.
+class PostQuantizePass : public impl::PostQuantizePassBase<PostQuantizePass> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PostQuantizePass)
+
+  explicit PostQuantizePass() = default;
+
+ private:
+  void runOnOperation() override;
+};
+
+// TODO: b/305815328 - Consider preserving leading and trailing QDQs for
+// ModifyIONodesPass in TFLite use cases.
+// Removes the back-to-back quantize and dequantize ops with volatile attribute.
+class RemoveVolatileQdqPattern
+    : public OpRewritePattern<mlir::quant::ir::DequantizeCastOp> {
+ public:
+  explicit RemoveVolatileQdqPattern(MLIRContext* context)
+      : OpRewritePattern<mlir::quant::ir::DequantizeCastOp>(context) {}
+
+  LogicalResult matchAndRewrite(mlir::quant::ir::DequantizeCastOp op,
+                                PatternRewriter& rewriter) const override {
+    auto input_op = op.getArg().getDefiningOp();
+    if (auto q =
+            llvm::dyn_cast_or_null<mlir::quant::ir::QuantizeCastOp>(input_op)) {
+      if (!q->getAttr(kVolatileOpAttrName)) return failure();
+
+      // If the quantize op is a requantize op, it is being used in other scale
+      // adjustments and should be kept. Instead, move dequantize op before the
+      // requantize op to remove the unnecessary requantize op.
+      if (const QuantizedType qtype =
+              QuantizedType::getQuantizedElementType(q.getArg().getType())) {
+        rewriter.setInsertionPoint(op);
+        rewriter.replaceOpWithNewOp<mlir::quant::ir::DequantizeCastOp>(
+            op, op.getResult().getType(), q.getArg());
+        return success();
+      }
+
+      op.replaceAllUsesWith(q.getArg());
+      return success();
+    }
+    return failure();
+  }
+};
+
+// Replaces constant and uniform_quantize ops with single quantized constant op.
+class QuantizeConstPattern
+    : public OpRewritePattern<mlir::stablehlo::UniformQuantizeOp> {
+ public:
+  explicit QuantizeConstPattern(MLIRContext* context)
+      : OpRewritePattern<mlir::stablehlo::UniformQuantizeOp>(context) {}
+
+  LogicalResult matchAndRewrite(mlir::stablehlo::UniformQuantizeOp op,
+                                PatternRewriter& rewriter) const override {
+    DenseFPElementsAttr attr;
+    if (matchPattern(op.getOperand(), m_Constant(&attr))) {
+      const Type qtype = op.getResult().getType();
+      ElementsAttr quantized_attr = Quantize(attr, qtype);
+      if (quantized_attr) {
+        rewriter.replaceOpWithNewOp<mlir::stablehlo::ConstantOp>(
+            op, qtype, quantized_attr);
+        return success();
+      }
+    }
+    return failure();
+  }
+};
+
+// Replaces quantfork.dcast with stablehlo.uniform_dequantize.
+class ConvertDequantizeCastToUniformDequantizePattern
+    : public OpRewritePattern<mlir::quant::ir::DequantizeCastOp> {
+ public:
+  explicit ConvertDequantizeCastToUniformDequantizePattern(MLIRContext* context)
+      : OpRewritePattern<mlir::quant::ir::DequantizeCastOp>(context) {}
+  LogicalResult matchAndRewrite(mlir::quant::ir::DequantizeCastOp dq_op,
+                                PatternRewriter& rewriter) const override {
+    rewriter.replaceOpWithNewOp<mlir::stablehlo::UniformDequantizeOp>(
+        dq_op, dq_op.getResult().getType(), dq_op.getArg());
+    return success();
+  }
+};
+
+// Replaces quantfork.qcast with stablehlo.uniform_quantize.
+class ConvertQuantizeCastToUniformQuantizePattern
+    : public OpRewritePattern<mlir::quant::ir::QuantizeCastOp> {
+ public:
+  explicit ConvertQuantizeCastToUniformQuantizePattern(MLIRContext* context)
+      : OpRewritePattern<mlir::quant::ir::QuantizeCastOp>(context) {}
+  LogicalResult matchAndRewrite(mlir::quant::ir::QuantizeCastOp q_op,
+                                PatternRewriter& rewriter) const override {
+    rewriter.replaceOpWithNewOp<mlir::stablehlo::UniformQuantizeOp>(
+        q_op, q_op.getResult().getType(), q_op.getArg());
+    return success();
+  }
+};
+
+void PostQuantizePass::runOnOperation() {
+  RewritePatternSet patterns(&getContext());
+  func::FuncOp func = getOperation();
+  MLIRContext* ctx = func.getContext();
+  // TODO: b/307463853 - Consider splitting passes for each pattern set.
+  patterns.add<FoldTrivalRequantizeOp<mlir::quant::ir::QuantizeCastOp>,
+               RemoveVolatileQdqPattern>(ctx);
+  if (failed(applyPatternsGreedily(func, std::move(patterns)))) {
+    signalPassFailure();
+  }
+
+  RewritePatternSet patterns_2(&getContext());
+  patterns_2
+      .add<QuantizeConstPattern, ConvertQuantizeCastToUniformQuantizePattern,
+           ConvertDequantizeCastToUniformDequantizePattern>(ctx);
+  if (failed(applyPatternsGreedily(func, std::move(patterns_2)))) {
+    signalPassFailure();
+  }
+}
+
+}  // namespace
+}  // namespace mlir::quant::stablehlo

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

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

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

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

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

@@ -0,0 +1,111 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <utility>
+
+#include "mlir/Dialect/Quant/IR/Quant.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h"  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/quantization_patterns.h"
+
+namespace mlir::quant::stablehlo {
+
+#define GEN_PASS_DEF_QUANTIZEPASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"
+
+namespace {
+
+// Base struct for quantization.
+template <typename ConcreteT,
+          typename RootOpT = mlir::quant::ir::DequantizeCastOp>
+struct StableHloQuantizationBase
+    : public StableHloQuantizationPattern<ConcreteT,
+                                          mlir::quant::ir::QuantizeCastOp,
+                                          mlir::quant::ir::DequantizeCastOp,
+                                          /*VerifierT=*/void, RootOpT> {
+  explicit StableHloQuantizationBase(MLIRContext* ctx)
+      : StableHloQuantizationPattern<ConcreteT, mlir::quant::ir::QuantizeCastOp,
+                                     mlir::quant::ir::DequantizeCastOp,
+                                     /*VerifierT=*/void, RootOpT>(ctx) {}
+
+  static bool AllowWeightOnlyQuantization(Operation& op) { return false; }
+};
+
+// Quantization rewrite pattern using DQ as the root op.
+struct StableHloQuantization
+    : public StableHloQuantizationBase<StableHloQuantization> {
+  explicit StableHloQuantization(MLIRContext* ctx)
+      : StableHloQuantizationBase<StableHloQuantization>(ctx) {}
+};
+
+// Quantization rewrite pattern using Q as the root op. This is for the
+// quantizable ops without floating-point operands.
+struct StableHloQuantizationReverse
+    : public StableHloQuantizationBase<StableHloQuantizationReverse,
+                                       mlir::quant::ir::QuantizeCastOp> {
+  explicit StableHloQuantizationReverse(MLIRContext* ctx)
+      : StableHloQuantizationBase<StableHloQuantizationReverse,
+                                  mlir::quant::ir::QuantizeCastOp>(ctx) {}
+};
+
+class QuantizePass : public impl::QuantizePassBase<QuantizePass> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(QuantizePass)
+
+  using impl::QuantizePassBase<QuantizePass>::QuantizePassBase;
+
+  explicit QuantizePass(const bool enable_per_channel_quantized_weight) {
+    enable_per_channel_quantized_weight_ = enable_per_channel_quantized_weight;
+  }
+
+ private:
+  void runOnOperation() override;
+};
+
+void QuantizePass::runOnOperation() {
+  ModuleOp module_op = getOperation();
+  MLIRContext& ctx = getContext();
+
+  RewritePatternSet patterns(&ctx);
+  patterns.add<StableHloQuantization, StableHloQuantizationReverse>(&ctx);
+
+  PopulateCommonQuantizationPatterns(ctx, patterns,
+                                     enable_per_channel_quantized_weight_);
+
+  // Quantize all quantizable ops, including ops that are not compute-heavy.
+  PopulateAllQuantizablePatterns(ctx, patterns);
+
+  if (failed(applyPatternsGreedily(module_op, std::move(patterns)))) {
+    // There are cases where no rewrites happen even if a pattern matches,
+    // causing this to result in a convergence failure. Consider this as a
+    // best-effort.
+    module_op.emitWarning("Failed to converge pattern at QuantizePass.");
+  }
+}
+
+}  // namespace
+
+}  // namespace mlir::quant::stablehlo

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

@@ -0,0 +1,114 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/Quant.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
+#include "tensorflow/compiler/mlir/quantization/tensorflow/cc/run_passes.h"
+#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"  // IWYU pragma: keep
+
+#define DEBUG_TYPE "quantize-composite-functions"
+
+namespace mlir::quant::stablehlo {
+
+#define GEN_PASS_DEF_QUANTIZECOMPOSITEFUNCTIONSPASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"
+
+namespace {
+
+using ::tensorflow::quantization::RunPassesOnModuleOp;
+
+class QuantizeCompositeFunctionsPass
+    : public impl::QuantizeCompositeFunctionsPassBase<
+          QuantizeCompositeFunctionsPass> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(QuantizeCompositeFunctionsPass)
+
+  using impl::QuantizeCompositeFunctionsPassBase<
+      QuantizeCompositeFunctionsPass>::QuantizeCompositeFunctionsPassBase;
+
+  explicit QuantizeCompositeFunctionsPass(
+      const bool enable_per_channel_quantized_weight) {
+    enable_per_channel_quantized_weight_ = enable_per_channel_quantized_weight;
+  }
+
+ private:
+  void runOnOperation() override;
+};
+
+void QuantizeCompositeFunctionsPass::runOnOperation() {
+  MLIRContext& ctx = getContext();
+
+  PassManager pm(&ctx);
+  // Intermediate output from QuantizePass will have quantized ops
+  // (XlaCallModuleOps) with quantized input and output types, which are not
+  // allowed in the TF dialect.
+  pm.enableVerifier(false);
+
+  PrepareQuantizePassOptions options;
+  options.enable_per_channel_quantized_weight_ =
+      enable_per_channel_quantized_weight_;
+  // Change this to user-given bit width once we have custom configuration.
+  options.bit_width_ = 8;
+
+  // Insert quantization parameters for weights for ops with `weight_only_ptq`
+  // attribute.
+  pm.addNestedPass<func::FuncOp>(createInsertWeightParamPass());
+
+  // PrepareQuantizePass uses SymbolTable to fetch relevant GEMM ops for
+  // determining quantization attributes. This requires module-level context.
+  pm.addPass(createPrepareQuantizePass(options));
+
+  QuantizePassOptions quantize_options;
+  quantize_options.enable_per_channel_quantized_weight_ =
+      enable_per_channel_quantized_weight_;
+
+  // QuantizePass modifies FuncOps referenced outside of its given scope
+  // and therefore requires a module-level context.
+  pm.addPass(createQuantizePass(quantize_options));
+  pm.addNestedPass<func::FuncOp>(createPostQuantizePass());
+
+  // Convert XlaCallModuleOps lifted but not quantized to func.call op.
+  // The reasons these ops are not quantized may be:
+  // 1. Disabled due to selective quantization.
+  // 2. Not supported, e.g. add op for server.
+  pm.addPass(createXlaCallModuleToCallPass());
+
+  // TODO: b/321729008 - move this implementation to quantization_patterns.cc.
+  if (merge_fusion_with_dequantize_) {
+    pm.addPass(createMergeFusionWithDequantizePass());
+  }
+
+  ModuleOp module_op = getOperation();
+  if (const absl::Status pm_run_status =
+          RunPassesOnModuleOp(mlir_dump_file_name_, pm, module_op);
+      !pm_run_status.ok()) {
+    signalPassFailure();
+  }
+}
+
+}  // namespace
+}  // namespace mlir::quant::stablehlo

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

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

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

@@ -0,0 +1,59 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include <utility>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Rewrite/FrozenRewritePatternSet.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo  // IWYU pragma: keep
+
+namespace mlir::quant::stablehlo {
+
+#define GEN_PASS_DEF_REMOVESHARDINGCUSTOMCALLPASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"
+
+// Include patterns generated from `remove_sharding_custom_call.td`.
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/remove_sharding_custom_call.inc"
+
+class RemoveShardingCustomCallPass
+    : public impl::RemoveShardingCustomCallPassBase<
+          RemoveShardingCustomCallPass> {
+ public:
+  using impl::RemoveShardingCustomCallPassBase<
+      RemoveShardingCustomCallPass>::RemoveShardingCustomCallPassBase;
+
+ private:
+  void runOnOperation() override;
+};
+
+void RemoveShardingCustomCallPass::runOnOperation() {
+  func::FuncOp func_op = getOperation();
+  MLIRContext& ctx = getContext();
+
+  RewritePatternSet patterns(&ctx);
+  populateWithGenerated(patterns);
+
+  FrozenRewritePatternSet frozen_patterns(std::move(patterns));
+  if (failed(applyPatternsGreedily(func_op, frozen_patterns))) {
+    func_op.emitWarning() << "Failed to converge "
+                          << RemoveShardingCustomCallPass::getArgumentName();
+  }
+}
+
+}  // namespace mlir::quant::stablehlo

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

@@ -0,0 +1,27 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+include "stablehlo/dialect/StablehloOps.td"
+
+class IsStringAttrOf<string value> : Constraint<
+  CPred<"::llvm::isa_and_nonnull<StringAttr>($_self) && llvm::cast<StringAttr>($_self).getValue() == \"" # value # "\"">,
+  "Is a string attribute whose value is \"" # value # "\""
+>;
+
+// Removes `stablehlo.custom_call @Sharding`. Assumes this call always accepts
+// one input. Other attributes are ignored.
+def RemoveShardingCustomCall : Pat<
+  (StableHLO_CustomCallOp (variadic $input0), $call_target_name, $_, $_, $_, $_, $_, $_, $_, $_),
+  (replaceWithValue $input0),
+  [(IsStringAttrOf<"Sharding">:$call_target_name)]>;

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

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

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

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

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

@@ -0,0 +1,40 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_PASSES_TESTING_PASSES_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_PASSES_TESTING_PASSES_H_
+
+#include "mlir/Pass/Pass.h"  // from @llvm-project  // IWYU pragma: keep
+
+namespace mlir::quant::stablehlo::testing {
+
+// Identifies predefined `QuantizationSpecs` for
+// `TestLiftQuantizableSpotsAsFunctionsWithQuantizationSpecsPass`. The pass
+// option argument is specified in line comments for each enum value.
+enum class TestQuantizationSpecs {
+  kEmpty,                         // empty
+  kDisableAllDotGeneral,          // disable-all-dot-general
+  kStaticRangePtqToAll,           // static-range-ptq-to-all
+  kStaticRangePtqToComputeHeavy,  // static-range-ptq-to-compute-heavy
+};
+
+// Adds generated pass default constructors or options definitions.
+#define GEN_PASS_DECL
+// Adds generated pass registration functions.
+#define GEN_PASS_REGISTRATION
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/testing/passes.h.inc"
+
+}  // namespace mlir::quant::stablehlo::testing
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_PASSES_TESTING_PASSES_H_

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

@@ -0,0 +1,94 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// Passes only used for testing purposes.
+
+include "mlir/Pass/PassBase.td"
+
+def TestPreCalibrationComponentPass : Pass<"stablehlo-test-pre-calibration-component", "mlir::ModuleOp"> {
+  let summary = "Test-only pass to test the PreCalibrationComponent.";
+  let description = [{
+    Runs the pre calibration passes for post-training quantization with default
+    configuration.
+  }];
+  let dependentDialects = [
+    "mlir::stablehlo::StablehloDialect", "mlir::TF::TensorFlowDialect",
+    "mlir::func::FuncDialect", "mlir::tf_executor::TensorFlowExecutorDialect",
+    "mlir::mhlo::MhloDialect", "mlir::vhlo::VhloDialect",
+  ];
+}
+
+def TestPostCalibrationComponentPass : Pass<"stablehlo-test-post-calibration-component", "mlir::ModuleOp"> {
+  let summary = "Test-only pass to test the PostCalibrationComponent.";
+  let description = [{
+    Runs the post-calibration passes for post-training quantization.
+  }];
+  let options = [
+    Option<"unpack_quantized_types_", "unpack-quantized-types", "bool",
+      /*default=*/"true", "Unpacks ops with uniform quantized types into "
+      "operations without uniform quantized types (mostly i8 or i32).">
+  ];
+  let dependentDialects = [
+    "mlir::stablehlo::StablehloDialect", "mlir::TF::TensorFlowDialect",
+    "mlir::func::FuncDialect", "mlir::mhlo::MhloDialect",
+    "mlir::quant::QuantDialect", "mlir::chlo::ChloDialect",
+    "mlir::vhlo::VhloDialect", "mlir::shape::ShapeDialect",
+    "mlir::quant::ir::TFQuantDialect",
+  ];
+}
+
+def TestTFToStablehloPass : Pass<"stablehlo-test-tf-to-stablehlo", "mlir::ModuleOp"> {
+  let summary = "Test-only pass to test TFToStablehloPasses.";
+  let description = [{
+    Runs the TFToStablehloPasses.
+  }];
+  let dependentDialects = [
+    "mlir::stablehlo::StablehloDialect", "mlir::TF::TensorFlowDialect",
+    "mlir::chlo::ChloDialect", "mlir::quant::QuantDialect",
+    "mlir::mhlo::MhloDialect", "mlir::shape::ShapeDialect",
+    "mlir::sparse_tensor::SparseTensorDialect", "mlir::ub::UBDialect",
+    "mlir::vhlo::VhloDialect",
+  ];
+}
+
+def TestLiftQuantizableSpotsAsFunctionsWithQuantizationSpecsPass :
+    Pass<"stablehlo-test-lift-quantizable-spots-as-functions-with-quantization-specs", "mlir::ModuleOp"> {
+  let summary = "Test-only pass for testing the LiftQuantizableSpotsAsFunctionsPass with a predefined QuantizationSpecs.";
+  let description = [{
+    This test-only pass is the same as `LiftQuantizableSpotsAsFunctionsPass` but
+    has predefined `QuantizationSpecs` to make FileCheck testing easier.
+  }];
+  let options = [
+    Option<"quantization_specs_", "quantization-specs",
+      "mlir::quant::stablehlo::testing::TestQuantizationSpecs",
+      /*default=*/"mlir::quant::stablehlo::testing::TestQuantizationSpecs::kEmpty",
+      "Sets one of the predefined `QuantizationSpecs` for testing.",
+      [{llvm::cl::values(
+        clEnumValN(mlir::quant::stablehlo::testing::TestQuantizationSpecs::kEmpty,
+          "empty", "Uses empty (default) QuantizationSpecs."),
+        clEnumValN(mlir::quant::stablehlo::testing::TestQuantizationSpecs::kDisableAllDotGeneral,
+          "disable-all-dot-general", "Disables all dot_general ops by matching lifted function names"),
+        clEnumValN(mlir::quant::stablehlo::testing::TestQuantizationSpecs::kStaticRangePtqToAll,
+          "static-range-ptq-to-all", "Applies `StaticRangePtq` to all quantizable units."),
+        clEnumValN(mlir::quant::stablehlo::testing::TestQuantizationSpecs::kStaticRangePtqToComputeHeavy,
+          "static-range-ptq-to-compute-heavy", "Applies `StaticRangePtq` to only compute heavy units.")
+      )}]>
+  ];
+  let dependentDialects = [
+      "mlir::func::FuncDialect",
+      "mlir::stablehlo::StablehloDialect",
+      "TF::TensorFlowDialect",
+  ];
+}

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

@@ -0,0 +1,140 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Pass/Pass.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/testing/passes.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"  // IWYU pragma: keep
+#include "tsl/platform/protobuf.h"  // IWYU pragma: keep
+
+namespace mlir::quant::stablehlo::testing {
+
+// NOLINTNEXTLINE - Automatically generated.
+#define GEN_PASS_DEF_TESTLIFTQUANTIZABLESPOTSASFUNCTIONSWITHQUANTIZATIONSPECSPASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/testing/passes.h.inc"
+
+namespace {
+
+using ::stablehlo::quantization::QuantizationSpecs;
+using ::tsl::protobuf::TextFormat;
+// NOLINTNEXTLINE(misc-include-cleaner) - Required for OSS.
+using ::tsl::protobuf::io::ArrayInputStream;
+
+// Empty (default) `QuantizationSpecs` proto.
+constexpr absl::string_view kSpecsEmpty = R"pb(specs
+                                               [])pb";
+
+// Configure `QuantizationSpecs` to disable quantization for all dot_general
+// quantizable units.
+constexpr absl::string_view kSpecsDisableAllDotGeneral =
+    R"pb(specs
+         [ {
+           matcher { function_name { regex: "composite_dot_general_.*" } }
+           method { no_quantization {} }
+         }])pb";
+
+// Configure `QuantizationSpecs` to apply `StaticRangePtq` to all quantizable
+// units.
+constexpr absl::string_view kSpecsStaticRangePtqToAll =
+    R"pb(specs
+         [ {
+           matcher { function_name { regex: ".*" } }
+           method { static_range_ptq {} }
+         }])pb";
+
+// Configure `QuantizationSpecs` to apply `StaticRangePtq` to compute heavy
+// units.
+constexpr absl::string_view kSpecsStaticRangePtqToComputeHeavy =
+    R"pb(specs
+         [ {
+           matcher { function_name { regex: "^.*(conv|dot|gather).*" } }
+           method { static_range_ptq {} }
+         }])pb";
+
+class TestLiftQuantizableSpotsAsFunctionsWithQuantizationSpecsPass
+    : public impl::
+          TestLiftQuantizableSpotsAsFunctionsWithQuantizationSpecsPassBase<
+              TestLiftQuantizableSpotsAsFunctionsWithQuantizationSpecsPass> {
+ public:
+  using impl::TestLiftQuantizableSpotsAsFunctionsWithQuantizationSpecsPassBase<
+      TestLiftQuantizableSpotsAsFunctionsWithQuantizationSpecsPass>::
+      TestLiftQuantizableSpotsAsFunctionsWithQuantizationSpecsPassBase;
+
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
+      TestLiftQuantizableSpotsAsFunctionsWithQuantizationSpecsPass)
+
+ private:
+  void runOnOperation() override;
+};
+
+// `TestQuantizationSpecs` -> predefined `QuantizationSpecs` textproto.
+absl::string_view GetQuantizationSpecsTextProto(
+    const TestQuantizationSpecs test_specs) {
+  switch (test_specs) {
+    case TestQuantizationSpecs::kEmpty:
+      return kSpecsEmpty;
+    case TestQuantizationSpecs::kDisableAllDotGeneral:
+      return kSpecsDisableAllDotGeneral;
+    case TestQuantizationSpecs::kStaticRangePtqToAll:
+      return kSpecsStaticRangePtqToAll;
+    case TestQuantizationSpecs::kStaticRangePtqToComputeHeavy:
+      return kSpecsStaticRangePtqToComputeHeavy;
+  }
+}
+
+// Parses a text proto into a `QuantizationSpecs` proto. Returns
+// `InvalidArgumentError` if `text_proto` is invalid.
+absl::StatusOr<QuantizationSpecs> ParseTextProto(
+    const absl::string_view text_proto) {
+  QuantizationSpecs quantization_specs;
+  TextFormat::Parser parser;
+  ArrayInputStream input_stream(text_proto.data(), text_proto.size());
+  if (parser.Parse(&input_stream, &quantization_specs)) {
+    return quantization_specs;
+  }
+  return absl::InvalidArgumentError("Could not parse text proto.");
+}
+
+void TestLiftQuantizableSpotsAsFunctionsWithQuantizationSpecsPass::
+    runOnOperation() {
+  PassManager pass_manager{&getContext()};
+
+  // Construct `QuantizationSpecs` from the pass option `quantization-specs`.
+  const absl::StatusOr<QuantizationSpecs> quantization_specs =
+      ParseTextProto(GetQuantizationSpecsTextProto(quantization_specs_));
+  if (!quantization_specs.ok()) {
+    signalPassFailure();
+    return;
+  }
+
+  pass_manager.addPass(
+      quant::stablehlo::CreateLiftQuantizableSpotsAsFunctionsPass(
+          *quantization_specs));
+
+  if (failed(pass_manager.run(getOperation()))) {
+    signalPassFailure();
+  }
+}
+
+}  // namespace
+}  // namespace mlir::quant::stablehlo::testing

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

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

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

@@ -0,0 +1,67 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "absl/status/statusor.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo  // IWYU pragma: keep
+#include "stablehlo/dialect/VhloOps.h"  // from @stablehlo  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/config.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/pre_calibration.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
+#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"  // IWYU pragma: keep
+#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"  // IWYU pragma: keep
+
+namespace mlir::quant::stablehlo::testing {
+
+#define GEN_PASS_DEF_TESTPRECALIBRATIONCOMPONENTPASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/testing/passes.h.inc"
+
+namespace {
+
+using ::stablehlo::quantization::ExpandPresets;
+using ::stablehlo::quantization::PopulateDefaults;
+using ::stablehlo::quantization::QuantizationConfig;
+
+class TestPreCalibrationComponentPass
+    : public impl::TestPreCalibrationComponentPassBase<
+          TestPreCalibrationComponentPass> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestPreCalibrationComponentPass)
+
+ private:
+  void runOnOperation() override;
+};
+
+void TestPreCalibrationComponentPass::runOnOperation() {
+  ModuleOp module_op = getOperation();
+  MLIRContext& ctx = getContext();
+
+  // Simply runs the PreCalibrationComponent with a default configuration.
+  PreCalibrationComponent component(&ctx);
+  QuantizationConfig quantization_config{};
+  quantization_config.mutable_static_range_ptq_preset();
+  quantization_config = ExpandPresets(PopulateDefaults(quantization_config));
+  if (!component.Run(module_op, quantization_config).ok()) {
+    signalPassFailure();
+  }
+}
+
+}  // namespace
+}  // namespace mlir::quant::stablehlo::testing

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

@@ -0,0 +1,70 @@
+/* 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 "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Dialect/Quant/IR/Quant.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Dialect/UB/IR/UBOps.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "stablehlo/dialect/ChloOps.h"  // from @stablehlo  // IWYU pragma: keep
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo  // IWYU pragma: keep
+#include "stablehlo/dialect/VhloOps.h"  // from @stablehlo  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
+#include "tensorflow/compiler/mlir/quantization/tensorflow/cc/run_passes.h"
+#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
+#include "tensorflow/compiler/mlir/quantization/tensorflow/quantize_preprocess.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"  // IWYU pragma: keep
+#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"  // IWYU pragma: keep
+
+namespace mlir::quant::stablehlo::testing {
+
+#define GEN_PASS_DEF_TESTTFTOSTABLEHLOPASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/testing/passes.h.inc"
+
+namespace {
+
+using ::tensorflow::quantization::AddTFToStablehloPasses;
+using ::tensorflow::quantization::RunPassesOnModuleOp;
+
+class TestTFToStablehloPass
+    : public impl::TestTFToStablehloPassBase<TestTFToStablehloPass> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestTFToStablehloPass)
+
+ private:
+  void runOnOperation() override;
+};
+
+void TestTFToStablehloPass::runOnOperation() {
+  ModuleOp module_op = getOperation();
+  MLIRContext* ctx = &getContext();
+  mlir::PassManager pm(ctx);
+
+  AddTFToStablehloPasses(pm);
+  if (!RunPassesOnModuleOp(
+           /*mlir_dump_file_name=*/"test_tf_to_stablehlo_pass", pm, module_op)
+           .ok()) {
+    return signalPassFailure();
+  }
+}
+
+}  // namespace
+}  // namespace mlir::quant::stablehlo::testing

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

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

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

@@ -0,0 +1,132 @@
+/* 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 "llvm/ADT/STLExtras.h"
+#include "llvm/Support/Casting.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/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/Region.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/xla_call_module_attrs.h"
+#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
+
+namespace mlir::quant::stablehlo {
+
+#define GEN_PASS_DEF_UNWRAPXLACALLMODULEOPPASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"
+
+namespace {
+
+// Unwraps XlaCallModule ops without quantizable trait that call function with
+// '_from_xla_call_module' trait.
+class UnwrapXlaCallModuleOpPass
+    : public impl::UnwrapXlaCallModuleOpPassBase<UnwrapXlaCallModuleOpPass> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(UnwrapXlaCallModuleOpPass)
+
+  explicit UnwrapXlaCallModuleOpPass() = default;
+
+ private:
+  void runOnOperation() override;
+};
+
+void UnwrapXlaCallModuleOp(TF::XlaCallModuleOp call_op,
+                           SymbolTable& symbol_table) {
+  // Do not inline lifted quantized functions used for fusing patterns.
+  // TODO - b/310539922: Remove reference to TF/TFL utils.
+  if (call_op->hasAttr(kQuantTraitAttrName)) {
+    return;
+  }
+
+  auto function_name = call_op
+                           ->getAttrOfType<FlatSymbolRefAttr>(
+                               TF::kStablehloEntryFunctionAttrName)
+                           .getValue();
+  func::FuncOp func_op = symbol_table.lookup<func::FuncOp>(function_name);
+
+  // We should not unwrap if the function is not from
+  // ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPass.
+  if (!func_op->hasAttr(TF::kFromXlaCallModuleAttrName)) {
+    return;
+  }
+
+  MLIRContext* context = call_op.getContext();
+  OpBuilder builder(context);
+  builder.setInsertionPointAfter(call_op);
+
+  IRMapping arg_mapper;
+  bool call_op_has_platform_index_arg = call_op.getPlatforms().size() > 1;
+  // Add an argument for platform_index. This allows for multiple platforms.
+  // TODO: b/310291615 - find a better way for multi-platform support.
+  if (call_op_has_platform_index_arg) {
+    arg_mapper.map(func_op.getArgument(0),
+                   mhlo::ConstantOp::create(builder, func_op.getLoc(),
+                                            builder.getI16IntegerAttr(0)));
+  }
+  for (auto [func_arg, operand] : llvm::zip_equal(
+           func_op.getArguments().take_back(call_op.getNumOperands()),
+           call_op.getOperands())) {
+    arg_mapper.map(func_arg, operand);
+  }
+
+  Region& function_body = func_op.getBody();
+  IRMapping new_op_mapper;
+  for (Operation& op : function_body.getOps()) {
+    if (llvm::isa<func::ReturnOp>(op)) {
+      for (auto [call_result, return_value] :
+           llvm::zip_equal(call_op.getResults(), op.getOperands())) {
+        Value new_result = new_op_mapper.lookup(return_value);
+
+        call_result.replaceAllUsesWith(new_result);
+      }
+      continue;
+    }
+
+    Operation& new_op = *builder.clone(op, arg_mapper);
+    for (auto [result, new_result] :
+         llvm::zip_equal(op.getResults(), new_op.getResults())) {
+      new_op_mapper.map(result, new_result);
+    }
+  }
+
+  call_op.erase();
+}
+
+void UnwrapXlaCallModuleOpPass::runOnOperation() {
+  ModuleOp module_op = getOperation();
+  SymbolTable symbol_table(module_op);
+
+  for (auto func_op : module_op.getOps<func::FuncOp>()) {
+    Region& function_body = func_op.getBody();
+
+    function_body.walk([&](TF::XlaCallModuleOp call_op) {
+      UnwrapXlaCallModuleOp(call_op, symbol_table);
+    });
+  }
+}
+
+}  // namespace
+
+}  // namespace mlir::quant::stablehlo

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

@@ -0,0 +1,84 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <utility>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+
+namespace mlir::quant::stablehlo {
+
+#define GEN_PASS_DEF_XLACALLMODULETOCALLPASS
+#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc"
+
+namespace {
+
+// Converts XlaCallModuleOps to func.call.
+class XlaCallModuleToCallPass
+    : public impl::XlaCallModuleToCallPassBase<XlaCallModuleToCallPass> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(XlaCallModuleToCallPass)
+
+  explicit XlaCallModuleToCallPass() = default;
+
+ private:
+  void runOnOperation() override;
+};
+
+// Converts XlaCallModuleOps to func.call.
+class XlaCallModuleOpToCallOp : public OpRewritePattern<TF::XlaCallModuleOp> {
+ public:
+  explicit XlaCallModuleOpToCallOp(MLIRContext* context)
+      : OpRewritePattern<TF::XlaCallModuleOp>(context) {}
+
+  LogicalResult matchAndRewrite(TF::XlaCallModuleOp op,
+                                PatternRewriter& rewriter) const override {
+    auto module_op = op->getParentOfType<ModuleOp>();
+    SymbolTable symbol_table(module_op);
+
+    auto entry_func_op = dyn_cast_or_null<func::FuncOp>(
+        symbol_table.lookup(GetEntryFunctionName(op)));
+    if (!entry_func_op) return failure();
+
+    // Replace the XlaCallModuleOp with a new CallOp.
+    rewriter.replaceOpWithNewOp<func::CallOp>(op, entry_func_op, op.getArgs());
+    return success();
+  }
+};
+
+void XlaCallModuleToCallPass::runOnOperation() {
+  ModuleOp module_op = getOperation();
+  MLIRContext* ctx = module_op.getContext();
+  RewritePatternSet patterns(&getContext());
+  patterns.add<XlaCallModuleOpToCallOp>(ctx);
+  if (failed(applyPatternsGreedily(module_op, std::move(patterns)))) {
+    signalPassFailure();
+  }
+}
+
+}  // namespace
+}  // namespace mlir::quant::stablehlo