[SOT] Store iterator holder when its a tensor (#67617)

SigureMo · web-flow · commit 455dd59c9596 · 2024-08-22T11:03:17.000+08:00
diff --git a/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py b/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py
@@ -82,6 +82,7 @@
     ContainerVariable,
     DictVariable,
     GlobalVariable,
+    IterVariable,
     ListVariable,
     MethodVariable,
     NullVariable,
@@ -1748,7 +1749,7 @@ def compile_return(self, ret_val):
         return Stop(state="Return")
 
     def get_compute_fn_and_update_changed_vars(
-        self, restore_names, stack, end_idx
+        self, restore_names, stack, end_idx, extra_store_vars
     ):
         """
         this function will:
@@ -1763,8 +1764,9 @@ def get_compute_fn_and_update_changed_vars(
             restore_names: the names used in resume functions.
             end_idx: instruction index where simulation get break.
             stack: current stack
+            extra_store_vars: for iterator, we need store the holder if it is a Tensor
         """
-        store_vars = list(OrderedSet(stack))
+        store_vars = list(OrderedSet(list(stack) + extra_store_vars))
         store_var_info = {var.id: None for var in stack}
 
         for name in restore_names:
@@ -1864,7 +1866,7 @@ def create_if_branch_fn(start_idx, input_var_names, is_pop_jump_branch):
         # 4. compile codes before if
         update_var_names = list(true_fn_read_names | false_fn_read_names)
         var_loader = self.get_compute_fn_and_update_changed_vars(
-            update_var_names, self.stack, cur_index
+            update_var_names, self.stack, cur_index, []
         )
 
         # 5. create if sturcture and call true_fn and false_fn
@@ -1964,7 +1966,7 @@ def create_resume_fn():
 
         # 3. compile sub graph before call
         var_loader = self.get_compute_fn_and_update_changed_vars(
-            read_names, self.stack, cur_index
+            read_names, self.stack, cur_index, []
         )
 
         # 4. recover stack
@@ -2130,8 +2132,17 @@ def create_after_loop_fn():
 
         # 5. compile sub graph before for-loop
         update_names = list(loop_body_read_names | after_loop_read_names)
+        extra_store_vars = (
+            [iterator]
+            if isinstance(iterator, IterVariable)
+            and isinstance(iterator.hold, TensorVariable)
+            else []
+        )
         var_loader = self.get_compute_fn_and_update_changed_vars(
-            update_names, self.stack, self.indexof(for_iter)
+            update_names,
+            self.stack,
+            self.indexof(for_iter),
+            extra_store_vars,
         )
 
         # 6. prepare a new loop and call loop body
diff --git a/test/sot/test_12_for_loop.py b/test/sot/test_12_for_loop.py
@@ -19,14 +19,14 @@
 
 import unittest
 
-from test_case_base import TestCaseBase
+from test_case_base import (
+    TestCaseBase,
+    test_instruction_translator_cache_context,
+)
 
 import paddle
 from paddle.jit import sot
 from paddle.jit.sot import symbolic_translate
-from paddle.jit.sot.opcode_translator.executor.executor_cache import (
-    OpcodeExecutorCache,
-)
 from paddle.jit.sot.utils import strict_mode_guard
 
 
@@ -259,9 +259,10 @@ def test_run(self):
             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(OpcodeExecutorCache().translate_count, 1)
+        with test_instruction_translator_cache_context() as ctx:
+            out = symbolic_translate(for_enumerate_cache)(func_list, x)
+            out = symbolic_translate(for_enumerate_cache)(func_list, x)
+            self.assertEqual(ctx.translate_count, 1)
 
 
 # after_loop_fn need zzz, and zzz is created as UndefinedVar when generating loop body
@@ -308,5 +309,17 @@ def test_list_comp_with_fallback(self):
         self.assert_results(comp_with_fallback, x)
 
 
+def for_arange(x):
+    for i in paddle.arange(0, 5):
+        x = x + i
+    return x
+
+
+class TestArange(TestCaseBase):
+    def test_arange(self):
+        x = paddle.to_tensor(1)
+        self.assert_results(for_arange, x)
+
+
 if __name__ == "__main__":
     unittest.main()
