[Compat][3.11] fix reset_offset and support jump breakgraph

sot/opcode_translator/instruction_utils/instruction_utils.py

@@ -5,7 +5,7 @@
 import sys
 from typing import TYPE_CHECKING, Any
 
-from .opcode_info import ABS_JUMP, ALL_JUMP, REL_JUMP
+from .opcode_info import ALL_JUMP, REL_JUMP
 
 if TYPE_CHECKING:
     import types
@@ -143,11 +143,19 @@ def reset_offset(instructions: list[Instruction]) -> None:
     Returns:
         None
     """
+    from ..executor.pycode_generator import get_instruction_size
+
+    if sys.version_info >= (3, 11):
+        current_offset = 0
+        for instr in instructions:
+            instr.offset = current_offset
+            current_offset += get_instruction_size(instr)
+        return
     for idx, instr in enumerate(instructions):
         instr.offset = idx * 2
 
 
-def relocate_jump_target(instuctions: list[Instruction]) -> None:
+def relocate_jump_target(instructions: list[Instruction]) -> None:
     """
     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.
@@ -159,23 +167,27 @@ def relocate_jump_target(instuctions: list[Instruction]) -> None:
         None
     """
     extended_arg = []
-    for instr in instuctions:
+    for instr in instructions:
         if instr.opname == "EXTENDED_ARG":
             extended_arg.append(instr)
             continue
 
         if instr.opname in ALL_JUMP:
+            assert instr.jump_to is not None
+            assert instr.offset is not None
             # if jump target has extended_arg, should jump to the first extended_arg opcode
             jump_target = (
                 instr.jump_to.offset
                 if instr.jump_to.first_ex_arg is None
                 else instr.jump_to.first_ex_arg.offset
             )
+            assert jump_target is not None
 
             if instr.opname in REL_JUMP:
                 new_arg = jump_target - instr.offset - 2
-            elif instr.opname in ABS_JUMP:
+            else:  # instr.opname in ABS_JUMP
                 new_arg = jump_target
+
             if sys.version_info >= (3, 10):
                 new_arg //= 2
 
@@ -192,7 +204,6 @@ def relocate_jump_target(instuctions: list[Instruction]) -> None:
                     extended_arg[0].arg += new_arg << 8
             else:
                 instr.arg = new_arg
-
         extended_arg.clear()
 
 

tests/test_11_jumps.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import sys
 import unittest
 
 from test_case_base import TestCaseBase
@@ -86,9 +85,6 @@ def test_simple(self):
         self.assert_results(pop_jump_if_not_none, None, a)
         self.assert_results(pop_jump_if_not_none, True, a)
 
-    @unittest.skipIf(
-        sys.version_info >= (3, 11), "Python 3.11+ not support breakgraph."
-    )
     def test_breakgraph(self):
         self.assert_results(pop_jump_if_false, true_tensor, a)
         self.assert_results(jump_if_false_or_pop, true_tensor, a)