[SOT][3.11] Combine PRECALL and CALL as a super instruction in simulation (#60280)

SigureMo · web-flow · commit 6eb0a84b8476 · 2023-12-25T11:41:54.000+08:00
diff --git a/python/paddle/jit/sot/opcode_translator/executor/function_graph.py b/python/paddle/jit/sot/opcode_translator/executor/function_graph.py
@@ -255,15 +255,16 @@ def load(self, var):
                 self._pycode_gen.gen_load(self._store_var_info[var.id])
 
         origin_instrs = get_instructions(self.pycode_gen._origin_code)
+        is_precall = origin_instrs[instr_idx].opname == "PRECALL"
+        current_idx = instr_idx
+        # skip CALL if current instr is PRECALL
+        next_idx = instr_idx + 1 + int(is_precall)
 
-        restore_instrs = origin_instrs[:instr_idx]
+        restore_instrs = origin_instrs[:current_idx]
         restore_instr_names = [
-            instr.opname for instr in restore_instrs[:instr_idx]
+            instr.opname for instr in restore_instrs[:current_idx]
         ]
-        # NOTE(SigureMo): Trailing KW_NAMES or PRECALL is no need to restore in Python 3.11+
-        if restore_instr_names[-1:] == ["PRECALL"]:
-            restore_instrs = restore_instrs[:-1]
-            restore_instr_names = restore_instr_names[:-1]
+        # NOTE(SigureMo): Trailing KW_NAMES is no need to restore in Python 3.11+
         if restore_instr_names[-1:] == ["KW_NAMES"]:
             restore_instrs = restore_instrs[:-1]
             restore_instr_names = restore_instr_names[:-1]
@@ -272,12 +273,12 @@ def load(self, var):
         nop = self.pycode_gen._add_instr("NOP")
 
         for instr in origin_instrs:
-            if instr.jump_to == origin_instrs[instr_idx]:
+            if instr.jump_to == origin_instrs[current_idx]:
                 instr.jump_to = nop
 
         self.pycode_gen.hooks.append(
             lambda: self.pycode_gen.extend_instrs(
-                iter(origin_instrs[instr_idx + 1 :])
+                iter(origin_instrs[next_idx:])
             )
         )
 
diff --git a/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py b/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py
@@ -558,8 +558,10 @@ def step(self, instr: Instruction):
             print(log_message)
             breakpoint()  # noqa: T100
 
-        with EventGuard(f"{instr.opname}", event_level=1):
-            return getattr(self, instr.opname)(instr)  # run single step.
+        opname = instr.opname if instr.opname != "PRECALL" else "PRECALL__CALL"
+        assert opname != "CALL", "CALL should fused with PRECALL"
+        with EventGuard(f"{opname}", event_level=1):
+            return getattr(self, opname)(instr)  # run single step.
 
     def indexof(self, instr: Instruction):
         """
@@ -1027,6 +1029,19 @@ def BUILD_MAP_UNPACK_WITH_CALL(self, instr: Instruction):
             )
         )
 
+    @call_break_graph_decorator(push_n=1)
+    def PRECALL__CALL(self, instr: Instruction):
+        """
+        presudo super-instruction for PRECALL + CALL
+        """
+        assert isinstance(instr.arg, int)
+        assert instr.opname == "PRECALL"
+        self.PRECALL(instr)
+        next_instr = self._instructions[self._lasti]
+        self._lasti += 1
+        assert next_instr.opname == "CALL"
+        self.CALL(next_instr)
+
     def PRECALL(self, instr: Instruction):
         assert isinstance(instr.arg, int)
         is_method_layout = not isinstance(
@@ -1045,7 +1060,6 @@ def KW_NAMES(self, instr: Instruction):
         assert isinstance(instr.arg, int)
         self._call_shape = self._co_consts[instr.arg].get_py_value()
 
-    @call_break_graph_decorator(push_n=1)
     def CALL(self, instr: Instruction):
         assert isinstance(instr.arg, int)
         assert instr.arg + 2 <= len(self.stack)
@@ -1660,19 +1674,24 @@ def _break_graph_in_call(
 
         """
         push_n = push_n(instr.arg) if callable(push_n) else push_n
+        is_precall = instr.opname == "PRECALL"
         index = self.indexof(instr)
+        # Use CALL instead of PRECALL to calculate the real stack effect
+        call_instr = self._instructions[index + int(is_precall)]
+        # skip CALL if current instr is PRECALL
+        next_index = index + 1 + int(is_precall)
         self.stack = origin_stack
 
         # gen call static fn opcode
 
-        resume_input_name = analysis_inputs(self._instructions, index + 1)
+        resume_input_name = analysis_inputs(self._instructions, next_index)
 
         var_loader = self.gen_compute_in_break_with_name_store(
-            resume_input_name, self.indexof(instr)
+            resume_input_name, index
         )
 
         # gen graph break call fn opcode
-        stack_effect = calc_stack_effect(instr)
+        stack_effect = calc_stack_effect(call_instr)
         pop_n = push_n - stack_effect
 
         for i, stack_arg in enumerate(self.stack):
@@ -1681,11 +1700,13 @@ def _break_graph_in_call(
         # gen call resume fn opcode
         # NOTE(SigureMo): In Python 3.11，we need generate KW_NAMES if the call shape is not None.
         self._graph.pycode_gen.gen_kw_names(self._call_shape)
-        self._graph.pycode_gen.extend_instrs([instr])
+        self._graph.pycode_gen.extend_instrs(
+            self._instructions[index:next_index]
+        )
         self.stack.pop_n(pop_n)
         stack_size = len(self.stack) + push_n
 
-        resume_fn, _ = self._create_resume_fn(index + 1, stack_size)
+        resume_fn, _ = self._create_resume_fn(next_index, stack_size)
 
         if resume_fn:
             self._graph.pycode_gen.gen_load_object(
diff --git a/test/sot/test_break_graph.py b/test/sot/test_break_graph.py
@@ -164,5 +164,26 @@ def test_break_graph_resume_pass_null(self):
         self.assert_results(break_graph_resume_pass_null, x, y)
 
 
+class MyLayer(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self.head = paddle.nn.Linear(3, 10)
+
+    def forward_features(self, x):
+        paddle.jit.sot.psdb.breakgraph()
+        return x
+
+    def forward(self, x):
+        x = self.forward_features(x)
+        return self.head(x)
+
+
+class TestBreakGraphInLayer(TestCaseBase):
+    def test_break_graph_in_layer(self):
+        x = paddle.rand([2, 3], dtype=paddle.float32)
+        net = MyLayer()
+        self.assert_results(net.forward, x)
+
+
 if __name__ == "__main__":
     unittest.main()
