[SOT][3.11] Do not restore KW_NAMES + PRECALL in 3.11 (#59017)

SigureMo · web-flow · commit 2b3edfc8b9ac · 2023-11-16T10:38:51.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
@@ -260,9 +260,17 @@ def load(self, var, allow_push_null=True):
                 # only restored vars in stack, so used gen_load to process global var
                 self._pycode_gen.gen_load(self._store_var_info[var])
 
-        origin_instr = get_instructions(self.pycode_gen._origin_code)
+        origin_instrs = get_instructions(self.pycode_gen._origin_code)
 
-        for instr in origin_instr[0:instr_idx]:
+        restore_instrs = origin_instrs[:instr_idx]
+        restore_instr_names = [
+            instr.opname for instr in restore_instrs[:instr_idx]
+        ]
+        # NOTE(SigureMo): Trailing KW_NAMES + PRECALL is no need to restore in Python 3.11+
+        if restore_instr_names[-2:] == ["KW_NAMES", "PRECALL"]:
+            restore_instrs = restore_instrs[:-2]
+
+        for instr in restore_instrs:
             if (
                 instr.opname == 'LOAD_FAST'
                 and instr.argval in self.pycode_gen._frame.f_locals.keys()
@@ -288,13 +296,13 @@ def load(self, var, allow_push_null=True):
 
         nop = self.pycode_gen._add_instr("NOP")
 
-        for instr in origin_instr:
-            if instr.jump_to == origin_instr[instr_idx]:
+        for instr in origin_instrs:
+            if instr.jump_to == origin_instrs[instr_idx]:
                 instr.jump_to = nop
 
         self.pycode_gen.hooks.append(
             lambda: self.pycode_gen.extend_instrs(
-                iter(origin_instr[instr_idx + 1 :])
+                iter(origin_instrs[instr_idx + 1 :])
             )
         )
 
diff --git a/test/sot/test_min_graph_size.py b/test/sot/test_min_graph_size.py
@@ -46,6 +46,16 @@ def case_call(x):
     return x
 
 
+def call_with_kwargs_inner(x):
+    return paddle.to_tensor(x.numpy())
+
+
+def call_with_kwargs(x):
+    y = call_with_kwargs_inner(x=x)
+    x += y
+    return x
+
+
 def case_all(x, vars):
     x = x + 1
     for y in vars:
@@ -82,6 +92,11 @@ def test_layer(self):
         layer = CustomLayer()
         self.assert_results(layer.forward, x)
 
+    @min_graph_size_guard(10)
+    def test_call_with_kwargs(self):
+        x = paddle.to_tensor(1)
+        self.assert_results(call_with_kwargs, x)
+
 
 if __name__ == "__main__":
     unittest.main()
