GH-130415: Improve the JIT's unneeded uop removal pass (GH-132333)

Author: brandtbucher

Lib/test/test_capi/test_opt.py

@@ -1283,7 +1283,7 @@ class Bar:
         load_attr_top = opnames.index("_LOAD_ATTR_NONDESCRIPTOR_WITH_VALUES", 0, call)
         load_attr_bottom = opnames.index("_LOAD_ATTR_NONDESCRIPTOR_WITH_VALUES", call)
         self.assertEqual(opnames[:load_attr_top].count("_GUARD_TYPE_VERSION"), 1)
-        self.assertEqual(opnames[call:load_attr_bottom].count("_CHECK_VALIDITY"), 1)
+        self.assertEqual(opnames[call:load_attr_bottom].count("_CHECK_VALIDITY"), 2)
 
     def test_guard_type_version_removed_escaping(self):
 
@@ -1306,7 +1306,7 @@ class Foo:
         load_attr_top = opnames.index("_LOAD_ATTR_NONDESCRIPTOR_WITH_VALUES", 0, call)
         load_attr_bottom = opnames.index("_LOAD_ATTR_NONDESCRIPTOR_WITH_VALUES", call)
         self.assertEqual(opnames[:load_attr_top].count("_GUARD_TYPE_VERSION"), 1)
-        self.assertEqual(opnames[call:load_attr_bottom].count("_CHECK_VALIDITY"), 1)
+        self.assertEqual(opnames[call:load_attr_bottom].count("_CHECK_VALIDITY"), 2)
 
     def test_guard_type_version_executor_invalidated(self):
         """
@@ -1601,7 +1601,7 @@ def testfunc(n):
         self.assertIsNotNone(ex)
         uops = get_opnames(ex)
         self.assertNotIn("_COMPARE_OP_INT", uops)
-        self.assertIn("_POP_TWO_LOAD_CONST_INLINE_BORROW", uops)
+        self.assertNotIn("_POP_TWO_LOAD_CONST_INLINE_BORROW", uops)
 
     def test_to_bool_bool_contains_op_set(self):
         """

Misc/NEWS.d/next/Core_and_Builtins/2025-04-09-14-05-54.gh-issue-130415.llQtUq.rst

@@ -0,0 +1,3 @@
+Improve the JIT's ability to remove unused constant and local variable
+loads, and fix an issue where deallocating unused values could cause JIT
+code to crash or behave incorrectly.

Python/optimizer_analysis.c

@@ -555,28 +555,47 @@ remove_unneeded_uops(_PyUOpInstruction *buffer, int buffer_size)
                 }
                 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:
             {
                 _PyUOpInstruction *last = &buffer[pc-1];
                 while (last->opcode == _NOP) {
                     last--;
                 }
-                if (last->opcode == _LOAD_CONST_INLINE  ||
-                    last->opcode == _LOAD_CONST_INLINE_BORROW ||
-                    last->opcode == _LOAD_FAST ||
-                    last->opcode == _LOAD_FAST_BORROW ||
-                    last->opcode == _COPY
-                ) {
-                    last->opcode = _NOP;
-                    buffer[pc].opcode = _NOP;
-                }
-                if (last->opcode == _REPLACE_WITH_TRUE) {
-                    last->opcode = _NOP;
+                switch (last->opcode) {
+                    case _POP_TWO_LOAD_CONST_INLINE_BORROW:
+                        last->opcode = _POP_TOP;
+                        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;
+                        }
                 }
-                break;
+                _Py_FALLTHROUGH;
             }
-            case _JUMP_TO_TOP:
-            case _EXIT_TRACE:
-                return pc + 1;
             default:
             {
                 /* _PUSH_FRAME doesn't escape or error, but it
@@ -591,7 +610,11 @@ remove_unneeded_uops(_PyUOpInstruction *buffer, int buffer_size)
                     buffer[last_set_ip].opcode = _SET_IP;
                     last_set_ip = -1;
                 }
+                break;
             }
+            case _JUMP_TO_TOP:
+            case _EXIT_TRACE:
+                return pc + 1;
         }
     }
     Py_UNREACHABLE();

Python/optimizer_bytecodes.c

@@ -912,6 +912,7 @@ dummy_func(void) {
     }
 
     op(_REPLACE_WITH_TRUE, (value -- res)) {
+        REPLACE_OP(this_instr, _POP_TOP_LOAD_CONST_INLINE_BORROW, 0, (uintptr_t)Py_True);
         res = sym_new_const(ctx, Py_True);
     }