// RUN: tac-opt-all-backends -tfl-device-transform-gpu %s -split-input-file -verify-diagnostics | FileCheck %s

func.func @pack(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<2x1xf32> {
  %0 = "tfl.pack"(%arg0, %arg1) {axis = 0 : i32, values_count = 2 : i32} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
  func.return %0 : tensor<2x1xf32>
}

// CHECK:   func @pack(%[[VAL_0:.*]]: tensor<1xf32>, %[[VAL_1:.*]]: tensor<1xf32>) -> tensor<2x1xf32> {

// `{ tac.device = "GPU", tac.inference_type = "FLOAT"}` to a function
// with the matching attributes. Assumed is that device type "CPU"
// is the only device that is allowed to call other devices. I.e. ancestors of a
// "CPU" `Operation` may only `Operations` without a device or other "CPU"
// `Operations`. Implied is that "CPU" ops may contain subgraphs of different

    // TODO(jpienaar): Remove this later.
    if (auto device = res->getAttrOfType<StringAttr>("device")) {
      if (!device.getValue().empty())
        result_id->setAttr("device", device);
      else
        result_id->setAttr("device", compile_device_op);
    } else if (compile_device_op) {
      result_id->setAttr("device", compile_device_op);
    }

    func.return %0 : tensor<2x100xf32>
  }
  func.func @func_0_GPU_FLOAT(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>, %arg2: tensor<100xf32>) -> tensor<100xf32> attributes {tac.cost = 4.000000e+01 : f32, tac.device = "GPU", tac.inference_type = "FLOAT", tac.interface_name = "func_0"} {

      task_and_device = {task, device};
    }
  }

  return topology;
}

// Determine execution devices when topology and device assignment are defined.
// With a topology device coordinate to task and device mapping, device
// assignment device coordinates can then be mapped to task and device for TPU
// devices. The device assignment array is also validated.
//

                                               Device* device,
                                               Tensor* output_tensor,
                                               StatusCallback done) const {
  if (!DeviceFactory::IsPluggableDevice(device->device_type())) {
    done(absl::UnimplementedError(
        "Same-device copies in PjRtDeviceContext is only implemented when "

  %0 = "tf.Const"() {device = "/device:CPU:0", value = dense<2> : tensor<i32>} : () -> tensor<i32>
  %1 = "tf.ReadVariableOp"(%arg) {device = "/device:CPU:0", dtype = i32} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
  %2 = "tf.Add"(%1, %0) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>

    %3 = "tf.LookupTableSizeV2"(%1) {device = ""} : (tensor<!tf_type.resource>) -> tensor<i64>
    %4 = "tf.AddV2"(%2, %3) {device = ""} : (tensor<i64>, tensor<i64>) -> tensor<i64>
    %5 = "tf.LookupTableFindV2"(%0, %arg0, %cst) {device = ""} : (tensor<!tf_type.resource>, tensor<?x!tf_type.string>, tensor<i64>) -> tensor<*xi64>

// CHECK: [[const:%.*]] = "tf.Const"() <{value = dense<0> : tensor<i32>}> {device = "/CPU:0"} : () -> tensor<i32>
// CHECK: [[x:%.*]] = "tf.AddV2"([[const]], [[const]]) {device = "/CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
// CHECK: "tf._TfrtSetResource"([[x]]) <{index = 0 : i64}> {device = "/CPU:0"} : (tensor<i32>) -> ()
// CHECK: [[const_1:%.*]] = "tf.Const"() <{value = dense<1> : tensor<i32>}> {device = "/CPU:0"} : () -> tensor<i32>

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