gh-146099: Optimize _GUARD_CODE_VERSION+IP via function version symbols

Include/internal/pycore_optimizer.h

@@ -312,6 +312,8 @@ 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,15 +36,16 @@ typedef enum _JitSymType {
     JIT_SYM_NON_NULL_TAG = 3,
     JIT_SYM_BOTTOM_TAG = 4,
     JIT_SYM_TYPE_VERSION_TAG = 5,
-    JIT_SYM_KNOWN_CLASS_TAG = 6,
-    JIT_SYM_KNOWN_VALUE_TAG = 7,
-    JIT_SYM_TUPLE_TAG = 8,
-    JIT_SYM_TRUTHINESS_TAG = 9,
-    JIT_SYM_COMPACT_INT = 10,
-    JIT_SYM_PREDICATE_TAG = 11,
-    JIT_SYM_RECORDED_VALUE_TAG = 12,
-    JIT_SYM_RECORDED_TYPE_TAG = 13,
-    JIT_SYM_RECORDED_GEN_FUNC_TAG = 14,
+    JIT_SYM_FUNC_VERSION_TAG = 6,
+    JIT_SYM_KNOWN_CLASS_TAG = 7,
+    JIT_SYM_KNOWN_VALUE_TAG = 8,
+    JIT_SYM_TUPLE_TAG = 9,
+    JIT_SYM_TRUTHINESS_TAG = 10,
+    JIT_SYM_COMPACT_INT = 11,
+    JIT_SYM_PREDICATE_TAG = 12,
+    JIT_SYM_RECORDED_VALUE_TAG = 13,
+    JIT_SYM_RECORDED_TYPE_TAG = 14,
+    JIT_SYM_RECORDED_GEN_FUNC_TAG = 15,
 } JitSymType;
 
 typedef struct _jit_opt_known_class {
@@ -58,6 +59,11 @@ typedef struct _jit_opt_known_version {
     uint32_t version;
 } JitOptKnownVersion;
 
+typedef struct _jit_opt_known_func_version {
+    uint8_t tag;
+    uint32_t func_version;
+} JitOptKnownFuncVersion;
+
 typedef struct _jit_opt_known_value {
     uint8_t tag;
     PyObject *value;
@@ -118,6 +124,7 @@ typedef union _jit_opt_symbol {
     JitOptKnownClass cls;
     JitOptKnownValue value;
     JitOptKnownVersion version;
+    JitOptKnownFuncVersion func_version;
     JitOptTuple tuple;
     JitOptTruthiness truthiness;
     JitOptCompactInt compact;
@@ -140,6 +147,7 @@ typedef struct _Py_UOpsAbstractFrame {
     int stack_len;
     int locals_len;
     bool caller; // We have made a call from this frame during the trace
+    JitOptRef callable;
     PyFunctionObject *func;
     PyCodeObject *code;
 

Lib/test/test_capi/test_opt.py

@@ -630,6 +630,8 @@ def dummy(x):
         self.assertIn("_PUSH_FRAME", uops)
         self.assertIn("_BINARY_OP_ADD_INT", uops)
         self.assertNotIn("_CHECK_PEP_523", uops)
+        self.assertNotIn("_GUARD_CODE_VERSION__PUSH_FRAME", uops)
+        self.assertNotIn("_GUARD_IP__PUSH_FRAME", uops)
 
     def test_int_type_propagate_through_range(self):
         def testfunc(n):
@@ -1540,8 +1542,9 @@ def testfunc(n):
         self.assertIsNotNone(ex)
         uops = get_opnames(ex)
         self.assertIn("_PUSH_FRAME", uops)
-        # Strength reduced version
-        self.assertIn("_CHECK_FUNCTION_VERSION_INLINE", 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)
         # Removed guard
         self.assertNotIn("_CHECK_FUNCTION_EXACT_ARGS", uops)

Python/optimizer_analysis.c

@@ -291,6 +291,8 @@ 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

@@ -857,12 +857,12 @@ dummy_func(void) {
     }
 
     op(_CHECK_FUNCTION_VERSION, (func_version/2, callable, self_or_null, unused[oparg] -- callable, self_or_null, unused[oparg])) {
-        if (sym_is_const(ctx, callable) && sym_matches_type(callable, &PyFunction_Type)) {
-            assert(PyFunction_Check(sym_get_const(ctx, callable)));
-            ADD_OP(_CHECK_FUNCTION_VERSION_INLINE, 0, func_version);
-            uop_buffer_last(&ctx->out_buffer)->operand1 = (uintptr_t)sym_get_const(ctx, callable);
+        if (sym_get_func_version(callable) == func_version) {
+            REPLACE_OP(this_instr, _NOP, 0, 0);
+        }
+        else {
+            sym_set_func_version(ctx, callable, func_version);
         }
-        sym_set_type(callable, &PyFunction_Type);
     }
 
     op(_CHECK_METHOD_VERSION, (func_version/2, callable, null, unused[oparg] -- callable, null, unused[oparg])) {
@@ -1767,12 +1767,46 @@ dummy_func(void) {
         sym_set_recorded_gen_func(nos, func);
     }
 
+    op(_GUARD_CODE_VERSION__PUSH_FRAME, (version/2 -- )) {
+        PyCodeObject *co = get_current_code_object(ctx);
+        if (co->co_version == version) {
+            _Py_BloomFilter_Add(dependencies, co);
+            if (sym_get_func_version(ctx->frame->callable) != 0) {
+                REPLACE_OP(this_instr, _NOP, 0, 0);
+            }
+        }
+        else {
+            ctx->done = true;
+        }
+    }
+
+    op(_GUARD_CODE_VERSION_RETURN_VALUE, (version/2 -- )) {
+        if (ctx->frame->caller) {
+            REPLACE_OP(this_instr, _NOP, 0, 0);
+        }
+    }
+
+    op(_GUARD_CODE_VERSION_YIELD_VALUE, (version/2 -- )) {
+        if (ctx->frame->caller) {
+            REPLACE_OP(this_instr, _NOP, 0, 0);
+        }
+    }
+
+    op(_GUARD_CODE_VERSION_RETURN_GENERATOR, (version/2 -- )) {
+        if (ctx->frame->caller) {
+            REPLACE_OP(this_instr, _NOP, 0, 0);
+        }
+    }
+
     op(_GUARD_IP__PUSH_FRAME, (ip/4 --)) {
         (void)ip;
         stack_pointer = sym_set_stack_depth((int)this_instr->operand1, stack_pointer);
-        // TO DO
-        // Normal function calls to known functions
-        // do not need an IP guard.
+        if (sym_get_func_version(ctx->frame->callable) != 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) {
+            REPLACE_OP(this_instr, _NOP, 0, 0);
+        }
     }
 
     op(_GUARD_IP_YIELD_VALUE, (ip/4 --)) {