// some special cases.
    if (auto func = llvm::dyn_cast<func::FuncOp>(op)) {
      AddRegionSideEffectsForOp(func.getBody(), op);
    } else if (auto call = llvm::dyn_cast<CallOpInterface>(op)) {
      func::FuncOp func_op = dyn_cast<func::FuncOp>(
          call.resolveCallable(&symbol_table_collection_));
      if (func_op) {

    op->erase();
  }
}
}  // namespace

std::unique_ptr<OperationPass<func::FuncOp>> CreateRaiseCustomOpsPass() {
  return std::make_unique<RaiseCustomOpsPass>();
}

// Creates an instance of the TensorFlow Lite dialect raise custom op pass.
std::unique_ptr<OperationPass<func::FuncOp>> CreateRaiseCustomOpsPass(
    const std::vector<std::string> &target_ops) {

  absl::StatusOr<ModuleOp> pre_calibration_result =
      component.Run(*module_op, quantization_config);

  EXPECT_THAT(pre_calibration_result, IsOk());

  SmallVector<func::FuncOp> func_ops;
  for (auto func_op : pre_calibration_result->getOps<func::FuncOp>()) {
    func_ops.push_back(func_op);
  }
  ASSERT_THAT(func_ops, SizeIs(2));
  EXPECT_THAT(func_ops, Contains(HasSymName("main")));

  body_yield_operand.replaceAllUsesWith(body_block_argument);
  candidate_op->erase();
  return true;
}

void RemoveUnusedWhileResultsPass::runOnOperation() {
  func::FuncOp func = getOperation();
  // Try to prune defining ops of unused results.
  func.walk([&](TF::WhileRegionOp while_op) {
    for (OpResult result : while_op.getResults()) {
      TryPruneResultDefiningOp(while_op, result);

    return definition->getNumOperands() == 0 &&
           definition->getNumResults() == 1;
  };

  // Check all generator aliases (ops or function argument) are on CPU.
  func::FuncOp func = iterator_op->getParentOfType<func::FuncOp>();
  return llvm::all_of(aliases, [&](Value alias) {
    // Ignore non-generator aliases.
    if (!is_generator(alias)) return true;

    StringAttr device;

    ModuleOp module = op->getParentOfType<ModuleOp>();
    SymbolTable symbol_table(module);

    const func::FuncOp ref_func =
        dyn_cast_or_null<func::FuncOp>(symbol_table.lookup(f_attr.getValue()));
    func::FuncOp new_ref_func = dyn_cast<func::FuncOp>(ref_func->clone());
    return symbol_table.insert(new_ref_func);
  }

  LogicalResult match(LiftedOpT op) const override {

                                builder.getContext()));
}

// Returns the aliasing argument number of a fucntion return value if it simply
// forwards the argument. Otherwise, returns -1.
int64_t FindAliasedInput(func::FuncOp func, int64_t return_index) {
  Value return_val = func.front().getTerminator()->getOperand(return_index);
  auto maybe_arg = mlir::dyn_cast<BlockArgument>(return_val);
  if (!maybe_arg) return -1;

  mlir::Builder builder(context);
  StringAttr resource_name_id = builder.getStringAttr(kResourceNameArgAttr);

  SmallSet<StringRef, 4> resource_names;

  for (func::FuncOp func : module.getOps<func::FuncOp>()) {
    for (int i = 0, e = func.getNumArguments(); i < e; ++i) {
      auto resource_arg =
          func.getArgAttrOfType<StringAttr>(i, kResourceNameArgAttr);
      if (!resource_arg) continue;

  void runOnOperation() override;
};

// 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

LogicalResult EncapsulateFuncAndSerialize(const std::string& module_name,
                                          func::FuncOp entry_func,
                                          std::string* serialized_func_module) {
  ModuleOp module = entry_func->getParentOfType<ModuleOp>();
  SymbolTable entry_module_table(module);
  llvm::SmallVector<func::FuncOp, 4> referenced({entry_func});

  // Create a new module to hold func and all referenced functions.