void runOnOperation() override;
};

LogicalResult HasNoNestedEntryFunctions(
    const llvm::SmallVector<func::FuncOp> &entry_funcs, SymbolTable &symtab) {
  auto calls_entry_functions = [&](SymbolUserOpInterface op) {
    llvm::SmallVector<func::FuncOp> callees;
    if (GetCallees(op, symtab, callees).failed()) {
      return false;
    }
    for (auto &callee : callees) {

// every host, a failure will be returned.
absl::StatusOr<llvm::SmallVector<llvm::SmallVector<ParsedDevice, 8>, 8>>
GetTPUDevices(ParsedDevices devices,
              llvm::ArrayRef<ParsedDevice> system_devices) {
  llvm::SmallVector<llvm::SmallVector<ParsedDevice, 8>, 8> tpu_devices;
  tpu_devices.reserve(system_devices.size());

// BlockArg or as a result of an Operation.
llvm::SmallVector<Value> AccumulateOperandsDefinedAbove(
    const llvm::SetVector<Operation*>& partition_ops);

// Similar to `AccumulateOperandsDefinedAbove()`, computes the Value(s) that are
// defined within a Subgraph and referenced in a descendant Operation. These
// Values(s) are to be returned by the new raised function.
llvm::SmallVector<Value> AccumulateResultsDefinedWithin(

  }
  return mlir::success();
}

llvm::SmallVector<llvm::SmallVector<int64_t, 4>, 4> GetMetadataArgumentMapping(
    const tpu::TPUCompileMetadataProto& metadata) {
  llvm::SmallVector<llvm::SmallVector<int64_t, 4>, 4> input_mappings(
      metadata.num_cores_per_replica(), llvm::SmallVector<int64_t, 4>());

  if (metadata.num_cores_per_replica() == 1) {

  SmallVector<int64_t, 8> perm_values;
  for (const auto& dim : perm_values_attr.getValues<APInt>())
    perm_values.push_back(dim.getSExtValue());

  // This should never happen unless the input graph is malformed.
  if (input_shape.size() != perm_values.size()) {
    return false;
  }

  SmallVector<int, 8> old_major_index_ordering;
  SmallVector<int, 8> new_major_index_ordering;

  bool rewrite = false;
  const int num_regions = parallel_execute.getNumRegions();
  llvm::SmallVector<Value, 4> results_to_remap;

  // Go through each region and find AssignVariableOps that can be moved into
  // the parallel_execute region. Result indices by region index are collected,
  // so they can be removed afterwards.
  llvm::SmallVector<llvm::SmallVector<int, 4>, 4> results_to_remove_by_region;
  results_to_remove_by_region.resize(num_regions);

                          Value &window_strides, Value &lhs_dilation,
                          Value &rhs_dilation, Value &feature_group_count,
                          int num_dims) {
  SmallVector<int32_t> lhs_dilation_values(num_dims - 2, 1);
  SmallVector<int32_t> stride_values, rhs_dilation_values;
  for (int64_t i : llvm::seq<int64_t>(1, num_dims - 1)) {
    stride_values.push_back(mlir::cast<IntegerAttr>(strides[i]).getInt());

      return failure();
    }

    SmallVector<int32_t, 4> begin, end, strides;
    SmallVector<int32_t, 4> padded_begin, padded_end, padded_strides;

    int num_input_dims = ranked_input_type.getRank();
    SmallVector<int32_t, 4> padding_begin(num_input_dims, 0);
    auto input_shape = ranked_input_type.getShape();
    SmallVector<int32_t, 4> padding_end(input_shape.begin(), input_shape.end());

  // pass the filter (returning true) will be included.
  const llvm::SmallVector<Operation*, 4>& DirectControlPredecessors(
      Operation* op) const;
  llvm::SmallVector<Operation*, 4> DirectControlPredecessors(
      Operation* op, llvm::function_ref<bool(Operation*)> filter) const;

  // pass the filter (returning true) will be included.
  const llvm::SmallVector<Operation*, 4>& DirectControlSuccessors(
      Operation* op) const;

bool IsEntryFunction(func::FuncOp func);

// Get all the entry functions in an MLIR module.
llvm::SmallVector<func::FuncOp> GetEntryFunctions(ModuleOp module);

// Get all the functions referenced in a symber user op and save them in
// `callees`.
LogicalResult GetCallees(SymbolUserOpInterface op, SymbolTable &symtab,
                         llvm::SmallVector<func::FuncOp> &callees);