saulshanabrook · May 24, 2024 · May 29, 2024 · May 30, 2024 · May 30, 2024 · Jun 10, 2024
diff --git a/Include/internal/pycore_optimizer.h b/Include/internal/pycore_optimizer.h
@@ -151,6 +151,7 @@ struct _Py_UopsSymbol {
     PyTypeObject *typ;  // Borrowed reference
     PyObject *const_val;  // Owned reference (!)
     unsigned int type_version; // currently stores type version
+    uint32_t function_version;
 };
 
 #define UOP_FORMAT_TARGET 0
@@ -234,10 +235,12 @@ extern _Py_UopsSymbol *_Py_uop_sym_new_null(_Py_UOpsContext *ctx);
 extern bool _Py_uop_sym_has_type(_Py_UopsSymbol *sym);
 extern bool _Py_uop_sym_matches_type(_Py_UopsSymbol *sym, PyTypeObject *typ);
 extern bool _Py_uop_sym_matches_type_version(_Py_UopsSymbol *sym, unsigned int version);
+extern bool _Py_uop_sym_matches_function_version(_Py_UopsSymbol *sym, uint32_t version);
 extern void _Py_uop_sym_set_null(_Py_UOpsContext *ctx, _Py_UopsSymbol *sym);
 extern void _Py_uop_sym_set_non_null(_Py_UOpsContext *ctx, _Py_UopsSymbol *sym);
 extern void _Py_uop_sym_set_type(_Py_UOpsContext *ctx, _Py_UopsSymbol *sym, PyTypeObject *typ);
 extern bool _Py_uop_sym_set_type_version(_Py_UOpsContext *ctx, _Py_UopsSymbol *sym, unsigned int version);
+extern void _Py_uop_sym_set_function_version(_Py_UOpsContext *ctx, _Py_UopsSymbol *sym, uint32_t version);
 extern void _Py_uop_sym_set_const(_Py_UOpsContext *ctx, _Py_UopsSymbol *sym, PyObject *const_val);
 extern bool _Py_uop_sym_is_bottom(_Py_UopsSymbol *sym);
 extern int _Py_uop_sym_truthiness(_Py_UopsSymbol *sym);

diff --git a/Lib/test/test_capi/test_opt.py b/Lib/test/test_capi/test_opt.py
@@ -1481,5 +1481,34 @@ def fn(a):
         fn(A())
 
 
