Allow keeping specialization enabled when specifying eval frame function

Author: DinoV

Doc/c-api/subinterpreters.rst

@@ -391,14 +391,31 @@ High-level APIs
    .. versionadded:: 3.9
 
 
-.. c:function:: void _PyInterpreterState_SetEvalFrameFunc(PyInterpreterState *interp, _PyFrameEvalFunction eval_frame)
+.. c:function:: void _PyInterpreterState_SetEvalFrameFunc(PyInterpreterState *interp, _PyFrameEvalFunction eval_frame, int allow_specialization)
 
    Set the frame evaluation function.
 
+   If *allow_specialization* is non-zero, the adaptive specializer will
+   continue to specialize bytecodes even though a custom eval frame function
+   is set. When *allow_specialization* is zero, setting a custom eval frame
+   disables specialization.
+
    See the :pep:`523` "Adding a frame evaluation API to CPython".
 
    .. versionadded:: 3.9
 
+   .. versionchanged:: 3.15
+      Added the *allow_specialization* parameter.
+
+
+.. c:function:: int _PyInterpreterState_IsSpecializationEnabled(PyInterpreterState *interp)
+
+   Return non-zero if adaptive specialization is enabled for the interpreter.
+   Specialization is enabled when no custom eval frame function is set, or
+   when one is set with *allow_specialization* enabled.
+
+   .. versionadded:: 3.15
+
 
 Low-level APIs
 --------------

Include/cpython/pystate.h

@@ -318,4 +318,7 @@ PyAPI_FUNC(_PyFrameEvalFunction) _PyInterpreterState_GetEvalFrameFunc(
     PyInterpreterState *interp);
 PyAPI_FUNC(void) _PyInterpreterState_SetEvalFrameFunc(
     PyInterpreterState *interp,
-    _PyFrameEvalFunction eval_frame);
+    _PyFrameEvalFunction eval_frame,
+    int allow_specialization);
+PyAPI_FUNC(int) _PyInterpreterState_IsSpecializationEnabled(
+    PyInterpreterState *interp);

Include/internal/pycore_interp_structs.h

