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

namespace mlir {
namespace TF {

llvm::StringRef GetDeviceOrEmpty(Operation *op) {
  if (auto device_attr = op->getAttrOfType<StringAttr>("device"))
    return device_attr.getValue();
  return llvm::StringRef();
}

// Returns resource handle value and id for resource op based on attributes. If

    Operation* op) {
  llvm::StringRef device;
  if (auto attr = op->getAttrOfType<mlir::StringAttr>("device")) {
    device = attr.getValue();
  }

  llvm::StringRef container;
  if (auto attr = op->getAttrOfType<mlir::StringAttr>("container")) {
    container = attr.getValue();
  }

  llvm::StringRef shared_name;
  if (auto attr = op->getAttrOfType<mlir::StringAttr>("shared_name")) {
    shared_name = attr.getValue();

// and writes X, Y as 32-bit integer attribute to `x`, `y`.
bool TFIntListIs1XY1(Operation *op, StringRef name, IntegerAttr *x,
                     IntegerAttr *y) {
  auto attr = op->getAttrOfType<ArrayAttr>(name);
  if (!attr) return false;

  auto elements = attr.getValue();
  if (elements.size() != 4 ||
      std::any_of(elements.begin(), elements.end(),

  auto num_cores_per_replica_attr = cluster_func->getAttrOfType<IntegerAttr>(
      tensorflow::kNumCoresPerReplicaAttr);
  if (!num_cores_per_replica_attr)
    return cluster_func.emitOpError(
        CreateMissingAttributeMsg(tensorflow::kNumCoresPerReplicaAttr));

  int num_cores_per_replica = num_cores_per_replica_attr.getInt();

  auto topology_attr =
      cluster_func->getAttrOfType<StringAttr>(tensorflow::kTopologyAttr);

namespace mlir {
namespace quant {

constexpr StringRef kTfEntryFunctionAttr = "tf.entry_function";

void AddEntryFunctionInput(StringRef input_name, func::FuncOp func_op) {
  auto entry_func_attr =
      func_op->getAttrOfType<DictionaryAttr>(kTfEntryFunctionAttr);
  if (!entry_func_attr) return;

  auto entry_func_attrs = SmallVector<NamedAttribute>(entry_func_attr.begin(),

std::unique_ptr<OpQuantSpec> GetStableHloOpQuantSpec(Operation* op) {
  auto spec = std::make_unique<OpQuantSpec>();
  if (auto call_op = dyn_cast_or_null<TF::XlaCallModuleOp>(op)) {
    auto entry_function =
        call_op->getAttrOfType<FlatSymbolRefAttr>("_entry_function");
    StringRef function_name = entry_function.getValue();
    if (!function_name.starts_with("composite_")) {
      return spec;
    }

    if (function_name.contains("conv")) {

  TF::ReorderOpResultUses(launch_op);
}

std::string GetDevice(Operation* op) {
  auto device_attr = op->getAttrOfType<StringAttr>("device");
  return device_attr ? device_attr.getValue().str() : "";
}

bool CanBeIgnoredInCluster(Operation* op) {
  auto device_attr = op->getAttrOfType<StringAttr>("device");
  return !device_attr || device_attr.getValue().empty();
}

// passes does not trigger additional logic due to presence of this attribute.
LogicalResult ExtractEmbeddingAttribute(
    Operation* op, llvm::StringMap<Operation*>* embedding_op_map) {
  auto embedding_attr = op->getAttrOfType<StringAttr>(kTPUEmbeddingAttr);
  if (!embedding_attr) return mlir::success();

  if (!embedding_op_map->insert({embedding_attr.getValue(), op}).second)
    return op->emitOpError(

namespace mlir {
namespace TF {

namespace {

const char kDeviceAttr[] = "device";
const char kDeviceGpu[] = "GPU";

std::optional<std::string> GetOpDevice(mlir::Operation *op) {
  mlir::StringAttr device = op->getAttrOfType<mlir::StringAttr>(kDeviceAttr);
  if (!device || device.getValue().empty()) {
    return std::nullopt;
  }
  tensorflow::DeviceNameUtils::ParsedName parsed_name;

      has_tensor_cores);
}

// Returns true if operation does not have explicit device placement that would
// prevent it from running on GPU device.
bool CanUseGpuDevice(Operation *op) {
  auto device_attr = op->getAttrOfType<StringAttr>("device");
  if (!device_attr || device_attr.getValue().empty()) return true;

  DeviceNameUtils::ParsedName device;
  if (!DeviceNameUtils::ParseFullName(device_attr.getValue().str(), &device))