// `container`, `name`, and `device` can be set. If an resource handle is tied
// to an instance of an operation (e.g. TensorFlow runtime operation caching),
// `op` can be set instead.
struct ResourceHandle {
  ResourceHandle(StringRef container, StringRef name, StringRef device,
                 Operation* op)
      : container(container), name(name), device(device), op(op) {}

  // CHECK-SAME: "TPU_REPLICATED_CORE_0"
  // CHECK: "tf.TPUCopyWithDynamicShape"
  func.func @valid_copy_op_in_replicated_host(
    %arg0: tensor<2048xi64> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"},
    %arg1: tensor<2048xi64> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"}) -> (tensor<2048xi32>, tensor<2048xi32>) {
    %cst = "tf.Const"() {value = dense<1024> : tensor<i32>} : () -> tensor<i32>

      // CHECK: %[[LAUNCH_OUTPUT:[0-9]*]] = "tf_device.launch_func"(%[[A_OUTPUT]]) <{device = "/device:test_device:0", func = @[[LAUNCH:.*]]}>
      %3 = "tf_device.launch"() ({
        %4 = "tf.B"(%2) : (tensor<?xi32>) -> tensor<?xi32>
        tf_device.return %4 : tensor<?xi32>
      }) {device = "/device:test_device:0"} : () -> tensor<?xi32>

      // CHECK: tf_executor.yield %[[LAUNCH_OUTPUT]]
      tf_executor.yield %3 : tensor<?xi32>
    }

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;

	if vendor == "" {
		return "", "", device, fmt.Errorf("unqualified device %q, missing vendor", device)
	}
	if class == "" {
		return "", "", device, fmt.Errorf("unqualified device %q, missing class", device)
	}
	if name == "" {
		return "", "", device, fmt.Errorf("unqualified device %q, missing device name", device)
	}

	if err := validateVendorName(vendor); err != nil {

        "Cannot get device metadata from non-PJRT device \"", device->name(),
        "\". GetMetadata must only be called on a device derived from "
        "PjRtBaseDevice. Either an internal bug has been triggered, or an "
        "XLA-specific op has been placed on the wrong device.");
  }
  return &pjrt_device->metadata_;
}

class DeviceInfoCache {
 public:
  bool IsGpu(DeviceId device) const { return is_gpu_[device.id()]; }
  bool IsCpu(DeviceId device) const { return is_cpu_[device.id()]; }

  absl::string_view GetNameFor(DeviceId device) const {
    return names_[device.id()];
  }

  absl::StatusOr<DeviceId> GetIdFor(absl::string_view name);

  LoggingDevice* device = new LoggingDevice;
  device->arrived_flag = arrived_flag;
  device->executed_flag = executed_flag;
  device->device_name = name;
  device->underlying_device = "/job:localhost/replica:0/task:0/device:CPU:0";
  device->strict_scope_placement = strict_scope_placement;
  TFE_RegisterCustomDevice(context, custom_device, name, device, status);
}

    %3:3 = tf_executor.Merge %2#0, %0#0 : tensor<*xi32> {device = "", N = 2, T = "tfdtype$DT_INT32"} loc("while/Merge")
    %4:2 = tf_executor.island(%3#2) wraps "tf.Const"() {device = "", dtype = "tfdtype$DT_INT32", value = dense<10> : tensor<i32>} : () -> tensor<i32> loc("while/Less/y")

    Op<TFRT_GPU_Dialect, mnemonic, traits> {
}

// TODO(b/260267885): We may add a device argument when we want to support
// GPU MIG.
def TransferToDeviceOp: Gpu_Op<"transfer_to_device"> {
  let summary = "Transfer a CPU tensor to device.";

  let description = [{
    Transfer a CPU tensor to device.

    Example:
      %device_tensor = gpurt.transfer_to_device %cpu_tensor
  }];