gh-119258: Eliminate Type Guards in Tier 2 Optimizer with Watcher (GH-119365)

Co-authored-by: parmeggiani <[email protected]>
Co-authored-by: dpdani <[email protected]>
Co-authored-by: blurb-it[bot] <43283697+blurb-it[bot]@users.noreply.github.com>
Co-authored-by: Brandt Bucher <[email protected]>
Co-authored-by: Ken Jin <[email protected]>

Author: 5 people

Include/internal/pycore_optimizer.h

@@ -33,6 +33,7 @@ struct _Py_UopsSymbol {
     int flags;  // 0 bits: Top; 2 or more bits: Bottom
     PyTypeObject *typ;  // Borrowed reference
     PyObject *const_val;  // Owned reference (!)
+    unsigned int type_version; // currently stores type version
 };
 
 #define UOP_FORMAT_TARGET 0
@@ -123,9 +124,11 @@ extern _Py_UopsSymbol *_Py_uop_sym_new_const(_Py_UOpsContext *ctx, PyObject *con
 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 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_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);
@@ -138,9 +141,9 @@ extern void _Py_uop_abstractcontext_fini(_Py_UOpsContext *ctx);
 extern _Py_UOpsAbstractFrame *_Py_uop_frame_new(
     _Py_UOpsContext *ctx,
     PyCodeObject *co,
-    _Py_UopsSymbol **localsplus_start,
-    int n_locals_already_filled,
-    int curr_stackentries);
+    int curr_stackentries,
+    _Py_UopsSymbol **args,
+    int arg_len);
 extern int _Py_uop_frame_pop(_Py_UOpsContext *ctx);
 
 PyAPI_FUNC(PyObject *) _Py_uop_symbols_test(PyObject *self, PyObject *ignored);

Include/internal/pycore_typeobject.h

@@ -63,6 +63,8 @@ typedef struct {
     PyObject *tp_weaklist;
 } managed_static_type_state;
 
