GH-131798: Narrow types more aggressively in the JIT (GH-134373)

6 files changed, 117 insertions, 50 deletions

diff --git a/Python/bytecodes.c b/Python/bytecodes.c
index 65f411f..a236702 100644
--- a/Python/bytecodes.c
+++ b/Python/bytecodes.c
@@ -344,6 +344,11 @@ dummy_func(
             PyStackRef_CLOSE(value);
         }
 
+        tier2 op(_POP_TWO, (nos, tos --)) {
+            PyStackRef_CLOSE(tos);
+            PyStackRef_CLOSE(nos);
+        }
+
         pure inst(PUSH_NULL, (-- res)) {
             res = PyStackRef_NULL;
         }
diff --git a/Python/executor_cases.c.h b/Python/executor_cases.c.h
index cc521bd..1c8239f 100644
--- a/Python/executor_cases.c.h
+++ b/Python/executor_cases.c.h
@@ -539,6 +539,24 @@
             break;
         }
 
+        case _POP_TWO: {
+            _PyStackRef tos;
+            _PyStackRef nos;
+            tos = stack_pointer[-1];
+            nos = stack_pointer[-2];
+            stack_pointer += -1;
+            assert(WITHIN_STACK_BOUNDS());
+            _PyFrame_SetStackPointer(frame, stack_pointer);
+            PyStackRef_CLOSE(tos);
+            stack_pointer = _PyFrame_GetStackPointer(frame);
+            stack_pointer += -1;
+            assert(WITHIN_STACK_BOUNDS());
+            _PyFrame_SetStackPointer(frame, stack_pointer);
+            PyStackRef_CLOSE(nos);
+            stack_pointer = _PyFrame_GetStackPointer(frame);
+            break;
+        }
+
         case _PUSH_NULL: {
             _PyStackRef res;
             res = PyStackRef_NULL;
diff --git a/Python/optimizer_analysis.c b/Python/optimizer_analysis.c
index 8b0bd1e..53ab289 100644
--- a/Python/optimizer_analysis.c
+++ b/Python/optimizer_analysis.c
@@ -523,6 +523,25 @@ error:
 
 }
 
+const uint16_t op_without_push[MAX_UOP_ID + 1] = {
+    [_COPY] = _NOP,
+    [_LOAD_CONST_INLINE] = _NOP,
+    [_LOAD_CONST_INLINE_BORROW] = _NOP,
+    [_LOAD_FAST] = _NOP,
+    [_LOAD_FAST_BORROW] = _NOP,
+    [_LOAD_SMALL_INT] = _NOP,
+    [_POP_TOP_LOAD_CONST_INLINE] = _POP_TOP,
+    [_POP_TOP_LOAD_CONST_INLINE_BORROW] = _POP_TOP,
+    [_POP_TWO_LOAD_CONST_INLINE_BORROW] = _POP_TWO,
+};
+
+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_LOAD_CONST_INLINE_BORROW] = _POP_TOP_LOAD_CONST_INLINE_BORROW,
+};
+
 
 static int
 remove_unneeded_uops(_PyUOpInstruction *buffer, int buffer_size)
@@ -551,50 +570,23 @@ remove_unneeded_uops(_PyUOpInstruction *buffer, int buffer_size)
                     buffer[pc].opcode = _NOP;
                 }
                 break;
