[SOT] Add BlockStackItem (#72474)

DrRyanHuang · web-flow · commit 29bf1f58f903 · 2025-04-27T18:38:20.000+08:00
diff --git a/python/paddle/jit/sot/opcode_translator/executor/virtual_frame.py b/python/paddle/jit/sot/opcode_translator/executor/virtual_frame.py
@@ -16,6 +16,7 @@
 
 import builtins
 import re
+from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, NamedTuple
 
 from ...utils import log
@@ -41,6 +42,7 @@
 
     from typing_extensions import TypeAlias
 
+    from ..instruction_utils import Instruction
     from .function_graph import FunctionGraph
     from .variables.callable import FunctionVariable
 
@@ -57,12 +59,22 @@ def validate_value(value):
     ), f"dangling variable {value} should not be pushed into stack."
 
 
+@dataclass
+class BlockStackItem:
+    # `PyTryBlock` in CPython source code
+    type: str
+    inst: Instruction
+    handler: Instruction
+    level: int
+
+
 class VirtualFrameState(NamedTuple):
     locals: dict[str, VariableBase]
     builtins: dict[str, VariableBase]
     cells: dict[str, VariableBase]
     lasti: int
     stack_data: list[VariableBase]
+    block_stack: list[BlockStackItem]
 
 
 class VirtualFrame:
@@ -74,6 +86,7 @@ class VirtualFrame:
     cells: dict[str, Any]
     lasti: int
     stack: VariableStack
+    block_stack: list[BlockStackItem]
 
     def __init__(self, code: types.CodeType):
         self.code = code
@@ -84,6 +97,7 @@ def __init__(self, code: types.CodeType):
         self.lasti = 0
         self.consts = []
         self.stack = VariableStack(validate_value_func=validate_value)
+        self.block_stack: list[BlockStackItem] = []
 
     @staticmethod
     def from_real_frame(frame: types.FrameType, graph: FunctionGraph):
@@ -213,6 +227,7 @@ def get_state(self):
             cells=self.cells.copy(),
             lasti=self.lasti,
             stack_data=list(self.stack._data),
+            block_stack=self.block_stack.copy(),
         )
 
     def restore_state(self, state: VirtualFrameState):
@@ -221,3 +236,4 @@ def restore_state(self, state: VirtualFrameState):
         self.cells = state.cells
         self.lasti = state.lasti
         self.stack._data = state.stack_data
+        self.block_stack = state.block_stack