@@ -899,6 +899,7 @@ struct _is {
     PyObject *builtins_copy;
     // Initialized to _PyEval_EvalFrameDefault().
     _PyFrameEvalFunction eval_frame;
+    int eval_frame_allow_specialization;
 
     PyFunction_WatchCallback func_watchers[FUNC_MAX_WATCHERS];
     // One bit is set for each non-NULL entry in func_watchers

Lib/test/test_capi/test_misc.py

@@ -2870,6 +2870,94 @@ def func():
         self.do_test(func, names)
 
 
+class Test_Pep523AllowSpecialization(unittest.TestCase):
+    """Tests for _PyInterpreterState_SetEvalFrameFunc with
+    allow_specialization=1."""
+
+    def test_is_specialization_enabled_default(self):
+        # With no custom eval frame, specialization should be enabled
+        self.assertTrue(_testinternalcapi.is_specialization_enabled())
+
+    def test_is_specialization_enabled_with_eval_frame(self):
+        # Setting eval frame with allow_specialization=0 disables specialization
+        try:
+            _testinternalcapi.set_eval_frame_record([])
+            self.assertFalse(_testinternalcapi.is_specialization_enabled())
+        finally:
+            _testinternalcapi.set_eval_frame_default()
+
+    def test_is_specialization_enabled_after_restore(self):
+        # Restoring the default eval frame re-enables specialization
+        try:
+            _testinternalcapi.set_eval_frame_record([])
+            self.assertFalse(_testinternalcapi.is_specialization_enabled())
+        finally:
+            _testinternalcapi.set_eval_frame_default()
+        self.assertTrue(_testinternalcapi.is_specialization_enabled())
+
+    def test_is_specialization_enabled_with_allow(self):
+        # Setting eval frame with allow_specialization=1 keeps it enabled
+        try:
+            _testinternalcapi.set_eval_frame_record_with_specialization([])
+            self.assertTrue(_testinternalcapi.is_specialization_enabled())
+        finally:
+            _testinternalcapi.set_eval_frame_default()
+
+    def test_allow_specialization_call(self):
+        # With allow_specialization, after enough iterations the CALL
+        # specialization should kick in and inline calls, so the eval frame
+        # function is NOT called for the inner function.
+        def inner(x=42):
+            pass
+        def func():
+            inner()
+
+        actual_calls = []
+        try:
+            _testinternalcapi.set_eval_frame_record_with_specialization(
+                actual_calls)
+            for i in range(SUFFICIENT_TO_DEOPT_AND_SPECIALIZE * 2):
+                func()
+        finally:
+            _testinternalcapi.set_eval_frame_default()
+
+        # With specialization enabled, calls to inner() can be inlined
+        # so "inner" should NOT appear in every call record.
+        # At minimum, "func" should appear for every call.
+        func_calls = [c for c in actual_calls if c == "func"]
+        self.assertEqual(len(func_calls),
+                         SUFFICIENT_TO_DEOPT_AND_SPECIALIZE * 2)
+        inner_calls = [c for c in actual_calls if c == "inner"]
+        # inner should appear fewer times than func because specialization
+        # inlines the call
+        self.assertLess(len(inner_calls), len(func_calls))
+
+    def test_no_specialization_call(self):
+        # Without allow_specialization, ALL calls go through the eval frame.
+        # This is the existing PEP 523 behavior.
+        def inner(x=42):
+            pass
+        def func():
+            inner()
+
+        # Pre-specialize
+        for _ in range(SUFFICIENT_TO_DEOPT_AND_SPECIALIZE):
+            func()
+
+        actual_calls = []
+        try:
+            _testinternalcapi.set_eval_frame_record(actual_calls)
+            for _ in range(SUFFICIENT_TO_DEOPT_AND_SPECIALIZE):
+                func()
+        finally:
+            _testinternalcapi.set_eval_frame_default()
+
+        # Without allow_specialization, every call including inner() goes
+        # through the eval frame
+        expected = ["func", "inner"] * SUFFICIENT_TO_DEOPT_AND_SPECIALIZE
+        self.assertEqual(actual_calls, expected)
+
+
 @unittest.skipUnless(support.Py_GIL_DISABLED, 'need Py_GIL_DISABLED')
 class TestPyThreadId(unittest.TestCase):
     def test_py_thread_id(self):

Modules/_testinternalcapi.c

@@ -929,7 +929,7 @@ static PyObject *
 set_eval_frame_default(PyObject *self, PyObject *Py_UNUSED(args))
 {
     module_state *state = get_module_state(self);
-    _PyInterpreterState_SetEvalFrameFunc(_PyInterpreterState_GET(), _PyEval_EvalFrameDefault);
+    _PyInterpreterState_SetEvalFrameFunc(_PyInterpreterState_GET(), _PyEval_EvalFrameDefault, 0);
     Py_CLEAR(state->record_list);
     Py_RETURN_NONE;
 }
@@ -961,7 +961,7 @@ set_eval_frame_record(PyObject *self, PyObject *list)
         return NULL;
     }
     Py_XSETREF(state->record_list, Py_NewRef(list));
-    _PyInterpreterState_SetEvalFrameFunc(_PyInterpreterState_GET(), record_eval);
+    _PyInterpreterState_SetEvalFrameFunc(_PyInterpreterState_GET(), record_eval, 0);
     Py_RETURN_NONE;
 }
 
@@ -998,10 +998,30 @@ get_eval_frame_stats(PyObject *self, PyObject *Py_UNUSED(args))
 static PyObject *
 set_eval_frame_interp(PyObject *self, PyObject *Py_UNUSED(args))
 {
-    _PyInterpreterState_SetEvalFrameFunc(_PyInterpreterState_GET(), Test_EvalFrame);
+    _PyInterpreterState_SetEvalFrameFunc(_PyInterpreterState_GET(), Test_EvalFrame, 0);
     Py_RETURN_NONE;
 }
 