-            case _POP_TOP:
-            case _POP_TOP_LOAD_CONST_INLINE:
-            case _POP_TOP_LOAD_CONST_INLINE_BORROW:
-            case _POP_TWO_LOAD_CONST_INLINE_BORROW:
-            optimize_pop_top_again:
+            default:
             {
-                _PyUOpInstruction *last = &buffer[pc-1];
-                while (last->opcode == _NOP) {
-                    last--;
-                }
-                switch (last->opcode) {
-                    case _POP_TWO_LOAD_CONST_INLINE_BORROW:
-                        last->opcode = _POP_TOP;
+                // Cancel out pushes and pops, repeatedly. So:
+                //     _LOAD_FAST + _POP_TWO_LOAD_CONST_INLINE_BORROW + _POP_TOP
+                // ...becomes:
+                //     _NOP + _POP_TOP + _NOP
+                while (op_without_pop[opcode]) {
+                    _PyUOpInstruction *last = &buffer[pc - 1];
+                    while (last->opcode == _NOP) {
+                        last--;
+                    }
+                    if (!op_without_push[last->opcode]) {
                         break;
-                    case _POP_TOP_LOAD_CONST_INLINE:
-                    case _POP_TOP_LOAD_CONST_INLINE_BORROW:
-                        last->opcode = _NOP;
-                        goto optimize_pop_top_again;
-                    case _COPY:
-                    case _LOAD_CONST_INLINE:
-                    case _LOAD_CONST_INLINE_BORROW:
-                    case _LOAD_FAST:
-                    case _LOAD_FAST_BORROW:
-                    case _LOAD_SMALL_INT:
-                        last->opcode = _NOP;
-                        if (opcode == _POP_TOP) {
-                            opcode = buffer[pc].opcode = _NOP;
-                        }
-                        else if (opcode == _POP_TOP_LOAD_CONST_INLINE) {
-                            opcode = buffer[pc].opcode = _LOAD_CONST_INLINE;
-                        }
-                        else if (opcode == _POP_TOP_LOAD_CONST_INLINE_BORROW) {
-                            opcode = buffer[pc].opcode = _LOAD_CONST_INLINE_BORROW;
-                        }
-                        else {
-                            assert(opcode == _POP_TWO_LOAD_CONST_INLINE_BORROW);
-                            opcode = buffer[pc].opcode = _POP_TOP_LOAD_CONST_INLINE_BORROW;
-                            goto optimize_pop_top_again;
-                        }
+                    }
+                    last->opcode = op_without_push[last->opcode];
+                    opcode = buffer[pc].opcode = op_without_pop[opcode];
                 }
-                _Py_FALLTHROUGH;
-            }
-            default:
-            {
                 /* _PUSH_FRAME doesn't escape or error, but it
                  * does need the IP for the return address */
                 bool needs_ip = opcode == _PUSH_FRAME;
diff --git a/Python/optimizer_bytecodes.c b/Python/optimizer_bytecodes.c
index 639b4b7..f12cd7b 100644
--- a/Python/optimizer_bytecodes.c
+++ b/Python/optimizer_bytecodes.c
@@ -118,6 +118,18 @@ dummy_func(void) {
         sym_set_type(left, &PyLong_Type);
     }
 
+    op(_CHECK_ATTR_CLASS, (type_version/2, owner -- owner)) {
+        PyObject *type = (PyObject *)_PyType_LookupByVersion(type_version);
+        if (type) {
+            if (type == sym_get_const(ctx, owner)) {
+                REPLACE_OP(this_instr, _NOP, 0, 0);
+            }
+            else {
+                sym_set_const(owner, type);
+            }
+        }
+    }
+
     op(_GUARD_TYPE_VERSION, (type_version/2, owner -- owner)) {
         assert(type_version);
         if (sym_matches_type_version(owner, type_version)) {
diff --git a/Python/optimizer_cases.c.h b/Python/optimizer_cases.c.h
index 0a539b2..602f5e2 100644
--- a/Python/optimizer_cases.c.h
+++ b/Python/optimizer_cases.c.h
@@ -102,6 +102,12 @@
             break;
         }
 
+        case _POP_TWO: {
+            stack_pointer += -2;
+            assert(WITHIN_STACK_BOUNDS());
+            break;
+        }
+
         case _PUSH_NULL: {
             JitOptSymbol *res;
             res = sym_new_null(ctx);
@@ -1259,6 +1265,18 @@
         }
 
         case _CHECK_ATTR_CLASS: {
+            JitOptSymbol *owner;
+            owner = stack_pointer[-1];
+            uint32_t type_version = (uint32_t)this_instr->operand0;
+            PyObject *type = (PyObject *)_PyType_LookupByVersion(type_version);
+            if (type) {
+                if (type == sym_get_const(ctx, owner)) {
+                    REPLACE_OP(this_instr, _NOP, 0, 0);
+                }
+                else {
+                    sym_set_const(owner, type);
+                }
+            }
             break;
         }
 
diff --git a/Python/optimizer_symbols.c b/Python/optimizer_symbols.c
index e8a4f87..2e619fa 100644
--- a/Python/optimizer_symbols.c
+++ b/Python/optimizer_symbols.c
@@ -200,6 +200,10 @@ _Py_uop_sym_set_type(JitOptContext *ctx, JitOptSymbol *sym, PyTypeObject *typ)
 bool
 _Py_uop_sym_set_type_version(JitOptContext *ctx, JitOptSymbol *sym, unsigned int version)
 {
+    PyTypeObject *type = _PyType_LookupByVersion(version);
+    if (type) {
+        _Py_uop_sym_set_type(ctx, sym, type);
+    }
     JitSymType tag = sym->tag;
     switch(tag) {
         case JIT_SYM_NULL_TAG:
@@ -215,18 +219,24 @@ _Py_uop_sym_set_type_version(JitOptContext *ctx, JitOptSymbol *sym, unsigned int
                 return true;
             }
         case JIT_SYM_KNOWN_VALUE_TAG:
-            Py_CLEAR(sym->value.value);
-            sym_set_bottom(ctx, sym);
-            return false;
+            if (Py_TYPE(sym->value.value)->tp_version_tag != version) {
+                Py_CLEAR(sym->value.value);
+                sym_set_bottom(ctx, sym);
+                return false;
+            };
+            return true;
         case JIT_SYM_TUPLE_TAG:
-            sym_set_bottom(ctx, sym);
-            return false;
+            if (PyTuple_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) {
-                return true;
+            if (sym->version.version != version) {
+                sym_set_bottom(ctx, sym);
+                return false;
             }
-            sym_set_bottom(ctx, sym);
-            return false;
+            return true;
         case JIT_SYM_BOTTOM_TAG:
             return false;
         case JIT_SYM_NON_NULL_TAG:
@@ -266,6 +276,18 @@ _Py_uop_sym_set_const(JitOptContext *ctx, JitOptSymbol *sym, PyObject *const_val
             }
             return;
         case JIT_SYM_TUPLE_TAG:
+            if (PyTuple_CheckExact(const_val)) {
+                Py_ssize_t len = _Py_uop_sym_tuple_length(sym);
+                if (len == PyTuple_GET_SIZE(const_val)) {
+                    for (Py_ssize_t i = 0; i < len; i++) {
+                        JitOptSymbol *sym_item = _Py_uop_sym_tuple_getitem(ctx, sym, i);
+                        PyObject *item = PyTuple_GET_ITEM(const_val, i);
+                        _Py_uop_sym_set_const(ctx, sym_item, item);
+                    }
+                    make_const(sym, const_val);
+                    return;
+                }
+            }
             sym_set_bottom(ctx, sym);
             return;
         case JIT_SYM_TYPE_VERSION_TAG: