// propagated through.
bool IsSupportedExecutorOp(Operation& op) {
  auto ops_have_same_device = [](Operation* lhs, Operation* rhs) {
    auto lhs_device_attr = lhs->getAttrOfType<StringAttr>(kDeviceAttr);
    auto rhs_device_attr = rhs->getAttrOfType<StringAttr>(kDeviceAttr);
    return (!lhs_device_attr && !rhs_device_attr) ||
           (lhs_device_attr && rhs_device_attr &&
            lhs_device_attr.getValue() == rhs_device_attr.getValue());
  };

#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
#include "tensorflow/core/util/device_name_utils.h"

namespace mlir {
namespace TFTPU {

namespace {

constexpr char kDeviceAttr[] = "device";
constexpr char kDeviceCPU[] = "CPU";
constexpr char kFuncDeviceAttr[] = "tf.device";

struct TPUDynamicLayoutPass
    : public TF::PerFunctionAggregateAnalysisConsumerPass<

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

#define DEBUG_TYPE "tf-resource-device-inference"

namespace mlir {
namespace TF {

namespace {
constexpr char kDeviceAttr[] = "device";
constexpr char kFuncDeviceAttr[] = "tf.device";

#define GEN_PASS_DEF_RESOURCEDEVICEINFERENCEPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

  // Both the device attribute and compile_device_type attribute should be
  // propagated to the reduce function call.
  reduce_call->setAttr(kDeviceAttr,
                       reduce_dataset->getAttrOfType<StringAttr>(kDeviceAttr));
  reduce_call->setAttr(
      TF::kCompileDeviceTypeAttr,
      reduce_dataset->getAttrOfType<StringAttr>(TF::kCompileDeviceTypeAttr));

  SmallVector<Value, 4> if_returns;

namespace TF {
namespace {

using DeviceNameUtils = ::tensorflow::DeviceNameUtils;
using ParsedName = ::tensorflow::DeviceNameUtils::ParsedName;

constexpr const char *kHostAttr = "host";
constexpr const char *kDeviceAttr = "device";
constexpr const char *kTFDeviceAttr = "tf.device";
// TODO(donglin): Handle the case where the address of localhost is different
// from /job:localhost/replica:0/task:0.

    auto range = llvm::make_filter_range(op->getUsers(), [&](Operation *user) {
      bool same_block = user->getBlock() == op->getBlock();
      bool same_device = op->getAttr(kDeviceAttr) == user->getAttr(kDeviceAttr);
      return same_block && same_device && state.IsMember(user);
    });
    return {range.begin(), range.end()};
  }

  // Users of block argument must be in the same block.

    }
  }
  return success();
}

namespace {
constexpr char kDeviceAttr[] = "device";
constexpr char kFuncDeviceAttr[] = "tf.device";
mlir::Attribute GetDeviceOfResource(mlir::func::FuncOp func,
                                    mlir::Value resource) {
  if (auto* resource_op = resource.getDefiningOp()) {
    return resource_op->getAttr(kDeviceAttr);
  } else {

  if (auto device_attr =
          func.getArgAttrOfType<mlir::StringAttr>(index, kDeviceAttr))
    *node_def->mutable_device() = device_attr.getValue().str();

  llvm::ArrayRef<mlir::NamedAttribute> func_arg_i_attrs =
      mlir::function_interface_impl::getArgAttrs(func, index);
  absl::flat_hash_set<absl::string_view> attrs_to_ignore = {kDeviceAttr,
                                                            kAliasingAttr};

  if (auto device_attr =
          func.getArgAttrOfType<mlir::StringAttr>(index, kDeviceAttr))
    *node_def->mutable_device() = device_attr.getValue().str();

  llvm::ArrayRef<mlir::NamedAttribute> func_arg_i_attrs =
      mlir::function_interface_impl::getArgAttrs(func, index);
  absl::flat_hash_set<absl::string_view> attrs_to_ignore = {kDeviceAttr,
                                                            kAliasingAttr};

    rewriter.replaceOp(op_to_replace, ValueRange({fused_op}));
    return success();
  }
};

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

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