bpo-46939: Specialize calls to Python classes

Include/cpython/object.h

@@ -234,6 +234,7 @@ struct _typeobject {
  * by code other than the specializer and interpreter. */
 struct _specialization_cache {
     PyObject *getitem;
+    PyObject *init;
 };
 
 /* The *real* layout of a type object when allocated on the heap */

Include/internal/pycore_frame.h

@@ -61,6 +61,7 @@ typedef struct _PyInterpreterFrame {
     PyFrameState f_state;  /* What state the frame is in */
     bool is_entry;  // Whether this is the "root" frame for the current _PyCFrame.
     bool is_generator;
+    PyObject *self; /* Strong reference consumed by __init__ frames to return self in RETURN_VALUE */
     PyObject *localsplus[1];
 } _PyInterpreterFrame;
 
@@ -118,6 +119,7 @@ _PyFrame_InitializeSpecials(
     frame->f_state = FRAME_CREATED;
     frame->is_entry = false;
     frame->is_generator = false;
+    frame->self = NULL;
 }
 
 /* Gets the pointer to the locals array

Lib/opcode.py

@@ -285,6 +285,7 @@ def jabs_op(name, op, entries=0):
     "PRECALL_NO_KW_METHOD_DESCRIPTOR_FAST",
     "PRECALL_BOUND_METHOD",
     "PRECALL_PYFUNC",
+    "PRECALL_PY_CLASS",
     "RESUME_QUICK",
     "STORE_ATTR_ADAPTIVE",
     "STORE_ATTR_INSTANCE_VALUE",

Misc/NEWS.d/next/Core and Builtins/2022-03-06-21-50-36.bpo-46939.nlQEGG.rst

@@ -0,0 +1,2 @@
+Calls to Python classes are now specialized. Creating objects from Python
+classes should now be faster. Patch by Ken Jin.

Python/ceval.c

@@ -1587,6 +1587,11 @@ pop_frame(PyThreadState *tstate, _PyInterpreterFrame *frame)
  */
 typedef struct {
     PyObject *kwnames;
+    /* __init__ is special because while it returns None, we need to return self
+      This tells CALL to pass the current self to the new frame (the __init__ frame).
+      Where it is eventually consumed by RETURN_VALUE.
+    */
+    bool init_pass_self;
 } CallShape;
 
 static inline bool
@@ -1618,6 +1623,7 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
     _PyCFrame cframe;
     CallShape call_shape;
     call_shape.kwnames = NULL; // Borrowed reference. Reset by CALL instructions.
+    call_shape.init_pass_self = 0;
 
     /* WARNING: Because the _PyCFrame lives on the C stack,
      * but can be accessed from a heap allocated object (tstate)
@@ -2386,6 +2392,18 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
 
         TARGET(RETURN_VALUE) {
             PyObject *retval = POP();
+            if (frame->self != NULL) {
+                if (Py_IsNone(retval)) {
+                    Py_SETREF(retval, frame->self);
+                    frame->self = NULL;
+                }
+                /* We need this to continue raising errors when bad-practice
+                   __init__s return their non-None values. This is later
+                   caught by the interpreter. */
+                else {
+                    Py_CLEAR(frame->self);
+                }
+            }
             assert(EMPTY());
             frame->f_state = FRAME_RETURNED;
             _PyFrame_SetStackPointer(frame, stack_pointer);
@@ -4617,6 +4635,45 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
             DISPATCH();
         }
 
+        TARGET(PRECALL_PY_CLASS) {
+            SpecializedCacheEntry *cache = GET_CACHE();
+            _PyAdaptiveEntry *cache0 = &cache[0].adaptive;
+            _PyCallCache *cache1 = &cache[-1].call;
+            int original_oparg = cache->adaptive.original_oparg;
+            int is_method = (PEEK(original_oparg + 2) != NULL);
+            DEOPT_IF(is_method, PRECALL);
+            PyObject *cls = PEEK(original_oparg + 1);
+            DEOPT_IF(!PyType_Check(cls), PRECALL);
+            PyTypeObject *cls_t = (PyTypeObject *)cls;
+            DEOPT_IF(cls_t->tp_version_tag != cache0->version, PRECALL);
+            assert(cls_t->tp_flags & Py_TPFLAGS_HEAPTYPE);
+            PyObject *init = ((PyHeapTypeObject *)cls_t)->_spec_cache.init;
+            assert(PyFunction_Check(init));
+            DEOPT_IF(((PyFunctionObject *)(init))->func_version != cache1->func_version, PRECALL);
+            DEOPT_IF(cls_t->tp_new != PyBaseObject_Type.tp_new, PRECALL);
+            STAT_INC(PRECALL, hit);
+
+            PyObject *args = _PyTuple_FromArray(&PEEK(original_oparg), original_oparg);
+            if (args == NULL) {
+                goto error;
+            }
+            PyObject *self = cls_t->tp_new(cls_t, args, call_shape.kwnames);
+            Py_DECREF(args);
+            if (self == NULL) {
+                goto error;
+            }
+            Py_INCREF(init);
+            PEEK(original_oparg+1) = self;
+            PEEK(original_oparg+2) = init;
+            Py_DECREF(cls);
+
+            /* For use in RETURN_VALUE later */
+            Py_INCREF(self);
+            assert(call_shape.init_pass_self == false);
+            call_shape.init_pass_self = true;
+            DISPATCH();
+        }
+
         TARGET(KW_NAMES) {
             assert(call_shape.kwnames == NULL);
             assert(oparg < PyTuple_GET_SIZE(consts));
@@ -4651,6 +4708,8 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
                 _PyFrame_SetStackPointer(frame, stack_pointer);
                 new_frame->previous = frame;
                 cframe.current_frame = frame = new_frame;
+                frame->self = call_shape.init_pass_self ? frame->localsplus[0] : NULL;
+                call_shape.init_pass_self = false;
                 CALL_STAT_INC(inlined_py_calls);
                 goto start_frame;
             }
