}
    callees.push_back(callee);
  }
  return success();
}

bool HasSingleBlock(func::FuncOp func) {
  return func->getNumRegions() == 1 && func.getBody().hasOneBlock();
}

  // Check if the island has only one block thats contain two ops, including
  // one resource op and one Yield op.
  auto island_op = llvm::dyn_cast_or_null<IslandOp>(op);
  if (!island_op || !island_op.getBody().hasOneBlock()) return nullptr;
  auto& island_block = island_op.getBody().front();
  if (++island_block.begin() != --island_block.end()) return nullptr;

  Operation* resource_op = &island_block.front();

  // The pattern we're matching is
  // ^block(arg0, arg1, ..., argN):
  //   r0, r1, ..., rN = func.call @foo(arg0, arg1, ..., argN)
  //   "tf.yield"(r0, r1, ..., rN)
  if (!region.hasOneBlock()) return std::nullopt;
  Block& block = region.front();
  if (std::distance(block.begin(), block.end()) != 2) return std::nullopt;
  TF::YieldOp yield =
      llvm::dyn_cast_or_null<TF::YieldOp>(block.getTerminator());

  auto function = dyn_cast_or_null<mlir::func::FuncOp>(
      SymbolTable::lookupSymbolIn(module, op.getReducer()));
  if (!function) return op.emitOpError() << "No reducer";
  if (!function.getBody().hasOneBlock())
    return op.emitOpError() << "reducer has more than one block";

  return success();
}

//===----------------------------------------------------------------------===//
// XlaVariadicSortOp

        start_indices_tensor);
    return success();
  };
};

template <typename ReturnOpType>
bool MatchTopKComparator(Region& comparator) {
  if (!comparator.hasOneBlock()) return false;
  Block& comparator_blk = comparator.front();
  using OpListType = llvm::iplist<Operation>;
  OpListType& operations = comparator_blk.getOperations();
  if (operations.size() != 2) return false;