From de4c33f5cd7ba6b192b512e1587077fc8a77186c Mon Sep 17 00:00:00 2001 From: reidenong Date: Thu, 22 Jan 2026 20:04:43 +0800 Subject: [PATCH 1/9] Add failing regression tests --- Lib/test/test_capi/test_opt.py | 42 ++++++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/Lib/test/test_capi/test_opt.py b/Lib/test/test_capi/test_opt.py index 7c33320e9f1785..bdf2df2bfd668e 100644 --- a/Lib/test/test_capi/test_opt.py +++ b/Lib/test/test_capi/test_opt.py @@ -890,6 +890,48 @@ def testfunc(n): self.assertLessEqual(len(guard_nos_unicode_count), 1) self.assertIn("_COMPARE_OP_STR", uops) + def test_compare_int_eq_narrows_to_constant(self): + def f(n): + def return_1(): + return 1 + + hits = 0 + v = return_1() + for _ in range(n): + if v == 1: + hits += v + 1 + return hits + + res, ex = self._run_with_optimizer(f, TIER2_THRESHOLD) + self.assertEqual(res, TIER2_THRESHOLD * 2) + self.assertIsNotNone(ex) + uops = get_opnames(ex) + + # v + 1 should be constant folded + self.assertLessEqual(count_ops(ex, "_BINARY_OP_ADD_INT"), 1) + + def test_compare_int_ne_narrows_to_constant(self): + def f(n): + def return_1(): + return 1 + + hits = 0 + v = return_1() + for _ in range(n): + if v != 1: + hits += 1000 + else: + hits += v + 1 + return hits + + res, ex = self._run_with_optimizer(f, TIER2_THRESHOLD) + self.assertEqual(res, TIER2_THRESHOLD * 2) + self.assertIsNotNone(ex) + uops = get_opnames(ex) + + # v + 1 should be constant folded + self.assertLessEqual(count_ops(ex, "_BINARY_OP_ADD_INT"), 1) + @unittest.skip("gh-139109 WIP") def test_combine_stack_space_checks_sequential(self): def dummy12(x): From ddd202a9c7f29e62e91ef63b2c203c61264a8e38 Mon Sep 17 00:00:00 2001 From: reidenong Date: Thu, 22 Jan 2026 21:12:01 +0800 Subject: [PATCH 2/9] Add constant narrowing for _COMPARE_OP_INT --- Include/internal/pycore_optimizer_types.h | 2 + Python/optimizer_bytecodes.c | 17 ++++++- Python/optimizer_cases.c.h | 14 +++++- Python/optimizer_symbols.c | 54 ++++++++++++++++++++++- 4 files changed, 84 insertions(+), 3 deletions(-) diff --git a/Include/internal/pycore_optimizer_types.h b/Include/internal/pycore_optimizer_types.h index 7e0dbddce2d6b8..e2323265798156 100644 --- a/Include/internal/pycore_optimizer_types.h +++ b/Include/internal/pycore_optimizer_types.h @@ -76,6 +76,8 @@ typedef struct { typedef enum { JIT_PRED_IS, JIT_PRED_IS_NOT, + JIT_PRED_EQ, + JIT_PRED_NE, } JitOptPredicateKind; typedef struct { diff --git a/Python/optimizer_bytecodes.c b/Python/optimizer_bytecodes.c index 0ccc788dff962d..ab7274be8591f9 100644 --- a/Python/optimizer_bytecodes.c +++ b/Python/optimizer_bytecodes.c @@ -514,7 +514,22 @@ dummy_func(void) { } op(_COMPARE_OP_INT, (left, right -- res, l, r)) { - res = sym_new_type(ctx, &PyBool_Type); + /* Comparison oparg masks */ + const int COMPARE_LT_MASK = 2; + const int COMPARE_GT_MASK = 4; + const int COMPARE_EQ_MASK = 8; + + int cmp_mask = oparg & (COMPARE_LT_MASK | COMPARE_GT_MASK | COMPARE_EQ_MASK); + + if (cmp_mask == COMPARE_EQ_MASK) { + res = sym_new_predicate(ctx, left, right, JIT_PRED_EQ); + } + else if (cmp_mask == (COMPARE_LT_MASK | COMPARE_GT_MASK)) { + res = sym_new_predicate(ctx, left, right, JIT_PRED_NE); + } + else { + res = sym_new_type(ctx, &PyBool_Type); + } l = left; r = right; REPLACE_OPCODE_IF_EVALUATES_PURE(left, right, res); diff --git a/Python/optimizer_cases.c.h b/Python/optimizer_cases.c.h index f62e15b987c0eb..8c9c975d84e66a 100644 --- a/Python/optimizer_cases.c.h +++ b/Python/optimizer_cases.c.h @@ -2167,7 +2167,19 @@ JitOptRef r; right = stack_pointer[-1]; left = stack_pointer[-2]; - res = sym_new_type(ctx, &PyBool_Type); + const int COMPARE_LT_MASK = 2; + const int COMPARE_GT_MASK = 4; + const int COMPARE_EQ_MASK = 8; + int cmp_mask = oparg & (COMPARE_LT_MASK | COMPARE_GT_MASK | COMPARE_EQ_MASK); + if (cmp_mask == COMPARE_EQ_MASK) { + res = sym_new_predicate(ctx, left, right, JIT_PRED_EQ); + } + else if (cmp_mask == (COMPARE_LT_MASK | COMPARE_GT_MASK)) { + res = sym_new_predicate(ctx, left, right, JIT_PRED_NE); + } + else { + res = sym_new_type(ctx, &PyBool_Type); + } l = left; r = right; if ( diff --git a/Python/optimizer_symbols.c b/Python/optimizer_symbols.c index a9640aaa5072c5..d599cf8a02cf80 100644 --- a/Python/optimizer_symbols.c +++ b/Python/optimizer_symbols.c @@ -875,9 +875,11 @@ _Py_uop_sym_apply_predicate_narrowing(JitOptContext *ctx, JitOptRef ref, bool br bool narrow = false; switch(pred.kind) { + case JIT_PRED_EQ: case JIT_PRED_IS: narrow = branch_is_true; break; + case JIT_PRED_NE: case JIT_PRED_IS_NOT: narrow = !branch_is_true; break; @@ -1300,7 +1302,7 @@ _Py_uop_symbols_test(PyObject *Py_UNUSED(self), PyObject *Py_UNUSED(ignored)) TEST_PREDICATE(_Py_uop_sym_is_const(ctx, subject), "predicate narrowing did not const-narrow subject (None)"); TEST_PREDICATE(_Py_uop_sym_get_const(ctx, subject) == Py_None, "predicate narrowing did not narrow subject to None"); - // Test narrowing subject to numerical constant + // Test narrowing subject to numerical constant from is comparison subject = _Py_uop_sym_new_unknown(ctx); PyObject *one_obj = PyLong_FromLong(1); JitOptRef const_one = _Py_uop_sym_new_const(ctx, one_obj); @@ -1314,6 +1316,56 @@ _Py_uop_symbols_test(PyObject *Py_UNUSED(self), PyObject *Py_UNUSED(ignored)) _Py_uop_sym_apply_predicate_narrowing(ctx, ref, true); TEST_PREDICATE(_Py_uop_sym_is_const(ctx, subject), "predicate narrowing did not const-narrow subject (1)"); TEST_PREDICATE(_Py_uop_sym_get_const(ctx, subject) == one_obj, "predicate narrowing did not narrow subject to 1"); + + // Test narrowing subject to numerical constant from EQ predicate + subject = _Py_uop_sym_new_unknown(ctx); + if (PyJitRef_IsNull(subject) || PyJitRef_IsNull(const_one)) { + goto fail; + } + ref = _Py_uop_sym_new_predicate(ctx, subject, const_one, JIT_PRED_EQ); + if (PyJitRef_IsNull(ref)) { + goto fail; + } + _Py_uop_sym_apply_predicate_narrowing(ctx, ref, true); + TEST_PREDICATE(_Py_uop_sym_is_const(ctx, subject), "predicate narrowing did not const-narrow subject (1)"); + TEST_PREDICATE(_Py_uop_sym_get_const(ctx, subject) == one_obj, "predicate narrowing did not narrow subject to 1"); + + // Resolving EQ predicate to False should not narrow subject + subject = _Py_uop_sym_new_unknown(ctx); + if (PyJitRef_IsNull(subject)) { + goto fail; + } + ref = _Py_uop_sym_new_predicate(ctx, subject, const_one, JIT_PRED_EQ); + if (PyJitRef_IsNull(ref)) { + goto fail; + } + _Py_uop_sym_apply_predicate_narrowing(ctx, ref, false); + TEST_PREDICATE(!_Py_uop_sym_is_const(ctx, subject), "predicate narrowing incorrectly narrowed subject (inverted/true)"); + + // Test narrowing subject to numerical constant from NE predicate + subject = _Py_uop_sym_new_unknown(ctx); + if (PyJitRef_IsNull(subject) || PyJitRef_IsNull(const_one)) { + goto fail; + } + ref = _Py_uop_sym_new_predicate(ctx, subject, const_one, JIT_PRED_NE); + if (PyJitRef_IsNull(ref)) { + goto fail; + } + _Py_uop_sym_apply_predicate_narrowing(ctx, ref, false); + TEST_PREDICATE(_Py_uop_sym_is_const(ctx, subject), "predicate narrowing did not const-narrow subject (1)"); + TEST_PREDICATE(_Py_uop_sym_get_const(ctx, subject) == one_obj, "predicate narrowing did not narrow subject to 1"); + + // Resolving NE predicate to true should not narrow subject + subject = _Py_uop_sym_new_unknown(ctx); + if (PyJitRef_IsNull(subject)) { + goto fail; + } + ref = _Py_uop_sym_new_predicate(ctx, subject, const_one, JIT_PRED_NE); + if (PyJitRef_IsNull(ref)) { + goto fail; + } + _Py_uop_sym_apply_predicate_narrowing(ctx, ref, true); + TEST_PREDICATE(!_Py_uop_sym_is_const(ctx, subject), "predicate narrowing incorrectly narrowed subject (inverted/true)"); val_big = PyNumber_Lshift(_PyLong_GetOne(), PyLong_FromLong(66)); if (val_big == NULL) { From 030bc286d3171182ddc2e5286af59274e55fefe9 Mon Sep 17 00:00:00 2001 From: reidenong Date: Thu, 22 Jan 2026 22:24:14 +0800 Subject: [PATCH 3/9] Refine unit test names --- Python/optimizer_bytecodes.c | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/Python/optimizer_bytecodes.c b/Python/optimizer_bytecodes.c index ab7274be8591f9..e27c91a30ab3ee 100644 --- a/Python/optimizer_bytecodes.c +++ b/Python/optimizer_bytecodes.c @@ -514,11 +514,6 @@ dummy_func(void) { } op(_COMPARE_OP_INT, (left, right -- res, l, r)) { - /* Comparison oparg masks */ - const int COMPARE_LT_MASK = 2; - const int COMPARE_GT_MASK = 4; - const int COMPARE_EQ_MASK = 8; - int cmp_mask = oparg & (COMPARE_LT_MASK | COMPARE_GT_MASK | COMPARE_EQ_MASK); if (cmp_mask == COMPARE_EQ_MASK) { @@ -536,7 +531,17 @@ dummy_func(void) { } op(_COMPARE_OP_FLOAT, (left, right -- res, l, r)) { - res = sym_new_type(ctx, &PyBool_Type); + int cmp_mask = oparg & (COMPARE_LT_MASK | COMPARE_GT_MASK | COMPARE_EQ_MASK); + + if (cmp_mask == COMPARE_EQ_MASK) { + res = sym_new_predicate(ctx, left, right, JIT_PRED_EQ); + } + else if (cmp_mask == (COMPARE_LT_MASK | COMPARE_GT_MASK)) { + res = sym_new_predicate(ctx, left, right, JIT_PRED_NE); + } + else { + res = sym_new_type(ctx, &PyBool_Type); + } l = left; r = right; REPLACE_OPCODE_IF_EVALUATES_PURE(left, right, res); From 7a03b220c8213b62f07e574e97650244f84c904d Mon Sep 17 00:00:00 2001 From: reidenong Date: Thu, 22 Jan 2026 22:39:00 +0800 Subject: [PATCH 4/9] Add failing unit tests for float EQ, NE --- Lib/test/test_capi/test_opt.py | 43 ++++++++++++++++++++++ Python/optimizer_symbols.c | 66 ++++++++++++++++++++++++++++++---- 2 files changed, 102 insertions(+), 7 deletions(-) diff --git a/Lib/test/test_capi/test_opt.py b/Lib/test/test_capi/test_opt.py index bdf2df2bfd668e..561351dab46995 100644 --- a/Lib/test/test_capi/test_opt.py +++ b/Lib/test/test_capi/test_opt.py @@ -932,6 +932,49 @@ def return_1(): # v + 1 should be constant folded self.assertLessEqual(count_ops(ex, "_BINARY_OP_ADD_INT"), 1) + def test_compare_float_eq_narrows_to_constant(self): + def f(n): + def return_tenth(): + return 0.1 + + hits = 0 + v = return_tenth() + for _ in range(n): + if v == 0.1: + if v == 0.1: + hits += 1 + return hits + + res, ex = self._run_with_optimizer(f, TIER2_THRESHOLD) + self.assertEqual(res, TIER2_THRESHOLD) + self.assertIsNotNone(ex) + uops = get_opnames(ex) + + self.assertLessEqual(count_ops(ex, "_COMPARE_OP_FLOAT"), 1) + + + def test_compare_float_ne_narrows_to_constant(self): + def f(n): + def return_tenth(): + return 0.1 + + hits = 0 + v = return_tenth() + for _ in range(n): + if v != 0.1: + hits += 1000 + else: + if v == 0.1: + hits += 1 + return hits + + res, ex = self._run_with_optimizer(f, TIER2_THRESHOLD) + self.assertEqual(res, TIER2_THRESHOLD) + self.assertIsNotNone(ex) + uops = get_opnames(ex) + + self.assertLessEqual(count_ops(ex, "_COMPARE_OP_FLOAT"), 1) + @unittest.skip("gh-139109 WIP") def test_combine_stack_space_checks_sequential(self): def dummy12(x): diff --git a/Python/optimizer_symbols.c b/Python/optimizer_symbols.c index d599cf8a02cf80..eb2895e91901e0 100644 --- a/Python/optimizer_symbols.c +++ b/Python/optimizer_symbols.c @@ -1306,7 +1306,7 @@ _Py_uop_symbols_test(PyObject *Py_UNUSED(self), PyObject *Py_UNUSED(ignored)) subject = _Py_uop_sym_new_unknown(ctx); PyObject *one_obj = PyLong_FromLong(1); JitOptRef const_one = _Py_uop_sym_new_const(ctx, one_obj); - if (PyJitRef_IsNull(subject) || PyJitRef_IsNull(const_one)) { + if (PyJitRef_IsNull(subject) || one_obj == NULL || PyJitRef_IsNull(const_one)) { goto fail; } ref = _Py_uop_sym_new_predicate(ctx, subject, const_one, JIT_PRED_IS); @@ -1317,9 +1317,9 @@ _Py_uop_symbols_test(PyObject *Py_UNUSED(self), PyObject *Py_UNUSED(ignored)) TEST_PREDICATE(_Py_uop_sym_is_const(ctx, subject), "predicate narrowing did not const-narrow subject (1)"); TEST_PREDICATE(_Py_uop_sym_get_const(ctx, subject) == one_obj, "predicate narrowing did not narrow subject to 1"); - // Test narrowing subject to numerical constant from EQ predicate + // Test narrowing subject to constant from EQ predicate for int subject = _Py_uop_sym_new_unknown(ctx); - if (PyJitRef_IsNull(subject) || PyJitRef_IsNull(const_one)) { + if (PyJitRef_IsNull(subject)) { goto fail; } ref = _Py_uop_sym_new_predicate(ctx, subject, const_one, JIT_PRED_EQ); @@ -1330,7 +1330,7 @@ _Py_uop_symbols_test(PyObject *Py_UNUSED(self), PyObject *Py_UNUSED(ignored)) TEST_PREDICATE(_Py_uop_sym_is_const(ctx, subject), "predicate narrowing did not const-narrow subject (1)"); TEST_PREDICATE(_Py_uop_sym_get_const(ctx, subject) == one_obj, "predicate narrowing did not narrow subject to 1"); - // Resolving EQ predicate to False should not narrow subject + // Resolving EQ predicate to False should not narrow subject for int subject = _Py_uop_sym_new_unknown(ctx); if (PyJitRef_IsNull(subject)) { goto fail; @@ -1342,9 +1342,9 @@ _Py_uop_symbols_test(PyObject *Py_UNUSED(self), PyObject *Py_UNUSED(ignored)) _Py_uop_sym_apply_predicate_narrowing(ctx, ref, false); TEST_PREDICATE(!_Py_uop_sym_is_const(ctx, subject), "predicate narrowing incorrectly narrowed subject (inverted/true)"); - // Test narrowing subject to numerical constant from NE predicate + // Test narrowing subject to constant from NE predicate for int subject = _Py_uop_sym_new_unknown(ctx); - if (PyJitRef_IsNull(subject) || PyJitRef_IsNull(const_one)) { + if (PyJitRef_IsNull(subject)) { goto fail; } ref = _Py_uop_sym_new_predicate(ctx, subject, const_one, JIT_PRED_NE); @@ -1355,7 +1355,7 @@ _Py_uop_symbols_test(PyObject *Py_UNUSED(self), PyObject *Py_UNUSED(ignored)) TEST_PREDICATE(_Py_uop_sym_is_const(ctx, subject), "predicate narrowing did not const-narrow subject (1)"); TEST_PREDICATE(_Py_uop_sym_get_const(ctx, subject) == one_obj, "predicate narrowing did not narrow subject to 1"); - // Resolving NE predicate to true should not narrow subject + // Resolving NE predicate to true should not narrow subject for int subject = _Py_uop_sym_new_unknown(ctx); if (PyJitRef_IsNull(subject)) { goto fail; @@ -1367,6 +1367,58 @@ _Py_uop_symbols_test(PyObject *Py_UNUSED(self), PyObject *Py_UNUSED(ignored)) _Py_uop_sym_apply_predicate_narrowing(ctx, ref, true); TEST_PREDICATE(!_Py_uop_sym_is_const(ctx, subject), "predicate narrowing incorrectly narrowed subject (inverted/true)"); + // Test narrowing subject to constant from EQ predicate for float + subject = _Py_uop_sym_new_unknown(ctx); + PyObject *float_tenth_obj = PyFloat_FromDouble(0.1); + JitOptRef const_float_tenth = _Py_uop_sym_new_const(ctx, float_tenth_obj); + if (PyJitRef_IsNull(subject) || float_tenth_obj == NULL || PyJitRef_IsNull(const_float_tenth)) { + goto fail; + } + ref = _Py_uop_sym_new_predicate(ctx, subject, const_float_tenth, JIT_PRED_EQ); + if (PyJitRef_IsNull(ref)) { + goto fail; + } + _Py_uop_sym_apply_predicate_narrowing(ctx, ref, true); + TEST_PREDICATE(_Py_uop_sym_is_const(ctx, subject), "predicate narrowing did not const-narrow subject (float)"); + TEST_PREDICATE(_Py_uop_sym_get_const(ctx, subject) == float_tenth_obj, "predicate narrowing did not narrow subject to 0.1"); + + // Resolving EQ predicate to False should not narrow subject for float + subject = _Py_uop_sym_new_unknown(ctx); + if (PyJitRef_IsNull(subject)) { + goto fail; + } + ref = _Py_uop_sym_new_predicate(ctx, subject, const_float_tenth, JIT_PRED_EQ); + if (PyJitRef_IsNull(ref)) { + goto fail; + } + _Py_uop_sym_apply_predicate_narrowing(ctx, ref, false); + TEST_PREDICATE(!_Py_uop_sym_is_const(ctx, subject), "predicate narrowing incorrectly narrowed subject (inverted/true)"); + + // Test narrowing subject to constant from NE predicate for float + subject = _Py_uop_sym_new_unknown(ctx); + if (PyJitRef_IsNull(subject)) { + goto fail; + } + ref = _Py_uop_sym_new_predicate(ctx, subject, const_float_tenth, JIT_PRED_NE); + if (PyJitRef_IsNull(ref)) { + goto fail; + } + _Py_uop_sym_apply_predicate_narrowing(ctx, ref, false); + TEST_PREDICATE(_Py_uop_sym_is_const(ctx, subject), "predicate narrowing did not const-narrow subject (float)"); + TEST_PREDICATE(_Py_uop_sym_get_const(ctx, subject) == float_tenth_obj, "predicate narrowing did not narrow subject to 0.1"); + + // Resolving NE predicate to true should not narrow subject for float + subject = _Py_uop_sym_new_unknown(ctx); + if (PyJitRef_IsNull(subject)) { + goto fail; + } + ref = _Py_uop_sym_new_predicate(ctx, subject, const_float_tenth, JIT_PRED_NE); + if (PyJitRef_IsNull(ref)) { + goto fail; + } + _Py_uop_sym_apply_predicate_narrowing(ctx, ref, true); + TEST_PREDICATE(!_Py_uop_sym_is_const(ctx, subject), "predicate narrowing incorrectly narrowed subject (inverted/true)"); + val_big = PyNumber_Lshift(_PyLong_GetOne(), PyLong_FromLong(66)); if (val_big == NULL) { goto fail; From eaf3f236375efa16363d24e67973f4faa32b2ed6 Mon Sep 17 00:00:00 2001 From: reidenong Date: Thu, 22 Jan 2026 22:39:45 +0800 Subject: [PATCH 5/9] Add constant narrowing for float EQ, NE --- Python/optimizer_analysis.c | 5 +++++ Python/optimizer_cases.c.h | 14 ++++++++++---- Python/optimizer_symbols.c | 2 +- 3 files changed, 16 insertions(+), 5 deletions(-) diff --git a/Python/optimizer_analysis.c b/Python/optimizer_analysis.c index e4e259a81b510f..9e012eb3f64878 100644 --- a/Python/optimizer_analysis.c +++ b/Python/optimizer_analysis.c @@ -250,6 +250,11 @@ add_op(JitOptContext *ctx, _PyUOpInstruction *this_instr, #define sym_new_predicate _Py_uop_sym_new_predicate #define sym_apply_predicate_narrowing _Py_uop_sym_apply_predicate_narrowing +/* Comparison oparg masks */ +#define COMPARE_LT_MASK 2 +#define COMPARE_GT_MASK 4 +#define COMPARE_EQ_MASK 8 + #define JUMP_TO_LABEL(label) goto label; static int diff --git a/Python/optimizer_cases.c.h b/Python/optimizer_cases.c.h index 8c9c975d84e66a..cc0c7294a2d990 100644 --- a/Python/optimizer_cases.c.h +++ b/Python/optimizer_cases.c.h @@ -2107,7 +2107,16 @@ JitOptRef r; right = stack_pointer[-1]; left = stack_pointer[-2]; - res = sym_new_type(ctx, &PyBool_Type); + int cmp_mask = oparg & (COMPARE_LT_MASK | COMPARE_GT_MASK | COMPARE_EQ_MASK); + if (cmp_mask == COMPARE_EQ_MASK) { + res = sym_new_predicate(ctx, left, right, JIT_PRED_EQ); + } + else if (cmp_mask == (COMPARE_LT_MASK | COMPARE_GT_MASK)) { + res = sym_new_predicate(ctx, left, right, JIT_PRED_NE); + } + else { + res = sym_new_type(ctx, &PyBool_Type); + } l = left; r = right; if ( @@ -2167,9 +2176,6 @@ JitOptRef r; right = stack_pointer[-1]; left = stack_pointer[-2]; - const int COMPARE_LT_MASK = 2; - const int COMPARE_GT_MASK = 4; - const int COMPARE_EQ_MASK = 8; int cmp_mask = oparg & (COMPARE_LT_MASK | COMPARE_GT_MASK | COMPARE_EQ_MASK); if (cmp_mask == COMPARE_EQ_MASK) { res = sym_new_predicate(ctx, left, right, JIT_PRED_EQ); diff --git a/Python/optimizer_symbols.c b/Python/optimizer_symbols.c index eb2895e91901e0..9fc66655b32f3d 100644 --- a/Python/optimizer_symbols.c +++ b/Python/optimizer_symbols.c @@ -1316,7 +1316,7 @@ _Py_uop_symbols_test(PyObject *Py_UNUSED(self), PyObject *Py_UNUSED(ignored)) _Py_uop_sym_apply_predicate_narrowing(ctx, ref, true); TEST_PREDICATE(_Py_uop_sym_is_const(ctx, subject), "predicate narrowing did not const-narrow subject (1)"); TEST_PREDICATE(_Py_uop_sym_get_const(ctx, subject) == one_obj, "predicate narrowing did not narrow subject to 1"); - + // Test narrowing subject to constant from EQ predicate for int subject = _Py_uop_sym_new_unknown(ctx); if (PyJitRef_IsNull(subject)) { From 2b5e3e0d93220dea4224e6a03f00fde7cd7b32d4 Mon Sep 17 00:00:00 2001 From: reidenong Date: Thu, 22 Jan 2026 22:45:56 +0800 Subject: [PATCH 6/9] Fix whitespace --- Lib/test/test_capi/test_opt.py | 1 - 1 file changed, 1 deletion(-) diff --git a/Lib/test/test_capi/test_opt.py b/Lib/test/test_capi/test_opt.py index 561351dab46995..fac086d84f63e9 100644 --- a/Lib/test/test_capi/test_opt.py +++ b/Lib/test/test_capi/test_opt.py @@ -952,7 +952,6 @@ def return_tenth(): self.assertLessEqual(count_ops(ex, "_COMPARE_OP_FLOAT"), 1) - def test_compare_float_ne_narrows_to_constant(self): def f(n): def return_tenth(): From 57ee7f37ad3904b16847d18d8918a851ef29a534 Mon Sep 17 00:00:00 2001 From: reidenong Date: Thu, 22 Jan 2026 22:50:35 +0800 Subject: [PATCH 7/9] Refine unit tests --- Lib/test/test_capi/test_opt.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/Lib/test/test_capi/test_opt.py b/Lib/test/test_capi/test_opt.py index fac086d84f63e9..15060e50788b94 100644 --- a/Lib/test/test_capi/test_opt.py +++ b/Lib/test/test_capi/test_opt.py @@ -899,16 +899,17 @@ def return_1(): v = return_1() for _ in range(n): if v == 1: - hits += v + 1 + if v == 1: + hits += 1 return hits res, ex = self._run_with_optimizer(f, TIER2_THRESHOLD) - self.assertEqual(res, TIER2_THRESHOLD * 2) + self.assertEqual(res, TIER2_THRESHOLD) self.assertIsNotNone(ex) uops = get_opnames(ex) - # v + 1 should be constant folded - self.assertLessEqual(count_ops(ex, "_BINARY_OP_ADD_INT"), 1) + # Constant narrowing allows constant folding for second comparison + self.assertLessEqual(count_ops(ex, "_COMPARE_OP_INT"), 1) def test_compare_int_ne_narrows_to_constant(self): def f(n): @@ -921,7 +922,8 @@ def return_1(): if v != 1: hits += 1000 else: - hits += v + 1 + if v == 1: + hits += v + 1 return hits res, ex = self._run_with_optimizer(f, TIER2_THRESHOLD) @@ -929,8 +931,8 @@ def return_1(): self.assertIsNotNone(ex) uops = get_opnames(ex) - # v + 1 should be constant folded - self.assertLessEqual(count_ops(ex, "_BINARY_OP_ADD_INT"), 1) + # Constant narrowing allows constant folding for second comparison + self.assertLessEqual(count_ops(ex, "_COMPARE_OP_INT"), 1) def test_compare_float_eq_narrows_to_constant(self): def f(n): @@ -950,6 +952,7 @@ def return_tenth(): self.assertIsNotNone(ex) uops = get_opnames(ex) + # Constant narrowing allows constant folding for second comparison self.assertLessEqual(count_ops(ex, "_COMPARE_OP_FLOAT"), 1) def test_compare_float_ne_narrows_to_constant(self): @@ -972,6 +975,7 @@ def return_tenth(): self.assertIsNotNone(ex) uops = get_opnames(ex) + # Constant narrowing allows constant folding for second comparison self.assertLessEqual(count_ops(ex, "_COMPARE_OP_FLOAT"), 1) @unittest.skip("gh-139109 WIP") From b8b217b9bf7fe0eb109bbe15f9d0b790ae100cde Mon Sep 17 00:00:00 2001 From: reidenong Date: Thu, 22 Jan 2026 23:08:30 +0800 Subject: [PATCH 8/9] Add failing unit tests for str comparison --- Lib/test/test_capi/test_opt.py | 44 ++++++++++++++++++++++++++++ Python/optimizer_symbols.c | 52 ++++++++++++++++++++++++++++++++++ 2 files changed, 96 insertions(+) diff --git a/Lib/test/test_capi/test_opt.py b/Lib/test/test_capi/test_opt.py index 15060e50788b94..e23ea81e7c4565 100644 --- a/Lib/test/test_capi/test_opt.py +++ b/Lib/test/test_capi/test_opt.py @@ -978,6 +978,50 @@ def return_tenth(): # Constant narrowing allows constant folding for second comparison self.assertLessEqual(count_ops(ex, "_COMPARE_OP_FLOAT"), 1) + def test_compare_str_eq_narrows_to_constant(self): + def f(n): + def return_hello(): + return "hello" + + hits = 0 + v = return_hello() + for _ in range(n): + if v == "hello": + if v == "hello": + hits += 1 + return hits + + res, ex = self._run_with_optimizer(f, TIER2_THRESHOLD) + self.assertEqual(res, TIER2_THRESHOLD) + self.assertIsNotNone(ex) + uops = get_opnames(ex) + + # Constant narrowing allows constant folding for second comparison + self.assertLessEqual(count_ops(ex, "_COMPARE_OP_STR"), 1) + + def test_compare_str_ne_narrows_to_constant(self): + def f(n): + def return_hello(): + return "hello" + + hits = 0 + v = return_hello() + for _ in range(n): + if v != "hello": + hits += 1000 + else: + if v == "hello": + hits += 1 + return hits + + res, ex = self._run_with_optimizer(f, TIER2_THRESHOLD) + self.assertEqual(res, TIER2_THRESHOLD) + self.assertIsNotNone(ex) + uops = get_opnames(ex) + + # Constant narrowing allows constant folding for second comparison + self.assertLessEqual(count_ops(ex, "_COMPARE_OP_STR"), 1) + @unittest.skip("gh-139109 WIP") def test_combine_stack_space_checks_sequential(self): def dummy12(x): diff --git a/Python/optimizer_symbols.c b/Python/optimizer_symbols.c index 9fc66655b32f3d..51cf6e189f0f49 100644 --- a/Python/optimizer_symbols.c +++ b/Python/optimizer_symbols.c @@ -1419,6 +1419,58 @@ _Py_uop_symbols_test(PyObject *Py_UNUSED(self), PyObject *Py_UNUSED(ignored)) _Py_uop_sym_apply_predicate_narrowing(ctx, ref, true); TEST_PREDICATE(!_Py_uop_sym_is_const(ctx, subject), "predicate narrowing incorrectly narrowed subject (inverted/true)"); + // Test narrowing subject to constant from EQ predicate for str + subject = _Py_uop_sym_new_unknown(ctx); + PyObject *str_hello_obj = PyUnicode_FromString("hello"); + JitOptRef const_str_hello = _Py_uop_sym_new_const(ctx, str_hello_obj); + if (PyJitRef_IsNull(subject) || str_hello_obj == NULL || PyJitRef_IsNull(const_str_hello)) { + goto fail; + } + ref = _Py_uop_sym_new_predicate(ctx, subject, const_str_hello, JIT_PRED_EQ); + if (PyJitRef_IsNull(ref)) { + goto fail; + } + _Py_uop_sym_apply_predicate_narrowing(ctx, ref, true); + TEST_PREDICATE(_Py_uop_sym_is_const(ctx, subject), "predicate narrowing did not const-narrow subject (str)"); + TEST_PREDICATE(_Py_uop_sym_get_const(ctx, subject) == str_hello_obj, "predicate narrowing did not narrow subject to hello"); + + // Resolving EQ predicate to False should not narrow subject for str + subject = _Py_uop_sym_new_unknown(ctx); + if (PyJitRef_IsNull(subject)) { + goto fail; + } + ref = _Py_uop_sym_new_predicate(ctx, subject, const_str_hello, JIT_PRED_EQ); + if (PyJitRef_IsNull(ref)) { + goto fail; + } + _Py_uop_sym_apply_predicate_narrowing(ctx, ref, false); + TEST_PREDICATE(!_Py_uop_sym_is_const(ctx, subject), "predicate narrowing incorrectly narrowed subject (inverted/true)"); + + // Test narrowing subject to constant from NE predicate for str + subject = _Py_uop_sym_new_unknown(ctx); + if (PyJitRef_IsNull(subject)) { + goto fail; + } + ref = _Py_uop_sym_new_predicate(ctx, subject, const_str_hello, JIT_PRED_NE); + if (PyJitRef_IsNull(ref)) { + goto fail; + } + _Py_uop_sym_apply_predicate_narrowing(ctx, ref, false); + TEST_PREDICATE(_Py_uop_sym_is_const(ctx, subject), "predicate narrowing did not const-narrow subject (str)"); + TEST_PREDICATE(_Py_uop_sym_get_const(ctx, subject) == str_hello_obj, "predicate narrowing did not narrow subject to hello"); + + // Resolving NE predicate to true should not narrow subject for str + subject = _Py_uop_sym_new_unknown(ctx); + if (PyJitRef_IsNull(subject)) { + goto fail; + } + ref = _Py_uop_sym_new_predicate(ctx, subject, const_str_hello, JIT_PRED_NE); + if (PyJitRef_IsNull(ref)) { + goto fail; + } + _Py_uop_sym_apply_predicate_narrowing(ctx, ref, true); + TEST_PREDICATE(!_Py_uop_sym_is_const(ctx, subject), "predicate narrowing incorrectly narrowed subject (inverted/true)"); + val_big = PyNumber_Lshift(_PyLong_GetOne(), PyLong_FromLong(66)); if (val_big == NULL) { goto fail; From 08c1464225b1af3fbea377b1062b8e9b8ef89758 Mon Sep 17 00:00:00 2001 From: reidenong Date: Thu, 22 Jan 2026 23:08:46 +0800 Subject: [PATCH 9/9] Add constant narrowing for str EQ, NE --- Python/optimizer_bytecodes.c | 12 +++++++++++- Python/optimizer_cases.c.h | 11 ++++++++++- 2 files changed, 21 insertions(+), 2 deletions(-) diff --git a/Python/optimizer_bytecodes.c b/Python/optimizer_bytecodes.c index e27c91a30ab3ee..2fdc0854cc27e6 100644 --- a/Python/optimizer_bytecodes.c +++ b/Python/optimizer_bytecodes.c @@ -548,7 +548,17 @@ dummy_func(void) { } op(_COMPARE_OP_STR, (left, right -- res, l, r)) { - res = sym_new_type(ctx, &PyBool_Type); + int cmp_mask = oparg & (COMPARE_LT_MASK | COMPARE_GT_MASK | COMPARE_EQ_MASK); + + if (cmp_mask == COMPARE_EQ_MASK) { + res = sym_new_predicate(ctx, left, right, JIT_PRED_EQ); + } + else if (cmp_mask == (COMPARE_LT_MASK | COMPARE_GT_MASK)) { + res = sym_new_predicate(ctx, left, right, JIT_PRED_NE); + } + else { + res = sym_new_type(ctx, &PyBool_Type); + } l = left; r = right; REPLACE_OPCODE_IF_EVALUATES_PURE(left, right, res); diff --git a/Python/optimizer_cases.c.h b/Python/optimizer_cases.c.h index cc0c7294a2d990..f9cb376d4c2ed3 100644 --- a/Python/optimizer_cases.c.h +++ b/Python/optimizer_cases.c.h @@ -2249,7 +2249,16 @@ JitOptRef r; right = stack_pointer[-1]; left = stack_pointer[-2]; - res = sym_new_type(ctx, &PyBool_Type); + int cmp_mask = oparg & (COMPARE_LT_MASK | COMPARE_GT_MASK | COMPARE_EQ_MASK); + if (cmp_mask == COMPARE_EQ_MASK) { + res = sym_new_predicate(ctx, left, right, JIT_PRED_EQ); + } + else if (cmp_mask == (COMPARE_LT_MASK | COMPARE_GT_MASK)) { + res = sym_new_predicate(ctx, left, right, JIT_PRED_NE); + } + else { + res = sym_new_type(ctx, &PyBool_Type); + } l = left; r = right; if (