// Test lowering of IfrtLoadVariableOp

// CHECK-LABEL: func @ifrt_load_variable_test
func.func @ifrt_load_variable_test() -> () {
  // CHECK: [[HANDLE:%.*]] = tf_mlrt.executeop()
  // CHECK-SAME:  VarHandleOp
  %0 = "tf.VarHandleOp"() {__op_key = 1: i32, device = "/device:CPU:0", container = "", shared_name = "variable"} : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>
  // CHECK-NEXT: "tf_mlrt.ifrt_load_variable"([[HANDLE]])

// Returns true if 'op' is TF op that accepts resource type, but is
// supported by TFLite.
bool IsSupportedTFLiteResourceOp(Operation* op) {
  return llvm::isa<TF::ReadVariableOp, TF::AssignVariableOp, TF::VarHandleOp,
                   TF::LookupTableFindV2Op, TF::LookupTableImportV2Op,
                   TF::LookupTableSizeV2Op>(op);
}

// Returns true if 'op' is TF/TFLite control flow op that can accept resource

func.func @resource_before_cluster() {
  // CHECK-NEXT: [[CONST:%.*]] = "tf.Const"
  %0 = "tf.Const"() {value = dense<1.000000e+00> : tensor<f32>} : () -> tensor<f32>

  // CHECK-NEXT: [[UNUSED_RESOURCE:%.*]] = "tf.VarHandleOp"
  // CHECK-NEXT: "tf.AssignAddVariableOp"([[UNUSED_RESOURCE]], [[CONST]])

  // CHECK:      "tf_device.cluster"
  // CHECK-NEXT:   "tf.NoOp"
  // CHECK-NEXT:   tf_device.return

      %out_0, %ctl_0 = tf_executor.island wraps "tf.Const"() {value = dense<"var_0"> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
      %out_1, %ctl_1 = tf_executor.island wraps "tf.VarHandleOp"() {container = "", shared_name = "var_0", device = "/device:CPU:0"} : () -> tensor<!tf_type.resource<tensor<2xf32>>>

      if (auto global = tf_saved_model::LookupBoundInputOfType<
              tf_saved_model::GlobalTensorOp>(func, i, syms)) {
        OpBuilder builder(func.getBody());
        auto dummy = builder.create<TF::VarHandleOp>(
            global.getLoc(), func.getArgument(i).getType(), "dummy", "dummy");
        func.getArgument(i).replaceAllUsesWith(dummy.getResult());
        argsToErase.set(i);
      }
    }

  TFE_DeleteContextOptions(opts);

  TFE_Op* var_op = TFE_NewOp(ctx, "VarHandleOp", status);
  TFE_OpSetAttrType(var_op, "dtype", TF_INT64);
  TFE_OpSetAttrShape(var_op, "shape", {}, 0, status);
  const TFE_OpAttrs* attributes = TFE_OpGetAttrs(var_op);

  TFE_Op* copy_op = TFE_NewOp(ctx, "VarHandleOp", status);
  TFE_OpSetAttrType(copy_op, "dtype", TF_FLOAT);
  TFE_OpAddAttrs(copy_op, attributes);

TFE_TensorHandle* TestVariable(TFE_Context* ctx, float value,
                               const tensorflow::string& device_name) {
  TF_Status* status = TF_NewStatus();
  // Create the variable handle.
  TFE_Op* op = TFE_NewOp(ctx, "VarHandleOp", status);
  if (TF_GetCode(status) != TF_OK) return nullptr;
  TFE_OpSetAttrType(op, "dtype", TF_FLOAT);
  TFE_OpSetAttrShape(op, "shape", {}, 0, status);
  TFE_OpSetAttrString(op, "container", "localhost", 0);

// Tests loop with two resource types, one of them being unique per iteration.
//
// Similar to above test but with one additional resource that is not unique per
// iteration (created by `tf.VarHandleOp`).

func.func @mixed_unique_resource_chain(%arg0: tensor<i32>, %arg1: tensor<f32>) {
  tf_executor.graph {

  "tf_device.cluster"() ({
    // CHECK:      "tf.VarHandleOp"
    // CHECK-SAME: _xla_outside_compilation
    "tf.VarHandleOp"() {allowed_devices = [], container = "", device = "", shared_name = ""} : () -> tensor<!tf_type.resource<tensor<!tf_type.string>>>
    tf_device.return

}

// CHECK-LABEL: test_expensive_ops
func.func @test_expensive_ops(%arg: tensor<?x512xf32>) -> tensor<?x512xf32> {
    // expected-remark@+1 {{Cost: 1}}
    %0 = "tf.VarHandleOp"() {allowed_devices = [], container = "", device = "/job:localhost/replica:0/task:0/device:CPU:0", shared_name = "var"} : () -> tensor<!tf_type.resource<tensor<512x512xf32>>>
    // expected-remark@+1 {{Cost: 2}}