GH-130415: Optimize constant comparison in JIT builds (GH-131489)

Include/internal/pycore_uop_ids.h

@@ -241,50 +241,51 @@ extern "C" {
 #define _POP_TOP POP_TOP
 #define _POP_TOP_LOAD_CONST_INLINE 447
 #define _POP_TOP_LOAD_CONST_INLINE_BORROW 448
+#define _POP_TWO_LOAD_CONST_INLINE_BORROW 449
 #define _PUSH_EXC_INFO PUSH_EXC_INFO
-#define _PUSH_FRAME 449
+#define _PUSH_FRAME 450
 #define _PUSH_NULL PUSH_NULL
-#define _PUSH_NULL_CONDITIONAL 450
-#define _PY_FRAME_GENERAL 451
-#define _PY_FRAME_KW 452
-#define _QUICKEN_RESUME 453
-#define _REPLACE_WITH_TRUE 454
+#define _PUSH_NULL_CONDITIONAL 451
+#define _PY_FRAME_GENERAL 452
+#define _PY_FRAME_KW 453
+#define _QUICKEN_RESUME 454
+#define _REPLACE_WITH_TRUE 455
 #define _RESUME_CHECK RESUME_CHECK
 #define _RETURN_GENERATOR RETURN_GENERATOR
 #define _RETURN_VALUE RETURN_VALUE
-#define _SAVE_RETURN_OFFSET 455
-#define _SEND 456
-#define _SEND_GEN_FRAME 457
+#define _SAVE_RETURN_OFFSET 456
+#define _SEND 457
+#define _SEND_GEN_FRAME 458
 #define _SETUP_ANNOTATIONS SETUP_ANNOTATIONS
 #define _SET_ADD SET_ADD
 #define _SET_FUNCTION_ATTRIBUTE SET_FUNCTION_ATTRIBUTE
 #define _SET_UPDATE SET_UPDATE
-#define _START_EXECUTOR 458
-#define _STORE_ATTR 459
-#define _STORE_ATTR_INSTANCE_VALUE 460
-#define _STORE_ATTR_SLOT 461
-#define _STORE_ATTR_WITH_HINT 462
+#define _START_EXECUTOR 459
+#define _STORE_ATTR 460
+#define _STORE_ATTR_INSTANCE_VALUE 461
+#define _STORE_ATTR_SLOT 462
+#define _STORE_ATTR_WITH_HINT 463
 #define _STORE_DEREF STORE_DEREF
-#define _STORE_FAST 463
-#define _STORE_FAST_0 464
-#define _STORE_FAST_1 465
-#define _STORE_FAST_2 466
-#define _STORE_FAST_3 467
-#define _STORE_FAST_4 468
-#define _STORE_FAST_5 469
-#define _STORE_FAST_6 470
-#define _STORE_FAST_7 471
+#define _STORE_FAST 464
+#define _STORE_FAST_0 465
+#define _STORE_FAST_1 466
+#define _STORE_FAST_2 467
+#define _STORE_FAST_3 468
+#define _STORE_FAST_4 469
+#define _STORE_FAST_5 470
+#define _STORE_FAST_6 471
+#define _STORE_FAST_7 472
 #define _STORE_FAST_LOAD_FAST STORE_FAST_LOAD_FAST
 #define _STORE_FAST_STORE_FAST STORE_FAST_STORE_FAST
 #define _STORE_GLOBAL STORE_GLOBAL
 #define _STORE_NAME STORE_NAME
-#define _STORE_SLICE 472
-#define _STORE_SUBSCR 473
+#define _STORE_SLICE 473
+#define _STORE_SUBSCR 474
 #define _STORE_SUBSCR_DICT STORE_SUBSCR_DICT
 #define _STORE_SUBSCR_LIST_INT STORE_SUBSCR_LIST_INT
 #define _SWAP SWAP
-#define _TIER2_RESUME_CHECK 474
-#define _TO_BOOL 475
+#define _TIER2_RESUME_CHECK 475
+#define _TO_BOOL 476
 #define _TO_BOOL_BOOL TO_BOOL_BOOL
 #define _TO_BOOL_INT TO_BOOL_INT
 #define _TO_BOOL_LIST TO_BOOL_LIST
@@ -294,13 +295,13 @@ extern "C" {
 #define _UNARY_NEGATIVE UNARY_NEGATIVE
 #define _UNARY_NOT UNARY_NOT
 #define _UNPACK_EX UNPACK_EX
-#define _UNPACK_SEQUENCE 476
+#define _UNPACK_SEQUENCE 477
 #define _UNPACK_SEQUENCE_LIST UNPACK_SEQUENCE_LIST
 #define _UNPACK_SEQUENCE_TUPLE UNPACK_SEQUENCE_TUPLE
 #define _UNPACK_SEQUENCE_TWO_TUPLE UNPACK_SEQUENCE_TWO_TUPLE
 #define _WITH_EXCEPT_START WITH_EXCEPT_START
 #define _YIELD_VALUE YIELD_VALUE
-#define MAX_UOP_ID 476
+#define MAX_UOP_ID 477
 
 #ifdef __cplusplus
 }

Include/internal/pycore_uop_metadata.h

@@ -273,6 +273,7 @@ const uint16_t _PyUop_Flags[MAX_UOP_ID+1] = {
     [_POP_TOP_LOAD_CONST_INLINE] = HAS_ESCAPES_FLAG | HAS_PURE_FLAG,
     [_LOAD_CONST_INLINE_BORROW] = HAS_PURE_FLAG,
     [_POP_TOP_LOAD_CONST_INLINE_BORROW] = HAS_ESCAPES_FLAG | HAS_PURE_FLAG,
+    [_POP_TWO_LOAD_CONST_INLINE_BORROW] = HAS_ESCAPES_FLAG | HAS_PURE_FLAG,
     [_CHECK_FUNCTION] = HAS_DEOPT_FLAG,
     [_START_EXECUTOR] = HAS_ESCAPES_FLAG,
     [_MAKE_WARM] = 0,
@@ -495,6 +496,7 @@ const char *const _PyOpcode_uop_name[MAX_UOP_ID+1] = {
     [_POP_TOP] = "_POP_TOP",
     [_POP_TOP_LOAD_CONST_INLINE] = "_POP_TOP_LOAD_CONST_INLINE",
     [_POP_TOP_LOAD_CONST_INLINE_BORROW] = "_POP_TOP_LOAD_CONST_INLINE_BORROW",
+    [_POP_TWO_LOAD_CONST_INLINE_BORROW] = "_POP_TWO_LOAD_CONST_INLINE_BORROW",
     [_PUSH_EXC_INFO] = "_PUSH_EXC_INFO",
     [_PUSH_FRAME] = "_PUSH_FRAME",
     [_PUSH_NULL] = "_PUSH_NULL",
@@ -1065,6 +1067,8 @@ int _PyUop_num_popped(int opcode, int oparg)
             return 0;
         case _POP_TOP_LOAD_CONST_INLINE_BORROW:
             return 1;
+        case _POP_TWO_LOAD_CONST_INLINE_BORROW:
+            return 2;
         case _CHECK_FUNCTION:
             return 0;
         case _START_EXECUTOR:

Lib/test/test_capi/test_opt.py

@@ -1531,7 +1531,7 @@ class TestUopsOptimization(unittest.TestCase):
         # But all of the appends we care about are still there:
         self.assertEqual(uops.count("_CALL_LIST_APPEND"), len("ABCDEFG"))
 
-def test_narrow_type_to_constant_str_empty(self):
+    def test_narrow_type_to_constant_str_empty(self):
         def f(n):
             trace = []
             for i in range(n):
@@ -1564,6 +1564,23 @@ def test_narrow_type_to_constant_str_empty(self):
         # But all of the appends we care about are still there:
         self.assertEqual(uops.count("_CALL_LIST_APPEND"), len("ABCDEFG"))
 
+    def test_compare_pop_two_load_const_inline_borrow(self):
+        def testfunc(n):
+            x = 0
+            for _ in range(n):
+                a = 10
+                b = 10
+                if a == b:
+                    x += 1
+            return x
+
+        res, ex = self._run_with_optimizer(testfunc, TIER2_THRESHOLD)
+        self.assertEqual(res, TIER2_THRESHOLD)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+        self.assertNotIn("_COMPARE_OP_INT", uops)
+        self.assertIn("_POP_TWO_LOAD_CONST_INLINE_BORROW", uops)
+
 def global_identity(x):
     return x
 

Misc/NEWS.d/next/Core_and_Builtins/2025-03-20-03-45-20.gh-issue-130415.VtJhoG.rst

@@ -0,0 +1 @@
+Optimize comparison of two constants in JIT builds

Python/bytecodes.c

@@ -5133,6 +5133,12 @@ dummy_func(
             value = PyStackRef_FromPyObjectImmortal(ptr);
         }
 
+        tier2 pure op(_POP_TWO_LOAD_CONST_INLINE_BORROW, (ptr/4, pop1, pop2 -- value)) {
+            PyStackRef_CLOSE(pop2);
+            PyStackRef_CLOSE(pop1);
+            value = PyStackRef_FromPyObjectImmortal(ptr);
+        }
+
         tier2 op(_CHECK_FUNCTION, (func_version/2 -- )) {
             assert(PyStackRef_FunctionCheck(frame->f_funcobj));
             PyFunctionObject *func = (PyFunctionObject *)PyStackRef_AsPyObjectBorrow(frame->f_funcobj);

Python/executor_cases.c.h

@@ -6925,6 +6925,30 @@
             break;
         }
 
+        case _POP_TWO_LOAD_CONST_INLINE_BORROW: {
+            _PyStackRef pop2;
+            _PyStackRef pop1;
+            _PyStackRef value;
+            pop2 = stack_pointer[-1];
+            pop1 = stack_pointer[-2];
+            PyObject *ptr = (PyObject *)CURRENT_OPERAND0();
+            stack_pointer += -1;
+            assert(WITHIN_STACK_BOUNDS());
+            _PyFrame_SetStackPointer(frame, stack_pointer);
+            PyStackRef_CLOSE(pop2);
+            stack_pointer = _PyFrame_GetStackPointer(frame);
+            stack_pointer += -1;
+            assert(WITHIN_STACK_BOUNDS());
+            _PyFrame_SetStackPointer(frame, stack_pointer);
+            PyStackRef_CLOSE(pop1);
+            stack_pointer = _PyFrame_GetStackPointer(frame);
+            value = PyStackRef_FromPyObjectImmortal(ptr);
+            stack_pointer[0] = value;
+            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));

Python/optimizer_bytecodes.c

@@ -446,7 +446,25 @@ dummy_func(void) {
     }
 
     op(_COMPARE_OP_INT, (left, right -- res)) {
-        res = sym_new_type(ctx, &PyBool_Type);
+        if (sym_is_const(ctx, left) && sym_is_const(ctx, right))
+        {
+            assert(PyLong_CheckExact(sym_get_const(ctx, left)));
+            assert(PyLong_CheckExact(sym_get_const(ctx, right)));
+            PyObject *tmp = PyObject_RichCompare(sym_get_const(ctx, left),
+                                                 sym_get_const(ctx, right),
+                                                 oparg >> 5);
+            if (tmp == NULL) {
+                goto error;
+            }
+            assert(PyBool_Check(tmp));
+            assert(_Py_IsImmortal(tmp));
+            REPLACE_OP(this_instr, _POP_TWO_LOAD_CONST_INLINE_BORROW, 0, (uintptr_t)tmp);
+            res = sym_new_const(ctx, tmp);
+            Py_DECREF(tmp);
+        }
+        else {
+            res = sym_new_type(ctx, &PyBool_Type);
+        }
     }
 
     op(_COMPARE_OP_FLOAT, (left, right -- res)) {

Python/optimizer_cases.c.h

@@ -1281,11 +1281,38 @@
         }
 
         case _COMPARE_OP_INT: {
+            JitOptSymbol *right;
+            JitOptSymbol *left;
             JitOptSymbol *res;
-            res = sym_new_type(ctx, &PyBool_Type);
-            stack_pointer[-2] = res;
-            stack_pointer += -1;
-            assert(WITHIN_STACK_BOUNDS());
+            right = stack_pointer[-1];
+            left = stack_pointer[-2];
+            if (sym_is_const(ctx, left) && sym_is_const(ctx, right))
+            {
+                assert(PyLong_CheckExact(sym_get_const(ctx, left)));
+                assert(PyLong_CheckExact(sym_get_const(ctx, right)));
+                stack_pointer += -2;
+                assert(WITHIN_STACK_BOUNDS());
+                PyObject *tmp = PyObject_RichCompare(sym_get_const(ctx, left),
+                    sym_get_const(ctx, right),
+                    oparg >> 5);
+                if (tmp == NULL) {
+                    goto error;
+                }
+                assert(PyBool_Check(tmp));
+                assert(_Py_IsImmortal(tmp));
+                REPLACE_OP(this_instr, _POP_TWO_LOAD_CONST_INLINE_BORROW, 0, (uintptr_t)tmp);
+                res = sym_new_const(ctx, tmp);
+                stack_pointer[0] = res;
+                stack_pointer += 1;
+                assert(WITHIN_STACK_BOUNDS());
+                Py_DECREF(tmp);
+            }
+            else {
+                res = sym_new_type(ctx, &PyBool_Type);
+                stack_pointer += -1;
+                assert(WITHIN_STACK_BOUNDS());
+            }
+            stack_pointer[-1] = res;
             break;
         }
 
@@ -2392,6 +2419,15 @@
             break;
         }
 
+        case _POP_TWO_LOAD_CONST_INLINE_BORROW: {
+            JitOptSymbol *value;
+            value = sym_new_not_null(ctx);
+            stack_pointer[-2] = value;
+            stack_pointer += -1;
+            assert(WITHIN_STACK_BOUNDS());
+            break;
+        }
+
         case _CHECK_FUNCTION: {
             break;
         }