+static PyObject *
+set_eval_frame_record_with_specialization(PyObject *self, PyObject *list)
+{
+    module_state *state = get_module_state(self);
+    if (!PyList_Check(list)) {
+        PyErr_SetString(PyExc_TypeError, "argument must be a list");
+        return NULL;
+    }
+    Py_XSETREF(state->record_list, Py_NewRef(list));
+    _PyInterpreterState_SetEvalFrameFunc(_PyInterpreterState_GET(), record_eval, 1);
+    Py_RETURN_NONE;
+}
+
+static PyObject *
+is_specialization_enabled(PyObject *self, PyObject *Py_UNUSED(args))
+{
+    return PyBool_FromLong(
+        _PyInterpreterState_IsSpecializationEnabled(_PyInterpreterState_GET()));
+}
+
 /*[clinic input]
 
 _testinternalcapi.compiler_cleandoc -> object
@@ -2863,6 +2883,9 @@ static PyMethodDef module_functions[] = {
     {"set_eval_frame_default", set_eval_frame_default, METH_NOARGS, NULL},
     {"set_eval_frame_interp", set_eval_frame_interp, METH_NOARGS, NULL},
     {"set_eval_frame_record", set_eval_frame_record, METH_O, NULL},
+    {"set_eval_frame_record_with_specialization",
+     set_eval_frame_record_with_specialization, METH_O, NULL},
+    {"is_specialization_enabled", is_specialization_enabled, METH_NOARGS, NULL},
     _TESTINTERNALCAPI_COMPILER_CLEANDOC_METHODDEF
     _TESTINTERNALCAPI_NEW_INSTRUCTION_SEQUENCE_METHODDEF
     _TESTINTERNALCAPI_COMPILER_CODEGEN_METHODDEF

Python/perf_trampoline.c

@@ -530,12 +530,12 @@ _PyPerfTrampoline_Init(int activate)
         code_watcher_id = -1;
     }
     if (!activate) {
-        _PyInterpreterState_SetEvalFrameFunc(tstate->interp, prev_eval_frame);
+        _PyInterpreterState_SetEvalFrameFunc(tstate->interp, prev_eval_frame, 0);
         perf_status = PERF_STATUS_NO_INIT;
     }
     else if (tstate->interp->eval_frame != py_trampoline_evaluator) {
         prev_eval_frame = _PyInterpreterState_GetEvalFrameFunc(tstate->interp);
-        _PyInterpreterState_SetEvalFrameFunc(tstate->interp, py_trampoline_evaluator);
+        _PyInterpreterState_SetEvalFrameFunc(tstate->interp, py_trampoline_evaluator, 0);
         extra_code_index = _PyEval_RequestCodeExtraIndex(NULL);
         if (extra_code_index == -1) {
             return -1;
@@ -568,7 +568,7 @@ _PyPerfTrampoline_Fini(void)
     }
     PyThreadState *tstate = _PyThreadState_GET();
     if (tstate->interp->eval_frame == py_trampoline_evaluator) {
-        _PyInterpreterState_SetEvalFrameFunc(tstate->interp, NULL);
+        _PyInterpreterState_SetEvalFrameFunc(tstate->interp, NULL, 0);
     }
     if (perf_status == PERF_STATUS_OK) {
         trampoline_api.free_state(trampoline_api.state);

Python/pystate.c

@@ -3010,12 +3010,14 @@ _PyInterpreterState_GetEvalFrameFunc(PyInterpreterState *interp)
 
 void
 _PyInterpreterState_SetEvalFrameFunc(PyInterpreterState *interp,
-                                     _PyFrameEvalFunction eval_frame)
+                                     _PyFrameEvalFunction eval_frame,
+                                     int allow_specialization)
 {
     if (eval_frame == _PyEval_EvalFrameDefault) {
         eval_frame = NULL;
     }
     if (eval_frame == interp->eval_frame) {
+        interp->eval_frame_allow_specialization = allow_specialization;
         return;
     }
 #ifdef _Py_TIER2
@@ -3026,9 +3028,17 @@ _PyInterpreterState_SetEvalFrameFunc(PyInterpreterState *interp,
     RARE_EVENT_INC(set_eval_frame_func);
     _PyEval_StopTheWorld(interp);
     interp->eval_frame = eval_frame;
+    interp->eval_frame_allow_specialization = allow_specialization;
     _PyEval_StartTheWorld(interp);
 }
 
+int
+_PyInterpreterState_IsSpecializationEnabled(PyInterpreterState *interp)
+{
+    return interp->eval_frame == NULL
+        || interp->eval_frame_allow_specialization;
+}
+
 
 const PyConfig*
 _PyInterpreterState_GetConfig(PyInterpreterState *interp)

Python/specialize.c

@@ -811,7 +811,7 @@ do_specialize_instance_load_attr(PyObject* owner, _Py_CODEUNIT* instr, PyObject*
                 return -1;
             }
             /* Don't specialize if PEP 523 is active */
