#include "mlir/IR/PatternMatch.h"  // from @llvm-project

namespace mlir {
namespace TF {

// Populates patterns to perform optimizations specific to tf.data operations.
void PopulateTFDataOptimizationPatterns(MLIRContext *context,
                                        RewritePatternSet *patterns);

}  // namespace TF
}  // namespace mlir

    rewriter.replaceOp(op, {fused_op.getHandle()});
    return failure();
  }
};

#include "tensorflow/compiler/mlir/tensorflow/transforms/generated_tf_data_optimization.inc"
}  // namespace

void PopulateTFDataOptimizationPatterns(MLIRContext *context,
                                        RewritePatternSet *patterns) {
  patterns->add<FuseParallelMapAndBatch>(context);
  populateWithGenerated(*patterns);
}

}  // namespace TF

struct TFDataOptimization
    : public impl::TFDataOptimizationPassBase<TFDataOptimization> {
  void runOnOperation() override {
    RewritePatternSet patterns(&getContext());
    mlir::TF::PopulateTFDataOptimizationPatterns(&getContext(), &patterns);

    (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
  }
};

}  // namespace