GH-91049: Introduce set vectorcall field API for PyFunctionObject (GH-92257)

Co-authored-by: Andrew Frost <adfrost@fb.com>
Co-authored-by: Itamar Ostricher <itamarost@gmail.com>

diff --git a/Doc/c-api/function.rst b/Doc/c-api/function.rst
index 56c18396d3221d51bc540c11f48ff377d6bb5d55..df88e85e518829ebb63219347fd9a8ebfda2ee51 100644 (file)
--- a/Doc/c-api/function.rst
+++ b/Doc/c-api/function.rst
@@ -83,6 +83,15 @@ There are a few functions specific to Python functions.
    Raises :exc:`SystemError` and returns ``-1`` on failure.
 
 
+.. c:function:: void PyFunction_SetVectorcall(PyFunctionObject *func, vectorcallfunc vectorcall)
+
+   Set the vectorcall field of a given function object *func*.
+
+   Warning: extensions using this API must preserve the behavior
+   of the unaltered (default) vectorcall function!
+
+   .. versionadded:: 3.12
+
 .. c:function:: PyObject* PyFunction_GetClosure(PyObject *op)
 
    Return the closure associated with the function object *op*. This can be ``NULL``

diff --git a/Doc/whatsnew/3.12.rst b/Doc/whatsnew/3.12.rst
index 4a6b9c0ba5a604df62998fee84afa1a7d88efc27..90355a73f080049070789e26e985f24ae8782be7 100644 (file)
--- a/Doc/whatsnew/3.12.rst
+++ b/Doc/whatsnew/3.12.rst
@@ -493,6 +493,10 @@ New Features
   functions in all running threads in addition to the calling one. (Contributed
   by Pablo Galindo in :gh:`93503`.)
 
+* Added new function :c:func:`PyFunction_SetVectorcall` to the C API
+  which sets the vectorcall field of a given :c:type:`PyFunctionObject`.
+  (Contributed by Andrew Frost in :gh:`92257`.)
+
 Porting to Python 3.12
 ----------------------
 

