message: Use OpRewritePattern instead of RewritePattern for ConvertFusedBatchNorm
diff: @@ -727,29 +727,28 @@ struct ConvertTFBroadcastTo : public RewritePattern {
  }
  };
  -struct ConvertFusedBatchNorm : public RewritePattern {
  +struct ConvertFusedBatchNorm : public OpRewritePattern<TF::FusedBatchNormOp> {
  explicit ConvertFusedBatchNorm(MLIRContext *context)
  - : RewritePattern(TF::FusedBatchNormOp::getOperationName(), 1, context) {}
  + : OpRewritePattern<TF::FusedBatchNormOp>(context) {}
  - LogicalResult matchAndRewrite(Operation *op,
  + LogicalResult matchAndRewrite(TF::FusedBatchNormOp tf_fused_batch_norm_op,
  PatternRewriter &rewriter) const override {
  - auto tf_fused_batch_norm_op = cast<TF::FusedBatchNormOp>(op);
  -
  - // FusedBatchNormV3 expects a 5th output reserve_space_3,
  - // but the output is unused; it doesn't matter what we pass there.
  - rewriter.replaceOpWithNewOp<TF::FusedBatchNormV3Op>(
  - op, tf_fused_batch_norm_op.y().getType(),
  - tf_fused_batch_norm_op.batch_mean().getType(),
  - tf_fused_batch_norm_op.batch_variance().getType(),
  - tf_fused_batch_norm_op.reserve_space_1().getType(),
  - tf_fused_batch_norm_op.reserve_space_2().getType(),
  - /*reserve_space_3=*/tf_fused_batch_norm_op.reserve_space_2().getType(),
  - tf_fused_batch_norm_op.x(), tf_fused_batch_norm_op.scale(),
  - tf_fused_batch_norm_op.offset(), tf_fused_batch_norm_op.mean(),
  - tf_fused_batch_norm_op.variance(), tf_fused_batch_norm_op.epsilon(),
  - tf_fused_batch_norm_op.exponential_avg_factor(),
  - tf_fused_batch_norm_op.data_format(),
  - tf_fused_batch_norm_op.is_training());
  + const auto &old_result_types = tf_fused_batch_norm_op.getResultTypes();
  + llvm::SmallVector<Type, 6> new_result_types(old_result_types.begin(),
  + old_result_types.end());
  + // reserve_space_3
  + new_result_types.push_back(
  + UnrankedTensorType::get(FloatType::getF32(rewriter.getContext())));
  +
  + OperationState new_state(tf_fused_batch_norm_op.getLoc(),
  + TF::FusedBatchNormV3Op::getOperationName(),
  + tf_fused_batch_norm_op.getOperands(),
  + new_result_types,
  + tf_fused_batch_norm_op.getAttrs());
  + Operation *tf_fused_batch_norm_op_v3 = rewriter.createOperation(new_state);
  +
  + rewriter.replaceOp(tf_fused_batch_norm_op,
  + tf_fused_batch_norm_op_v3->getResults());
  return success();
  }
  };
  