[SOT][3.12] Support LOAD_ATTR opcode and skip PRECALL_CALL check in Python 3.12 (#61305)

gouzil · web-flow · commit 56b6f63331f1 · 2024-02-18T14:31:34.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
@@ -561,8 +561,10 @@ def step(self, instr: Instruction):
             print(log_message)
             breakpoint()  # noqa: T100
 
-        opname = instr.opname if instr.opname != "PRECALL" else "PRECALL__CALL"
-        assert opname != "CALL", "CALL should fused with PRECALL"
+        opname = instr.opname
+        if sys.version_info < (3, 12):
+            opname = opname if opname != "PRECALL" else "PRECALL__CALL"
+            assert opname != "CALL", "CALL should fused with PRECALL"
         with EventGuard(f"{opname}", event_level=2):
             return getattr(self, opname)(instr)  # run single step.
 
@@ -717,7 +719,14 @@ def NOP(self, instr: Instruction):
 
     @call_break_graph_decorator(push_n=1)
     def LOAD_ATTR(self, instr: Instruction):
-        attr_name = self._code.co_names[instr.arg]
+        if sys.version_info >= (3, 12):
+            assert isinstance(instr.arg, int)
+            attr_name = self._code.co_names[instr.arg >> 1]
+            if instr.arg & 1:
+                self.load_method(attr_name)
+                return
+        else:
+            attr_name = self._code.co_names[instr.arg]
         attr_name_var = ConstantVariable.wrap_literal(attr_name, self._graph)
         obj = self.stack.pop()
         self.stack.push(
@@ -779,8 +788,7 @@ def LOAD_GLOBAL(self, instr: Instruction):
             raise InnerError(f"{name} not in globals and builtins")
         self.stack.push(value)
 
-    def LOAD_METHOD(self, instr: Instruction):
-        method_name = self._code.co_names[instr.arg]
+    def load_method(self, method_name):
         method_name_var = ConstantVariable.wrap_literal(
             method_name, self._graph
         )
@@ -802,6 +810,10 @@ def LOAD_METHOD(self, instr: Instruction):
             self.stack.push(NullVariable())
             self.stack.push(method)
 
+    def LOAD_METHOD(self, instr: Instruction):
+        method_name = self._code.co_names[instr.arg]
+        self.load_method(method_name)
+
     @call_break_graph_decorator(push_n=0)
     def STORE_ATTR(self, instr: Instruction):
         obj = self.stack.pop()
diff --git a/python/paddle/jit/sot/opcode_translator/executor/pycode_generator.py b/python/paddle/jit/sot/opcode_translator/executor/pycode_generator.py
@@ -878,7 +878,8 @@ def gen_unpack_sequence(self, count):
 
     def gen_call_function(self, argc=0):
         if sys.version_info >= (3, 11):
-            self._add_instr("PRECALL", arg=argc, argval=argc)
+            if sys.version_info >= (3, 11) and sys.version_info < (3, 12):
+                self._add_instr("PRECALL", arg=argc, argval=argc)
             self._add_instr("CALL", arg=argc, argval=argc)
         else:
             self._add_instr("CALL_FUNCTION", arg=argc, argval=argc)
@@ -891,7 +892,8 @@ def gen_call_function_ex(self, has_kwargs):
 
     def gen_call_method(self, argc=0):
         if sys.version_info >= (3, 11):
-            self._add_instr("PRECALL", arg=argc, argval=argc)
+            if sys.version_info >= (3, 11) and sys.version_info < (3, 12):
+                self._add_instr("PRECALL", arg=argc, argval=argc)
             self._add_instr("CALL", arg=argc, argval=argc)
         else:
             self._add_instr("CALL_METHOD", arg=argc, argval=argc)
diff --git a/python/paddle/jit/sot/opcode_translator/instruction_utils/opcode_info.py b/python/paddle/jit/sot/opcode_translator/instruction_utils/opcode_info.py
@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from __future__ import annotations
 
 import sys
 from enum import Enum
@@ -41,18 +42,46 @@ class PopJumpCond(Enum):
     NOT_NONE = "NOT_NONE"
 
 
-# Cache for some opcodes, it's for Python 3.11+
-# https://github.com/python/cpython/blob/3.11/Include/internal/pycore_opcode.h#L41-L53
-PYOPCODE_CACHE_SIZE = {
-    "BINARY_SUBSCR": 4,
-    "STORE_SUBSCR": 1,
-    "UNPACK_SEQUENCE": 1,
-    "STORE_ATTR": 4,
-    "LOAD_ATTR": 4,
-    "COMPARE_OP": 2,
-    "LOAD_GLOBAL": 5,
-    "BINARY_OP": 1,
-    "LOAD_METHOD": 10,
-    "PRECALL": 1,
-    "CALL": 4,
-}
+def get_pyopcode_cache_size() -> dict[str, int]:
+    if sys.version_info >= (3, 11) and sys.version_info < (3, 12):
+        # Cache for some opcodes, it's for Python 3.11+
+        # https://github.com/python/cpython/blob/3.11/Include/internal/pycore_opcode.h#L41-L53
+        return {
+            "BINARY_SUBSCR": 4,
+            "STORE_SUBSCR": 1,
+            "UNPACK_SEQUENCE": 1,
+            "STORE_ATTR": 4,
+            "LOAD_ATTR": 4,
+            "COMPARE_OP": 2,
+            "LOAD_GLOBAL": 5,
+            "BINARY_OP": 1,
+            "LOAD_METHOD": 10,
+            "PRECALL": 1,
+            "CALL": 4,
+        }
+    elif sys.version_info >= (3, 12) and sys.version_info < (3, 13):
+        # Cache for some opcodes, it's for Python 3.12+
+        # https://github.com/python/cpython/blob/3.12/Include/internal/pycore_opcode.h#L34-L47
+        return {
+            "BINARY_SUBSCR": 1,
+            "STORE_SUBSCR": 1,
+            "UNPACK_SEQUENCE": 1,
+            "FOR_ITER": 1,
+            "STORE_ATTR": 4,
+            "LOAD_ATTR": 9,
+            "COMPARE_OP": 1,
+            "LOAD_GLOBAL": 4,
+            "BINARY_OP": 1,
+            "SEND": 1,
+            "LOAD_SUPER_ATTR": 1,
+            "CALL": 3,
+        }
+    elif sys.version_info >= (3, 13):
+        raise NotImplementedError(
+            f"Need to supplement cache operation code, for Python {sys.version_info}"
+        )
+    else:
+        return {}
+
+
+PYOPCODE_CACHE_SIZE = get_pyopcode_cache_size()
diff --git a/test/sot/skip_files_py312 b/test/sot/skip_files_py312
@@ -1,19 +1,14 @@
 ./test_01_basic.py
-./test_02_store_inplace.py
 ./test_03_tuple.py
 ./test_04_list.py
 ./test_05_dict.py
-./test_06_call_function.py
 ./test_07_unpack.py
-./test_08_rot.py
 ./test_09_f_string.py
 ./test_10_build_unpack.py
 ./test_11_jumps.py
 ./test_12_for_loop.py
-./test_13_make_function.py
 ./test_14_operators.py
 ./test_15_slice.py
-./test_16_paddle_api.py
 ./test_17_paddle_layer.py
 ./test_18_tensor_method.py
 ./test_19_closure.py
@@ -26,26 +21,14 @@
 ./test_builtin_map.py
 ./test_builtin_range.py
 ./test_builtin_zip.py
-./test_call_ast.py
-./test_call_object.py
-./test_case_base.py
 ./test_constant_graph.py
-./test_delete_fast.py
-./test_dtype.py
 ./test_dup_top.py
 ./test_enumerate.py
-./test_execution_base.py
-./test_guard_outputs.py
 ./test_guard_user_defined_fn.py
 ./test_inplace_api.py
-./test_instruction_translator_cache.py
 ./test_min_graph_size.py
-./test_model_switch_training.py
-./test_multiple_args.py
-./test_numpy.py
 ./test_numpy_var_if.py
 ./test_output_restoration.py
-./test_segment_linear.py
 ./test_side_effects.py
 ./test_simulate_initialize.py
 ./test_sir_rollback.py
@@ -57,6 +40,3 @@
 ./test_specialization.py
 ./test_str_format.py
 ./test_tensor_dtype_in_guard.py
-./test_tensor_slice.py
-./test_trace_list_arg.py
-./test_unsupport_function.py
