[SOT][3.13] Support non-breakgraph for loop (#69436)

SigureMo · web-flow · commit 7584a89db7a5 · 2024-11-17T13:31:13.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
@@ -1809,7 +1809,12 @@ def FOR_ITER(self, instr):
             self._lasti = self.indexof(instr.jump_to)
             if sys.version_info >= (3, 12):
                 assert self._instructions[self._lasti].opname == "END_FOR"
-                self._lasti += 1
+                # NOTE(SigureMo): From Python 3.12, END_FOR indicates the end of the for loop.
+                # But it's never executed, so we need to skip it.
+                # In Python 3.12, it equivalent to POP_TOP + POP_TOP (in one instruction)
+                # In Python 3.13, it equivalent to POP_TOP, and it common with other POP_TOP
+                skip_n_instrs = 2 if sys.version_info >= (3, 13) else 1
+                self._lasti += skip_n_instrs
         except BreakGraphError as e:
             log(3, f"[BreakGraph] FOR_ITER sim for loop failed for: {e}\n")
             if backup_iter_idx:
@@ -2241,7 +2246,8 @@ def create_after_loop_fn():
             # skip resume END_FOR in python3.12
             if sys.version_info >= (3, 12):
                 assert origin_instrs[loop_body_end_idx].opname == "END_FOR"
-                resume_fn_end_idx += 1
+                skip_n_instrs = 2 if sys.version_info >= (3, 13) else 1
+                resume_fn_end_idx += skip_n_instrs
 
             resume_fn_creator.set_inputs(
                 after_loop_fn_inputs, stack_size=len(self.stack) - 1
@@ -2322,6 +2328,8 @@ def create_after_loop_fn():
 
         if sys.version_info >= (3, 12):
             end_for = self._graph.pycode_gen.add_instr("END_FOR")
+            if sys.version_info >= (3, 13):
+                self._graph.pycode_gen.gen_pop_top()
 
         nop = self._graph.pycode_gen.add_instr("NOP")
 
@@ -2411,6 +2419,8 @@ def create_inline_call_fn():
 
             if sys.version_info >= (3, 12):
                 end_for = pycode_gen.add_instr("END_FOR")
+                if sys.version_info >= (3, 13):
+                    self._graph.pycode_gen.gen_pop_top()
             nop_for_break = pycode_gen.add_instr("NOP")
 
             # 2.4. relocate jumps
diff --git a/python/paddle/jit/sot/opcode_translator/executor/opcode_inline_executor.py b/python/paddle/jit/sot/opcode_translator/executor/opcode_inline_executor.py
@@ -319,7 +319,8 @@ def FOR_ITER(self, instr: Instruction):
                 self._lasti = self.indexof(instr.jump_to)
                 if sys.version_info >= (3, 12):
                     assert self._instructions[self._lasti].opname == "END_FOR"
-                    self._lasti += 1
+                    skip_n_instrs = 2 if sys.version_info >= (3, 13) else 1
+                    self._lasti += skip_n_instrs
 
         else:
             self._graph.remove_global_guarded_variable(iterator)
diff --git a/python/paddle/jit/sot/opcode_translator/executor/pycode_generator.py b/python/paddle/jit/sot/opcode_translator/executor/pycode_generator.py
@@ -43,12 +43,12 @@
     apply_instr_pass,
     calc_stack_effect,
     gen_instr,
+    get_instruction_size,
     instrs_info,
     modify_instrs,
     modify_vars,
 )
 from ..instruction_utils.opcode_info import (
-    PYOPCODE_CACHE_SIZE,
     UNCONDITIONAL_JUMP,
     JumpDirection,
     PopJumpCond,
@@ -217,13 +217,6 @@ def to_byte(num):
     return num
 
 
-def get_instruction_size(instr: Instruction) -> int:
-    cache_size = 0
-    if sys.version_info >= (3, 11):
-        cache_size = PYOPCODE_CACHE_SIZE.get(instr.opname, 0)
-    return 2 * (cache_size + 1)
-
-
 def create_linetable_calculator(firstlineno: int):
     """
     Creates a line table calculator function.
@@ -933,6 +926,11 @@ def gen_pop_jump(
         direction: JumpDirection = JumpDirection.FORWARD,
         suffix: PopJumpCond = PopJumpCond.NONE,
     ) -> Instruction:
+        if sys.version_info >= (3, 13) and suffix in [
+            PopJumpCond.TRUE,
+            PopJumpCond.FALSE,
+        ]:
+            self.add_instr("TO_BOOL")
         if sys.version_info >= (3, 11) and sys.version_info < (3, 12):
             return self.add_instr(
                 f"POP_JUMP_{direction.value}_IF_{suffix.value}", jump_to=jump_to
diff --git a/python/paddle/jit/sot/opcode_translator/instruction_utils/__init__.py b/python/paddle/jit/sot/opcode_translator/instruction_utils/__init__.py
@@ -20,6 +20,7 @@
     calc_stack_effect,
     convert_instruction,
     gen_instr,
+    get_instruction_size,
     get_instructions,
     instrs_info,
     modify_extended_args,
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
@@ -57,6 +57,13 @@ def __eq__(self, instr):
         return id(self) == id(instr)
 
 
+def get_instruction_size(instr: Instruction) -> int:
+    cache_size = 0
+    if sys.version_info >= (3, 11):
+        cache_size = PYOPCODE_CACHE_SIZE.get(instr.opname, 0)
+    return 2 * (cache_size + 1)
+
+
 def gen_instr(name, arg=None, argval=None, gened=True, jump_to=None):
     return Instruction(
         opcode=dis.opmap[name],
@@ -207,8 +214,10 @@ def modify_instrs(instructions: list[Instruction]) -> None:
     modify_completed = False
     while not modify_completed:
         reset_offset(instructions)
-        relocate_jump_target(instructions)
-        modify_completed = modify_extended_args(instructions)
+        has_inverted_jump = relocate_jump_target(instructions)
+        modify_completed = (
+            modify_extended_args(instructions) and not has_inverted_jump
+        )
 
 
 def reset_offset(instructions: list[Instruction]) -> None:
@@ -233,18 +242,21 @@ def reset_offset(instructions: list[Instruction]) -> None:
         instr.offset = idx * 2
 
 
-def correct_jump_direction(instr: Instruction, arg: int) -> Instruction:
+def correct_jump_direction(
+    instr: Instruction, arg: int
+) -> tuple[Instruction, bool]:
     """
     Corrects the jump direction of the given instruction.
     NOTE(zrr1999): In Python 3.11, JUMP_ABSOLUTE is removed, so python generates JUMP_FORWARD or JUMP_BACKWARD instead,
     but in for loop breakgraph, we reuse JUMP_BACKWARD to jump forward, so we need to change it to JUMP_FORWARD.
 
     Args:
         instr (Instruction): The instruction to be corrected.
+        invert_jump (bool): Whether to invert the jump direction.
     """
     if instr.opname in ABS_JUMP:
         instr.arg = arg
-        return instr
+        return instr, False
     elif instr.opname in REL_JUMP:
         if arg < 0:
             if instr.opname in REL_BWD_JUMP:
@@ -260,14 +272,16 @@ def correct_jump_direction(instr: Instruction, arg: int) -> Instruction:
                 instr.opname = backward_op_name
                 instr.opcode = dis.opmap[backward_op_name]
             instr.arg = -arg
+            invert_jump = True
         else:
             instr.arg = arg
-        return instr
+            invert_jump = False
+        return instr, invert_jump
     else:
         raise ValueError(f"unknown jump type: {instr.opname}")
 
 
-def relocate_jump_target(instructions: list[Instruction]) -> None:
+def relocate_jump_target(instructions: list[Instruction]) -> bool:
     """
     If a jump instruction is found, this function will adjust the jump targets based on the presence of EXTENDED_ARG instructions.
     If an EXTENDED_ARG instruction exists for the jump target, use its offset as the new target.
@@ -276,8 +290,9 @@ def relocate_jump_target(instructions: list[Instruction]) -> None:
         instructions (list): The list of Instruction objects representing bytecode instructions.
 
     Returns:
-        None
+        bool: True if the jump direction is inverted, False otherwise.
     """
+    has_inverted_jump = False
     extended_arg = []
     for instr in instructions:
         if instr.opname == "EXTENDED_ARG":
@@ -305,7 +320,8 @@ def relocate_jump_target(instructions: list[Instruction]) -> None:
 
             if sys.version_info >= (3, 10):
                 new_arg //= 2
-            correct_jump_direction(instr, new_arg)
+            _, invert_jump = correct_jump_direction(instr, new_arg)
+            has_inverted_jump = has_inverted_jump or invert_jump
             assert instr.arg is not None
             if extended_arg:
                 instr.arg &= 0xFF
@@ -319,6 +335,7 @@ def relocate_jump_target(instructions: list[Instruction]) -> None:
                 if new_arg > 0:
                     extended_arg[0].arg += new_arg << 8
         extended_arg.clear()
+    return has_inverted_jump
 
 
 def modify_extended_args(instructions: list[Instruction]) -> bool:
diff --git a/test/sot/skip_files_py313 b/test/sot/skip_files_py313
@@ -5,12 +5,10 @@ test/sot/test_14_operators.py
 test/sot/test_15_slice.py
 test/sot/test_17_paddle_layer.py
 test/sot/test_19_closure.py
-test/sot/test_21_global.py
 test/sot/test_analysis_inputs.py
 test/sot/test_break_graph.py
 test/sot/test_builtin_bool.py
 test/sot/test_builtin_dispatch.py
-test/sot/test_builtin_map.py
 test/sot/test_builtin_range.py
 test/sot/test_builtin_zip.py
 test/sot/test_dup_top.py
