Add function symbol (with version) to optimizer symbols

markshannon · markshannon · commit d28150d27e4f · 2025-08-15T17:31:12.000+01:00
diff --git a/Include/internal/pycore_optimizer.h b/Include/internal/pycore_optimizer.h
@@ -181,6 +181,7 @@ typedef enum _JitSymType {
     JIT_SYM_TUPLE_TAG = 8,
     JIT_SYM_TRUTHINESS_TAG = 9,
     JIT_SYM_COMPACT_INT = 10,
+    JIT_SYM_FUNCTION_TAG = 11,
 } JitSymType;
 
 typedef struct _jit_opt_known_class {
@@ -217,6 +218,14 @@ typedef struct {
     uint8_t tag;
 } JitOptCompactInt;
 
+typedef struct {
+    uint8_t tag;
+    bool value_is_code;
+    uint32_t version;
+    PyObject *function;
+} JitOptFunction;
+
+
 typedef union _jit_opt_symbol {
     uint8_t tag;
     JitOptKnownClass cls;
@@ -225,6 +234,7 @@ typedef union _jit_opt_symbol {
     JitOptTuple tuple;
     JitOptTruthiness truthiness;
     JitOptCompactInt compact;
+    JitOptFunction function;
 } JitOptSymbol;
 
 
@@ -276,6 +286,9 @@ struct _Py_UOpsAbstractFrame {
     // Max stacklen
     int stack_len;
     int locals_len;
+    bool function_checked;
+    bool builtins_watched;
+    bool globals_watched;
 
     JitOptRef *stack_pointer;
     JitOptRef *stack;
@@ -298,6 +311,7 @@ typedef struct _JitOptContext {
     _Py_UOpsAbstractFrame *frame;
     _Py_UOpsAbstractFrame frames[MAX_ABSTRACT_FRAME_DEPTH];
     int curr_frame_depth;
+    uint32_t prechecked_function_version;
 
     // Arena for the symbolic types.
     ty_arena t_arena;
@@ -339,6 +353,8 @@ extern JitOptRef _Py_uop_sym_new_truthiness(JitOptContext *ctx, JitOptRef value,
 extern bool _Py_uop_sym_is_compact_int(JitOptRef sym);
 extern JitOptRef _Py_uop_sym_new_compact_int(JitOptContext *ctx);
 extern void _Py_uop_sym_set_compact_int(JitOptContext *ctx,  JitOptRef sym);
+extern void _Py_uop_sym_set_function_version(JitOptContext *ctx, JitOptRef ref, uint32_t version);
+extern uint32_t _Py_uop_sym_get_function_version(JitOptRef ref);
 
 extern void _Py_uop_abstractcontext_init(JitOptContext *ctx);
 extern void _Py_uop_abstractcontext_fini(JitOptContext *ctx);
diff --git a/Python/optimizer_analysis.c b/Python/optimizer_analysis.c
@@ -344,6 +344,8 @@ remove_globals(_PyInterpreterFrame *frame, _PyUOpInstruction *buffer,
 #define sym_set_type_version(SYM, VERSION) _Py_uop_sym_set_type_version(ctx, SYM, VERSION)
 #define sym_set_const(SYM, CNST) _Py_uop_sym_set_const(ctx, SYM, CNST)
 #define sym_set_compact_int(SYM) _Py_uop_sym_set_compact_int(ctx, SYM)
+#define sym_get_function_version(SYM) _Py_uop_sym_get_function_version(SYM)
+#define sym_set_function_version(SYM, VER) _Py_uop_sym_set_function_version(ctx, SYM, VER)
 #define sym_is_bottom _Py_uop_sym_is_bottom
 #define sym_truthiness _Py_uop_sym_truthiness
 #define frame_new _Py_uop_frame_new
diff --git a/Python/optimizer_bytecodes.c b/Python/optimizer_bytecodes.c
@@ -28,6 +28,8 @@ typedef struct _Py_UOpsAbstractFrame _Py_UOpsAbstractFrame;
 #define sym_set_type_version(SYM, VERSION) _Py_uop_sym_set_type_version(ctx, SYM, VERSION)
 #define sym_set_const(SYM, CNST) _Py_uop_sym_set_const(ctx, SYM, CNST)
 #define sym_set_compact_int(SYM) _Py_uop_sym_set_compact_int(ctx, SYM)
+#define sym_get_function_version(SYM) _Py_uop_sym_get_function_version(SYM)
+#define sym_set_function_version(SYM, VER) _Py_uop_sym_set_function_version(ctx, SYM, VER)
 #define sym_is_bottom _Py_uop_sym_is_bottom
 #define frame_new _Py_uop_frame_new
 #define frame_pop _Py_uop_frame_pop
@@ -663,12 +665,18 @@ 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)));
-            REPLACE_OP(this_instr, _CHECK_FUNCTION_VERSION_INLINE, 0, func_version);
-            this_instr->operand1 = (uintptr_t)sym_get_const(ctx, callable);
+        assert(func_version != 0);
+        if (sym_get_function_version(callable) == func_version) {
+            REPLACE_OP(this_instr, NOP, 0, 0);
+        }
+        else {
+            if (sym_is_const(ctx, callable) && sym_matches_type(callable, &PyFunction_Type)) {
+                assert(PyFunction_Check(sym_get_const(ctx, callable)));
+                REPLACE_OP(this_instr, _CHECK_FUNCTION_VERSION_INLINE, 0, func_version);
+                this_instr->operand1 = (uintptr_t)sym_get_const(ctx, callable);
+            }
+            sym_set_function_version(callable, func_version);
         }
-        sym_set_type(callable, &PyFunction_Type);
     }
 
     op(_CHECK_METHOD_VERSION, (func_version/2, callable, null, unused[oparg] -- callable, null, unused[oparg])) {
diff --git a/Python/optimizer_cases.c.h b/Python/optimizer_cases.c.h
diff --git a/Python/optimizer_symbols.c b/Python/optimizer_symbols.c
@@ -121,6 +121,12 @@ static void make_const(JitOptSymbol *sym, PyObject *val)
 static inline void
 sym_set_bottom(JitOptContext *ctx, JitOptSymbol *sym)
 {
+    if (sym->tag == JIT_SYM_KNOWN_VALUE_TAG) {
+        Py_CLEAR(sym->value.value);
+    }
+    if (sym->tag == JIT_SYM_FUNCTION_TAG) {
+        Py_CLEAR(sym->function.function);
+    }
     sym->tag = JIT_SYM_BOTTOM_TAG;
     ctx->done = true;
     ctx->contradiction = true;
@@ -143,19 +149,24 @@ bool
 _Py_uop_sym_is_const(JitOptContext *ctx, JitOptRef ref)
 {
     JitOptSymbol *sym = PyJitRef_Unwrap(ref);
-    if (sym->tag == JIT_SYM_KNOWN_VALUE_TAG) {
-        return true;
-    }
-    if (sym->tag == JIT_SYM_TRUTHINESS_TAG) {
-        JitOptSymbol *value = allocation_base(ctx) + sym->truthiness.value;
-        int truthiness = _Py_uop_sym_truthiness(ctx, PyJitRef_Wrap(value));
-        if (truthiness < 0) {
-            return false;
+    switch(sym->tag) {
+        case JIT_SYM_KNOWN_VALUE_TAG:
+            return true;
+        case JIT_SYM_TRUTHINESS_TAG:
+        {
+            JitOptSymbol *value = allocation_base(ctx) + sym->truthiness.value;
+            int truthiness = _Py_uop_sym_truthiness(ctx, PyJitRef_Wrap(value));
+            if (truthiness < 0) {
+                return false;
+            }
+            make_const(sym, (truthiness ^ sym->truthiness.invert) ? Py_True : Py_False);
+            return true;
         }
-        make_const(sym, (truthiness ^ sym->truthiness.invert) ? Py_True : Py_False);
-        return true;
+        case JIT_SYM_FUNCTION_TAG:
+            return sym->function.function != NULL;
+        default:
+            return false;
     }
-    return false;
 }
 
 bool
@@ -182,6 +193,9 @@ _Py_uop_sym_get_const(JitOptContext *ctx, JitOptRef ref)
         make_const(sym, res);
         return res;
     }
+    if (sym->tag == JIT_SYM_FUNCTION_TAG) {
+        return sym->function.function;
+    }
     return NULL;
 }
 
@@ -269,6 +283,11 @@ _Py_uop_sym_set_type(JitOptContext *ctx, JitOptRef ref, PyTypeObject *typ)
                 sym_set_bottom(ctx, sym);
             }
             return;
+        case JIT_SYM_FUNCTION_TAG:
+            if (typ != &PyFunction_Type) {
+                sym_set_bottom(ctx, sym);
+            }
+            return;
     }
 }
 
@@ -307,6 +326,12 @@ _Py_uop_sym_set_type_version(JitOptContext *ctx, JitOptRef ref, unsigned int ver
                 return false;
             };
             return true;
+        case JIT_SYM_FUNCTION_TAG:
+            if (PyFunction_Type.tp_version_tag != version) {
+                sym_set_bottom(ctx, sym);
+                return false;
+            };
+            return true;
         case JIT_SYM_TYPE_VERSION_TAG:
             if (sym->version.version != version) {
                 sym_set_bottom(ctx, sym);
@@ -424,6 +449,14 @@ _Py_uop_sym_set_const(JitOptContext *ctx, JitOptRef ref, PyObject *const_val)
                 sym_set_bottom(ctx, sym);
             }
             return;
+        case JIT_SYM_FUNCTION_TAG:
+            if (PyFunction_Check(const_val)) {
+                sym->function.function = Py_NewRef(const_val);
+                if (((PyFunctionObject *)const_val)->func_version != sym->function.version) {
+                    sym_set_bottom(ctx, sym);
+                }
+            }
+            return;
     }
 }
 
