gh-130415: Narrowing to constants in branches involving is comparisons with a constant (GH-143895)

This commit is contained in:
reiden 2026-01-22 17:37:45 +08:00 committed by GitHub
parent 6181b69970
commit 0b08438ea6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 241 additions and 20 deletions

View File

@ -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);

View File

@ -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

View File

@ -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):

View File

@ -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;

View File

@ -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);

View File

@ -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);

View File

@ -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) {