gh-130415: Narrow types to constants in branches involving specialized comparisons with a constant (GH-144150)

Author: reidenong

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 {

Lib/test/test_capi/test_opt.py

@@ -890,6 +890,138 @@ 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:
+                    if v == 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)
+
+        # 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):
+            def return_1():
+                return 1
+
+            hits = 0
+            v = return_1()
+            for _ in range(n):
+                if v != 1:
+                    hits += 1000
+                else:
+                    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)
+
+        # 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):
+            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)
+
+        # 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):
+        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)
+
+        # 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):

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

Python/optimizer_bytecodes.c

@@ -521,21 +521,51 @@ dummy_func(void) {
     }
 
     op(_COMPARE_OP_INT, (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);
     }
 
     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);
     }
 
     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);