gh-130415: Narrowing to constants in branches involving `is` comparisons with a const...

diff --git a/Include/internal/pycore_optimizer.h b/Include/internal/pycore_optimizer.h
index 0592221f15226e4ae819eb1494e0ce35d20ebe67..fbe403b492d5ac964c654197b9b2a3447dec4e4c 100644 (file)
--- a/Include/internal/pycore_optimizer.h
+++ b/Include/internal/pycore_optimizer.h
@@ -205,6 +205,8 @@ extern JitOptRef _Py_uop_sym_new_truthiness(JitOptContext *ctx, JitOptRef value,
 extern bool _Py_uop_sym_is_compact_int(JitOptRef sym);
 extern JitOptRef _Py_uop_sym_new_compact_int(JitOptContext *ctx);
 extern void _Py_uop_sym_set_compact_int(JitOptContext *ctx,  JitOptRef sym);
+extern JitOptRef _Py_uop_sym_new_predicate(JitOptContext *ctx, JitOptRef lhs_ref, JitOptRef rhs_ref, JitOptPredicateKind kind);
+extern void _Py_uop_sym_apply_predicate_narrowing(JitOptContext *ctx, JitOptRef sym, bool branch_is_true);
 
 extern void _Py_uop_abstractcontext_init(JitOptContext *ctx);
 extern void _Py_uop_abstractcontext_fini(JitOptContext *ctx);

diff --git a/Include/internal/pycore_optimizer_types.h b/Include/internal/pycore_optimizer_types.h
index 6501ce869c1425266440ad44a77c951ef7fa56f9..7e0dbddce2d6b80b35d376befd597edb6d4e1763 100644 (file)
--- a/Include/internal/pycore_optimizer_types.h
+++ b/Include/internal/pycore_optimizer_types.h
@@ -40,6 +40,7 @@ typedef enum _JitSymType {
     JIT_SYM_TUPLE_TAG = 8,
     JIT_SYM_TRUTHINESS_TAG = 9,
     JIT_SYM_COMPACT_INT = 10,
+    JIT_SYM_PREDICATE_TAG = 11,
 } JitSymType;
 
 typedef struct _jit_opt_known_class {
@@ -72,6 +73,18 @@ typedef struct {
     uint16_t value;
 } JitOptTruthiness;
 
+typedef enum {
+    JIT_PRED_IS,
+    JIT_PRED_IS_NOT,
+} JitOptPredicateKind;
+
+typedef struct {
+    uint8_t tag;
+    uint8_t kind;
+    uint16_t lhs;
+    uint16_t rhs;
+} JitOptPredicate;
+
 typedef struct {
     uint8_t tag;
 } JitOptCompactInt;
@@ -84,6 +97,7 @@ typedef union _jit_opt_symbol {
     JitOptTuple tuple;
     JitOptTruthiness truthiness;
     JitOptCompactInt compact;
+    JitOptPredicate predicate;
 } JitOptSymbol;
 
 // This mimics the _PyStackRef API

diff --git a/Lib/test/test_capi/test_opt.py b/Lib/test/test_capi/test_opt.py
index 307eac6ee51756d114a78e404ad85bc0575f84ca..7c33320e9f1785690ae2e79ba78440b21bdd1468 100644 (file)
--- a/Lib/test/test_capi/test_opt.py
+++ b/Lib/test/test_capi/test_opt.py
@@ -3551,6 +3551,46 @@ class TestUopsOptimization(unittest.TestCase):
         self.assertIn("_POP_TOP_NOP", uops)
         self.assertLessEqual(count_ops(ex, "_POP_TOP"), 2)
 
+    def test_is_true_narrows_to_constant(self):
+        def f(n):
+            def return_true():
+                return True
+
+            hits = 0
+            v = return_true()
+            for i in range(n):
+                if v is True:
+                    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.assertNotIn("_BINARY_OP", uops)
+
+    def test_is_false_narrows_to_constant(self):
+        def f(n):
+            def return_false():
+                return False
+
+            hits = 0
+            v = return_false()
+            for i in range(n):
+                if v is False:
+                    hits += v + 1
+            return hits
+
+        res, ex = self._run_with_optimizer(f, TIER2_THRESHOLD)
+        self.assertEqual(res, TIER2_THRESHOLD)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+
+        # v + 1 should be constant folded
+        self.assertNotIn("_BINARY_OP", uops)
+
     def test_for_iter_gen_frame(self):
         def f(n):
             for i in range(n):

diff --git a/Python/optimizer_analysis.c b/Python/optimizer_analysis.c
index 01fd24564f8c2bacc904d750f3d875c66a7fb455..e4e259a81b510f6a8358848f20fc4c1422bbb36a 100644 (file)
--- a/Python/optimizer_analysis.c
+++ b/Python/optimizer_analysis.c
@@ -247,6 +247,8 @@ add_op(JitOptContext *ctx, _PyUOpInstruction *this_instr,
 #define sym_is_compact_int _Py_uop_sym_is_compact_int
 #define sym_new_compact_int _Py_uop_sym_new_compact_int
 #define sym_new_truthiness _Py_uop_sym_new_truthiness
+#define sym_new_predicate _Py_uop_sym_new_predicate
+#define sym_apply_predicate_narrowing _Py_uop_sym_apply_predicate_narrowing
 
 #define JUMP_TO_LABEL(label) goto label;
 

diff --git a/Python/optimizer_bytecodes.c b/Python/optimizer_bytecodes.c
index 366094d939a3966e08d39b8edd34b2f34e7e752b..0ccc788dff962d901f94d4b7a9389ff3221f3937 100644 (file)
--- a/Python/optimizer_bytecodes.c
+++ b/Python/optimizer_bytecodes.c
@@ -38,6 +38,8 @@ typedef struct _Py_UOpsAbstractFrame _Py_UOpsAbstractFrame;
 #define sym_new_compact_int _Py_uop_sym_new_compact_int
 #define sym_is_compact_int _Py_uop_sym_is_compact_int
 #define sym_new_truthiness _Py_uop_sym_new_truthiness
+#define sym_new_predicate _Py_uop_sym_new_predicate
+#define sym_apply_predicate_narrowing _Py_uop_sym_apply_predicate_narrowing
 
 extern int
 optimize_to_bool(
@@ -533,7 +535,7 @@ dummy_func(void) {
     }
 
     op(_IS_OP, (left, right -- b, l, r)) {
-        b = sym_new_type(ctx, &PyBool_Type);
+        b = sym_new_predicate(ctx, left, right, (oparg ? JIT_PRED_IS_NOT : JIT_PRED_IS));
         l = left;
         r = right;
     }
@@ -1173,6 +1175,8 @@ dummy_func(void) {
     }
 
     op(_GUARD_IS_TRUE_POP, (flag -- )) {
+        sym_apply_predicate_narrowing(ctx, flag, true);
+
         if (sym_is_const(ctx, flag)) {
             PyObject *value = sym_get_const(ctx, flag);
             assert(value != NULL);
@@ -1191,6 +1195,8 @@ dummy_func(void) {
     }
 
     op(_GUARD_IS_FALSE_POP, (flag -- )) {
+        sym_apply_predicate_narrowing(ctx, flag, false);
+
         if (sym_is_const(ctx, flag)) {
             PyObject *value = sym_get_const(ctx, flag);
             assert(value != NULL);

diff --git a/Python/optimizer_cases.c.h b/Python/optimizer_cases.c.h
index 2cc53937925a24f9607a0bd95212ff6223dd7347..f62e15b987c0eb717aa7653c69c54c9f25a7b0f0 100644 (file)
--- a/Python/optimizer_cases.c.h
+++ b/Python/optimizer_cases.c.h
@@ -2293,7 +2293,7 @@
             JitOptRef r;
             right = stack_pointer[-1];
             left = stack_pointer[-2];
-            b = sym_new_type(ctx, &PyBool_Type);
+            b = sym_new_predicate(ctx, left, right, (oparg ? JIT_PRED_IS_NOT : JIT_PRED_IS));
             l = left;
             r = right;
             CHECK_STACK_BOUNDS(1);
@@ -3715,6 +3715,7 @@
         case _GUARD_IS_TRUE_POP: {
             JitOptRef flag;
             flag = stack_pointer[-1];
+            sym_apply_predicate_narrowing(ctx, flag, true);
             if (sym_is_const(ctx, flag)) {
                 PyObject *value = sym_get_const(ctx, flag);
                 assert(value != NULL);
@@ -3739,6 +3740,7 @@
         case _GUARD_IS_FALSE_POP: {
             JitOptRef flag;
             flag = stack_pointer[-1];
+            sym_apply_predicate_narrowing(ctx, flag, false);
             if (sym_is_const(ctx, flag)) {
                 PyObject *value = sym_get_const(ctx, flag);
                 assert(value != NULL);

diff --git a/Python/optimizer_symbols.c b/Python/optimizer_symbols.c
index 5f5086d33b5c4c1ab824d6f2ae9e1a40ca097aae..a9640aaa5072c5b7e3dfdd5c61acaa7747ef3155 100644 (file)
--- a/Python/optimizer_symbols.c
+++ b/Python/optimizer_symbols.c
@@ -25,24 +25,24 @@ state represents no information, and the BOTTOM state represents contradictory
 information. Though symbols logically progress through all intermediate nodes,
 we often skip in-between states for convenience:
 
-   UNKNOWN
-   |     |
-NULL     |
-|        |                <- Anything below this level is an object.
-|        NON_NULL-+
-|          |      |       <- Anything below this level has a known type version.
-|    TYPE_VERSION |
-|    |            |       <- Anything below this level has a known type.
-|    KNOWN_CLASS  |
-|    |  |  |   |  |
-|    |  | INT* |  |
-|    |  |  |   |  |       <- Anything below this level has a known truthiness.
-|    |  |  |   |  TRUTHINESS
-|    |  |  |   |  |
-| TUPLE |  |   |  |
-|    |  |  |   |  |       <- Anything below this level is a known constant.
-|    KNOWN_VALUE--+
-|    |                    <- Anything below this level is unreachable.
+   UNKNOWN-------------------+
+   |     |                   |
+NULL     |                   |
+|        |                   |   <- Anything below this level is an object.
+|        NON_NULL-+          |
+|          |      |          |   <- Anything below this level has a known type version.
+|    TYPE_VERSION |          |
+|    |            |          |   <- Anything below this level has a known type.
+|    KNOWN_CLASS  |          |
+|    |  |  |   |  |          PREDICATE
+|    |  | INT* |  |          |
+|    |  |  |   |  |          |   <- Anything below this level has a known truthiness.
+|    |  |  |   |  TRUTHINESS |
+|    |  |  |   |  |          |
+| TUPLE |  |   |  |          |
+|    |  |  |   |  |          |   <- Anything below this level is a known constant.
+|    KNOWN_VALUE--+----------+
+|    |                           <- Anything below this level is unreachable.
 BOTTOM
 
 For example, after guarding that the type of an UNKNOWN local is int, we can
@@ -309,6 +309,7 @@ _Py_uop_sym_set_type(JitOptContext *ctx, JitOptRef ref, PyTypeObject *typ)
             sym->cls.version = 0;
             sym->cls.type = typ;
             return;
+        case JIT_SYM_PREDICATE_TAG:
         case JIT_SYM_TRUTHINESS_TAG:
             if (typ != &PyBool_Type) {
                 sym_set_bottom(ctx, sym);
@@ -370,6 +371,7 @@ _Py_uop_sym_set_type_version(JitOptContext *ctx, JitOptRef ref, unsigned int ver
             sym->tag = JIT_SYM_TYPE_VERSION_TAG;
             sym->version.version = version;
             return true;
+        case JIT_SYM_PREDICATE_TAG:
         case JIT_SYM_TRUTHINESS_TAG:
             if (version != PyBool_Type.tp_version_tag) {
                 sym_set_bottom(ctx, sym);
@@ -436,6 +438,13 @@ _Py_uop_sym_set_const(JitOptContext *ctx, JitOptRef ref, PyObject *const_val)
         case JIT_SYM_UNKNOWN_TAG:
             make_const(sym, const_val);
             return;
+        case JIT_SYM_PREDICATE_TAG:
+            if (!PyBool_Check(const_val)) {
+                sym_set_bottom(ctx, sym);
+                return;
+            }
+            make_const(sym, const_val);
+            return;
         case JIT_SYM_TRUTHINESS_TAG:
             if (!PyBool_Check(const_val) ||
                 (_Py_uop_sym_is_const(ctx, ref) &&
@@ -589,6 +598,7 @@ _Py_uop_sym_get_type(JitOptRef ref)
             return _PyType_LookupByVersion(sym->version.version);
         case JIT_SYM_TUPLE_TAG:
             return &PyTuple_Type;
+        case JIT_SYM_PREDICATE_TAG:
         case JIT_SYM_TRUTHINESS_TAG:
             return &PyBool_Type;
         case JIT_SYM_COMPACT_INT:
@@ -617,6 +627,7 @@ _Py_uop_sym_get_type_version(JitOptRef ref)
             return Py_TYPE(sym->value.value)->tp_version_tag;
         case JIT_SYM_TUPLE_TAG:
             return PyTuple_Type.tp_version_tag;
+        case JIT_SYM_PREDICATE_TAG:
         case JIT_SYM_TRUTHINESS_TAG:
             return PyBool_Type.tp_version_tag;
         case JIT_SYM_COMPACT_INT:
@@ -810,6 +821,7 @@ _Py_uop_sym_set_compact_int(JitOptContext *ctx, JitOptRef ref)
             }
             return;
         case JIT_SYM_TUPLE_TAG:
+        case JIT_SYM_PREDICATE_TAG:
         case JIT_SYM_TRUTHINESS_TAG:
             sym_set_bottom(ctx, sym);
             return;
@@ -823,6 +835,70 @@ _Py_uop_sym_set_compact_int(JitOptContext *ctx, JitOptRef ref)
     }
 }
 
+JitOptRef
+_Py_uop_sym_new_predicate(JitOptContext *ctx, JitOptRef lhs_ref, JitOptRef rhs_ref, JitOptPredicateKind kind)
+{
+    JitOptSymbol *lhs = PyJitRef_Unwrap(lhs_ref);
+    JitOptSymbol *rhs = PyJitRef_Unwrap(rhs_ref);
+
+    JitOptSymbol *res = sym_new(ctx);
+    if (res == NULL) {
+        return out_of_space_ref(ctx);
+    }
+
+    res->tag = JIT_SYM_PREDICATE_TAG;
+    res->predicate.kind = kind;
+    res->predicate.lhs = (uint16_t)(lhs - allocation_base(ctx));
+    res->predicate.rhs = (uint16_t)(rhs - allocation_base(ctx));
+
+    return PyJitRef_Wrap(res);
+}
+
+void
+_Py_uop_sym_apply_predicate_narrowing(JitOptContext *ctx, JitOptRef ref, bool branch_is_true)
+{
+    JitOptSymbol *sym = PyJitRef_Unwrap(ref);
+    if (sym->tag != JIT_SYM_PREDICATE_TAG) {
+        return;
+    }
+
+    JitOptPredicate pred = sym->predicate;
+
+    JitOptRef lhs_ref = PyJitRef_Wrap(allocation_base(ctx) + pred.lhs);
+    JitOptRef rhs_ref = PyJitRef_Wrap(allocation_base(ctx) + pred.rhs);
+
+    bool lhs_is_const = _Py_uop_sym_is_const(ctx, lhs_ref);
+    bool rhs_is_const = _Py_uop_sym_is_const(ctx, rhs_ref);
+    if (!lhs_is_const && !rhs_is_const) {
+        return;
+    }
+
+    bool narrow = false;
+    switch(pred.kind) {
+        case JIT_PRED_IS:
+            narrow = branch_is_true;
+            break;
+        case JIT_PRED_IS_NOT:
+            narrow = !branch_is_true;
+            break;
+        default:
+            return;
+    }
+    if (!narrow) {
+        return;
+    }
+
+    JitOptRef subject_ref = lhs_is_const ? rhs_ref : lhs_ref;
+    JitOptRef const_ref = lhs_is_const ? lhs_ref : rhs_ref;
+
+    PyObject *const_val = _Py_uop_sym_get_const(ctx, const_ref);
+    if (const_val == NULL) {
+        return;
+    }
+    _Py_uop_sym_set_const(ctx, subject_ref, const_val);
+    assert(_Py_uop_sym_is_const(ctx, subject_ref));
+}
+
 JitOptRef
 _Py_uop_sym_new_truthiness(JitOptContext *ctx, JitOptRef ref, bool truthy)
 {
@@ -1159,6 +1235,85 @@ _Py_uop_symbols_test(PyObject *Py_UNUSED(self), PyObject *Py_UNUSED(ignored))
     TEST_PREDICATE(_Py_uop_sym_is_const(ctx, value) == true, "value is not constant");
     TEST_PREDICATE(_Py_uop_sym_get_const(ctx, value) == Py_True, "value is not True");
 
+    // Resolving predicate result to True should narrow subject to True
+    JitOptRef subject = _Py_uop_sym_new_unknown(ctx);
+    JitOptRef const_true = _Py_uop_sym_new_const(ctx, Py_True);
+    if (PyJitRef_IsNull(subject) || PyJitRef_IsNull(const_true)) {
+        goto fail;
+    }
+    ref = _Py_uop_sym_new_predicate(ctx, subject, const_true, JIT_PRED_IS);
+    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");
+    TEST_PREDICATE(_Py_uop_sym_get_const(ctx, subject) == Py_True, "predicate narrowing did not narrow subject to True");
+
+    // Resolving predicate result 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_true, JIT_PRED_IS);
+    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");
+
+    // Resolving inverted predicate to False should narrow subject to True
+    subject = _Py_uop_sym_new_unknown(ctx);
+    if (PyJitRef_IsNull(subject)) {
+        goto fail;
+    }
+    ref = _Py_uop_sym_new_predicate(ctx, subject, const_true, JIT_PRED_IS_NOT);
+    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 (inverted) did not const-narrow subject");
+    TEST_PREDICATE(_Py_uop_sym_get_const(ctx, subject) == Py_True, "predicate narrowing (inverted) did not narrow subject to True");
+
+    // Resolving inverted 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_true, JIT_PRED_IS_NOT);
+    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)");
+
+    // Test narrowing subject to None
+    subject = _Py_uop_sym_new_unknown(ctx);
+    JitOptRef const_none = _Py_uop_sym_new_const(ctx, Py_None);
+    if (PyJitRef_IsNull(subject) || PyJitRef_IsNull(const_none)) {
+        goto fail;
+    }
+    ref = _Py_uop_sym_new_predicate(ctx, subject, const_none, JIT_PRED_IS);
+    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 (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
+    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)) {
+        goto fail;
+    }
+    ref = _Py_uop_sym_new_predicate(ctx, subject, const_one, JIT_PRED_IS);
+    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");
 
     val_big = PyNumber_Lshift(_PyLong_GetOne(), PyLong_FromLong(66));
     if (val_big == NULL) {