+#define TYPE_VERSION_CACHE_SIZE (1<<12)  /* Must be a power of 2 */
+
 struct types_state {
     /* Used to set PyTypeObject.tp_version_tag.
        It starts at _Py_MAX_GLOBAL_TYPE_VERSION_TAG + 1,
@@ -118,6 +120,12 @@ struct types_state {
         managed_static_type_state initialized[_Py_MAX_MANAGED_STATIC_EXT_TYPES];
     } for_extensions;
     PyMutex mutex;
+
+    // Borrowed references to type objects whose
+    // tp_version_tag % TYPE_VERSION_CACHE_SIZE
+    // once was equal to the index in the table.
+    // They are cleared when the type object is deallocated.
+    PyTypeObject *type_version_cache[TYPE_VERSION_CACHE_SIZE];
 };
 
 
@@ -230,6 +238,9 @@ extern void _PyType_SetFlags(PyTypeObject *self, unsigned long mask,
 extern void _PyType_SetFlagsRecursive(PyTypeObject *self, unsigned long mask,
                                       unsigned long flags);
 
+extern unsigned int _PyType_GetVersionForCurrentState(PyTypeObject *tp);
+PyAPI_FUNC(void) _PyType_SetVersion(PyTypeObject *tp, unsigned int version);
+PyTypeObject *_PyType_LookupByVersion(unsigned int version);
 
 #ifdef __cplusplus
 }

Lib/test/test_capi/test_opt.py

@@ -1333,6 +1333,153 @@ def test_modified_local_is_seen_by_optimized_code(self):
         self.assertIs(type(s), float)
         self.assertEqual(s, 1024.0)
 
+    def test_guard_type_version_removed(self):
+        def thing(a):
+            x = 0
+            for _ in range(100):
+                x += a.attr
+                x += a.attr
+            return x
+
+        class Foo:
+            attr = 1
+
+        res, ex = self._run_with_optimizer(thing, Foo())
+        opnames = list(iter_opnames(ex))
+        self.assertIsNotNone(ex)
+        self.assertEqual(res, 200)
+        guard_type_version_count = opnames.count("_GUARD_TYPE_VERSION")
+        self.assertEqual(guard_type_version_count, 1)
+
+    def test_guard_type_version_removed_inlined(self):
+        """
+        Verify that the guard type version if we have an inlined function
+        """
+
+        def fn():
+            pass
+
+        def thing(a):
+            x = 0
+            for _ in range(100):
+                x += a.attr
+                fn()
+                x += a.attr
+            return x
+
+        class Foo:
+            attr = 1
+
+        res, ex = self._run_with_optimizer(thing, Foo())
+        opnames = list(iter_opnames(ex))
+        self.assertIsNotNone(ex)
+        self.assertEqual(res, 200)
+        guard_type_version_count = opnames.count("_GUARD_TYPE_VERSION")
+        self.assertEqual(guard_type_version_count, 1)
+
+    def test_guard_type_version_not_removed(self):
+        """
+        Verify that the guard type version is not removed if we modify the class
+        """
+
+        def thing(a):
+            x = 0
+            for i in range(100):
+                x += a.attr
+                # for the first 90 iterations we set the attribute on this dummy function which shouldn't
+                # trigger the type watcher
+                # then after 90  it should trigger it and stop optimizing
+                # Note that the code needs to be in this weird form so it's optimized inline without any control flow
+                setattr((Foo, Bar)[i < 90], "attr", 2)
+                x += a.attr
+            return x
+
+        class Foo:
+            attr = 1
+
+        class Bar:
+            pass
+
+        res, ex = self._run_with_optimizer(thing, Foo())
+        opnames = list(iter_opnames(ex))
+
+        self.assertIsNotNone(ex)
+        self.assertEqual(res, 219)
+        guard_type_version_count = opnames.count("_GUARD_TYPE_VERSION")
+        self.assertEqual(guard_type_version_count, 2)
+
+
+    @unittest.expectedFailure
+    def test_guard_type_version_not_removed_escaping(self):
+        """
+        Verify that the guard type version is not removed if have an escaping function
+        """
+
+        def thing(a):
+            x = 0
+            for i in range(100):
+                x += a.attr
+                # eval should be escaping and so should cause optimization to stop and preserve both type versions
+                eval("None")
+                x += a.attr
+            return x
+
+        class Foo:
+            attr = 1
+        res, ex = self._run_with_optimizer(thing, Foo())
+        opnames = list(iter_opnames(ex))
+        self.assertIsNotNone(ex)
+        self.assertEqual(res, 200)
+        guard_type_version_count = opnames.count("_GUARD_TYPE_VERSION")
+        # Note: This will actually be 1 for noe
+        # https://github.com/python/cpython/pull/119365#discussion_r1626220129
+        self.assertEqual(guard_type_version_count, 2)
+
+
+    def test_guard_type_version_executor_invalidated(self):
+        """
+        Verify that the executor is invalided on a type change.
+        """
+
+        def thing(a):
+            x = 0
+            for i in range(100):
+                x += a.attr
+                x += a.attr
+            return x
+
+        class Foo:
+            attr = 1
+
+        res, ex = self._run_with_optimizer(thing, Foo())
+        self.assertEqual(res, 200)
+        self.assertIsNotNone(ex)
+        self.assertEqual(list(iter_opnames(ex)).count("_GUARD_TYPE_VERSION"), 1)
+        self.assertTrue(ex.is_valid())
+        Foo.attr = 0
+        self.assertFalse(ex.is_valid())
+
+    def test_type_version_doesnt_segfault(self):
+        """
+        Tests that setting a type version doesn't cause a segfault when later looking at the stack.
+        """
+
+        # Minimized from mdp.py benchmark
+
+        class A:
+            def __init__(self):
+                self.attr = {}
+
+            def method(self, arg):
+                self.attr[arg] = None
+
+        def fn(a):
+            for _ in range(100):
+                (_ for _ in [])
+                (_ for _ in [a.method(None)])
+
+        fn(A())
+
 
 if __name__ == "__main__":
     unittest.main()

Lib/test/test_capi/test_watchers.py

@@ -282,8 +282,10 @@ class C: pass
             self.watch(wid, C)
             with catch_unraisable_exception() as cm:
                 C.foo = "bar"
-                self.assertEqual(cm.unraisable.err_msg,
-                    f"Exception ignored in type watcher callback #0 for {C!r}")
+                self.assertEqual(
+                    cm.unraisable.err_msg,
+                    f"Exception ignored in type watcher callback #1 for {C!r}",
+                )
                 self.assertIs(cm.unraisable.object, None)
                 self.assertEqual(str(cm.unraisable.exc_value), "boom!")
             self.assert_events([])

Lib/test/test_type_cache.py

@@ -10,8 +10,9 @@
 
 # Skip this test if the _testcapi module isn't available.
 _testcapi = import_helper.import_module("_testcapi")
+_testinternalcapi = import_helper.import_module("_testinternalcapi")
 type_get_version = _testcapi.type_get_version
-type_assign_specific_version_unsafe = _testcapi.type_assign_specific_version_unsafe
+type_assign_specific_version_unsafe = _testinternalcapi.type_assign_specific_version_unsafe
 type_assign_version = _testcapi.type_assign_version
 type_modified = _testcapi.type_modified
 

Misc/NEWS.d/next/Core and Builtins/2024-05-23-20-17-37.gh-issue-119258.wZFIpt.rst

@@ -0,0 +1,3 @@
+Eliminate type version guards in the tier two interpreter.
+
+Note that setting the ``tp_version_tag`` manually (which has never been supported) may result in crashes.

Modules/_testcapimodule.c

@@ -2403,21 +2403,6 @@ type_modified(PyObject *self, PyObject *type)
     Py_RETURN_NONE;
 }
 
-// Circumvents standard version assignment machinery - use with caution and only on
-// short-lived heap types
-static PyObject *
-type_assign_specific_version_unsafe(PyObject *self, PyObject *args)
-{
-    PyTypeObject *type;
-    unsigned int version;
-    if (!PyArg_ParseTuple(args, "Oi:type_assign_specific_version_unsafe", &type, &version)) {
-        return NULL;
-    }
-    assert(!PyType_HasFeature(type, Py_TPFLAGS_IMMUTABLETYPE));
-    type->tp_version_tag = version;
-    type->tp_flags |= Py_TPFLAGS_VALID_VERSION_TAG;
-    Py_RETURN_NONE;
-}
 
 static PyObject *
 type_assign_version(PyObject *self, PyObject *type)
@@ -3427,8 +3412,6 @@ static PyMethodDef TestMethods[] = {
     {"test_py_is_funcs", test_py_is_funcs, METH_NOARGS},
     {"type_get_version", type_get_version, METH_O, PyDoc_STR("type->tp_version_tag")},
     {"type_modified", type_modified, METH_O, PyDoc_STR("PyType_Modified")},
-    {"type_assign_specific_version_unsafe", type_assign_specific_version_unsafe, METH_VARARGS,
-     PyDoc_STR("forcefully assign type->tp_version_tag")},
     {"type_assign_version", type_assign_version, METH_O, PyDoc_STR("PyUnstable_Type_AssignVersionTag")},
     {"type_get_tp_bases", type_get_tp_bases, METH_O},
     {"type_get_tp_mro", type_get_tp_mro, METH_O},

Modules/_testinternalcapi.c

@@ -2002,6 +2002,22 @@ has_inline_values(PyObject *self, PyObject *obj)
 }
 
 
+// Circumvents standard version assignment machinery - use with caution and only on
+// short-lived heap types
+static PyObject *
+type_assign_specific_version_unsafe(PyObject *self, PyObject *args)
+{
+    PyTypeObject *type;
+    unsigned int version;
+    if (!PyArg_ParseTuple(args, "Oi:type_assign_specific_version_unsafe", &type, &version)) {
+        return NULL;
+    }
+    assert(!PyType_HasFeature(type, Py_TPFLAGS_IMMUTABLETYPE));
+    _PyType_SetVersion(type, version);
+    type->tp_flags |= Py_TPFLAGS_VALID_VERSION_TAG;
+    Py_RETURN_NONE;
+}
+
 /*[clinic input]
 gh_119213_getargs
 
@@ -2102,6 +2118,9 @@ static PyMethodDef module_functions[] = {
     {"get_rare_event_counters", get_rare_event_counters, METH_NOARGS},
     {"reset_rare_event_counters", reset_rare_event_counters, METH_NOARGS},
     {"has_inline_values", has_inline_values, METH_O},
+    {"type_assign_specific_version_unsafe", type_assign_specific_version_unsafe, METH_VARARGS,
+     PyDoc_STR("forcefully assign type->tp_version_tag")},
+
 #ifdef Py_GIL_DISABLED
     {"py_thread_id", get_py_thread_id, METH_NOARGS},
 #endif