dq.getArg());
      dq.getResult().replaceAllUsesWith(dcast);
      if (auto extra_attr = op->getAttr(mlir::quant::kVolatileOpAttrName)) {
        dcast->setAttr(mlir::quant::kVolatileOpAttrName, extra_attr);
      }
      dq.erase();
    } else if (auto q = llvm::dyn_cast<quantfork::QuantizeCastOp>(op)) {
      auto out_type = q.getResult().getType();

                                PatternRewriter& rewriter) const override {
    auto input_op = op.getArg().getDefiningOp();
    if (auto q = llvm::dyn_cast_or_null<quantfork::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

                                PatternRewriter& rewriter) const override {
    auto input_op = op.getArg().getDefiningOp();
    if (auto q = llvm::dyn_cast_or_null<quantfork::QuantizeCastOp>(input_op)) {
      if (!q->getAttr(kVolatileOpAttrName)) return failure();

      if (remove_volatile_ops_type == kPreserveInputsAndOutputs) {
        // Don't remove leading and trailing QDQ for PTQ workflow, so the io
        // modifying lib can work correctly.

    if (!preceding_qcast) return failure();

    auto new_qcast = rewriter.create<quantfork::QuantizeCastOp>(
        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:

    if (!preceding_qcast) return failure();

    auto new_qcast = rewriter.create<quantfork::QuantizeCastOp>(
        q_op.getLoc(), q_op.getType(), preceding_qcast.getArg());
    new_qcast->setAttr(kVolatileOpAttrName, rewriter.getUnitAttr());
    q_op->replaceAllUsesWith(new_qcast);
    return success();
  }
};

bool PrepareQuantizePass::ContainsQuantizeOps(func::FuncOp func) {
  for (const auto& op : func.getOps()) {

    auto input_op = op.getInput().getDefiningOp();
    if (auto q = llvm::dyn_cast_or_null<QuantizeOp>(input_op)) {
      if (!q->getAttr(mlir::quant::kVolatileOpAttrName)) return failure();

      if (remove_volatile_ops_type == kPreserveInputsAndOutputs) {
        // Don't remove leading and trailing QDQ for PTQ workflow, so the io
        // modifying lib can work correctly.

namespace quant {

// A unit attribute can be attached to the quantize/dequantize ops which are
// added by the quantization passes. These ops can be removed erased without
// losing accuracy.
inline constexpr char kVolatileOpAttrName[] = "volatile";

// Following attributes are used to mark ops that are not quantizable during
// debug model generation process for whole-model verify mode. If these

  // quantization pass. These ops can be removed without losing original
  // program accuracy.
  // TODO: b/323478683 - Make the attribute being part of op definition.
  quantize->setAttr(kVolatileOpAttrName, builder_.getUnitAttr());

  // `original_result` has a use to `quantize`, so this will replace that use
  // by the result of `dequantize`. Remember to reset that use afterwards
  value.replaceAllUsesWith(dequantize);