gh-146261: JIT: protect against function version changes (#146300)

Co-authored-by: Kumar Aditya <kumaraditya@python.org>

Author: Fidget-Spinner

Include/internal/pycore_optimizer.h

@@ -424,8 +424,6 @@ extern PyCodeObject *_Py_uop_sym_get_probable_func_code(JitOptRef sym);
 extern PyObject *_Py_uop_sym_get_probable_value(JitOptRef sym);
 extern PyTypeObject *_Py_uop_sym_get_probable_type(JitOptRef sym);
 extern JitOptRef *_Py_uop_sym_set_stack_depth(JitOptContext *ctx, int stack_depth, JitOptRef *current_sp);
-extern uint32_t _Py_uop_sym_get_func_version(JitOptRef ref);
-bool _Py_uop_sym_set_func_version(JitOptContext *ctx, JitOptRef ref, uint32_t version);
 
 extern void _Py_uop_abstractcontext_init(JitOptContext *ctx, _PyBloomFilter *dependencies);
 extern void _Py_uop_abstractcontext_fini(JitOptContext *ctx);

Include/internal/pycore_optimizer_types.h

@@ -36,7 +36,6 @@ typedef enum _JitSymType {
     JIT_SYM_NON_NULL_TAG = 3,
     JIT_SYM_BOTTOM_TAG = 4,
     JIT_SYM_TYPE_VERSION_TAG = 5,
-    JIT_SYM_FUNC_VERSION_TAG = 6,
     JIT_SYM_KNOWN_CLASS_TAG = 7,
     JIT_SYM_KNOWN_VALUE_TAG = 8,
     JIT_SYM_TUPLE_TAG = 9,

Lib/test/test_capi/test_opt.py

@@ -1622,9 +1622,8 @@ def testfunc(n):
         self.assertEqual(uops.count("_PUSH_FRAME"), 2)
         # Type version propagation: one guard covers both method lookups
         self.assertEqual(uops.count("_GUARD_TYPE_VERSION"), 1)
-        # Function checks eliminated (type info resolves the callable)
-        self.assertNotIn("_CHECK_FUNCTION_VERSION", uops)
-        self.assertNotIn("_CHECK_FUNCTION_EXACT_ARGS", uops)
+        # Function checks cannot be eliminated for safety reasons.
+        self.assertIn("_CHECK_FUNCTION_VERSION", uops)
 
     def test_method_chain_guard_elimination(self):
         """
@@ -1669,10 +1668,7 @@ def testfunc(n):
         self.assertIsNotNone(ex)
         uops = get_opnames(ex)
         self.assertIn("_PUSH_FRAME", uops)
-        # Both should be not present, as this is a call
-        # to a simple function with a known function version.
-        self.assertNotIn("_CHECK_FUNCTION_VERSION_INLINE", uops)
-        self.assertNotIn("_CHECK_FUNCTION_VERSION", uops)
+        self.assertIn("_CHECK_FUNCTION_VERSION", uops)
         # Removed guard
         self.assertNotIn("_CHECK_FUNCTION_EXACT_ARGS", uops)
 
@@ -5178,6 +5174,27 @@ def g():
         PYTHON_JIT="1", PYTHON_JIT_STRESS="1")
         self.assertEqual(result[0].rc, 0, result)
 
+    def test_func_version_guarded_on_change(self):
+        def testfunc(n):
+            for i in range(n):
+                # Only works on functions promoted to constants
+                global_identity_code_will_be_modified(i)
+
+        testfunc(TIER2_THRESHOLD)
+
+        ex = get_first_executor(testfunc)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+        self.assertIn("_PUSH_FRAME", uops)
+        self.assertIn("_CHECK_FUNCTION_VERSION", uops)
+
+        global_identity_code_will_be_modified.__code__ = (lambda a: 0xdeadead).__code__
+        _testinternalcapi.clear_executor_deletion_list()
+        ex = get_first_executor(testfunc)
+        self.assertIsNone(ex)
+        # JItted code should've deopted.
+        self.assertEqual(global_identity_code_will_be_modified(None), 0xdeadead)
+
     def test_call_super(self):
         class A:
             def method1(self):
@@ -5224,6 +5241,9 @@ def testfunc(n):
 def global_identity(x):
     return x
 
+def global_identity_code_will_be_modified(x):
+    return x
+
 class TestObject:
     def test(self, *args, **kwargs):
         return args[0]

Objects/funcobject.c

@@ -8,6 +8,7 @@
 #include "pycore_modsupport.h"    // _PyArg_NoKeywords()
 #include "pycore_object.h"        // _PyObject_GC_UNTRACK()
 #include "pycore_object_deferred.h" // _PyObject_SetDeferredRefcount()
+#include "pycore_optimizer.h"       // _Py_Executors_InvalidateDependency()
 #include "pycore_pyerrors.h"      // _PyErr_Occurred()
 #include "pycore_setobject.h"     // _PySet_NextEntry()
 #include "pycore_stats.h"
@@ -63,6 +64,13 @@ handle_func_event(PyFunction_WatchEvent event, PyFunctionObject *func,
         case PyFunction_EVENT_MODIFY_DEFAULTS:
         case PyFunction_EVENT_MODIFY_KWDEFAULTS:
         case PyFunction_EVENT_MODIFY_QUALNAME:
+#if _Py_TIER2
+            // Note: we only invalidate JIT code if a function version changes.
+            // Not when the function is deallocated.
+            // Function deallocation occurs frequently (think: lambdas),
+            // so we want to minimize dependency invalidation there.
+            _Py_Executors_InvalidateDependency(interp, func, 1);
+#endif
             RARE_EVENT_INTERP_INC(interp, func_modification);
             break;
         default:

Python/optimizer_analysis.c

@@ -283,8 +283,6 @@ add_op(JitOptContext *ctx, _PyUOpInstruction *this_instr,
 #define sym_get_probable_func_code _Py_uop_sym_get_probable_func_code
 #define sym_get_probable_value _Py_uop_sym_get_probable_value
 #define sym_set_stack_depth(DEPTH, SP) _Py_uop_sym_set_stack_depth(ctx, DEPTH, SP)
-#define sym_get_func_version _Py_uop_sym_get_func_version
-#define sym_set_func_version _Py_uop_sym_set_func_version
 
 /* Comparison oparg masks */
 #define COMPARE_LT_MASK 2

Python/optimizer_bytecodes.c

@@ -1001,12 +1001,15 @@ dummy_func(void) {
     }
 
     op(_CHECK_FUNCTION_VERSION, (func_version/2, callable, self_or_null, unused[oparg] -- callable, self_or_null, unused[oparg])) {
-        if (sym_get_func_version(callable) == func_version) {
-            REPLACE_OP(this_instr, _NOP, 0, 0);
-        }
-        else {
-            sym_set_func_version(ctx, callable, func_version);
+        PyObject *func = sym_get_probable_value(callable);
+        if (func == NULL || !PyFunction_Check(func) || ((PyFunctionObject *)func)->func_version != func_version) {
+            ctx->contradiction = true;
+            ctx->done = true;
+            break;
         }
+        // Guarded on this, so it can be promoted.
+        sym_set_const(callable, func);
+        _Py_BloomFilter_Add(dependencies, func);
     }
 
     op(_CHECK_METHOD_VERSION, (func_version/2, callable, null, unused[oparg] -- callable, null, unused[oparg])) {
@@ -2229,7 +2232,8 @@ dummy_func(void) {
         if (co->co_version == version) {
             _Py_BloomFilter_Add(dependencies, co);
             // Functions derive their version from code objects.
-            if (sym_get_func_version(ctx->frame->callable) == version) {
+            PyFunctionObject *func = (PyFunctionObject *)sym_get_const(ctx, ctx->frame->callable);
+            if (func != NULL && func->func_version == version) {
                 REPLACE_OP(this_instr, _NOP, 0, 0);
             }
         }
@@ -2262,7 +2266,8 @@ dummy_func(void) {
     op(_GUARD_IP__PUSH_FRAME, (ip/4 --)) {
         (void)ip;
         stack_pointer = sym_set_stack_depth((int)this_instr->operand1, stack_pointer);
-        if (sym_get_func_version(ctx->frame->callable) != 0 &&
+        PyFunctionObject *func = (PyFunctionObject *)sym_get_const(ctx, ctx->frame->callable);
+        if (func != NULL && func->func_version != 0 &&
             // We can remove this guard for simple function call targets.
             (((PyCodeObject *)ctx->frame->func->func_code)->co_flags &
                 (CO_GENERATOR | CO_COROUTINE | CO_ASYNC_GENERATOR)) == 0) {