#include <utility>

#include "llvm/ADT/StringExtras.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project

namespace mlir {
namespace quant {

constexpr StringRef kTfEntryFunctionAttr = "tf.entry_function";

void AddEntryFunctionInput(StringRef input_name, func::FuncOp func_op) {
  auto entry_func_attr =
      func_op->getAttrOfType<DictionaryAttr>(kTfEntryFunctionAttr);
  if (!entry_func_attr) return;

  const StringSet<> noinline_functions = GetNoinlineFunctionsSet();

  func::FuncOp func_op = getOperation();
  Builder builder(&getContext());

  // Adds the `tf._noinline = true` attribute to the function if the name
  // matches.
  if (noinline_functions.contains(func_op.getSymName())) {
    func_op->setAttr(kTfNoinlineAttr, builder.getBoolAttr(true));
    LLVM_DEBUG(llvm::dbgs()

  Block &nested_block = nested_module->getRegion(0).front();
  for (func::FuncOp func_op :
       llvm::make_early_inc_range(nested_block.getOps<func::FuncOp>())) {
    if (!symbol_table.lookupSymbolIn(getOperation(), func_op.getName())) {
      nested_block.getOperations().remove(func_op.getOperation());
      symbol_table.insert(func_op.getOperation());
    }
  }
  nested_module->erase();
}

}  // namespace

  // Return if the function contains only non-side-effecting ops or TensorArray
  // ops.
  bool HasAtMostTensorArrayEffect(mlir::func::FuncOp func_op) const {
    return set_.count(func_op) > 0;
  }

 private:
  llvm::DenseSet<mlir::func::FuncOp> set_;
};

}  // namespace tfrt_compiler
}  // namespace tensorflow

class CostAnalysis {
 public:
  explicit CostAnalysis(
      mlir::func::FuncOp func_op,
      const tfrt_stub::CostRecorder* cost_recorder = nullptr) {
    cost_recorder_ = cost_recorder;
    AnalyzeArguments(func_op);
    AnalyzeBlock(&func_op.front());
  }

  int64_t GetCost(mlir::Operation* op) const;

 private:
  void AnalyzeArguments(mlir::func::FuncOp func_op);
  void AnalyzeBlock(mlir::Block* block);

    return "Add remarks based on cost analysis for testing purpose.";
  }
  void runOnOperation() override {
    const auto& cost_analysis = getAnalysis<CostAnalysis>();

    auto func_op = getOperation();
    for (auto& op : func_op.front()) {
      op.emitRemark() << "Cost: " << cost_analysis.GetCost(&op);
    }
  }
};

static mlir::PassRegistration<TestCostAnalysis> pass;

}  // namespace tfrt_compiler

    }
    return failure();
  }
};

namespace {

bool IsShloMainFuncOp(func::FuncOp func_op) {
  if (func_op == nullptr) {
    return false;
  }

  if (!func_op.getSymName().contains(kStablehloModuleDefaultEntryFuncName)) {
    return false;
  }

  if (func_op.getSymVisibility() == "nested" ||
      func_op.getSymVisibility() == "private") {
    return false;
  }

  return true;
}

};

void RemoveShardingCustomCallPass::runOnOperation() {
  func::FuncOp func_op = getOperation();
  MLIRContext& ctx = getContext();

  RewritePatternSet patterns(&ctx);
  populateWithGenerated(patterns);

  FrozenRewritePatternSet frozen_patterns(std::move(patterns));
  if (failed(applyPatternsAndFoldGreedily(func_op, frozen_patterns))) {
    func_op.emitWarning() << "Failed to converge "

};

// Folds tf.IdentityOp and tf.IdentityNOp if op device and the argument devices
// from the defining ops match.
void TensorDeviceCopyConversionPass::runOnOperation() {
  func::FuncOp func_op = getOperation();

  auto should_fold_op_func = [&func_op](const Value &arg,
                                        const StringAttr &op_device) {
    // In TFRT TPU, tensor transfer is handled specifically by D2H and

  return op->getName().getStringRef().str() == kClusterFuncOpName;
}

void VerifyInputDialectToExecutorPass::runOnOperation() {
  Operation* func_op = getOperation();

  auto walk_result = func_op->walk([&](Operation* op) {
    if (!tensorflow::tf2xla::internal::IsInBridgeAcceptableDialects(op)) {
      std::string error = "op is in dialect " +