Fix cache miss when enumerate obj as resume fn param

sot/opcode_translator/executor/function_graph.py

@@ -283,7 +283,12 @@ def _put_inner(self, var):
         )
 
     def add_global_guarded_variable(self, variable: VariableBase):
-        self._global_guarded_variables.append(variable)
+        if variable not in self._global_guarded_variables:
+            self._global_guarded_variables.append(variable)
+
+    def remove_global_guarded_variable(self, variable: VariableBase):
+        if variable in self._global_guarded_variables:
+            self._global_guarded_variables.remove(variable)
 
     def _find_tensor_outputs(
         self, outputs: list[VariableBase]

sot/opcode_translator/executor/opcode_executor.py

@@ -1758,6 +1758,7 @@ def FOR_ITER(self, instr):
             if not isinstance(
                 iterator, (SequenceIterVariable, DictIterVariable)
             ):
+                self._graph.remove_global_guarded_variable(iterator)
                 raise BreakGraphError()
             backup_iter_idx = iterator.idx
             self._inline_call_for_loop(iterator, instr)

sot/opcode_translator/executor/opcode_inline_executor.py

@@ -305,4 +305,5 @@ def FOR_ITER(self, instr: Instruction):
                 self._lasti = self.indexof(instr.jump_to)
 
         else:
+            self._graph.remove_global_guarded_variable(iterator)
             raise BreakGraphError("For loop fallback.")

tests/test_12_for_loop.py

@@ -10,6 +10,9 @@
 
 import paddle
 from sot import symbolic_translate
+from sot.opcode_translator.executor.opcode_executor import (
+    InstructionTranslatorCache,
+)
 
 
 def gener():
@@ -157,6 +160,27 @@ def test_list_comp(self):
         self.assert_results(run_list_comp, x)
 
 
+def for_enumerate_cache(func_list, x):
+    out = None
+    for idx, func in enumerate(func_list):
+        out = func(x[idx])
+    return out
+
+
+class TestEnumerateCache(TestCaseBase):
+    def test_run(self):
+        func_list = [
+            paddle.nn.Linear(10, 10),
+        ]
+        x = [
+            paddle.randn([5, 10]),
+        ]
+
+        out = symbolic_translate(for_enumerate_cache)(func_list, x)
+        out = symbolic_translate(for_enumerate_cache)(func_list, x)
+        self.assert_nest_match(InstructionTranslatorCache().translate_count, 4)
+
+
 if __name__ == "__main__":
     with strict_mode_guard(0 if sys.version_info >= (3, 10) else 1):
         unittest.main()