// * 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<mlir::stablehlo::ConvolutionOp>::OpRewritePattern;

  LogicalResult match(mlir::stablehlo::ConvolutionOp op) const override {
    // Handles 2D convolutions only.
    if (!HasRankOf(op.getOperand(0), /*rank=*/4) ||

  LogicalResult matchAndRewrite(mhlo::MulOp mul_op,
                                PatternRewriter &rewriter) const override {
    // Variables for capturing values and attributes used while creating ops.
    mhlo::ConvolutionOp conv_op;
    Operation *bcast_or_const_op;
    shape::ShapeOfOp shape_of_op;
    mhlo::ConstantOp filter;
    mhlo::ConstantOp multiplier;
    mlir::ElementsAttr filter_value, mul_value;

//   * Not a depthwise convolution.
class RewriteQuantizedConvolutionOp
    : public OpRewritePattern<stablehlo::ConvolutionOp> {
 public:
  using OpRewritePattern<stablehlo::ConvolutionOp>::OpRewritePattern;
  LogicalResult match(stablehlo::ConvolutionOp op) const override {
    const bool has_i32_output =
        IsI32F32UniformQuantizedPerAxisType(GetElementType(op.getResult()));

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)

  }
}

template <int num_spatial_dims>
class ConvertNdConvOp : public OpConversionPattern<mhlo::ConvolutionOp> {
 public:
  using OpConversionPattern::OpConversionPattern;

  LogicalResult matchAndRewrite(
      mhlo::ConvolutionOp conv_op, OpAdaptor adaptor,
      ConversionPatternRewriter& rewriter) const final {
    SetDefaultConvAttributes(conv_op, rewriter);

// ```
class ComposeUniformQuantizedConvolutionOp
    : public OpRewritePattern<stablehlo::ConvolutionOp> {
 public:
  using OpRewritePattern<stablehlo::ConvolutionOp>::OpRewritePattern;

  LogicalResult match(stablehlo::ConvolutionOp op) const final {
    // Verify operands' types.
    for (Type operand_type : op.getOperandTypes()) {
      if (Type element_type =

    if (!quantization_method.has_static_range_ptq()) {
      return failure();
    }
    return MatchGemmStyleOp<ConvolutionOp>(entry_func_op);
  }

  void rewrite(func::FuncOp entry_func_op, const Method& quantization_method,
               PatternRewriter& rewriter) const override {
    RewriteGemmStyleOp<ConvolutionOp>(
        entry_func_op, rewriter,
        enable_per_channel_quantized_weight_ &&

#include "tsl/platform/status_matchers.h"

namespace mlir::quant {
namespace {

using ::mlir::stablehlo::AddOp;
using ::mlir::stablehlo::ConstantOp;
using ::mlir::stablehlo::ConvolutionOp;
using ::mlir::stablehlo::DotGeneralOp;
using ::mlir::stablehlo::SubtractOp;
using ::testing::ElementsAreArray;
using ::testing::Eq;
using ::testing::IsEmpty;
using ::testing::IsNull;
using ::testing::NotNull;

    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());

}

def NchwConvolutionToNhwcPass : Pass<"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