// This pattern checks if an i8 arith::ConstantOp tensor has all values within
// the INT4 range, i.e. [-8,7] and converts it into i4 if so. This assumes that
// the input is sign-extended two's complement.
class CheckRangeAndConvertI8ToI4 : public OpRewritePattern<arith::ConstantOp> {
 public:
  using OpRewritePattern<arith::ConstantOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(arith::ConstantOp op,

    // 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;
    mlir::DenseIntElementsAttr broadcast_dims;

    // Match and capture values/attributes.
    Value lhs = mul_op.getLhs();

  Value GetZeroPointsValue() { return call_op_.getOperand(2); }

  stablehlo::ConstantOp GetZeroPointsConstantOp() {
    return cast<stablehlo::ConstantOp>(GetZeroPointsValue().getDefiningOp());
  }

  stablehlo::ConstantOp GetInverseScalesConstantOp() {
    return cast<stablehlo::ConstantOp>(GetInverseScalesValue().getDefiningOp());
  }

  ElementsAttr GetZeroPointsValueAttr() {

class PrepareDRQQuantizableOp : public OpRewritePattern<arith::ConstantOp> {
 public:
  explicit PrepareDRQQuantizableOp(MLIRContext* context,
                                   const quant::QuantizationSpecs& quant_specs,
                                   OpSet op_set,
                                   bool enable_per_channel_quantization)
      : OpRewritePattern<arith::ConstantOp>(context),
        quant_specs_(quant_specs),

    // Cannot use the pattern rewriter because the driver applies folders by
    // default.
    module.walk([&](mhlo::ConstantOp const_op) {
      UnfoldSplatConstant(&op_builder, const_op);
    });
  }

 private:
  void UnfoldSplatConstant(mlir::OpBuilder* op_builder,
                           mhlo::ConstantOp const_op) const {
    auto splat_elements_attr =
        mlir::dyn_cast<SplatElementsAttr>(const_op.getValue());

  auto attr = CreateTypedAttr(dense_type, value);

  return rewriter->create<arith::ConstantOp>(loc, dense_type,
                                             cast<TypedAttr>(*attr));
}

absl::StatusOr<arith::ConstantOp> CreateConstOpWithSingleValue(
    PatternRewriter* rewriter, Location loc, ShapedType shaped_type,
    int value) {
  ShapedType scalar_type =

template <typename Op, typename ElementType = Type, typename ValType,
          typename Convert>
static Attribute BinaryFolder(Op *op) {
  auto lhs_op = op->getLhs().template getDefiningOp<mhlo::ConstantOp>();
  auto rhs_op = op->getRhs().template getDefiningOp<mhlo::ConstantOp>();
  if (!lhs_op || !lhs_op) return {};

  auto lhs = dyn_cast_or_null<DenseElementsAttr>(lhs_op.getValue());
  auto rhs = dyn_cast_or_null<DenseElementsAttr>(rhs_op.getValue());

      .replaceAllUsesWith(cloned_const_op->getResult(0));
}

bool IsConstOrQConstInt(Operation* op) {
  if (!llvm::isa<TFL::ConstOp, TFL::QConstOp, arith::ConstantOp>(op))
    return false;

  if (auto arith_const_op = dyn_cast_or_null<arith::ConstantOp>(op)) {
    // arith ConstOp path.
    auto type =
        mlir::cast<ShapedType>(arith_const_op.getType()).getElementType();

namespace mlir {
namespace TFL {

// Returns a Constant op with a single value.
absl::StatusOr<arith::ConstantOp> CreateConstOpWithSingleValue(
    PatternRewriter* rewriter, Location loc, ShapedType shaped_type, int value);

// Returns a Constant op with a splat vector value.
absl::StatusOr<arith::ConstantOp> CreateConstOpWithVectorValue(
    PatternRewriter* rewriter, Location loc, ShapedType shaped_type, int value);

  static std::optional<FoldAdaptor> Create(Operation* operation) {
    auto foldable_opr = [](Value val) -> bool {
      return !llvm::isa<BlockArgument>(val) &&
             llvm::isa<stablehlo::ConstantOp>(val.getDefiningOp());
    };
    if (!llvm::all_of(operation->getOperands(), foldable_opr)) {
      return std::nullopt;
    }
    return FoldAdaptor(operation);
  }