-            if (_PyInterpreterState_GET()->eval_frame) {
+            if (!_PyInterpreterState_IsSpecializationEnabled(_PyInterpreterState_GET())) {
                 SPECIALIZATION_FAIL(LOAD_ATTR, SPEC_FAIL_OTHER);
                 return -1;
             }
@@ -890,7 +890,7 @@ do_specialize_instance_load_attr(PyObject* owner, _Py_CODEUNIT* instr, PyObject*
                 return -1;
             }
             /* Don't specialize if PEP 523 is active */
-            if (_PyInterpreterState_GET()->eval_frame) {
+            if (!_PyInterpreterState_IsSpecializationEnabled(_PyInterpreterState_GET())) {
                 SPECIALIZATION_FAIL(LOAD_ATTR, SPEC_FAIL_OTHER);
                 return -1;
             }
@@ -1697,7 +1697,7 @@ specialize_py_call(PyFunctionObject *func, _Py_CODEUNIT *instr, int nargs,
     PyCodeObject *code = (PyCodeObject *)func->func_code;
     int kind = function_kind(code);
     /* Don't specialize if PEP 523 is active */
-    if (_PyInterpreterState_GET()->eval_frame) {
+    if (!_PyInterpreterState_IsSpecializationEnabled(_PyInterpreterState_GET())) {
         SPECIALIZATION_FAIL(CALL, SPEC_FAIL_CALL_PEP_523);
         return -1;
     }
@@ -1740,7 +1740,7 @@ specialize_py_call_kw(PyFunctionObject *func, _Py_CODEUNIT *instr, int nargs,
     PyCodeObject *code = (PyCodeObject *)func->func_code;
     int kind = function_kind(code);
     /* Don't specialize if PEP 523 is active */
-    if (_PyInterpreterState_GET()->eval_frame) {
+    if (!_PyInterpreterState_IsSpecializationEnabled(_PyInterpreterState_GET())) {
         SPECIALIZATION_FAIL(CALL, SPEC_FAIL_CALL_PEP_523);
         return -1;
     }
@@ -2003,7 +2003,7 @@ binary_op_fail_kind(int oparg, PyObject *lhs, PyObject *rhs)
                     return SPEC_FAIL_WRONG_NUMBER_ARGUMENTS;
                 }
 
-                if (_PyInterpreterState_GET()->eval_frame) {
+                if (!_PyInterpreterState_IsSpecializationEnabled(_PyInterpreterState_GET())) {
                     /* Don't specialize if PEP 523 is active */
                     Py_DECREF(descriptor);
                     return SPEC_FAIL_OTHER;
@@ -2312,7 +2312,7 @@ _Py_Specialize_BinaryOp(_PyStackRef lhs_st, _PyStackRef rhs_st, _Py_CODEUNIT *in
                 PyHeapTypeObject *ht = (PyHeapTypeObject *)container_type;
                 if (kind == SIMPLE_FUNCTION &&
                     fcode->co_argcount == 2 &&
-                    !_PyInterpreterState_GET()->eval_frame && /* Don't specialize if PEP 523 is active */
+                    _PyInterpreterState_IsSpecializationEnabled(_PyInterpreterState_GET()) && /* Don't specialize if PEP 523 is active */
                     _PyType_CacheGetItemForSpecialization(ht, descriptor, (uint32_t)tp_version))
                 {
                     specialize(instr, BINARY_OP_SUBSCR_GETITEM);
@@ -2570,7 +2570,7 @@ _Py_Specialize_ForIter(_PyStackRef iter, _PyStackRef null_or_index, _Py_CODEUNIT
                 instr[oparg + INLINE_CACHE_ENTRIES_FOR_ITER + 1].op.code == INSTRUMENTED_END_FOR
             );
             /* Don't specialize if PEP 523 is active */
-            if (_PyInterpreterState_GET()->eval_frame) {
+            if (!_PyInterpreterState_IsSpecializationEnabled(_PyInterpreterState_GET())) {
                 goto failure;
             }
             specialize(instr, FOR_ITER_GEN);
@@ -2609,7 +2609,7 @@ _Py_Specialize_Send(_PyStackRef receiver_st, _Py_CODEUNIT *instr)
     PyTypeObject *tp = Py_TYPE(receiver);
     if (tp == &PyGen_Type || tp == &PyCoro_Type) {
         /* Don't specialize if PEP 523 is active */
-        if (_PyInterpreterState_GET()->eval_frame) {
+        if (!_PyInterpreterState_IsSpecializationEnabled(_PyInterpreterState_GET())) {
             SPECIALIZATION_FAIL(SEND, SPEC_FAIL_OTHER);
             goto failure;
         }
@@ -2632,7 +2632,7 @@ _Py_Specialize_CallFunctionEx(_PyStackRef func_st, _Py_CODEUNIT *instr)
 
     if (Py_TYPE(func) == &PyFunction_Type &&
         ((PyFunctionObject *)func)->vectorcall == _PyFunction_Vectorcall) {
-        if (_PyInterpreterState_GET()->eval_frame) {
+        if (!_PyInterpreterState_IsSpecializationEnabled(_PyInterpreterState_GET())) {
             goto failure;
         }
         specialize(instr, CALL_EX_PY);