namespace mlir {

class Operation;

namespace TF {

class RuntimeDevices;

// Returns true if at least one GPU device is available at runtime.
bool CanUseGpuDevice(const RuntimeDevices &devices);

// Returns true if all of the GPUs available at runtime support TensorCores
// (NVIDIA compute capability >= 7.0).
bool CanUseTensorCores(const RuntimeDevices &devices);

TEST(DeviceUtilTest, GetDevicesFromOpNoDevicesAttribute) {
  mlir::MLIRContext context;
  mlir::OwningOpRef<mlir::ModuleOp> module_ref =
      mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));

  mlir::TF::RuntimeDevices devices;
  EXPECT_TRUE(mlir::succeeded(GetDevicesFromOp(*module_ref, &devices)));
}

TEST(DeviceUtilTest, GetDevicesFromOpBadDevicesAttributeType) {
  mlir::MLIRContext context;

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_structs.h"

#include <optional>

namespace mlir {
namespace TF {

void RuntimeDevices::AddDevice(const ParsedName& device) {
  device_names_.push_back(device);
}

void RuntimeDevices::AddGpuDevice(const ParsedName& device,
                                  const GpuDeviceMetadata& metadata) {
  device_names_.push_back(device);

// Returns true if at least one GPU device is available at runtime.
bool CanUseGpuDevice(const RuntimeDevices &devices) {
  return llvm::any_of(devices.device_names(), IsGpuDevice);
}

// Returns true if all of the GPUs available at runtime support TensorCores
// (NVIDIA compute capability >= 7.0).
bool CanUseTensorCores(const RuntimeDevices &devices) {
  auto has_tensor_cores = [&](const DeviceNameUtils::ParsedName &device) {

// `tf.devices` and remove this function.
mlir::LogicalResult GetDevicesFromOp(mlir::Operation* op,
                                     mlir::ArrayAttr array_attr,
                                     mlir::TF::RuntimeDevices* devices) {
  DeviceNameUtils::ParsedName device;

  for (const auto& kv : llvm::enumerate(array_attr)) {
    const int idx = kv.index();

    auto string_attr = mlir::dyn_cast<mlir::StringAttr>(kv.value());

// failure if can't parse device metadata from the attribute.
mlir::LogicalResult GetDevicesFromOp(mlir::Operation* op,
                                     mlir::TF::RuntimeDevices* devices);

// Parses a device string and returns its ordinal (id). This will return an
// error if the device string is invalid or has no id.
mlir::LogicalResult GetDeviceOrdinalFromDeviceString(mlir::Location loc,

// Tensorflow devices available at runtime with corresponding metadata if it is
// available. It's completely valid to have a device without any metadata
// attached to it.
class RuntimeDevices {
  using DeviceNameUtils = ::tensorflow::DeviceNameUtils;
  using ParsedName = ::tensorflow::DeviceNameUtils::ParsedName;

 public:
  // Adds a device with and empty metadata. Device can be of any type.

      "SmallVector<unsigned, 4>", "GetLayoutDependentResults", (ins)
    >,
    InterfaceMethod<
      [{Returns the optimal data layout based on the available devices.}],
      "StringRef", "GetOptimalLayout", (ins "const RuntimeDevices&":$devices)
    >,
    InterfaceMethod<
      [{Updates operation attributes and operands to account for the updated
        data format. If data format is not supported, must return failure.}],

  if (!tf_dialect) {
    getOperation().emitError() << "'tf' dialect is not registered";
    return signalPassFailure();
  }

  // Get the number of devices per host.
  int device_num = 0;
  mlir::TF::RuntimeDevices devices;
  if (failed(tensorflow::GetDevicesFromOp(
          getOperation()->getParentOfType<ModuleOp>(), &devices)))
    return signalPassFailure();
  for (const auto& device_name : devices.device_names()) {

    bool pass_host_device, ModuleOp module,
    std::function<WalkResult(Operation*, tf_device::ClusterOp,
                             std::optional<std::string>)>
        callback) {
  mlir::TF::RuntimeDevices devices;
  if (failed(tensorflow::GetDevicesFromOp(module, &devices))) return failure();
  const CallGraph call_graph(module);
  // symbol_table caches callees in the CallGraph.
  SymbolTableCollection symbol_table;