[SOT][3.12] replace POP_JUMP_{BACKWARD,FORWARD}_IF_{TRUE,FALSE} to POP_JUMP_IF_{TRUE,FALSE}

python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py

@@ -640,10 +640,6 @@ def _rot_top_n(self, n: int):
     def POP_TOP(self, instr: Instruction):
         self.stack.pop()
 
-    def END_FOR(self, instr: Instruction):
-        self.POP_TOP(instr)
-        self.POP_TOP(instr)
-
     def PUSH_NULL(self, instr: Instruction):
         self.stack.push(NullVariable())
 
@@ -1690,8 +1686,9 @@ def FOR_ITER(self, instr):
 
             self._inline_call_for_loop(iterator, instr)
             self._lasti = self.indexof(instr.jump_to)
-            next_instr = self._instructions[self._lasti]
-            self._lasti += int(next_instr.opname == 'END_FOR')
+            if sys.version_info >= (3, 12):
+                assert self._instructions[self._lasti].opname == "END_FOR"
+                self._lasti += 1
         except BreakGraphError as e:
             log(3, f"[BreakGraph] FOR_ITER sim for loop failed for: {e}\n")
             if backup_iter_idx:
@@ -2067,10 +2064,8 @@ def create_after_loop_fn():
             resume_fn_end_idx = loop_body_end_idx
 
             # skip resume END_FOR in python3.12
-            if (
-                sys.version_info >= (3, 12)
-                and origin_instrs[loop_body_end_idx].opname == "END_FOR"
-            ):
+            if sys.version_info >= (3, 12):
+                assert origin_instrs[loop_body_end_idx].opname == "END_FOR"
                 resume_fn_end_idx += 1
 
             pycode_gen.set_function_inputs(

python/paddle/jit/sot/opcode_translator/executor/opcode_inline_executor.py

@@ -17,6 +17,7 @@
 import contextlib
 import inspect
 import re
+import sys
 from typing import TYPE_CHECKING
 
 from ...profiler import event_register
@@ -316,8 +317,9 @@ def FOR_ITER(self, instr: Instruction):
                 self.stack.pop()
                 assert isinstance(instr.jump_to, Instruction)
                 self._lasti = self.indexof(instr.jump_to)
-                next_instr = self._instructions[self._lasti]
-                self._lasti += int(next_instr.opname == 'END_FOR')
+                if sys.version_info >= (3, 12):
+                    assert self._instructions[self._lasti].opname == "END_FOR"
+                    self._lasti += 1
 
         else:
             self._graph.remove_global_guarded_variable(iterator)

python/paddle/jit/sot/opcode_translator/instruction_utils/instruction_utils.py

@@ -21,7 +21,13 @@
 from typing import TYPE_CHECKING, Any
 
 from ...utils import InnerError
-from .opcode_info import ABS_JUMP, ALL_JUMP, REL_BWD_JUMP, REL_JUMP
+from .opcode_info import (
+    ABS_JUMP,
+    ALL_JUMP,
+    PYOPCODE_CACHE_SIZE,
+    REL_BWD_JUMP,
+    REL_JUMP,
+)
 
 if TYPE_CHECKING:
     import types
@@ -239,6 +245,9 @@ def relocate_jump_target(instructions: list[Instruction]) -> None:
             if instr.opname in ABS_JUMP:
                 new_arg = jump_target
             else:  # instr.opname in REL_JUMP
+                if sys.version_info >= (3, 12):
+                    cache = PYOPCODE_CACHE_SIZE.get(instr.opname, 0)
+                    jump_target -= 2 * cache
                 new_arg = jump_target - instr.offset - 2
                 if instr.opname in REL_BWD_JUMP:
                     new_arg = -new_arg
@@ -336,19 +345,6 @@ def modify_vars(instructions: list[Instruction], code_options):
                     instrs.argval in namemap
                 ), f"`{instrs.argval}` not in {namemap}"
                 instrs.arg = namemap.index(instrs.argval)
-        elif instrs.opname == "FOR_ITER":
-            if sys.version_info >= (3, 12):
-                assert instrs.jump_to is not None
-                assert instrs.arg is not None
-                assert instrs.offset is not None
-                instrs.arg -= 1
-                # END_FOR offset = (((FOR_ITER arg + 2) << 1) + FOR_ITER offset)
-                if instrs.jump_to.offset != (
-                    ((instrs.arg + 2) << 1) + instrs.offset
-                ):
-                    raise InnerError(
-                        'FOR_ITER jump_to offset is not equal to "(((FOR_ITER arg + 2) << 1) + FOR_ITER offset)" '
-                    )
 
 
 def calc_offset_from_bytecode_offset(

python/paddle/jit/sot/opcode_translator/instruction_utils/opcode_info.py

@@ -45,7 +45,7 @@ class PopJumpCond(Enum):
     NOT_NONE = "NOT_NONE"
 
 
-def get_pyopcode_cache_size() -> dict[str, int]:
+def _get_pyopcode_cache_size() -> dict[str, int]:
     if sys.version_info >= (3, 11) and sys.version_info < (3, 12):
         # Cache for some opcodes, it's for Python 3.11+
         # https://github.com/python/cpython/blob/3.11/Include/internal/pycore_opcode.h#L41-L53
@@ -87,4 +87,4 @@ def get_pyopcode_cache_size() -> dict[str, int]:
         return {}
 
 
-PYOPCODE_CACHE_SIZE = get_pyopcode_cache_size()
+PYOPCODE_CACHE_SIZE = _get_pyopcode_cache_size()