// if the given operation is the only user.
Operation* GetInputOpWithOneUse(Operation* op, int opr_num) {
  if (opr_num >= op->getNumOperands()) return nullptr;
  auto opr = op->getOperand(opr_num);
  if (llvm::isa<BlockArgument>(opr)) return nullptr;
  auto* res = opr.getDefiningOp();
  if (!res->hasOneUse()) return nullptr;
  return res;
}

// Checks if the given operand of given operation refers to a splat constant
// with given val.

      return failure();
    }

    Operation* operand_op = operand.getDefiningOp();
    if (operand_op == nullptr) {
      // When `QuantizeOpT`'s operand does not have a defining op, it means it
      // is a `BlockArgument`. The pattern does not match if there is no op to
      // quantize.
      return failure();
    }

    if (operand_op->hasTrait<OpTrait::ConstantLike>()) {

  return new_device;
}

// Propagates devices from function arguments.
void PropagateDevicesFromArguments(
    func::FuncOp func,
    llvm::DenseMap<Value, llvm::StringRef>& value_to_device) {
  for (BlockArgument& arg : func.getArguments()) {
    auto arg_device_attr =
        func.getArgAttrOfType<StringAttr>(arg.getArgNumber(), kFuncDeviceAttr);
    if (!arg_device_attr || arg_device_attr.getValue().empty() ||

    constexpr char kShardingAttr[] = "mhlo.sharding";
    if (auto sharding =
            func.getArgAttrOfType<mlir::StringAttr>(i, kShardingAttr)) {
      if (!sharding.getValue().empty()) {
        BlockArgument arg = func.getArgument(i);
        // TODO(hinsu): Instead of setting both 'sharding' and '_XlaSharding'
        // attributes, only set the 'sharding' attribute. Both attributes are

  // the result cannot be backtracked to a region argument, returns
  // std::nullopt.
  std::optional<int> GetArg(int result_index) const {
    if (auto arg = mlir::dyn_cast<BlockArgument>(GetValue(result_index)))
      if (arg.getParentBlock() == &region_->front()) return arg.getArgNumber();
    return std::nullopt;
  }

 private:
  friend class detail::BacktrackAnalysis;

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