@@ -4762,6 +4821,8 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
             _PyFrame_SetStackPointer(frame, stack_pointer);
             new_frame->previous = frame;
             frame = cframe.current_frame = new_frame;
+            frame->self = call_shape.init_pass_self ? frame->localsplus[0] : NULL;
+            call_shape.init_pass_self = false;
             goto start_frame;
         }
 
@@ -4803,6 +4864,8 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
             _PyFrame_SetStackPointer(frame, stack_pointer);
             new_frame->previous = frame;
             frame = cframe.current_frame = new_frame;
+            frame->self = call_shape.init_pass_self ? frame->localsplus[0] : NULL;
+            call_shape.init_pass_self = false;
             goto start_frame;
         }
 
@@ -5247,6 +5310,10 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
             assert(PyTuple_CheckExact(callargs));
 
             result = do_call_core(tstate, func, callargs, kwargs, cframe.use_tracing);
+            if (call_shape.init_pass_self) {
+                Py_SETREF(result, PyTuple_GET_ITEM(callargs, 0));
+                call_shape.init_pass_self = false;
+            }
             Py_DECREF(func);
             Py_DECREF(callargs);
             Py_XDECREF(kwargs);

Python/specialize.c

@@ -591,15 +591,16 @@ initial_counter_value(void) {
 #define SPEC_FAIL_CALL_BAD_CALL_FLAGS 17
 #define SPEC_FAIL_CALL_CLASS 18
 #define SPEC_FAIL_CALL_PYTHON_CLASS 19
-#define SPEC_FAIL_CALL_METHOD_DESCRIPTOR 20
-#define SPEC_FAIL_CALL_BOUND_METHOD 21
-#define SPEC_FAIL_CALL_STR 22
-#define SPEC_FAIL_CALL_CLASS_NO_VECTORCALL 23
-#define SPEC_FAIL_CALL_CLASS_MUTABLE 24
-#define SPEC_FAIL_CALL_KWNAMES 25
-#define SPEC_FAIL_CALL_METHOD_WRAPPER 26
-#define SPEC_FAIL_CALL_OPERATOR_WRAPPER 27
-#define SPEC_FAIL_CALL_PYFUNCTION 28
+#define SPEC_FAIL_CALL_PYTHON_CLASS_NON_PY_INIT 20
+#define SPEC_FAIL_CALL_METHOD_DESCRIPTOR 21
+#define SPEC_FAIL_CALL_BOUND_METHOD 22
+#define SPEC_FAIL_CALL_STR 23
+#define SPEC_FAIL_CALL_CLASS_NO_VECTORCALL 24
+#define SPEC_FAIL_CALL_CLASS_MUTABLE 25
+#define SPEC_FAIL_CALL_KWNAMES 26
+#define SPEC_FAIL_CALL_METHOD_WRAPPER 27
+#define SPEC_FAIL_CALL_OPERATOR_WRAPPER 28
+#define SPEC_FAIL_CALL_PYFUNCTION 29
 
 /* COMPARE_OP */
 #define SPEC_FAIL_COMPARE_OP_DIFFERENT_TYPES 12
@@ -1523,7 +1524,33 @@ specialize_class_call(
     assert(_Py_OPCODE(*instr) == PRECALL_ADAPTIVE);
     PyTypeObject *tp = _PyType_CAST(callable);
     if (tp->tp_new == PyBaseObject_Type.tp_new) {
-        SPECIALIZATION_FAIL(PRECALL, SPEC_FAIL_CALL_PYTHON_CLASS);
+        _PyAdaptiveEntry *cache0 = &cache[0].adaptive;
+        _PyCallCache *cache1 = &cache[-1].call;
+        PyObject *descriptor = _PyType_Lookup(tp, &_Py_ID(__init__));
+        if (descriptor && Py_TYPE(descriptor) == &PyFunction_Type) {
+            if (!(tp->tp_flags & Py_TPFLAGS_HEAPTYPE)) {
+                return -1;
+            }
+            PyFunctionObject *func = (PyFunctionObject *)descriptor;
+            PyCodeObject *fcode = (PyCodeObject *)func->func_code;
+            int kind = function_kind(fcode);
+            if (kind != SIMPLE_FUNCTION) {
+                SPECIALIZATION_FAIL(PRECALL, kind);
+                return -1;
+            }
+            assert(tp->tp_version_tag != 0);
+            cache0->version = tp->tp_version_tag;
+            int version = _PyFunction_GetVersionForCurrentState(func);
+            if (version == 0 || version != (uint16_t)version) {
+                SPECIALIZATION_FAIL(PRECALL, SPEC_FAIL_OUT_OF_VERSIONS);
+                return -1;
+            }
+            cache1->func_version = version;
+            ((PyHeapTypeObject *)tp)->_spec_cache.init = descriptor;
+            *instr = _Py_MAKECODEUNIT(PRECALL_PY_CLASS, _Py_OPARG(*instr));
+            return 0;
+        }
+        SPECIALIZATION_FAIL(PRECALL, SPEC_FAIL_CALL_PYTHON_CLASS_NON_PY_INIT);
         return -1;
     }
     if (tp->tp_flags & Py_TPFLAGS_IMMUTABLETYPE) {