[SOT] Fix for breakgraph load same variable with different name (#68938)

SigureMo · web-flow · commit 65126f558a5c · 2024-10-25T21:24:35.000+08:00
diff --git a/python/paddle/jit/sot/opcode_translator/executor/function_graph.py b/python/paddle/jit/sot/opcode_translator/executor/function_graph.py
@@ -184,7 +184,9 @@ def load(self, var):
         elif isinstance(var, NullVariable):
             var.reconstruct(self._pycode_gen)
         else:
-            self._pycode_gen.gen_load(self._store_var_info[var.id])
+            # NOTE: One variable may have multiple names, we can
+            # use any name to load it.
+            self._pycode_gen.gen_load(self._store_var_info[var.id][0])
 
 
 class FunctionGraph:
@@ -367,13 +369,18 @@ def _restore_origin_opcode(self, stack_vars, store_var_info, instr_idx):
         # here is not update changed values, it just give names to stack vars
         # and want keep same interface as _build_compile_fn_with_name_store
         for var in stack_vars[::-1]:
-            if store_var_info[var.id] is None:
-                store_var_info[var.id] = name_gen.next()
-                self.pycode_gen.gen_store_fast(store_var_info[var.id])
+            if not store_var_info[var.id]:
+                name = name_gen.next()
+                store_var_info[var.id].append(name)
+                self.pycode_gen.gen_store_fast(name)
             else:
-                self.pycode_gen.gen_store(
-                    store_var_info[var.id], self.pycode_gen._origin_code
-                )
+                all_names = store_var_info[var.id]
+                for _ in range(len(all_names) - 1):
+                    self.pycode_gen.gen_dup_top()
+                for name in all_names:
+                    self.pycode_gen.gen_store(
+                        name, self.pycode_gen._origin_code
+                    )
 
         return VariableLoader(store_var_info, self.pycode_gen)
 
@@ -391,13 +398,18 @@ def _build_compile_fn_with_name_store(
         name_gen = NameGenerator("___compile_fn_saved_")
 
         for var in to_store_vars[::-1]:
-            if store_var_info[var.id] is None:
-                store_var_info[var.id] = name_gen.next()
-                self.pycode_gen.gen_store_fast(store_var_info[var.id])
+            if not store_var_info[var.id]:
+                name = name_gen.next()
+                store_var_info[var.id].append(name)
+                self.pycode_gen.gen_store_fast(name)
             else:
-                self.pycode_gen.gen_store(
-                    store_var_info[var.id], self.pycode_gen._origin_code
-                )
+                all_names = store_var_info[var.id]
+                for _ in range(len(all_names) - 1):
+                    self.pycode_gen.gen_dup_top()
+                for name in all_names:
+                    self.pycode_gen.gen_store(
+                        name, self.pycode_gen._origin_code
+                    )
 
         return VariableLoader(store_var_info, self.pycode_gen)
 
diff --git a/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py b/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py
@@ -1775,15 +1775,16 @@ def get_compute_fn_and_update_changed_vars(
             extra_store_vars: for iterator, we need store the holder if it is a Tensor
         """
         store_vars = list(OrderedSet(list(stack) + extra_store_vars))
-        store_var_info = {var.id: None for var in stack}
+        store_var_info = {var.id: [] for var in stack}
 
         for name in restore_names:
             _var = self.get_var(name, allow_undefined=True)
             if _var is SotUndefinedVar():
                 continue
             if _var not in stack:
                 store_vars.append(_var)
-            store_var_info[_var.id] = name
+            store_var_info.setdefault(_var.id, [])
+            store_var_info[_var.id].append(name)
 
         compile_graph_result = self._graph.compile_graph(*store_vars)
         graph_fn, _ = compile_graph_result
@@ -2169,6 +2170,7 @@ def create_after_loop_fn():
         )
 
         for name in loop_body_inputs[:-1]:
+            # var_loader.load(self.get_var(name))
             self._graph.pycode_gen.gen_load(name)
 
         # this is the _break_flag
diff --git a/python/paddle/jit/sot/opcode_translator/executor/pycode_generator.py b/python/paddle/jit/sot/opcode_translator/executor/pycode_generator.py
@@ -935,6 +935,11 @@ def gen_shift_n(self, s: int, n: int):
                     "shift_n is not supported before python3.11"
                 )
 
+    def gen_dup_top(self):
+        if sys.version_info >= (3, 11):
+            return self.add_instr("COPY", arg=0)
+        return self.add_instr("DUP_TOP")
+
     def gen_swap(self, n):
         if sys.version_info >= (3, 11):
             self.add_instr("SWAP", arg=n)
diff --git a/test/sot/test_12_for_loop.py b/test/sot/test_12_for_loop.py
@@ -321,5 +321,24 @@ def test_arange(self):
         self.assert_results(for_arange, x)
 
 
+def for_break_with_load_same_consts(x: paddle.Tensor):
+    y = None
+    z = None
+    for i in [1, 2, 3]:
+        if y is None:
+            y = i
+        if z is None:
+            z = i
+        x += y + z
+        sot.psdb.breakgraph()
+    return x
+
+
+class TestForBreakWithLoadSameConsts(TestCaseBase):
+    def test_for_break_with_load_same_consts(self):
+        x = paddle.to_tensor(1)
+        self.assert_results(for_break_with_load_same_consts, x)
+
+
 if __name__ == "__main__":
     unittest.main()
