SmallVector<Value> ConstantFoldOpIfPossible(Operation* op);

// This pattern tries to constant-fold the quantizable operands of supported
// TF operations.
struct ConstantFoldQuantizableOperands : public RewritePattern {
 public:
  explicit ConstantFoldQuantizableOperands(MLIRContext* context)
      : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}
  LogicalResult matchAndRewrite(Operation* op,

  const auto test_func =
      module_op_ref->lookupSymbol<func::FuncOp>("test_fold_constant");
  ASSERT_THAT(test_func, NotNull());

  RewritePatternSet patterns(ctx_.get());
  patterns.add<ConstantFoldQuantizableOperands>(ctx_.get());
  EXPECT_TRUE(
      succeeded(applyPatternsAndFoldGreedily(test_func, std::move(patterns))));

  auto depthwise_conv_op =
      FindOperationOfType<TF::DepthwiseConv2dNativeOp>(test_func);

  OpBuilder builder(op);
  SmallVector<Value> results;
  if (failed(FoldOperation(builder, op, results))) {
    return op->getResults();
  }
  return results;
}

LogicalResult ConstantFoldQuantizableOperands::matchAndRewrite(
    Operation* op, PatternRewriter& rewriter) const {
  absl::flat_hash_set<int> quantizable_operands = GetQuantizableOperands(op);
  if (quantizable_operands.empty()) return failure();

  // with a constant operand to a preceding affine operation.
  RewritePatternSet patterns(ctx);
  populateWithGenerated(patterns);
  patterns.add<RemoveIdentity, ConstantFoldQuantizableOperands>(ctx);
  if (op_set_ != OpSet::XLA) {
    // Convert Einsum into BatchMatMul for non-XLA opsets.
    // For the uniform opset, it is requested to maintain the BatchMatmul logic.