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

6 files changed, 227 insertions, 48 deletions

diff --git a/Python/bytecodes.c b/Python/bytecodes.c
index a2367026cde..652bda9c182 100644
--- a/Python/bytecodes.c
+++ b/Python/bytecodes.c
@@ -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);
diff --git a/Python/executor_cases.c.h b/Python/executor_cases.c.h
index 1c8239f38ee..fcde31a3012 100644
--- a/Python/executor_cases.c.h
+++ b/Python/executor_cases.c.h
@@ -7105,6 +7105,36 @@
             break;
         }
 
+        case _LOAD_CONST_UNDER_INLINE: {
+            _PyStackRef old;
+            _PyStackRef value;
+            _PyStackRef new;
+            old = stack_pointer[-1];
+            PyObject *ptr = (PyObject *)CURRENT_OPERAND0();
+            new = old;
+            value = PyStackRef_FromPyObjectNew(ptr);
+            stack_pointer[-1] = value;
+            stack_pointer[0] = new;
+            stack_pointer += 1;
+            assert(WITHIN_STACK_BOUNDS());
+            break;
+        }
+
+        case _LOAD_CONST_UNDER_INLINE_BORROW: {
+            _PyStackRef old;
+            _PyStackRef value;
+            _PyStackRef new;
+            old = stack_pointer[-1];
+            PyObject *ptr = (PyObject *)CURRENT_OPERAND0();
+            new = old;
+            value = PyStackRef_FromPyObjectBorrow(ptr);
+            stack_pointer[-1] = value;
+            stack_pointer[0] = new;
+            stack_pointer += 1;
+            assert(WITHIN_STACK_BOUNDS());
+            break;
+        }
+
         case _CHECK_FUNCTION: {
             uint32_t func_version = (uint32_t)CURRENT_OPERAND0();
             assert(PyStackRef_FunctionCheck(frame->f_funcobj));
diff --git a/Python/optimizer_analysis.c b/Python/optimizer_analysis.c
index 53ab289b75c..5c50228a13b 100644
--- a/Python/optimizer_analysis.c
+++ b/Python/optimizer_analysis.c
@@ -375,6 +375,23 @@ eliminate_pop_guard(_PyUOpInstruction *this_instr, bool exit)
     }
 }
 
+static JitOptSymbol *
+lookup_attr(JitOptContext *ctx, _PyUOpInstruction *this_instr,
+            PyTypeObject *type, PyObject *name, uint16_t immortal,
+            uint16_t mortal)
+{
+    // The cached value may be dead, so we need to do the lookup again... :(
+    if (type && PyType_Check(type)) {
+        PyObject *lookup = _PyType_Lookup(type, name);
+        if (lookup) {
+            int opcode = _Py_IsImmortal(lookup) ? immortal : mortal;
+            REPLACE_OP(this_instr, opcode, 0, (uintptr_t)lookup);
+            return sym_new_const(ctx, lookup);
+        }
+    }
+    return sym_new_not_null(ctx);
+}
+
 /* _PUSH_FRAME/_RETURN_VALUE's operand can be 0, a PyFunctionObject *, or a
  * PyCodeObject *. Retrieve the code object if possible.
  */
