#include "tensorflow/compiler/mlir/lite/transforms/passes.h.inc"

// Module pass to optimize TensorFlow functional ops.
struct OptimizeFunctionalOpsPass
    : public impl::OptimizeFunctionalOpsPassBase<OptimizeFunctionalOpsPass> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(OptimizeFunctionalOpsPass)

  void runOnOperation() override;
};

// Updates function return type of the given functions to match the terminator

// Creates an instance of the TensorFlow Lite dialect SplitMergedOperandsPass.
std::unique_ptr<OperationPass<func::FuncOp>> CreateSplitMergedOperandsPass();

// Creates an instance of the TensorFlow Lite dialect OptimizeFunctionalOpsPass.
std::unique_ptr<OperationPass<ModuleOp>> CreateOptimizeFunctionalOpsPass();

std::unique_ptr<OperationPass<func::FuncOp>> CreateModifyIONodesPass(
    mlir::Type input_type, mlir::Type output_type);

  let summary = "Optimize FC with BatchMatmul within the TensorFlow Lite dialect";
  let constructor = "CreateOptimizeBatchMatmulPass()";
  let dependentDialects = ["TFL::TensorFlowLiteDialect"];
}

def OptimizeFunctionalOpsPass : Pass<"tfl-optimize-functional-ops", "mlir::ModuleOp"> {
  let summary = "Optimize TensorFlow functional op";
  let constructor = "CreateOptimizeFunctionalOpsPass()";
}