inline constexpr absl::string_view kDeviceAssignmentAttr = "device_assignment";

// A TPU device for execution alongside its associated host CPU device.
struct TPUDeviceAndHost {
  TPUDeviceAndHost() = default;
  TPUDeviceAndHost(llvm::StringRef device, llvm::StringRef host)
      : device(device), host(host) {}

  std::string device;
  std::string host;
};

"909", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg7: tensor<*x!tf_type.resource<tensor<25001x64xf32>>> {tf._user_specified_name = "911", tf.device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0"}, %arg8: tensor<*x!tf_type.resource<tensor<25001x64xf32>>> {tf._user_specified_name = "913", tf.device = "/job:tpu_host_worker/replica:0/task:1/device:CPU:0"}, %arg9: tensor<*x!tf_type.resource<tensor<25001x64xf32>>> {tf._user_specified_name = "915", tf.device = "/job:t...

  %0 = "tf.HashTableV2"() {container = "", device = "", key_dtype = !tf_type.string, shared_name = "hash_table_1dd4fef4-646d-491f-a3a8-bf5334f45813", use_node_name_sharing = false, value_dtype = i64} : () -> tensor<!tf_type.resource>
  "tf.LookupTableImportV2"(%0, %cst, %cst_1) {device = ""} : (tensor<!tf_type.resource>, tensor<3x!tf_type.string>, tensor<3xi64>) -> ()

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;

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

    %0 = tfl.add %arg0, %arg1 {fused_activation_function = "RELU6", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<1xf32>
    %1 = tfl.mul %0, %arg2 {fused_activation_function = "RELU6", tac.device = "GPU", tac.inference_type = "FLOAT"} : tensor<1xf32>
    func.return %1 : tensor<1xf32>

    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
  }];

  %1 = "tfl.mul"(%0, %arg2) {tac.device = "GPU", fused_activation_function = "RELU6", tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
  %2 = "tfl.add"(%arg0, %arg3) {tac.device = "GPU", fused_activation_function = "RELU6", tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>

                             StringPiece tensor_name, Device* device,
                             Tensor* cpu_tensor, StatusCallback done) override;

  void CopyTensorInSameDevice(const Tensor* input_tensor, Device* device,
                              Tensor* output_tensor,
                              StatusCallback done) const override {
    done(errors::Internal("device->device copy not implemented."));
  }

 private:

    %cst = "tf.Const"() <{value = dense<[0, 1]> : tensor<2xi32>}> {device = ""} : () -> tensor<2xi32>
    %cst_0 = "tf.Const"() <{value = dense<1.000000e+01> : tensor<f32>}> {device = ""} : () -> tensor<f32>
    %0 = "tf.Sum"(%arg0, %cst) <{keep_dims = false}> {device = ""} : (tensor<1x4xf32>, tensor<2xi32>) -> tensor<f32>
    %1 = "tf.Greater"(%0, %cst_0) {device = ""} : (tensor<f32>, tensor<f32>) -> tensor<i1>