@@ -527,6 +544,8 @@ const uint16_t op_without_push[MAX_UOP_ID + 1] = {
     [_COPY] = _NOP,
     [_LOAD_CONST_INLINE] = _NOP,
     [_LOAD_CONST_INLINE_BORROW] = _NOP,
+    [_LOAD_CONST_UNDER_INLINE] = _POP_TOP_LOAD_CONST_INLINE,
+    [_LOAD_CONST_UNDER_INLINE_BORROW] = _POP_TOP_LOAD_CONST_INLINE_BORROW,
     [_LOAD_FAST] = _NOP,
     [_LOAD_FAST_BORROW] = _NOP,
     [_LOAD_SMALL_INT] = _NOP,
@@ -535,10 +554,16 @@ const uint16_t op_without_push[MAX_UOP_ID + 1] = {
     [_POP_TWO_LOAD_CONST_INLINE_BORROW] = _POP_TWO,
 };
 
+const bool op_skip[MAX_UOP_ID + 1] = {
+    [_NOP] = true,
+    [_CHECK_VALIDITY] = true,
+};
+
 const uint16_t op_without_pop[MAX_UOP_ID + 1] = {
     [_POP_TOP] = _NOP,
     [_POP_TOP_LOAD_CONST_INLINE] = _LOAD_CONST_INLINE,
     [_POP_TOP_LOAD_CONST_INLINE_BORROW] = _LOAD_CONST_INLINE_BORROW,
+    [_POP_TWO] = _POP_TOP,
     [_POP_TWO_LOAD_CONST_INLINE_BORROW] = _POP_TOP_LOAD_CONST_INLINE_BORROW,
 };
 
@@ -578,7 +603,7 @@ remove_unneeded_uops(_PyUOpInstruction *buffer, int buffer_size)
                 //     _NOP + _POP_TOP + _NOP
                 while (op_without_pop[opcode]) {
                     _PyUOpInstruction *last = &buffer[pc - 1];
-                    while (last->opcode == _NOP) {
+                    while (op_skip[last->opcode]) {
                         last--;
                     }
                     if (!op_without_push[last->opcode]) {
@@ -586,6 +611,10 @@ remove_unneeded_uops(_PyUOpInstruction *buffer, int buffer_size)
                     }
                     last->opcode = op_without_push[last->opcode];
                     opcode = buffer[pc].opcode = op_without_pop[opcode];
+                    if (op_without_pop[last->opcode]) {
+                        opcode = last->opcode;
+                        pc = last - buffer;
+                    }
                 }
                 /* _PUSH_FRAME doesn't escape or error, but it
                  * does need the IP for the return address */
diff --git a/Python/optimizer_bytecodes.c b/Python/optimizer_bytecodes.c
index b9ebd8678e0..0b6bbd133d6 100644
--- a/Python/optimizer_bytecodes.c
+++ b/Python/optimizer_bytecodes.c
@@ -522,7 +522,7 @@ dummy_func(void) {
     }
 
     op(_LOAD_CONST, (-- value)) {
-        PyObject *val = PyTuple_GET_ITEM(co->co_consts, this_instr->oparg);
+        PyObject *val = PyTuple_GET_ITEM(co->co_consts, oparg);
         REPLACE_OP(this_instr, _LOAD_CONST_INLINE_BORROW, 0, (uintptr_t)val);
         value = sym_new_const(ctx, val);
     }
@@ -608,7 +608,7 @@ dummy_func(void) {
     op(_LOAD_ATTR, (owner -- attr, self_or_null[oparg&1])) {
         (void)owner;
         attr = sym_new_not_null(ctx);
-        if (oparg &1) {
+        if (oparg & 1) {
             self_or_null[0] = sym_new_unknown(ctx);
         }
     }
@@ -624,25 +624,59 @@ dummy_func(void) {
     }
 
     op(_LOAD_ATTR_CLASS, (descr/4, owner -- attr)) {
-        attr = sym_new_not_null(ctx);
         (void)descr;
+        PyTypeObject *type = (PyTypeObject *)sym_get_const(ctx, owner);
+        PyObject *name = PyTuple_GET_ITEM(co->co_names, oparg >> 1);
+        attr = lookup_attr(ctx, this_instr, type, name,
+                           _POP_TOP_LOAD_CONST_INLINE_BORROW,
+                           _POP_TOP_LOAD_CONST_INLINE);
+    }
+
+    op(_LOAD_ATTR_NONDESCRIPTOR_WITH_VALUES, (descr/4, owner -- attr)) {
+        (void)descr;
+        PyTypeObject *type = sym_get_type(owner);
+        PyObject *name = PyTuple_GET_ITEM(co->co_names, oparg >> 1);
+        attr = lookup_attr(ctx, this_instr, type, name,
+                           _POP_TOP_LOAD_CONST_INLINE_BORROW,
+                           _POP_TOP_LOAD_CONST_INLINE);
+    }
+
+    op(_LOAD_ATTR_NONDESCRIPTOR_NO_DICT, (descr/4, owner -- attr)) {
+        (void)descr;
+        PyTypeObject *type = sym_get_type(owner);
+        PyObject *name = PyTuple_GET_ITEM(co->co_names, oparg >> 1);
+        attr = lookup_attr(ctx, this_instr, type, name,
+                           _POP_TOP_LOAD_CONST_INLINE_BORROW,
+                           _POP_TOP_LOAD_CONST_INLINE);
     }
 
     op(_LOAD_ATTR_METHOD_WITH_VALUES, (descr/4, owner -- attr, self)) {
         (void)descr;
-        attr = sym_new_not_null(ctx);
+        PyTypeObject *type = sym_get_type(owner);
+        PyObject *name = PyTuple_GET_ITEM(co->co_names, oparg >> 1);
+        attr = lookup_attr(ctx, this_instr, type, name,
+                           _LOAD_CONST_UNDER_INLINE_BORROW,
+                           _LOAD_CONST_UNDER_INLINE);
         self = owner;
     }
 
     op(_LOAD_ATTR_METHOD_NO_DICT, (descr/4, owner -- attr, self)) {
         (void)descr;
-        attr = sym_new_not_null(ctx);
+        PyTypeObject *type = sym_get_type(owner);
+        PyObject *name = PyTuple_GET_ITEM(co->co_names, oparg >> 1);
+        attr = lookup_attr(ctx, this_instr, type, name,
+                           _LOAD_CONST_UNDER_INLINE_BORROW,
+                           _LOAD_CONST_UNDER_INLINE);
         self = owner;
     }
 
     op(_LOAD_ATTR_METHOD_LAZY_DICT, (descr/4, owner -- attr, self)) {
         (void)descr;
-        attr = sym_new_not_null(ctx);
+        PyTypeObject *type = sym_get_type(owner);
+        PyObject *name = PyTuple_GET_ITEM(co->co_names, oparg >> 1);
+        attr = lookup_attr(ctx, this_instr, type, name,
+                           _LOAD_CONST_UNDER_INLINE_BORROW,
+                           _LOAD_CONST_UNDER_INLINE);
         self = owner;
     }
 
diff --git a/Python/optimizer_cases.c.h b/Python/optimizer_cases.c.h
index 0ba45e1f58f..5a9fcf3b1b6 100644
--- a/Python/optimizer_cases.c.h
+++ b/Python/optimizer_cases.c.h
@@ -68,7 +68,7 @@
 
         case _LOAD_CONST: {
             JitOptSymbol *value;
-            PyObject *val = PyTuple_GET_ITEM(co->co_consts, this_instr->oparg);
+            PyObject *val = PyTuple_GET_ITEM(co->co_consts, oparg);
             REPLACE_OP(this_instr, _LOAD_CONST_INLINE_BORROW, 0, (uintptr_t)val);
             value = sym_new_const(ctx, val);
             stack_pointer[0] = value;
@@ -1174,7 +1174,7 @@
             self_or_null = &stack_pointer[0];
             (void)owner;
             attr = sym_new_not_null(ctx);
-            if (oparg &1) {
+            if (oparg & 1) {
                 self_or_null[0] = sym_new_unknown(ctx);
             }
             stack_pointer[-1] = attr;
@@ -1284,10 +1284,16 @@
         }
 
         case _LOAD_ATTR_CLASS: {
+            JitOptSymbol *owner;
             JitOptSymbol *attr;
+            owner = stack_pointer[-1];
             PyObject *descr = (PyObject *)this_instr->operand0;
-            attr = sym_new_not_null(ctx);
             (void)descr;
+            PyTypeObject *type = (PyTypeObject *)sym_get_const(ctx, owner);
+            PyObject *name = PyTuple_GET_ITEM(co->co_names, oparg >> 1);
+            attr = lookup_attr(ctx, this_instr, type, name,
+                           _POP_TOP_LOAD_CONST_INLINE_BORROW,
+                           _POP_TOP_LOAD_CONST_INLINE);
             stack_pointer[-1] = attr;
             break;
         }
@@ -1701,7 +1707,11 @@
             owner = stack_pointer[-1];
             PyObject *descr = (PyObject *)this_instr->operand0;
             (void)descr;
-            attr = sym_new_not_null(ctx);
+            PyTypeObject *type = sym_get_type(owner);
+            PyObject *name = PyTuple_GET_ITEM(co->co_names, oparg >> 1);
+            attr = lookup_attr(ctx, this_instr, type, name,
+                           _LOAD_CONST_UNDER_INLINE_BORROW,
+                           _LOAD_CONST_UNDER_INLINE);
             self = owner;
             stack_pointer[-1] = attr;
             stack_pointer[0] = self;
@@ -1717,7 +1727,11 @@
             owner = stack_pointer[-1];
             PyObject *descr = (PyObject *)this_instr->operand0;
             (void)descr;
-            attr = sym_new_not_null(ctx);
+            PyTypeObject *type = sym_get_type(owner);
+            PyObject *name = PyTuple_GET_ITEM(co->co_names, oparg >> 1);
+            attr = lookup_attr(ctx, this_instr, type, name,
+                           _LOAD_CONST_UNDER_INLINE_BORROW,
+                           _LOAD_CONST_UNDER_INLINE);
             self = owner;
             stack_pointer[-1] = attr;
             stack_pointer[0] = self;
@@ -1727,15 +1741,31 @@
         }
 
         case _LOAD_ATTR_NONDESCRIPTOR_WITH_VALUES: {
+            JitOptSymbol *owner;
             JitOptSymbol *attr;
-            attr = sym_new_not_null(ctx);
+            owner = stack_pointer[-1];
+            PyObject *descr = (PyObject *)this_instr->operand0;
+            (void)descr;
+            PyTypeObject *type = sym_get_type(owner);
+            PyObject *name = PyTuple_GET_ITEM(co->co_names, oparg >> 1);
+            attr = lookup_attr(ctx, this_instr, type, name,
+                           _POP_TOP_LOAD_CONST_INLINE_BORROW,
+                           _POP_TOP_LOAD_CONST_INLINE);
             stack_pointer[-1] = attr;
             break;
         }
 
         case _LOAD_ATTR_NONDESCRIPTOR_NO_DICT: {
+            JitOptSymbol *owner;
             JitOptSymbol *attr;
-            attr = sym_new_not_null(ctx);
+            owner = stack_pointer[-1];
+            PyObject *descr = (PyObject *)this_instr->operand0;
+            (void)descr;
+            PyTypeObject *type = sym_get_type(owner);
+            PyObject *name = PyTuple_GET_ITEM(co->co_names, oparg >> 1);
+            attr = lookup_attr(ctx, this_instr, type, name,
+                           _POP_TOP_LOAD_CONST_INLINE_BORROW,
+                           _POP_TOP_LOAD_CONST_INLINE);
             stack_pointer[-1] = attr;
             break;
         }
@@ -1751,7 +1781,11 @@
             owner = stack_pointer[-1];
             PyObject *descr = (PyObject *)this_instr->operand0;
             (void)descr;
-            attr = sym_new_not_null(ctx);
+            PyTypeObject *type = sym_get_type(owner);
+            PyObject *name = PyTuple_GET_ITEM(co->co_names, oparg >> 1);
+            attr = lookup_attr(ctx, this_instr, type, name,
+                           _LOAD_CONST_UNDER_INLINE_BORROW,
+                           _LOAD_CONST_UNDER_INLINE);
             self = owner;
             stack_pointer[-1] = attr;
             stack_pointer[0] = self;
@@ -2594,6 +2628,30 @@
             break;
         }
 
+        case _LOAD_CONST_UNDER_INLINE: {
+            JitOptSymbol *value;
+            JitOptSymbol *new;
+            value = sym_new_not_null(ctx);
+            new = sym_new_not_null(ctx);
+            stack_pointer[-1] = value;
+            stack_pointer[0] = new;
+            stack_pointer += 1;
+            assert(WITHIN_STACK_BOUNDS());
+            break;
+        }
+
+        case _LOAD_CONST_UNDER_INLINE_BORROW: {
+            JitOptSymbol *value;
+            JitOptSymbol *new;
+            value = sym_new_not_null(ctx);
+            new = sym_new_not_null(ctx);
+            stack_pointer[-1] = value;
+            stack_pointer[0] = new;
+            stack_pointer += 1;
+            assert(WITHIN_STACK_BOUNDS());
+            break;
+        }
+
         case _CHECK_FUNCTION: {
             break;
         }
diff --git a/Python/optimizer_symbols.c b/Python/optimizer_symbols.c
index 2e619fa6f99..25de5d83166 100644
--- a/Python/optimizer_symbols.c
+++ b/Python/optimizer_symbols.c
@@ -13,22 +13,46 @@
 #include <stdint.h>
 #include <stddef.h>
 
-/* Symbols
-   =======
-
-   See the diagram at
-   https://github.com/faster-cpython/ideas/blob/main/3.13/redundancy_eliminator.md
-
-   We represent the nodes in the diagram as follows
-   (the flag bits are only defined in optimizer_symbols.c):
-   - Top: no flag bits, typ and const_val are NULL.
-   - NULL: IS_NULL flag set, type and const_val NULL.
-   - Not NULL: NOT_NULL flag set, type and const_val NULL.
-   - None/not None: not used. (None could be represented as any other constant.)
-   - Known type: NOT_NULL flag set and typ set; const_val is NULL.
-   - Known constant: NOT_NULL flag set, type set, const_val set.
-   - Bottom: IS_NULL and NOT_NULL flags set, type and const_val NULL.
- */
+/*
+
+Symbols
+=======
+
+https://github.com/faster-cpython/ideas/blob/main/3.13/redundancy_eliminator.md
+
+Logically, all symbols begin as UNKNOWN, and can transition downwards along the
+edges of the lattice, but *never* upwards (see the diagram below). The UNKNOWN
+state represents no information, and the BOTTOM state represents contradictory
+information. Though symbols logically progress through all intermediate nodes,
+we often skip in-between states for convenience:
+
+   UNKNOWN
+   |     |
+NULL     |
+|        |                <- Anything below this level is an object.
+|        NON_NULL
+|        |      |         <- Anything below this level has a known type version.
+| TYPE_VERSION  |
+| |             |         <- Anything below this level has a known type.
+| KNOWN_CLASS   |
+| |         |   |         <- Anything below this level has a known truthiness.
+| |         |  TRUTHINESS
+| |         |  |
+| TUPLE     |  |
+|     |     |  |          <- Anything below this level is a known constant.
+|    KNOWN_VALUE
+|    |                    <- Anything below this level is unreachable.
+BOTTOM
+
+For example, after guarding that the type of an UNKNOWN local is int, we can
+narrow the symbol to KNOWN_CLASS (logically progressing though NON_NULL and
+TYPE_VERSION to get there). Later, we may learn that it is falsey based on the
+result of a truth test, which would allow us to narrow the symbol to KNOWN_VALUE
+(with a value of integer zero). If at any point we encounter a float guard on
+the same symbol, that would be a contradiction, and the symbol would be set to
+BOTTOM (indicating that the code is unreachable).
+
+*/
 
 #ifdef Py_DEBUG
 static inline int get_lltrace(void) {
@@ -420,7 +444,6 @@ _Py_uop_sym_get_type(JitOptSymbol *sym)
     JitSymType tag = sym->tag;
     switch(tag) {
         case JIT_SYM_NULL_TAG:
-        case JIT_SYM_TYPE_VERSION_TAG:
         case JIT_SYM_BOTTOM_TAG:
         case JIT_SYM_NON_NULL_TAG:
         case JIT_SYM_UNKNOWN_TAG:
@@ -429,6 +452,8 @@ _Py_uop_sym_get_type(JitOptSymbol *sym)
             return sym->cls.type;
         case JIT_SYM_KNOWN_VALUE_TAG:
             return Py_TYPE(sym->value.value);
+        case JIT_SYM_TYPE_VERSION_TAG:
+            return _PyType_LookupByVersion(sym->version.version);
         case JIT_SYM_TUPLE_TAG:
             return &PyTuple_Type;
         case JIT_SYM_TRUTHINESS_TAG:
@@ -464,21 +489,7 @@ _Py_uop_sym_get_type_version(JitOptSymbol *sym)
 bool
 _Py_uop_sym_has_type(JitOptSymbol *sym)
 {
-    JitSymType tag = sym->tag;
-    switch(tag) {
-        case JIT_SYM_NULL_TAG:
-        case JIT_SYM_TYPE_VERSION_TAG:
-        case JIT_SYM_BOTTOM_TAG:
-        case JIT_SYM_NON_NULL_TAG:
-        case JIT_SYM_UNKNOWN_TAG:
-            return false;
-        case JIT_SYM_KNOWN_CLASS_TAG:
-        case JIT_SYM_KNOWN_VALUE_TAG:
-        case JIT_SYM_TUPLE_TAG:
-        case JIT_SYM_TRUTHINESS_TAG:
-            return true;
-    }
-    Py_UNREACHABLE();
+    return _Py_uop_sym_get_type(sym) != NULL;
 }
 
 bool
@@ -576,7 +587,7 @@ _Py_uop_sym_tuple_getitem(JitOptContext *ctx, JitOptSymbol *sym, int item)
     else if (sym->tag == JIT_SYM_TUPLE_TAG && item < sym->tuple.length) {
         return allocation_base(ctx) + sym->tuple.items[item];
     }
-    return _Py_uop_sym_new_unknown(ctx);
+    return _Py_uop_sym_new_not_null(ctx);
 }
 
 int
@@ -863,6 +874,11 @@ _Py_uop_symbols_test(PyObject *Py_UNUSED(self), PyObject *Py_UNUSED(ignored))
         _Py_uop_sym_get_const(ctx, _Py_uop_sym_tuple_getitem(ctx, sym, 1)) == val_43,
         "tuple item does not match value used to create tuple"
     );
+    sym = _Py_uop_sym_new_type(ctx, &PyTuple_Type);
+    TEST_PREDICATE(
+        _Py_uop_sym_is_not_null(_Py_uop_sym_tuple_getitem(ctx, sym, 42)),
+        "Unknown tuple item is not narrowed to non-NULL"
+    );
     JitOptSymbol *value = _Py_uop_sym_new_type(ctx, &PyBool_Type);
     sym = _Py_uop_sym_new_truthiness(ctx, value, false);
     TEST_PREDICATE(_Py_uop_sym_matches_type(sym, &PyBool_Type), "truthiness is not boolean");