diff --git a/Include/cpython/funcobject.h b/Include/cpython/funcobject.h
index 05f76656c4bf1d738a90768ff9668baf189bda75..dd8f20b2c20b39bbd42a3a67e4c14a4a6c37aeb8 100644 (file)
--- a/Include/cpython/funcobject.h
+++ b/Include/cpython/funcobject.h
@@ -48,7 +48,8 @@ typedef struct {
      *     defaults
      *     kwdefaults (only if the object changes, not the contents of the dict)
      *     code
-     *     annotations */
+     *     annotations
+     *     vectorcall function pointer */
     uint32_t func_version;
 
     /* Invariant:
@@ -69,6 +70,7 @@ PyAPI_FUNC(PyObject *) PyFunction_GetGlobals(PyObject *);
 PyAPI_FUNC(PyObject *) PyFunction_GetModule(PyObject *);
 PyAPI_FUNC(PyObject *) PyFunction_GetDefaults(PyObject *);
 PyAPI_FUNC(int) PyFunction_SetDefaults(PyObject *, PyObject *);
+PyAPI_FUNC(void) PyFunction_SetVectorcall(PyFunctionObject *, vectorcallfunc);
 PyAPI_FUNC(PyObject *) PyFunction_GetKwDefaults(PyObject *);
 PyAPI_FUNC(int) PyFunction_SetKwDefaults(PyObject *, PyObject *);
 PyAPI_FUNC(PyObject *) PyFunction_GetClosure(PyObject *);

diff --git a/Lib/test/test_call.py b/Lib/test/test_call.py
index c00de27b265d2734d403212089bd03dd15071cdd..c1a386228ff0d023978c6bd0771c16820d2d90a3 100644 (file)
--- a/Lib/test/test_call.py
+++ b/Lib/test/test_call.py
@@ -580,6 +580,9 @@ def testfunction_kw(self, *, kw):
     return self
 
 
+QUICKENING_WARMUP_DELAY = 8
+
+
 class TestPEP590(unittest.TestCase):
 
     def test_method_descriptor_flag(self):
@@ -760,6 +763,54 @@ class TestPEP590(unittest.TestCase):
                 self.assertEqual(expected, meth(*args1, **kwargs))
                 self.assertEqual(expected, wrapped(*args, **kwargs))
 
+    def test_setvectorcall(self):
+        from _testcapi import function_setvectorcall
+        def f(num): return num + 1
+        assert_equal = self.assertEqual
+        num = 10
+        assert_equal(11, f(num))
+        function_setvectorcall(f)
+        # make sure specializer is triggered by running > 50 times
+        for _ in range(10 * QUICKENING_WARMUP_DELAY):
+            assert_equal("overridden", f(num))
+
+    def test_setvectorcall_load_attr_specialization_skip(self):
+        from _testcapi import function_setvectorcall
+
+        class X:
+            def __getattribute__(self, attr):
+                return attr
+
+        assert_equal = self.assertEqual
+        x = X()
+        assert_equal("a", x.a)
+        function_setvectorcall(X.__getattribute__)
+        # make sure specialization doesn't trigger
+        # when vectorcall is overridden
+        for _ in range(QUICKENING_WARMUP_DELAY):
+            assert_equal("overridden", x.a)
+
+    def test_setvectorcall_load_attr_specialization_deopt(self):
+        from _testcapi import function_setvectorcall
+
+        class X:
+            def __getattribute__(self, attr):
+                return attr
+
+        def get_a(x):
+            return x.a
+
+        assert_equal = self.assertEqual
+        x = X()
+        # trigger LOAD_ATTR_GETATTRIBUTE_OVERRIDDEN specialization
+        for _ in range(QUICKENING_WARMUP_DELAY):
+            assert_equal("a", get_a(x))
+        function_setvectorcall(X.__getattribute__)
+        # make sure specialized LOAD_ATTR_GETATTRIBUTE_OVERRIDDEN
+        # gets deopted due to overridden vectorcall
+        for _ in range(QUICKENING_WARMUP_DELAY):
+            assert_equal("overridden", get_a(x))
+
     @requires_limited_api
     def test_vectorcall_limited(self):
         from _testcapi import pyobject_vectorcall

diff --git a/Misc/NEWS.d/next/C API/2022-05-03-19-35-37.gh-issue-92193.61VoFL.rst b/Misc/NEWS.d/next/C API/2022-05-03-19-35-37.gh-issue-92193.61VoFL.rst
new file mode 100644 (file)
index 0000000..e0755bb
--- /dev/null
+++ b/Misc/NEWS.d/next/C API/2022-05-03-19-35-37.gh-issue-92193.61VoFL.rst
@@ -0,0 +1,5 @@
+Add new function :c:func:`PyFunction_SetVectorcall` to the C API
+which sets the vectorcall field of a given :c:type:`PyFunctionObject`.
+
+Warning: extensions using this API must preserve the behavior
+of the unaltered function!

diff --git a/Modules/_testcapi/vectorcall.c b/Modules/_testcapi/vectorcall.c
index 626706eafab5e8ba9f44ef4878e124d2e66a3e5f..e9c863a75704583472d5ab4d49d16716f3d92256 100644 (file)
--- a/Modules/_testcapi/vectorcall.c
+++ b/Modules/_testcapi/vectorcall.c
@@ -102,6 +102,24 @@ test_pyobject_vectorcall(PyObject *self, PyObject *args)
     return PyObject_Vectorcall(func, stack, nargs, kwnames);
 }
 
+static PyObject *
+override_vectorcall(PyObject *callable, PyObject *const *args, size_t nargsf,
+                    PyObject *kwnames)
+{
+    return PyUnicode_FromString("overridden");
+}
+
+static PyObject *
+function_setvectorcall(PyObject *self, PyObject *func)
+{
+    if (!PyFunction_Check(func)) {
+        PyErr_SetString(PyExc_TypeError, "'func' must be a function");
+        return NULL;
+    }
+    PyFunction_SetVectorcall((PyFunctionObject *)func, (vectorcallfunc)override_vectorcall);
+    Py_RETURN_NONE;
+}
+
 static PyObject *
 test_pyvectorcall_call(PyObject *self, PyObject *args)
 {
@@ -244,6 +262,7 @@ static PyMethodDef TestMethods[] = {
     {"pyobject_fastcall", test_pyobject_fastcall, METH_VARARGS},
     {"pyobject_fastcalldict", test_pyobject_fastcalldict, METH_VARARGS},
     {"pyobject_vectorcall", test_pyobject_vectorcall, METH_VARARGS},
+    {"function_setvectorcall", function_setvectorcall, METH_O},
     {"pyvectorcall_call", test_pyvectorcall_call, METH_VARARGS},
     _TESTCAPI_MAKE_VECTORCALL_CLASS_METHODDEF
     _TESTCAPI_HAS_VECTORCALL_FLAG_METHODDEF

diff --git a/Objects/funcobject.c b/Objects/funcobject.c
index 32b4155c03e6ae01a06822386a8406a7cd613a55..7f257a9986987b00fd1768666cced47bb03afd4d 100644 (file)
--- a/Objects/funcobject.c
+++ b/Objects/funcobject.c
@@ -134,6 +134,9 @@ uint32_t _PyFunction_GetVersionForCurrentState(PyFunctionObject *func)
     if (func->func_version != 0) {
         return func->func_version;
     }
+    if (func->vectorcall != _PyFunction_Vectorcall) {
+        return 0;
+    }
     if (next_func_version == 0) {
         return 0;
     }
@@ -209,6 +212,14 @@ PyFunction_SetDefaults(PyObject *op, PyObject *defaults)
     return 0;
 }
 
+void
+PyFunction_SetVectorcall(PyFunctionObject *func, vectorcallfunc vectorcall)
+{
+    assert(func != NULL);
+    func->func_version = 0;
+    func->vectorcall = vectorcall;
+}
+
 PyObject *
 PyFunction_GetKwDefaults(PyObject *op)
 {

diff --git a/Python/ceval.c b/Python/ceval.c
index b61cc0852ed96a677f711c1575f266171b7a04d3..8891d6cfa07a153e1342e5d3e25c131b432a8e1c 100644 (file)
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -3126,8 +3126,11 @@ handle_eval_breaker:
             PyObject *getattribute = read_obj(cache->descr);
             assert(Py_IS_TYPE(getattribute, &PyFunction_Type));
             PyFunctionObject *f = (PyFunctionObject *)getattribute;
+            uint32_t func_version = read_u32(cache->keys_version);
+            assert(func_version != 0);
+            DEOPT_IF(f->func_version != func_version, LOAD_ATTR);
             PyCodeObject *code = (PyCodeObject *)f->func_code;
-            DEOPT_IF(((PyCodeObject *)f->func_code)->co_argcount != 2, LOAD_ATTR);
+            assert(code->co_argcount == 2);
             DEOPT_IF(!_PyThreadState_HasStackSpace(tstate, code->co_framesize), CALL);
             STAT_INC(LOAD_ATTR, hit);
 
@@ -4133,7 +4136,10 @@ handle_eval_breaker:
             function = PEEK(total_args + 1);
             int positional_args = total_args - KWNAMES_LEN();
             // Check if the call can be inlined or not
-            if (Py_TYPE(function) == &PyFunction_Type && tstate->interp->eval_frame == NULL) {
+            if (Py_TYPE(function) == &PyFunction_Type &&
+                tstate->interp->eval_frame == NULL &&
+                ((PyFunctionObject *)function)->vectorcall == _PyFunction_Vectorcall)
+            {
                 int code_flags = ((PyCodeObject*)PyFunction_GET_CODE(function))->co_flags;
                 PyObject *locals = code_flags & CO_OPTIMIZED ? NULL : Py_NewRef(PyFunction_GET_GLOBALS(function));
                 STACK_SHRINK(total_args);

diff --git a/Python/specialize.c b/Python/specialize.c
index 93f1d289b3acb7c5c562dc8077cf8a23e94fd6a3..b7c321e4878b981435716a9deff73bc263f56f6a 100644 (file)
--- a/Python/specialize.c
+++ b/Python/specialize.c
@@ -837,6 +837,11 @@ _Py_Specialize_LoadAttr(PyObject *owner, _Py_CODEUNIT *instr, PyObject *name)
             if (!function_check_args(descr, 2, LOAD_ATTR)) {
                 goto fail;
             }
+            uint32_t version = function_get_version(descr, LOAD_ATTR);
+            if (version == 0) {
+                goto fail;
+            }
+            write_u32(lm_cache->keys_version, version);
             /* borrowed */
             write_obj(lm_cache->descr, descr);
             write_u32(lm_cache->type_version, type->tp_version_tag);