[lite] Update analyze-variables pass to allow resource to be passed t…

…o TFLite control flow ops (If, While, CallOnce) these pass the resources for other ops inside the functions, if any of the other ops that use the resources are not supported we shouldn't enable variables. This should be handled since we inspect all ops in the graph.

PiperOrigin-RevId: 399222824
Change-Id: I0546d8abc9f72d3d9f0247e72910b6bfaf24f339

tensorflow/compiler/mlir/lite/tests/analyze-variables.mlir

@@ -57,7 +57,7 @@ module {
 
 // -----
 
-// CHECK: module attributes {tfl._legalize_tfl_variables = false}
+// CHECK: module attributes {tfl._legalize_tfl_variables = true}
 module {
   func @main() -> tensor<i32> {
     %0 = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor<*x!tf_type.resource<tensor<*xi32>>>

tensorflow/compiler/mlir/lite/transforms/analyze_variables.cc

@@ -33,9 +33,11 @@ bool IsSupportedTFLiteResourceOp(Operation* op) {
                    TF::LookupTableSizeV2Op>(op);
 }
 
-// Returns true if 'op' is a TFLite control flow operation.
-bool IsTFLiteControlFlowOp(Operation* op) {
-  return llvm::isa<TFL::WhileOp, TFL::IfOp, TFL::CallOnceOp>(op);
+// Returns true if 'op' is TF/TFLite control flow op that can accept resource
+// type. Usually these ops are just pass through, they call another subgraph and
+// pass the operands to.
+bool IsSupportedTFLiteControlFlow(Operation* op) {
+  return llvm::isa<TFL::WhileOp, TFL::IfOp, TF::IfOp, TFL::CallOnceOp>(op);
 }
 }  // namespace
 
@@ -67,15 +69,10 @@ class AnalyzeVariablesPass
     module.walk([&](Operation* op) {
       // Skip ops that are supported natively by TFLite.
       if (IsSupportedTFLiteResourceOp(op)) return WalkResult::advance();
+      if (IsSupportedTFLiteControlFlow(op)) return WalkResult::advance();
 
       // Check for ops that are legalized to TFLite.
       if (op->getDialect()->getNamespace() == "tfl") {
-        // TODO(b/189370197): Enable control flow ops after updating
-        // checks to handle them.
-        if (IsTFLiteControlFlowOp(op)) {
-          legalize_to_tfl = false;
-          return WalkResult::interrupt();
-        }
         return WalkResult::advance();
       }
       // Check for ops that are not legalized to TFLite.