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());

  auto cluster = builder.create<mlir::tf_device::ClusterOp>(
      mlir::UnknownLoc::get(&context), result_types);
  cluster->setAttr(kDeviceAssignmentAttr, builder.getArrayAttr({}));

  mlir::TF::RuntimeDevices devices;
  std::string host_device;
  EXPECT_TRUE(mlir::failed(
      GetHostDeviceOutsideComputation(devices, cluster, &host_device)));
}

// 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,

}

bool HasTPUDevice(const mlir::TF::RuntimeDevices& devices) {
  for (const auto& device : devices.device_names()) {
    if (device.has_type && device.type == "TPU") return true;
  }
  return false;
}

mlir::LogicalResult GetHostDeviceOutsideCompilationInGenericPipeline(
    mlir::TF::RuntimeDevices devices, std::string* host_device) {
  for (const auto& device : devices.device_names()) {

// Returns true if the devices list contain any TPU devices
bool HasTPUDevice(const mlir::TF::RuntimeDevices& devices);

// Returns the host device used for outside compilation in generic pipeline.
mlir::LogicalResult GetHostDeviceOutsideCompilationInGenericPipeline(
    mlir::TF::RuntimeDevices devices, std::string* host_device);

// Parses XLA compilation and execution devices from a tf_device.cluster and

// 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.