GH-130415: Use boolean guards to narrow types to values in the JIT (GH-130659)

Author: brandtbucher

Include/internal/pycore_optimizer.h

@@ -172,6 +172,7 @@ typedef enum _JitSymType {
     JIT_SYM_KNOWN_CLASS_TAG = 6,
     JIT_SYM_KNOWN_VALUE_TAG = 7,
     JIT_SYM_TUPLE_TAG = 8,
+    JIT_SYM_TRUTHINESS_TAG = 9,
 } JitSymType;
 
 typedef struct _jit_opt_known_class {
@@ -198,12 +199,19 @@ typedef struct _jit_opt_tuple {
     uint16_t items[MAX_SYMBOLIC_TUPLE_SIZE];
 } JitOptTuple;
 
+typedef struct {
+    uint8_t tag;
+    bool not;
+    uint16_t value;
+} JitOptTruthiness;
+
 typedef union _jit_opt_symbol {
     uint8_t tag;
     JitOptKnownClass cls;
     JitOptKnownValue value;
     JitOptKnownVersion version;
     JitOptTuple tuple;
+    JitOptTruthiness truthiness;
 } JitOptSymbol;
 
 
@@ -245,8 +253,8 @@ typedef struct _JitOptContext {
 
 extern bool _Py_uop_sym_is_null(JitOptSymbol *sym);
 extern bool _Py_uop_sym_is_not_null(JitOptSymbol *sym);
-extern bool _Py_uop_sym_is_const(JitOptSymbol *sym);
-extern PyObject *_Py_uop_sym_get_const(JitOptSymbol *sym);
+extern bool _Py_uop_sym_is_const(JitOptContext *ctx, JitOptSymbol *sym);
+extern PyObject *_Py_uop_sym_get_const(JitOptContext *ctx, JitOptSymbol *sym);
 extern JitOptSymbol *_Py_uop_sym_new_unknown(JitOptContext *ctx);
 extern JitOptSymbol *_Py_uop_sym_new_not_null(JitOptContext *ctx);
 extern JitOptSymbol *_Py_uop_sym_new_type(
@@ -262,12 +270,13 @@ extern void _Py_uop_sym_set_type(JitOptContext *ctx, JitOptSymbol *sym, PyTypeOb
 extern bool _Py_uop_sym_set_type_version(JitOptContext *ctx, JitOptSymbol *sym, unsigned int version);
 extern void _Py_uop_sym_set_const(JitOptContext *ctx, JitOptSymbol *sym, PyObject *const_val);
 extern bool _Py_uop_sym_is_bottom(JitOptSymbol *sym);
-extern int _Py_uop_sym_truthiness(JitOptSymbol *sym);
+extern int _Py_uop_sym_truthiness(JitOptContext *ctx, JitOptSymbol *sym);
 extern PyTypeObject *_Py_uop_sym_get_type(JitOptSymbol *sym);
 extern bool _Py_uop_sym_is_immortal(JitOptSymbol *sym);
 extern JitOptSymbol *_Py_uop_sym_new_tuple(JitOptContext *ctx, int size, JitOptSymbol **args);
 extern JitOptSymbol *_Py_uop_sym_tuple_getitem(JitOptContext *ctx, JitOptSymbol *sym, int item);
 extern int _Py_uop_sym_tuple_length(JitOptSymbol *sym);
+extern JitOptSymbol *_Py_uop_sym_new_truthiness(JitOptContext *ctx, JitOptSymbol *value, bool truthy);
 
 extern void _Py_uop_abstractcontext_init(JitOptContext *ctx);
 extern void _Py_uop_abstractcontext_fini(JitOptContext *ctx);

Lib/test/test_capi/test_opt.py

@@ -1437,6 +1437,68 @@ def crash_addition():
 
         crash_addition()
 
+    def test_narrow_type_to_constant_bool_false(self):
+        def f(n):
+            trace = []
+            for i in range(n):
+                # false is always False, but we can only prove that it's a bool:
+                false = i == TIER2_THRESHOLD
+                trace.append("A")
+                if not false:  # Kept.
+                    trace.append("B")
+                    if not false:  # Removed!
+                        trace.append("C")
+                    trace.append("D")
+                    if false:  # Removed!
+                        trace.append("X")
+                    trace.append("E")
+                trace.append("F")
+                if false:  # Removed!
+                    trace.append("X")
+                trace.append("G")
+            return trace
+
+        trace, ex = self._run_with_optimizer(f, TIER2_THRESHOLD)
+        self.assertEqual(trace, list("ABCDEFG") * TIER2_THRESHOLD)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+        # Only one guard remains:
+        self.assertEqual(uops.count("_GUARD_IS_FALSE_POP"), 1)
+        self.assertEqual(uops.count("_GUARD_IS_TRUE_POP"), 0)
+        # 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_bool_true(self):
+        def f(n):
+            trace = []
+            for i in range(n):
+                # true always True, but we can only prove that it's a bool:
+                true = i != TIER2_THRESHOLD
+                trace.append("A")
+                if true:  # Kept.
+                    trace.append("B")
+                    if not true:  # Removed!
+                        trace.append("X")
+                    trace.append("C")
+                    if true:  # Removed!
+                        trace.append("D")
+                    trace.append("E")
+                trace.append("F")
+                if not true:  # Removed!
+                    trace.append("X")
+                trace.append("G")
+            return trace
+
+        trace, ex = self._run_with_optimizer(f, TIER2_THRESHOLD)
+        self.assertEqual(trace, list("ABCDEFG") * TIER2_THRESHOLD)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+        # Only one guard remains:
+        self.assertEqual(uops.count("_GUARD_IS_FALSE_POP"), 0)
+        self.assertEqual(uops.count("_GUARD_IS_TRUE_POP"), 1)
+        # But all of the appends we care about are still there:
+        self.assertEqual(uops.count("_CALL_LIST_APPEND"), len("ABCDEFG"))
+
 
 def global_identity(x):
     return x

Misc/NEWS.d/next/Core_and_Builtins/2025-02-27-17-05-05.gh-issue-130415.iijvfW.rst

@@ -0,0 +1,2 @@
+Improve the experimental JIT's ability to narrow boolean values based on the
+results of truthiness tests.

Python/optimizer_analysis.c

@@ -136,26 +136,6 @@ incorrect_keys(_PyUOpInstruction *inst, PyObject *obj)
     return 0;
 }
 
-static int
-check_next_uop(_PyUOpInstruction *buffer, int size, int pc, uint16_t expected)
-{
-    if (pc + 1 >= size) {
-        DPRINTF(1, "Cannot rewrite %s at pc %d: buffer too small\n",
-                _PyOpcode_uop_name[buffer[pc].opcode], pc);
-        return 0;
-    }
-    uint16_t next_opcode = buffer[pc + 1].opcode;
-    if (next_opcode != expected) {
-        DPRINTF(1,
-                "Cannot rewrite %s at pc %d: unexpected next opcode %s, "
-                "expected %s\n",
-                _PyOpcode_uop_name[buffer[pc].opcode], pc,
-                _PyOpcode_uop_name[next_opcode], _PyOpcode_uop_name[expected]);
-        return 0;
-    }
-    return 1;
-}
-
 /* Returns 1 if successfully optimized
  *         0 if the trace is not suitable for optimization (yet)
  *        -1 if there was an error. */
@@ -363,6 +343,7 @@ remove_globals(_PyInterpreterFrame *frame, _PyUOpInstruction *buffer,
 #define sym_tuple_getitem _Py_uop_sym_tuple_getitem
 #define sym_tuple_length _Py_uop_sym_tuple_length
 #define sym_is_immortal _Py_uop_sym_is_immortal
+#define sym_new_truthiness _Py_uop_sym_new_truthiness
 
 static int
 optimize_to_bool(
@@ -376,7 +357,7 @@ optimize_to_bool(
         *result_ptr = value;
         return 1;
     }
-    int truthiness = sym_truthiness(value);
+    int truthiness = sym_truthiness(ctx, value);
     if (truthiness >= 0) {
         PyObject *load = truthiness ? Py_True : Py_False;
         REPLACE_OP(this_instr, _POP_TOP_LOAD_CONST_INLINE_BORROW, 0, (uintptr_t)load);