[SOT] Modify DELETE_FAST instruction arg (#67936)

SigureMo · web-flow · commit d32ae25e4bed · 2024-09-03T18:37:53.000+08:00
diff --git a/python/paddle/jit/sot/opcode_translator/instruction_utils/instruction_utils.py b/python/paddle/jit/sot/opcode_translator/instruction_utils/instruction_utils.py
@@ -323,11 +323,15 @@ def bind_ex_arg_with_instr(ex_arg, instr):
 
 
 def modify_vars(instructions: list[Instruction], code_options):
-    co_names = code_options['co_names']
     co_varnames = code_options['co_varnames']
     co_freevars = code_options['co_freevars']
     for instrs in instructions:
-        if instrs.opname in ['LOAD_FAST', 'LOAD_FAST_CHECK', 'STORE_FAST']:
+        if instrs.opname in [
+            'LOAD_FAST',
+            'LOAD_FAST_CHECK',
+            'STORE_FAST',
+            'DELETE_FAST',
+        ]:
             assert (
                 instrs.argval in co_varnames
             ), f"`{instrs.argval}` not in {co_varnames}"
diff --git a/test/sot/test_delete_fast.py b/test/sot/test_delete_fast.py
@@ -19,6 +19,7 @@
 from test_case_base import TestCaseBase
 
 import paddle
+from paddle.jit.sot.psdb import breakgraph
 
 
 def test_delete_fast(a):
@@ -28,11 +29,33 @@ def test_delete_fast(a):
     return a
 
 
+def inner_call_with_breakgraph(a):
+    breakgraph()
+    return a + 1
+
+
+def test_delete_fast_with_breakgraph(a):
+    a = a + 2
+    t1 = a * 3
+    out1 = inner_call_with_breakgraph(a + t1)
+    t2 = a - 4
+    t3 = a / 5
+    out2 = inner_call_with_breakgraph(t1 + t2)
+    del t1, t2, t3
+    return a + out1 + out2
+
+
 class TestDeleteFast(TestCaseBase):
     def test_simple(self):
         a = paddle.to_tensor(1)
         self.assert_results(test_delete_fast, a)
 
 
+class TestDeleteFastWithBreakGraph(TestCaseBase):
+    def test_delete_fast_with_break_graph(self):
+        a = paddle.to_tensor(1)
+        self.assert_results(test_delete_fast_with_breakgraph, a)
+
+
 if __name__ == "__main__":
     unittest.main()
