[SOT][3.13] Support CALL_KW in python 3.13 (#69431)

SigureMo · web-flow · commit a8aca59f2397 · 2024-11-16T12:00:58.000+08:00
diff --git a/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py b/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py
@@ -370,7 +370,7 @@ def __init__(self, code: types.CodeType, graph: FunctionGraph):
         self.guard_fn = None
         self._name = "Executor"
         self._call_shape: tuple[str, ...] | None = (
-            None  # store kwnames for Python 3.11+
+            None  # store kwnames for Python 3.11 and 3.12
         )
         self._prepare_virtual_env()
         self.stop_state = None
@@ -1158,12 +1158,15 @@ def KW_NAMES(self, instr: Instruction):
         assert isinstance(instr.arg, int)
         self._call_shape = self._co_consts[instr.arg].get_py_value()
 
-    def call(self, instr: Instruction):
-        assert isinstance(instr.arg, int)
-        assert instr.arg + 2 <= len(self.stack)
-        is_method = not isinstance(self.stack.peek[instr.arg + 2], NullVariable)
-        total_args = instr.arg + int(is_method)
-        kwnames = self._call_shape if self._call_shape is not None else []
+    def call_impl(
+        self,
+        num_args: int,
+        kwnames_getter: Callable[[OpcodeExecutor], tuple[str, ...]],
+    ):
+        kwnames = kwnames_getter(self)
+        assert num_args + 2 <= len(self.stack)
+        is_method = not isinstance(self.stack.peek[num_args + 2], NullVariable)
+        total_args = num_args + int(is_method)
         n_kwargs = len(kwnames)
         n_positional_args = total_args - n_kwargs
         kwargs_list = self.stack.pop_n(n_kwargs)
@@ -1176,12 +1179,33 @@ def call(self, instr: Instruction):
         self.stack.push(fn(*args, **kwargs))
         self._call_shape = None
 
+    def call(self, instr: Instruction):
+        assert isinstance(instr.arg, int)
+        self.call_impl(
+            instr.arg,
+            lambda exe: (
+                exe._call_shape if exe._call_shape is not None else ()
+            ),
+        )
+
     CALL = (
         call_break_graph_decorator(push_n=1)(call)
         if sys.version_info >= (3, 12)
         else call
     )
 
+    @call_break_graph_decorator(push_n=1)
+    def CALL_KW(self, instr: Instruction):
+        assert isinstance(instr.arg, int)
+
+        def get_kwnames(exe: OpcodeExecutor):
+            kwnames_var = exe.stack.pop()
+            assert isinstance(kwnames_var, TupleVariable)
+            kwnames = kwnames_var.get_py_value()
+            return kwnames
+
+        self.call_impl(instr.arg, get_kwnames)
+
     @call_break_graph_decorator(push_n=1)
     def CALL_FUNCTION(self, instr: Instruction):
         assert isinstance(instr.arg, int)
@@ -2018,7 +2042,7 @@ def create_resume_fn():
             var_loader.load(stack_arg)
 
         # 5. run the break CALL with origin python
-        # NOTE(SigureMo): In Python 3.11，we need generate KW_NAMES if the call shape is not None.
+        # NOTE(SigureMo): In Python 3.11 and 3.12，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(
             self._instructions[cur_index:next_index]
diff --git a/python/paddle/jit/sot/opcode_translator/executor/pycode_generator.py b/python/paddle/jit/sot/opcode_translator/executor/pycode_generator.py
@@ -821,8 +821,10 @@ def gen_call_method(self, argc=0):
     def gen_kw_names(self, kw_names: tuple[str, ...] | None):
         if kw_names is None:
             return
-        if sys.version_info < (3, 11):
-            raise InnerError("gen_kw_names is not supported before python3.11")
+        if sys.version_info < (3, 11) or sys.version_info >= (3, 13):
+            raise InnerError(
+                "gen_kw_names is only supported in Python 3.11 and 3.12"
+            )
         if kw_names not in self._code_options["co_consts"]:
             self._code_options["co_consts"].append(kw_names)
         idx = self._code_options["co_consts"].index(kw_names)
diff --git a/test/sot/skip_files_py313 b/test/sot/skip_files_py313
@@ -1,6 +1,4 @@
-test/sot/test_03_tuple.py
 test/sot/test_04_list.py
-test/sot/test_06_call_function.py
 test/sot/test_11_jumps.py
 test/sot/test_12_for_loop.py
 test/sot/test_13_make_function.py
@@ -12,15 +10,13 @@ test/sot/test_18_tensor_method.py
 test/sot/test_19_closure.py
 test/sot/test_21_global.py
 test/sot/test_analysis_inputs.py
-test/sot/test_binary_operator_tracker.py
 test/sot/test_break_graph.py
 test/sot/test_builtin_bool.py
 test/sot/test_builtin_dispatch.py
 test/sot/test_builtin_map.py
 test/sot/test_builtin_range.py
 test/sot/test_builtin_zip.py
 test/sot/test_call_object.py
-test/sot/test_constant_graph.py
 test/sot/test_dtype.py
 test/sot/test_dup_top.py
 test/sot/test_enumerate.py
@@ -29,14 +25,12 @@ test/sot/test_inplace_api.py
 test/sot/test_min_graph_size.py
 test/sot/test_model_switch_training.py
 test/sot/test_numpy.py
-test/sot/test_numpy_var_if.py
+test/sot/test_resume_cache.py
 test/sot/test_side_effects.py
 test/sot/test_simulate_initialize.py
 test/sot/test_sot_cost_model.py
 test/sot/test_sot_dynamic_shape.py
-test/sot/test_sot_exception.py
 test/sot/test_sot_export.py
 test/sot/test_sot_resnet.py
 test/sot/test_sot_resnet50_backward.py
 test/sot/test_str_format.py
-test/sot/test_tensor_dtype_in_guard.py
