[lite] Add new types for ReadVariable and AssignVariables ops in MLIR converter for TFLite.

This is in progress and currently experimental.

PiperOrigin-RevId: 354678544
Change-Id: I6f32122e90fa126042a1ad4383bae6a80009d716

tensorflow/compiler/mlir/lite/converter_gen.cc

@@ -241,9 +241,15 @@ static void EmitGetBuiltinOpCode(const std::vector<Record *> &defs,
 
   for (const auto *def : defs) {
     StringRef op_name = def->getName().drop_front(4);
+    auto operator_name = GetOperatorName(*def);
+    // TODO(b/149099381): Remove this part after kernels are added as
+    // builtin op.
+    if (operator_name == "ASSIGN_VARIABLE" ||
+        operator_name == "READ_VARIABLE") {
+      continue;
+    }
     os << "  if (isa<mlir::TFL::" << op_name << ">(op))\n"
-       << "    return tflite::BuiltinOperator_" << GetOperatorName(*def)
+       << "    return tflite::BuiltinOperator_" << operator_name << ";\n";
-       << ";\n";
   }
 
   os << "  return llvm::None;\n"

tensorflow/compiler/mlir/lite/ir/tfl_ops.td

@@ -4556,4 +4556,41 @@ the dimension is padded with zeros.
   );
 }
 
+
+def TFL_AssignVariableOp : TFL_Op<"assign_variable", []> {
+  let summary = "Assigns a new value to a variable.";
+
+  let description = [{
+Any ReadVariableOp with a control dependency on this op is guaranteed to return
+this value or a subsequent newer value of the variable.
+  }];
+
+  let arguments = (ins
+    // TODO(b/149099381): Remove integer IDs after adding the new variable
+    // handle type.
+    TFL_TensorOf<[I32]>:$resource_id,
+    // TODO(b/149099381): Support other types too.
+    TFL_TensorOf<[F32]>:$value
+  );
+
+  let results = (outs);
+}
+
+def TFL_ReadVariableOp : TFL_Op<"read_variable", []> {
+  let summary = "Reads variable value.";
+
+  let description = [{
+Read variable data identified by 'resource_id'.
+  }];
+
+  let arguments = (ins
+    // TODO(b/149099381): Remove integer IDs after adding the new variable
+    // handle type.
+    TFL_TensorOf<[I32]>:$resource_id
+  );
+
+  // TODO(b/149099381): Support other types too.
+  let results = (outs TFL_TensorOf<[F32]>:$result);
+}
+
 #endif // TFL_OPS