@@ -543,7 +576,8 @@ _Py_uop_sym_get_type(JitOptRef ref)
             return &PyBool_Type;
         case JIT_SYM_COMPACT_INT:
             return &PyLong_Type;
-
+        case JIT_SYM_FUNCTION_TAG:
+            return &PyFunction_Type;
     }
     Py_UNREACHABLE();
 }
@@ -571,6 +605,8 @@ _Py_uop_sym_get_type_version(JitOptRef ref)
             return PyBool_Type.tp_version_tag;
         case JIT_SYM_COMPACT_INT:
             return PyLong_Type.tp_version_tag;
+        case JIT_SYM_FUNCTION_TAG:
+            return PyFunction_Type.tp_version_tag;
     }
     Py_UNREACHABLE();
 }
@@ -606,6 +642,8 @@ _Py_uop_sym_truthiness(JitOptContext *ctx, JitOptRef ref)
         case JIT_SYM_UNKNOWN_TAG:
         case JIT_SYM_COMPACT_INT:
             return -1;
+        case JIT_SYM_FUNCTION_TAG:
+            return 1;
         case JIT_SYM_KNOWN_CLASS_TAG:
             /* TODO :
              * Instances of some classes are always
@@ -761,6 +799,7 @@ _Py_uop_sym_set_compact_int(JitOptContext *ctx, JitOptRef ref)
             return;
         case JIT_SYM_TUPLE_TAG:
         case JIT_SYM_TRUTHINESS_TAG:
+        case JIT_SYM_FUNCTION_TAG:
             sym_set_bottom(ctx, sym);
             return;
         case JIT_SYM_BOTTOM_TAG:
@@ -773,6 +812,72 @@ _Py_uop_sym_set_compact_int(JitOptContext *ctx, JitOptRef ref)
     }
 }
 
+void
+_Py_uop_sym_set_function_version(JitOptContext *ctx, JitOptRef ref, uint32_t version)
+{
+    JitOptSymbol *sym = PyJitRef_Unwrap(ref);
+    JitSymType tag = sym->tag;
+    switch(tag) {
+        case JIT_SYM_TUPLE_TAG:
+        case JIT_SYM_TRUTHINESS_TAG:
+        case JIT_SYM_BOTTOM_TAG:
+        case JIT_SYM_COMPACT_INT:
+        case JIT_SYM_NULL_TAG:
+            sym_set_bottom(ctx, sym);
+            return;
+        case JIT_SYM_KNOWN_CLASS_TAG:
+            if (sym->cls.type == &PyFunction_Type) {
+                sym->tag = JIT_SYM_FUNCTION_TAG;
+                sym->function.function = NULL;
+                sym->function.version = version;
+            } else {
+                sym_set_bottom(ctx, sym);
+            }
+            return;
+        case JIT_SYM_TYPE_VERSION_TAG:
+            if (sym->version.version == PyLong_Type.tp_version_tag) {
+                sym->tag = JIT_SYM_COMPACT_INT;
+            }
+            else {
+                sym_set_bottom(ctx, sym);
+            }
+            return;
+        case JIT_SYM_KNOWN_VALUE_TAG:
+            if (PyFunction_Check(sym->value.value)) {
+                sym->tag = JIT_SYM_FUNCTION_TAG;
+                sym->function.function = sym->value.value;
+                sym->function.version = version;
+            }
+            else {
+                Py_CLEAR(sym->value.value);
+                sym_set_bottom(ctx, sym);
+            }
+            return;
+        case JIT_SYM_NON_NULL_TAG:
+        case JIT_SYM_UNKNOWN_TAG:
+            sym->tag = JIT_SYM_FUNCTION_TAG;
+            sym->function.function = NULL;
+            sym->function.version = version;
+            return;
+        case JIT_SYM_FUNCTION_TAG:
+            if (sym->function.version != version) {
+                Py_CLEAR(sym->function.function);
+                sym_set_bottom(ctx, sym);
+            }
+            return;
+    }
+}
+
+uint32_t
+_Py_uop_sym_get_function_version(JitOptRef ref)
+{
+    JitOptSymbol *sym = PyJitRef_Unwrap(ref);
+    if (sym->tag == JIT_SYM_FUNCTION_TAG) {
+        return sym->function.function;
+    }
+    return 0;
+}
+
 JitOptRef
 _Py_uop_sym_new_truthiness(JitOptContext *ctx, JitOptRef ref, bool truthy)
 {
@@ -827,6 +932,9 @@ _Py_uop_frame_new(
     frame->locals = ctx->n_consumed;
     frame->stack = frame->locals + co->co_nlocalsplus;
     frame->stack_pointer = frame->stack + curr_stackentries;
+    frame->function_checked = false;
+    frame->builtins_watched = false;
+    frame->globals_watched = false;
     ctx->n_consumed = ctx->n_consumed + (co->co_nlocalsplus + co->co_stacksize);
     if (ctx->n_consumed >= ctx->limit) {
         ctx->done = true;
@@ -860,13 +968,17 @@ _Py_uop_abstractcontext_fini(JitOptContext *ctx)
     if (ctx == NULL) {
         return;
     }
+    ctx->prechecked_function_version = 0;
     ctx->curr_frame_depth = 0;
     int tys = ctx->t_arena.ty_curr_number;
     for (int i = 0; i < tys; i++) {
         JitOptSymbol *sym = &ctx->t_arena.arena[i];
         if (sym->tag == JIT_SYM_KNOWN_VALUE_TAG) {
             Py_CLEAR(sym->value.value);
         }
+        if (sym->tag == JIT_SYM_FUNCTION_TAG) {
+            Py_CLEAR(sym->function.function);
+        }
     }
 }
 
