[SOT][3.13] Copy opcode TO_BOOL when necessary (#69700)

GoldenStain · web-flow · commit c6b195e08f88 · 2024-11-27T03:28:03.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
@@ -49,7 +49,12 @@
     calc_stack_effect,
     get_instructions,
 )
-from ..instruction_utils.opcode_info import RETURN, JumpDirection, PopJumpCond
+from ..instruction_utils.opcode_info import (
+    NEED_TO_BOOL,
+    RETURN,
+    JumpDirection,
+    PopJumpCond,
+)
 from .dispatch_functions import (
     operator_BAD,
     operator_exception_match,
@@ -1344,11 +1349,9 @@ def COMPARE_OP(self, instr: Instruction):
     def TO_BOOL(self, instr: Instruction):
         # we don't do anything in TO_BOOL, we simply check if the bytecode is legal
         next_instr = self._instructions[self._lasti]
-        assert next_instr.opname in [
-            'POP_JUMP_IF_TRUE',
-            'POP_JUMP_IF_FALSE',
-            'UNARY_NOT',
-        ], f"The bytecode is illegal! The opcode following TO_BOOL must be in ['POP_JUMP_IF_TRUE', 'POP_JUMP_IF_FALSE', 'UNARY_NOT'], the next instuction now is {next_instr.opname}"
+        assert (
+            next_instr.opname in NEED_TO_BOOL
+        ), f"The bytecode is illegal! The opcode following TO_BOOL must be in ['POP_JUMP_IF_TRUE', 'POP_JUMP_IF_FALSE', 'UNARY_NOT'], the next instuction now is {next_instr.opname}"
 
     @call_break_graph_decorator(push_n=1)
     def IS_OP(self, instr: Instruction):
@@ -2038,6 +2041,7 @@ def _break_graph_when_if(self, result: TensorVariable, instr: Instruction):
 
         # 1. analyse info
         cur_index = self.indexof(instr)
+        prefix_opname = self._instructions[cur_index - 1].opname
         true_fn_start_index = cur_index + 1
         false_fn_start_index = self.indexof(instr.jump_to)
         stack_size_after_if = len(self.stack) - 1
@@ -2112,6 +2116,11 @@ def create_if_branch_fn(start_idx, input_var_names, is_pop_jump_branch):
 
         # 5. create if sturcture and call true_fn and false_fn
         var_loader.load(result)
+
+        # in 3.13, we have to copy the original 'TO_BOOL' to make the generated bytecode valid.
+        if sys.version_info >= (3, 13) and prefix_opname == "TO_BOOL":
+            self._graph.pycode_gen.add_instr('TO_BOOL')
+
         if_code = self._graph.pycode_gen.add_instr(instr.opname)
 
         assert true_fn is not None
@@ -2174,6 +2183,7 @@ def _break_graph_when_call(
         push_n = push_n(instr.arg) if callable(push_n) else push_n
         is_precall = instr.opname == "PRECALL"
         cur_index = self.indexof(instr)
+        prefix_opname = self._instructions[cur_index - 1].opname
         # Use CALL instead of PRECALL to calculate the real stack effect
         call_instr = self._instructions[cur_index + int(is_precall)]
         # skip CALL if current instr is PRECALL
@@ -2223,6 +2233,10 @@ def create_resume_fn():
         # 5. run the break CALL with origin python
         # NOTE(SigureMo): In Python 3.11 and 3.12，we need generate KW_NAMES if the call shape is not None.
         self._graph.pycode_gen.gen_kw_names(self._call_shape)
+        # in 3.13, We have to copy the original 'TO_BOOL' to make the generated bytecode valid.
+        if sys.version_info >= (3, 13) and prefix_opname == 'TO_BOOL':
+            self._graph.pycode_gen.add_instr('TO_BOOL')
+
         self._graph.pycode_gen.extend_instrs(
             self._instructions[cur_index:next_index]
         )
@@ -2243,10 +2257,8 @@ def create_resume_fn():
             self._graph.pycode_gen.gen_call_function(
                 argc=resume_fn.__code__.co_argcount,
             )
-
         # gen RETURN_VALUE
         self._graph.pycode_gen.gen_return()
-
         self.new_code = self._graph.pycode_gen.gen_pycode()
         self.guard_fn = self._graph.guard_fn
 
diff --git a/python/paddle/jit/sot/opcode_translator/instruction_utils/opcode_info.py b/python/paddle/jit/sot/opcode_translator/instruction_utils/opcode_info.py
@@ -23,6 +23,7 @@
 ABS_JUMP = {opcode.opname[x] for x in opcode.hasjabs}
 HAS_LOCAL = {opcode.opname[x] for x in opcode.haslocal}
 HAS_FREE = {opcode.opname[x] for x in opcode.hasfree}
+NEED_TO_BOOL = {"UNARY_NOT", "POP_JUMP_IF_FALSE", "POP_JUMP_IF_TRUE"}
 ALL_JUMP = REL_JUMP | ABS_JUMP
 UNCONDITIONAL_JUMP = {"JUMP_ABSOLUTE", "JUMP_FORWARD"}
 if sys.version_info >= (3, 11):
diff --git a/test/sot/skip_files_py313 b/test/sot/skip_files_py313
@@ -1,8 +1,5 @@
-test/sot/test_11_jumps.py
-test/sot/test_12_for_loop.py
 test/sot/test_19_closure.py
 test/sot/test_break_graph.py
 test/sot/test_min_graph_size.py
 test/sot/test_numpy.py
-test/sot/test_resume_cache.py
 test/sot/test_simulate_initialize.py
diff --git a/test/sot/test_break_graph.py b/test/sot/test_break_graph.py
@@ -200,5 +200,15 @@ def test_break_graph_when_call_generator_function(self):
         self.assert_results(break_graph_call_generator_function, [x, y])
 
 
+def unary_not_break_graph(x):
+    return not x
+
+
+class TestUnaryNot(TestCaseBase):
+    def test_unary_not_break_graph(self):
+        x = paddle.to_tensor(0)
+        self.assert_results(unary_not_break_graph, x)
+
+
 if __name__ == "__main__":
     unittest.main()
