// 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> {

#define GEN_PASS_DECL_LAYOUTASSIGNMENTPASS
#define GEN_PASS_DECL_LEGALIZEHLOTOTFPASS
#define GEN_PASS_DECL_LEGALIZETFGTOTFPASS
#define GEN_PASS_DECL_LOCALIZEVARHANDLESPASS
#define GEN_PASS_DECL_LOWERQUANTIZEDPASS
#define GEN_PASS_DECL_MARKINPUTOUTPUTALIASESPASS
#define GEN_PASS_DECL_MATERIALIZEPASSTHROUGHOP
#define GEN_PASS_DECL_MERGECONTROLFLOWPASS
#define GEN_PASS_DECL_MOVETRANSPOSESPASS
#define GEN_PASS_DECL_ORDERBYDIALECTPASS

//     OnTFDeviceView(params->device->struct_size);
//     params->device = { SP_DEVICE_STRUCT_SIZE };
//     params->device->device_handle = get_my_device_handle(device->ordinal);
//     params->device->ordinal = params->ordinal;
//     ...
//   }
//
//   void destroy_device(const SP_Platform* platform, SP_Device* device) {
//     delete_my_device_handle(device->device_handle);
//   }
//
//   void SE_InitPlugin(

#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"

// A pass that propagates device assignment of resources on a module. It

limitations under the License.
==============================================================================*/

#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V1_COMPILE_MLIR_UTIL_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V1_COMPILE_MLIR_UTIL_H_

#include <memory>

#include "absl/base/attributes.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/StringRef.h"

      parse_options);
}

void CreateOp::print(OpAsmPrinter &p) {
  CreateOp op = *this;
  p << "(" << op.getInCh() << ") key("
    << op->getAttrOfType<mlir::IntegerAttr>("op_key").getInt() << ") device("
    << op->getAttr("device") << ") " << op->getAttr("op_name") << "()";

  fallback_common::PrintExecuteOpCommon(p, op);
  fallback_common::PrintExecuteOpFuncAttribute(p, op);

          %c_9 = tf_executor.island wraps "tf.InitializeTableFromTextFileV2"(%o, %arg1) <{delimiter = "\09", key_index = -2 : i64, value_index = -1 : i64, vocab_size = -1 : i64}> {_has_manual_control_dependencies = true, device = ""} : (tensor<!tf_type.resource>, tensor<!tf_type.string>) -> ()

  if (cc_ctx->op_device_context() == nullptr) {  // CPU Device
    status->status = absl::FailedPreconditionError(
        "Accessing device stream is not supported for a CPU device.");
    return nullptr;
  } else if (!cc_ctx->op_device_context()->IsPluggableDevice()) {
    status->status = absl::FailedPreconditionError(
        "Accessing device stream is only supported for pluggable devices.");
    return nullptr;

FunctionDef SquareFunc() {
  return FunctionDefHelper::Define(
      // Function Name
      "SquareFunc",
      // Args
      {"x: float"},
      // Returns
      {"y: float"},
      // Attr def
      {},
      // Nodes
      {{/*ret=*/{"y"},
        /*op=*/"Square",
        /*arg=*/{"x"},
        /*attr=*/{{"T", DT_FLOAT}},
        /*dep=*/{},
        /*device=*/"",
        /*name=*/"square"}});
}

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