// During cloning, let it know that the source function's argument
    // corresponds to the main function's newly created argument when cloning
    // ops from src -> main.
    BlockArgument main_arg = main_func_op.getArgument(main_arg_idx);
    mapper.map(src_arg, main_arg);
  }

  return mapper;
}

// Copies ops from `src_func_op` to `main_body` except for the FetchOps. Returns

    func::FuncOp func, const llvm::SmallSet<int, 4> &tensor_list_args) {
  // `indexes` stores the argument index of tensorlists which size may get
  // updated in the function.
  llvm::SmallSet<int, 4> indexes;
  for (BlockArgument &arg : func.getArguments()) {
    if (tensor_list_args.contains(arg.getArgNumber())) {
      for (const mlir::OpOperand &use : arg.getUses()) {
        mlir::Operation *op = use.getOwner();

    }
  } else {
    info.signature_change = true;
    for (auto& entry : callee_map) {
      auto buffer_arg = mlir::dyn_cast<BlockArgument>(entry.getFirst());
      if (!buffer_arg) continue;
      info.buffer_arg_to_size_arg[buffer_arg.getArgNumber()] =
          mlir::cast<BlockArgument>(entry.getSecond().size).getArgNumber();
    }
    if (lowered_callee != callee) {
      // Add the clone with a new name.

  int result_idx = 0;
  llvm::SmallVector<Value> call_op_returns;
  for (auto func_op : module_op.getOps<func::FuncOp>()) {
    if (!ShouldIncludeInMainFunction(func_op)) continue;

    llvm::ArrayRef<BlockArgument> new_args = llvm::ArrayRef(
        main_func.getArguments().begin() + arg_idx, func_op.getNumArguments());
    arg_idx += func_op.getNumArguments();
    llvm::ArrayRef<Type> new_types = llvm::ArrayRef(

    LLVM_DEBUG(llvm::dbgs() << GemmStyleOp::getOperationName()
                            << " op must have ranked tensor type.\n");
    return failure();
  }

  MutableArrayRef<BlockArgument> operands =
      entry_func_op.getBody().getArguments();
  // Function must have input, filter, and optionally bias.
  if (operands.size() != 2 && operands.size() != 3) {

        operand_not_in_launch);
  }

  // Handle pass through block arguments.
  for (OpOperand& operand :
       original_launch_op.GetBody().getTerminator()->getOpOperands()) {
    if (mlir::isa<BlockArgument>(operand.get())) {
      original_launch_op.getResult(operand.getOperandNumber())
          .replaceAllUsesWith(operand.get());
    }
  }
}

  mlir::Block& block = main.getRegion().front();
  xla::XlaBuilder builder("main");

  // Create xla_params.
  std::vector<xla::XlaOp> xla_params;
  for (mlir::BlockArgument& arg : block.getArguments()) {
    auto num = arg.getArgNumber();
    xla::Shape shape = xla::TypeToShape(arg.getType());
    xla::XlaOp argop =
        xla::Parameter(&builder, num, shape, absl::StrCat("Arg_", num));

  auto builder = mlir::OpBuilder(main_fn.getBody());
  std::vector<int> args_to_erase;
  for (int idx = 0; idx < args.size(); idx++) {
    const XlaArgument& xla_arg = args[idx];
    mlir::BlockArgument mlir_arg = main_fn.getArgument(idx);
    if (xla_arg.kind == XlaArgument::kResource) {
      mlir::Type element_type;
      if (xla_arg.type == DT_INVALID) {
        return errors::Unimplemented(absl::StrCat(

    // Additionally remove arguments from the outputs, as it provides liveness
    // throughout (functions as an invisible op above the very first op that
    // returns the arguments).
    for (const BlockArgument arg : main_func.getArguments()) {
      outputs.remove(arg);
    }
  }

  ReplaceStablehloOpsWithXlaCallModuleOp(
      SmallVector<Value>(inputs.begin(), inputs.end()),

        // Check device matching for the node defining the resource.
        if (!IsResourceMergeable(resource_attr, device_attr)) continue;
      } else {
        auto resource_arg = resource.dyn_cast<BlockArgument>();
        assert(resource_arg);
        if (resource_arg.getOwner() != &func.front()) continue;
        // Check device matching for the argument defining the resource.