pythonGH-131798: Optimize cached class attributes and methods in the JIT (pythonGH-134403)

Author: brandtbucher

Include/internal/pycore_backoff.h

@@ -95,8 +95,10 @@ backoff_counter_triggers(_Py_BackoffCounter counter)
     return counter.value_and_backoff < UNREACHABLE_BACKOFF;
 }
 
-/* Initial JUMP_BACKWARD counter.
- * This determines when we create a trace for a loop. */
+// Initial JUMP_BACKWARD counter.
+// Must be larger than ADAPTIVE_COOLDOWN_VALUE, otherwise when JIT code is
+// invalidated we may construct a new trace before the bytecode has properly
+// re-specialized:
 #define JUMP_BACKWARD_INITIAL_VALUE 4095
 #define JUMP_BACKWARD_INITIAL_BACKOFF 12
 static inline _Py_BackoffCounter

Include/internal/pycore_code.h

@@ -451,6 +451,9 @@ write_location_entry_start(uint8_t *ptr, int code, int length)
 #define ADAPTIVE_COOLDOWN_BACKOFF 0
 
 // Can't assert this in pycore_backoff.h because of header order dependencies
+#if JUMP_BACKWARD_INITIAL_VALUE <= ADAPTIVE_COOLDOWN_VALUE
+#  error  "JIT threshold value should be larger than adaptive cooldown value"
+#endif
 #if SIDE_EXIT_INITIAL_VALUE <= ADAPTIVE_COOLDOWN_VALUE
 #  error  "Cold exit value should be larger than adaptive cooldown value"
 #endif

Include/internal/pycore_uop_ids.h

@@ -1280,8 +1280,8 @@ class Bar:
         self.assertIsNotNone(ex)
         self.assertEqual(res, TIER2_THRESHOLD * 6 + 1)
         call = opnames.index("_CALL_BUILTIN_FAST")
-        load_attr_top = opnames.index("_LOAD_ATTR_NONDESCRIPTOR_WITH_VALUES", 0, call)
-        load_attr_bottom = opnames.index("_LOAD_ATTR_NONDESCRIPTOR_WITH_VALUES", call)
+        load_attr_top = opnames.index("_POP_TOP_LOAD_CONST_INLINE_BORROW", 0, call)
+        load_attr_bottom = opnames.index("_POP_TOP_LOAD_CONST_INLINE_BORROW", call)
         self.assertEqual(opnames[:load_attr_top].count("_GUARD_TYPE_VERSION"), 1)
         self.assertEqual(opnames[call:load_attr_bottom].count("_CHECK_VALIDITY"), 2)
 
@@ -1303,8 +1303,8 @@ class Foo:
         self.assertIsNotNone(ex)
         self.assertEqual(res, TIER2_THRESHOLD * 2)
         call = opnames.index("_CALL_BUILTIN_FAST_WITH_KEYWORDS")
-        load_attr_top = opnames.index("_LOAD_ATTR_NONDESCRIPTOR_WITH_VALUES", 0, call)
-        load_attr_bottom = opnames.index("_LOAD_ATTR_NONDESCRIPTOR_WITH_VALUES", call)
+        load_attr_top = opnames.index("_POP_TOP_LOAD_CONST_INLINE_BORROW", 0, call)
+        load_attr_bottom = opnames.index("_POP_TOP_LOAD_CONST_INLINE_BORROW", call)
         self.assertEqual(opnames[:load_attr_top].count("_GUARD_TYPE_VERSION"), 1)
         self.assertEqual(opnames[call:load_attr_bottom].count("_CHECK_VALIDITY"), 2)
 
@@ -2169,6 +2169,45 @@ def testfunc(n):
         self.assertNotIn("_LOAD_SMALL_INT", uops)
         self.assertIn("_LOAD_CONST_INLINE_BORROW", uops)
 
+    def test_cached_attributes(self):
+        class C:
+            A = 1
+            def m(self):
+                return 1
+        class D:
+            __slots__ = ()
+            A = 1
+            def m(self):
+                return 1
+        class E(Exception):
+            def m(self):
+                return 1
+        def f(n):
+            x = 0
+            c = C()
+            d = D()
+            e = E()
+            for _ in range(n):
+                x += C.A  # _LOAD_ATTR_CLASS
+                x += c.A  # _LOAD_ATTR_NONDESCRIPTOR_WITH_VALUES
+                x += d.A  # _LOAD_ATTR_NONDESCRIPTOR_NO_DICT
+                x += c.m()  # _LOAD_ATTR_METHOD_WITH_VALUES
+                x += d.m()  # _LOAD_ATTR_METHOD_NO_DICT
+                x += e.m()  # _LOAD_ATTR_METHOD_LAZY_DICT
+            return x
+
+        res, ex = self._run_with_optimizer(f, TIER2_THRESHOLD)
+        self.assertEqual(res, 6 * TIER2_THRESHOLD)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+        self.assertNotIn("_LOAD_ATTR_CLASS", uops)
+        self.assertNotIn("_LOAD_ATTR_NONDESCRIPTOR_WITH_VALUES", uops)
+        self.assertNotIn("_LOAD_ATTR_NONDESCRIPTOR_NO_DICT", uops)
+        self.assertNotIn("_LOAD_ATTR_METHOD_WITH_VALUES", uops)
+        self.assertNotIn("_LOAD_ATTR_METHOD_NO_DICT", uops)
+        self.assertNotIn("_LOAD_ATTR_METHOD_LAZY_DICT", uops)
+
+
 def global_identity(x):
     return x
 

Include/internal/pycore_uop_metadata.h

@@ -0,0 +1,2 @@
+Improve the JIT's ability to optimize away cached class attribute and method
+loads.

Lib/test/test_capi/test_opt.py

@@ -5307,6 +5307,18 @@ dummy_func(
             value = PyStackRef_FromPyObjectBorrow(ptr);
         }
 
+        tier2 op(_LOAD_CONST_UNDER_INLINE, (ptr/4, old -- value, new)) {
+            new = old;
+            DEAD(old);
+            value = PyStackRef_FromPyObjectNew(ptr);
+        }
+
+        tier2 op(_LOAD_CONST_UNDER_INLINE_BORROW, (ptr/4, old -- value, new)) {
+            new = old;
+            DEAD(old);
+            value = PyStackRef_FromPyObjectBorrow(ptr);
+        }
+
         tier2 op(_CHECK_FUNCTION, (func_version/2 -- )) {
             assert(PyStackRef_FunctionCheck(frame->f_funcobj));
             PyFunctionObject *func = (PyFunctionObject *)PyStackRef_AsPyObjectBorrow(frame->f_funcobj);