+    def test_guard_function_version_removed(self):
+        def thing(f):
+            x = 0
+            for _ in range(100):
+                x += f() + f()
+            return x
+
+        res, ex = self._run_with_optimizer(thing, lambda: 1)
+        opnames = list(iter_opnames(ex))
+        self.assertIsNotNone(ex)
+        self.assertEqual(res, 200)
+        self.assertEqual(opnames.count("_CHECK_FUNCTION_VERSION"), 1)
+
+    def test_guard_function_version_invalidated(self):
+        def thing(f):
+            x = 0
+            for _ in range(100):
+                x += f()
+                x += f()
+            return x
+
+        def fn():
+            return 1
+        res, ex = self._run_with_optimizer(thing, fn)
+        self.assertTrue(ex.is_valid())
+        fn.__code__ = fn.__code__.replace(co_name="new_name")
+        self.assertFalse(ex.is_valid())
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/Objects/funcobject.c b/Objects/funcobject.c
@@ -70,7 +70,7 @@ PyFunction_AddWatcher(PyFunction_WatchCallback callback)
 {
     PyInterpreterState *interp = _PyInterpreterState_GET();
     assert(interp->_initialized);
-    for (int i = 0; i < FUNC_MAX_WATCHERS; i++) {
+    for (int i = 1; i < FUNC_MAX_WATCHERS; i++) {
         if (interp->func_watchers[i] == NULL) {
             interp->func_watchers[i] = callback;
             interp->active_func_watchers |= (1 << i);
@@ -346,11 +346,11 @@ _PyFunction_LookupByVersion(uint32_t version, PyObject **p_code)
     if (slot->code) {
         assert(PyCode_Check(slot->code));
         PyCodeObject *code = (PyCodeObject *)slot->code;
-        if (code->co_version == version) {
+        if (code->co_version == version && p_code) {
             *p_code = slot->code;
         }
     }
-    else {
+    else if (p_code) {
         *p_code = NULL;
     }
     if (slot->func && slot->func->func_version == version) {

@@ -79,7 +79,21 @@ increment_mutations(PyObject* dict) {
 #define BUILTINS_WATCHER_ID 0
 #define GLOBALS_WATCHER_ID  1
 #define TYPE_WATCHER_ID  0
+#define FUNCTION_WATCHER_ID 0
 
+static int
+function_watcher_callback(
+  PyFunction_WatchEvent event,
+  PyFunctionObject *func,
+  PyObject *new_value)
+{
+    // Don't invalidate if we are destroying funciton
+    if (event == PyFunction_EVENT_DESTROY) {
+        return 0;
+    }
+    _Py_Executors_InvalidateDependency(_PyInterpreterState_GET(), func, 1);
+    return 0;
+}
 static int
 globals_watcher_callback(PyDict_WatchEvent event, PyObject* dict,
                          PyObject* key, PyObject* new_value)
@@ -178,6 +192,10 @@ remove_globals(_PyInterpreterFrame *frame, _PyUOpInstruction *buffer,
     if (interp->type_watchers[TYPE_WATCHER_ID] == NULL) {
         interp->type_watchers[TYPE_WATCHER_ID] = type_watcher_callback;
     }
+    if (interp->func_watchers[GLOBALS_WATCHER_ID] == NULL) {
+        interp->active_func_watchers |= 0x1;
+        interp->func_watchers[FUNCTION_WATCHER_ID] = function_watcher_callback;
+    }
     for (int pc = 0; pc < buffer_size; pc++) {
         _PyUOpInstruction *inst = &buffer[pc];
         int opcode = inst->opcode;
@@ -327,10 +345,12 @@ remove_globals(_PyInterpreterFrame *frame, _PyUOpInstruction *buffer,
 #define sym_get_type _Py_uop_sym_get_type
 #define sym_matches_type _Py_uop_sym_matches_type
 #define sym_matches_type_version _Py_uop_sym_matches_type_version
+#define sym_matches_function_version _Py_uop_sym_matches_function_version
 #define sym_set_null(SYM) _Py_uop_sym_set_null(ctx, SYM)
 #define sym_set_non_null(SYM) _Py_uop_sym_set_non_null(ctx, SYM)
 #define sym_set_type(SYM, TYPE) _Py_uop_sym_set_type(ctx, SYM, TYPE)
 #define sym_set_type_version(SYM, VERSION) _Py_uop_sym_set_type_version(ctx, SYM, VERSION)
+#define sym_set_function_version(SYM, VERSION) _Py_uop_sym_set_function_version(ctx, SYM, VERSION)
 #define sym_set_const(SYM, CNST) _Py_uop_sym_set_const(ctx, SYM, CNST)
 #define sym_is_bottom _Py_uop_sym_is_bottom
 #define sym_truthiness _Py_uop_sym_truthiness

@@ -22,12 +22,14 @@ typedef struct _Py_UOpsAbstractFrame _Py_UOpsAbstractFrame;
 #define sym_new_null _Py_uop_sym_new_null
 #define sym_matches_type _Py_uop_sym_matches_type
 #define sym_matches_type_version _Py_uop_sym_matches_type_version
+#define sym_matches_function_version _Py_uop_sym_matches_function_version
 #define sym_get_type _Py_uop_sym_get_type
 #define sym_has_type _Py_uop_sym_has_type
 #define sym_set_null(SYM) _Py_uop_sym_set_null(ctx, SYM)
 #define sym_set_non_null(SYM) _Py_uop_sym_set_non_null(ctx, SYM)
 #define sym_set_type(SYM, TYPE) _Py_uop_sym_set_type(ctx, SYM, TYPE)
 #define sym_set_type_version(SYM, VERSION) _Py_uop_sym_set_type_version(ctx, SYM, VERSION)
+#define sym_set_function_version(SYM, VERSION) _Py_uop_sym_set_function_version(ctx, SYM, VERSION)
 #define sym_set_const(SYM, CNST) _Py_uop_sym_set_const(ctx, SYM, CNST)
 #define sym_is_bottom _Py_uop_sym_is_bottom
 #define frame_new _Py_uop_frame_new
@@ -97,6 +99,21 @@ dummy_func(void) {
         res = sym_new_null(ctx);
     }
 
+    op(_CHECK_FUNCTION_VERSION, (func_version/2, callable, unused, unused[oparg] -- callable, unused, unused[oparg])) {
+        if (sym_matches_function_version(callable, func_version)) {
+            REPLACE_OP(this_instr, _NOP, 0, 0);
+        } else {
+            PyFunctionObject *function = _PyFunction_LookupByVersion(func_version, NULL);
+            // if the type is null, it was not found in the cache (there was a conflict)
+            // with the key, in which case we can't trust the version
+            if (function) {
+                sym_set_function_version(callable, func_version);
+                _Py_BloomFilter_Add(dependencies, function);
+            }
+
+        }
+    }
+
     op(_GUARD_BOTH_INT, (left, right -- left, right)) {
         if (sym_matches_type(left, &PyLong_Type)) {
             if (sym_matches_type(right, &PyLong_Type)) {

diff --git a/Python/optimizer_cases.c.h b/Python/optimizer_cases.c.h
@@ -53,6 +53,7 @@ static _Py_UopsSymbol NO_SPACE_SYMBOL = {
     .typ = NULL,
     .const_val = NULL,
     .type_version = 0,
+    .function_version = 0,
 };
 
 _Py_UopsSymbol *
@@ -77,6 +78,7 @@ sym_new(_Py_UOpsContext *ctx)
     self->typ = NULL;
     self->const_val = NULL;
     self->type_version = 0;
+    self->function_version = 0;
 
     return self;
 }
@@ -165,6 +167,24 @@ _Py_uop_sym_set_type_version(_Py_UOpsContext *ctx, _Py_UopsSymbol *sym, unsigned
     return true;
 }
 
+void
+_Py_uop_sym_set_function_version(_Py_UOpsContext *ctx, _Py_UopsSymbol *sym, uint32_t version)
+{
+    sym->function_version = version;
+}
+
+uint32_t
+_Py_uop_sym_get_function_version(_Py_UopsSymbol *sym)
+{
+    return sym->function_version;
+}
+
+bool
+_Py_uop_sym_matches_function_version(_Py_UopsSymbol *sym, uint32_t version)
+{
+    return _Py_uop_sym_get_function_version(sym) == version;
+}
+
 void
 _Py_uop_sym_set_const(_Py_UOpsContext *ctx, _Py_UopsSymbol *sym, PyObject *const_val)
 {