namespace mlir {
namespace TFTPU {

// Creates a pass that rewrites `tf_device.launch_func` on TPUs into TPU runtime
// ops.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> CreateTPURewritePass(
    llvm::StringRef module_name = llvm::StringRef());

// Creates a pass that adds ops which perform formatting on variables at
// run-time according to compilation result.
std::unique_ptr<OperationPass<ModuleOp>>

}

}  // namespace

void AddTPULowerClusterToRuntimeOpsPassPipeline(OpPassManager& pm,
                                                llvm::StringRef module_name) {
  pm.addPass(mlir::TFTPU::CreateTPURewritePass(module_name));
  pm.addPass(mlir::createSymbolDCEPass());
  pm.addNestedPass<FuncOp>(
      mlir::TFDevice::CreateReplicateInvariantOpHoistingPass());